@@ -57,18 +57,16 @@ struct PacketizeJointMatrix {
5757
5858 /* ! @brief Performs a coalesced non-vectorized load when the current block is
5959 * not internal.
60- * @tparam trans Whether the source matrix is transposed or not.
6160 * @tparam internal True if the current block is internal and no bounds
6261 * checking is required.
63- * @tparam ld The leading dimension of the destination memory.
6462 */
6563
66- template <bool trans, bool internal, int ld , typename SrcPointerType ,
67- typename DestPointerType, typename EdgePredicate>
64+ template <bool internal, typename SrcPointerType , typename DestPointerType ,
65+ typename EdgePredicate>
6866 static PORTBLAS_INLINE typename std::enable_if<!internal>::type load (
6967 const bool in_range, SrcPointerType src, DestPointerType dest,
7068 EdgePredicate) {
71- value_t val = in_range ? *( src) : value_t {0 };
69+ value_t val = in_range ? *src : value_t {0 };
7270 using address_t = cl::sycl::access::address_space;
7371 if constexpr (std::is_same<cl::sycl::multi_ptr<cl::sycl::half,
7472 address_t ::local_space>,
@@ -79,93 +77,96 @@ struct PacketizeJointMatrix {
7977 cl::sycl::ext::oneapi::bfloat16,
8078 address_t ::local_space>,
8179 DestPointerType>::value) {
82- using dtype = cl::sycl::ext::oneapi::bfloat16 ;
83- *dest = static_cast <dtype> (val);
80+ using namespace cl ::sycl::ext::oneapi;
81+ *dest = bfloat16 (val);
8482 } else {
8583 using namespace cl ::sycl::ext::oneapi::experimental::matrix;
8684 *dest = round_to_tf32 (val);
8785 }
8886 }
87+
8988 /* ! @brief Performs a vectorised load using sycl::vec::load when the current
9089 * block is internal. In the case where k < the
9190 * number of elements being loaded then edge loads will be element wise with
9291 * additional bounds checking.
93- * @tparam trans Whether the source matrix is transposed or not.
9492 * @tparam internal True if the current block is internal and no bounds
9593 * checking is required.
96- * @tparam ld The leading dimension of the destination memory. * /
97- template <bool trans, bool internal, index_t ld , typename SrcPointerType ,
98- typename DestPointerType, typename EdgePredicate>
94+ */
95+ template <bool internal, typename SrcPointerType , typename DestPointerType ,
96+ typename EdgePredicate>
9997 static PORTBLAS_INLINE typename std::enable_if<internal>::type load (
10098 const bool in_range, SrcPointerType src, DestPointerType dest,
10199 EdgePredicate edge_in_range) {
102100 PacketType packet{};
103101
102+ using address_t = cl::sycl::access::address_space;
104103 if (in_range) {
105- using address_t = cl::sycl::access::address_space;
106104 packet.template load <address_t ::global_space>(
107105 0 , cl::sycl::multi_ptr<const value_t , address_t ::global_space>(src));
106+ store (packet, dest);
108107 } else {
108+ // avoid writing to variable, instead directly write to
109+ // shared local memory to avoid race condition experienced
110+ // with release compiler.
109111#pragma unroll
110- for (index_t i = 0 ; i < packet_size; i++) {
111- reinterpret_cast <value_t *>(&packet)[i] =
112- edge_in_range (i) ? *(src + i) : value_t {0 };
113- }
114- }
115- store<trans, ld>(packet, dest);
116- }
117- /* ! @brief Store a vector packet into local memory when the source is
118- * transposed. This will untranspose the elements individually when storing so
119- * the data in local memory is always consistent.
120- * @tparam trans Whether the source matrix is transposed or not.
121- * @tparam ld The leading dimension of the destination memory.*/
122- template <bool trans, index_t ld, typename DestPointerType>
123- static PORTBLAS_INLINE typename std::enable_if<trans>::type store (
124- PacketType &packet, DestPointerType dest) {
125- using address_t = cl::sycl::access::address_space;
126- #pragma unroll
127- for (index_t i = 0 ; i < packet_size; i++) {
128- value_t val = reinterpret_cast <value_t *>(&packet)[i];
129- if constexpr (std::is_same<cl::sycl::multi_ptr<cl::sycl::half,
130- address_t ::local_space>,
131- DestPointerType>::value) {
132- using dtype = cl::sycl::half;
133- *(dest + ld * i) = static_cast <dtype>(val);
134- } else if constexpr (std::is_same<cl::sycl::multi_ptr<
135- cl::sycl::ext::oneapi::bfloat16,
136- address_t ::local_space>,
137- DestPointerType>::value) {
138- using dtype = cl::sycl::ext::oneapi::bfloat16;
139- *(dest + ld * i) = static_cast <dtype>(val);
140- } else {
141- using namespace cl ::sycl::ext::oneapi::experimental::matrix;
142- *(dest + ld * i) = round_to_tf32 (val);
112+ for (index_t i = 0 ; i < packet_size; i++, dest++, src++) {
113+ if constexpr (std::is_same<cl::sycl::multi_ptr<cl::sycl::half,
114+ address_t ::local_space>,
115+ DestPointerType>::value) {
116+ using dtype = cl::sycl::half;
117+ *dest = static_cast <dtype>(edge_in_range (i) ? *src : 0 );
118+ } else if constexpr (std::is_same<cl::sycl::multi_ptr<
119+ cl::sycl::ext::oneapi::bfloat16,
120+ address_t ::local_space>,
121+ DestPointerType>::value) {
122+ using namespace cl ::sycl::ext::oneapi;
123+ *dest = bfloat16 (edge_in_range (i) ? *src : 0 .f );
124+ } else {
125+ using namespace cl ::sycl::ext::oneapi::experimental::matrix;
126+ *dest = edge_in_range (i) ? round_to_tf32 (*src) : 0 .f ;
127+ }
143128 }
144129 }
145130 }
146131
147- /* ! @brief Store a vector packet into local memory when the source is not
148- * transposed. This will use sycl::vec::store function.
149- * @tparam trans Whether the source matrix is transposed or not.
150- * @tparam ld The leading dimension of the destination memory.*/
151- template <bool trans, int ld, typename DestPointerType>
152- static PORTBLAS_INLINE typename std::enable_if<!trans>::type store (
153- PacketType &packet, DestPointerType dest) {
132+ /* ! @brief Store a vector packet into local memory. This will use
133+ * sycl::vec::store function.
134+ */
135+ template <typename DestPointerType>
136+ static PORTBLAS_INLINE void store (PacketType &packet, DestPointerType dest) {
154137 using address_t = cl::sycl::access::address_space;
155138 if constexpr (std::is_same<cl::sycl::multi_ptr<cl::sycl::half,
156139 address_t ::local_space>,
157140 DestPointerType>::value) {
158141 using dtype = cl::sycl::half;
159- *dest = static_cast <dtype>(packet[0 ]);
142+ cl::sycl::vec<dtype, vector_size> new_vec{};
143+ for (index_t i = 0 ; i < packet_size; i++) {
144+ reinterpret_cast <dtype *>(&new_vec)[i] =
145+ static_cast <dtype>(reinterpret_cast <value_t *>(&packet)[i]);
146+ }
147+ new_vec.template store <address_t ::local_space>(
148+ 0 , cl::sycl::multi_ptr<dtype, address_t ::local_space>(dest));
160149 } else if constexpr (std::is_same<cl::sycl::multi_ptr<
161150 cl::sycl::ext::oneapi::bfloat16,
162151 address_t ::local_space>,
163152 DestPointerType>::value) {
164- using dtype = cl::sycl::ext::oneapi::bfloat16;
165- *dest = static_cast <dtype>(packet[0 ]);
153+ // sycl::vec doesn't accept bfloat16 as a valid input type
154+ // so we need to write the packet elements individually to
155+ // the shared memory.
156+ using namespace cl ::sycl::ext::oneapi;
157+ for (index_t i = 0 ; i < packet_size; i++, dest++) {
158+ *dest = bfloat16 (reinterpret_cast <value_t *>(&packet)[i]);
159+ }
166160 } else {
167161 using namespace cl ::sycl::ext::oneapi::experimental::matrix;
168- *dest = round_to_tf32 (packet[0 ]);
162+ using dtype = float ;
163+ cl::sycl::vec<dtype, vector_size> new_vec;
164+ for (index_t i = 0 ; i < packet_size; i++) {
165+ reinterpret_cast <dtype *>(&new_vec)[i] =
166+ round_to_tf32 (reinterpret_cast <value_t *>(&packet)[i]);
167+ }
168+ new_vec.template store <address_t ::local_space>(
169+ 0 , cl::sycl::multi_ptr<dtype, address_t ::local_space>(dest));
169170 }
170171 }
171172};
0 commit comments