refactor buffer load

composable_kernel/include/utility/amd_buffer_addressing.hpp

@@ -6,6 +6,17 @@
 
 namespace ck {
 
+template <typename T>
+union BufferResource
+{
+    // 128 bit SGPRs to supply buffer resource in buffer instructions
+    // https://rocm-documentation.readthedocs.io/en/latest/GCN_ISA_Manuals/testdocbook.html#vector-memory-buffer-instructions
+    int32x4_t data;
+    T* address[2];
+    int32_t range[4];
+    int32_t config[4];
+};
+
 __device__ float __llvm_amdgcn_buffer_load_f32(int32x4_t srsrc,
                                                index_t vindex,
                                                index_t offset,

composable_kernel/include/utility/amd_buffer_addressing_v2.hpp

@@ -6,27 +6,27 @@
 namespace ck {
 
 template <typename T>
-union BufferResource
+union BufferResource_v2
 {
     // 128 bit SGPRs to supply buffer resource in buffer instructions
     // https://rocm-documentation.readthedocs.io/en/latest/GCN_ISA_Manuals/testdocbook.html#vector-memory-buffer-instructions
     int32x4_t data;
-    T* address[2];
-    int32_t range[4];
-    int32_t config[4];
+    StaticallyIndexedArray<T*, 2> address;
+    StaticallyIndexedArray<int32_t, 4> range;
+    StaticallyIndexedArray<int32_t, 4> config;
 };
 
 template <typename T>
 __device__ int32x4_t make_wave_buffer_resource(T* p_wave, index_t data_space_size)
 {
-    BufferResource<T> wave_buffer_resource;
+    BufferResource_v2<T> wave_buffer_resource;
 
     // wavewise base address (64 bit)
-    wave_buffer_resource.address[0] = const_cast<remove_cv_t<T>*>(p_wave);
+    wave_buffer_resource.address(Number<0>{}) = const_cast<remove_cv_t<T>*>(p_wave);
     // wavewise range (32 bit)
-    wave_buffer_resource.range[2] = data_space_size * sizeof(T);
+    wave_buffer_resource.range(Number<2>{}) = data_space_size * sizeof(T);
     // wavewise setting (32 bit)
-    wave_buffer_resource.config[3] = CK_BUFFER_RESOURCE_3RD_DWORD;
+    wave_buffer_resource.config(Number<3>{}) = CK_BUFFER_RESOURCE_3RD_DWORD;
 
     return wave_buffer_resource.data;
 }
@@ -37,6 +37,19 @@ __llvm_amdgcn_raw_buffer_load_i8(int32x4_t srsrc,
                                  index_t voffset,
                                  index_t soffset,
                                  index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.i8");
+
+__device__ int8x2_t
+__llvm_amdgcn_raw_buffer_load_i8x2(int32x4_t srsrc,
+                                   index_t voffset,
+                                   index_t soffset,
+                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v2i8");
+
+__device__ int8x4_t
+__llvm_amdgcn_raw_buffer_load_i8x4(int32x4_t srsrc,
+                                   index_t voffset,
+                                   index_t soffset,
+                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v4i8");
+
 __device__ int16_t
 __llvm_amdgcn_raw_buffer_load_i16(int32x4_t srsrc,
                                   index_t voffset,
@@ -86,6 +99,20 @@ __llvm_amdgcn_raw_buffer_store_i8(int8_t vdata,
                                   index_t soffset,
                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.i8");
 
+__device__ void
+__llvm_amdgcn_raw_buffer_store_i8x2(int8x2_t vdata,
+                                    int32x4_t rsrc,
+                                    index_t voffset,
+                                    index_t soffset,
+                                    index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v2i8");
+
+__device__ void
+__llvm_amdgcn_raw_buffer_store_i8x4(int8x4_t vdata,
+                                    int32x4_t rsrc,
+                                    index_t voffset,
+                                    index_t soffset,
+                                    index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v4i8");
+
 __device__ void
 __llvm_amdgcn_raw_buffer_store_i16(int16_t vdata,
                                    int32x4_t rsrc,
@@ -142,6 +169,7 @@ amd_buffer_load_impl_v2(int32x4_t src_wave_buffer_resource,
                         index_t src_wave_addr_offset)
 {
     static_assert((is_same<T, float>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
+                      (is_same<T, int8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
                       (is_same<T, int32_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)),
                   "wrong! not implemented");
 
@@ -169,8 +197,11 @@ amd_buffer_load_impl_v2(int32x4_t src_wave_buffer_resource,
             tmp.Vectors(Number<4>{})(Number<0>{}) = __llvm_amdgcn_raw_buffer_load_fp32x4(
                 src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
 
-            tmp.Vectors(Number<4>{})(Number<1>{}) = __llvm_amdgcn_raw_buffer_load_fp32x4(
-                src_wave_buffer_resource, src_thread_addr_offset, 4 * sizeof(float), 0);
+            tmp.Vectors(Number<4>{})(Number<1>{}) =
+                __llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
+                                                     src_thread_addr_offset,
+                                                     src_wave_addr_offset + 4 * sizeof(float),
+                                                     0);
 
             return tmp.Vector();
         }
@@ -199,8 +230,44 @@ amd_buffer_load_impl_v2(int32x4_t src_wave_buffer_resource,
             tmp.Vectors(Number<4>{})(Number<0>{}) = __llvm_amdgcn_raw_buffer_load_i32x4(
                 src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
 
-            tmp.Vectors(Number<4>{})(Number<1>{}) = __llvm_amdgcn_raw_buffer_load_i32x4(
-                src_wave_buffer_resource, src_thread_addr_offset, 4 * sizeof(int32_t), 0);
+            tmp.Vectors(Number<4>{})(Number<1>{}) =
+                __llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
+                                                    src_thread_addr_offset,
+                                                    src_wave_addr_offset + 4 * sizeof(int32_t),
+                                                    0);
+
+            return tmp.Vector();
+        }
+    }
+    else if constexpr(is_same<T, int8_t>::value)
+    {
+        if constexpr(N == 1)
+        {
+            return __llvm_amdgcn_raw_buffer_load_i8(
+                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
+        }
+        else if constexpr(N == 2)
+        {
+            return __llvm_amdgcn_raw_buffer_load_i8x2(
+                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
+        }
+        else if constexpr(N == 4)
+        {
+            return __llvm_amdgcn_raw_buffer_load_i8x4(
+                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
+        }
+        else if constexpr(N == 8)
+        {
+            vector_type<int8_t, 8> tmp;
+
+            tmp.Vectors(Number<4>{})(Number<0>{}) = __llvm_amdgcn_raw_buffer_load_i8x4(
+                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
+
+            tmp.Vectors(Number<4>{})(Number<1>{}) =
+                __llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
+                                                   src_thread_addr_offset,
+                                                   src_wave_addr_offset + 4 * sizeof(int8_t),
+                                                   0);
 
             return tmp.Vector();
         }

composable_kernel/include/utility/float_type.amd.hpp.in

@@ -274,6 +274,7 @@ using int32x8_t = typename vector_type<int32_t, 8>::type;
 // i8
 // hack for int8x4_t, because compiler does not have native support for int8x4_t
 // int8x4_t is defined as int32_t
+using int8x2_t = typename vector_type<int8_t, 2>::type;
 using int8x4_t = typename vector_type<int8_t, 4>::type;
 
 // data type conversion