[hipblaslt] modify the batch offset value as in elements.

Author: shbae

projects/hipblaslt/clients/common/include/testing_matmul_batch_offset.hpp

@@ -62,7 +62,7 @@ void testing_matmul_batch_offset_impl(const Arguments& arg)
 
     int32_t batch_count = arg.batch_count;
 
-    // Batch offsets (from YAML - assumed to be in BYTES)
+    // Batch offsets (from YAML - in ELEMENTS)
     int64_t offset_a = arg.batch_offset_a;
     int64_t offset_b = arg.batch_offset_b;
     int64_t offset_c = arg.batch_offset_c;
@@ -86,22 +86,16 @@ void testing_matmul_batch_offset_impl(const Arguments& arg)
     size_t size_C_sub = size_t(ldc) * size_t(N);
     size_t size_D_sub = size_t(ldd) * size_t(N);
 
-    // Convert byte offsets to element offsets
-    // Note: offsets must be aligned to element size
-    size_t offset_a_elem = offset_a / sizeof(Ti);
-    size_t offset_b_elem = offset_b / sizeof(Ti);
-    size_t offset_c_elem = offset_c / sizeof(To);
-    size_t offset_d_elem = offset_d / sizeof(To);
-
     // Full buffer sizes: [offset padding] + [matrix data]
     // The offset padding comes FIRST because:
     // - Pointer array points to buffer BASE
     // - Kernel adds offset to BASE to get actual data location
     // - So actual data is at BASE + offset
-    size_t size_A_full = offset_a_elem + size_A_sub;
-    size_t size_B_full = offset_b_elem + size_B_sub;
-    size_t size_C_full = offset_c_elem + size_C_sub;
-    size_t size_D_full = offset_d_elem + size_D_sub;
+    // Offsets are already in elements (the API takes element offsets).
+    size_t size_A_full = offset_a + size_A_sub;
+    size_t size_B_full = offset_b + size_B_sub;
+    size_t size_C_full = offset_c + size_C_sub;
+    size_t size_D_full = offset_d + size_D_sub;
 
     // Allocate host memory for full buffers
     host_vector<Ti> h_A_full(size_A_full * batch_count);
@@ -116,9 +110,9 @@ void testing_matmul_batch_offset_impl(const Arguments& arg)
     for(int b = 0; b < batch_count; b++)
     {
         // Initialize each batch's sub-matrix at the offset location
-        Ti* A_batch = h_A_full.data() + b * size_A_full + offset_a_elem;
-        Ti* B_batch = h_B_full.data() + b * size_B_full + offset_b_elem;
-        To* C_batch = h_C_full.data() + b * size_C_full + offset_c_elem;
+        Ti* A_batch = h_A_full.data() + b * size_A_full + offset_a;
+        Ti* B_batch = h_B_full.data() + b * size_B_full + offset_b;
+        To* C_batch = h_C_full.data() + b * size_C_full + offset_c;
 
         // Simple initialization: A and B with small integers
         for(int64_t j = 0; j < A_col; j++)
@@ -308,9 +302,9 @@ void testing_matmul_batch_offset_impl(const Arguments& arg)
     for(int b = 0; b < batch_count; b++)
     {
         // Get pointers to sub-matrices (with element offset applied)
-        Ti* A_sub = h_A_full.data() + b * size_A_full + offset_a_elem;
-        Ti* B_sub = h_B_full.data() + b * size_B_full + offset_b_elem;
-        To* C_sub = h_C_full.data() + b * size_C_full + offset_c_elem;
+        Ti* A_sub = h_A_full.data() + b * size_A_full + offset_a;
+        Ti* B_sub = h_B_full.data() + b * size_B_full + offset_b;
+        To* C_sub = h_C_full.data() + b * size_C_full + offset_c;
         To* D_sub = h_D_gold.data() + b * size_D_sub;
 
         // Simple GEMM: D = alpha * A * B + beta * C
@@ -346,7 +340,7 @@ void testing_matmul_batch_offset_impl(const Arguments& arg)
     for(int b = 0; b < batch_count; b++)
     {
         // GPU result is at (base + offset) within each batch's buffer
-        To* result_gpu = h_D_full.data() + b * size_D_full + offset_d_elem;
+        To* result_gpu = h_D_full.data() + b * size_D_full + offset_d;
         To* result_cpu = h_D_gold.data() + b * size_D_sub;
 
         for(size_t i = 0; i < size_D_sub; i++)

projects/hipblaslt/library/include/hipblaslt/hipblaslt.h

@@ -177,7 +177,7 @@ typedef enum {
   /** Matrix Offset.
    *
    * For ``General Batched GEMM``, we can support for users to access a sub-matrix of
-   * the original matrix by adding an ``offset`` value from the base address.
+   * the original matrix by adding an ``offset`` value (in elements) from the base address.
    * Note that for non-batched or Strided Batch GEMM case, we can directly apply
    * the offset value by using the strided-offset value.
    */

projects/hipblaslt/library/src/amd_detail/include/auxiliary.hpp

@@ -158,6 +158,22 @@ constexpr const char* hip_datatype_to_string(hipDataType type)
     return "invalid";
 }
 
+// Returns true for sub-byte MX-style data types (fp6/fp4).
+// Used to reject features that require byte-addressable elements.
+HIPBLASLT_EXPORT
+constexpr bool hip_datatype_is_mxtype(hipDataType type)
+{
+    switch(type)
+    {
+    case HIP_R_6F_E2M3:
+    case HIP_R_6F_E3M2:
+    case HIP_R_4F_E2M1:
+        return true;
+    default:
+        return false;
+    }
+}
+
 // return precision string for hipDataType
 HIPBLASLT_EXPORT
 constexpr const char* hipblas_computetype_to_string(hipblasComputeType_t type)

projects/hipblaslt/library/src/amd_detail/rocblaslt/src/include/rocblaslt_mat_utils.hpp

@@ -229,6 +229,19 @@ inline rocblaslt_status validateMatmulArgs(int64_t                       m,
         return rocblaslt_status_invalid_size;
     }
 
+    // Batch offsets are expressed in elements and converted to bytes; sub-byte
+    // (MX) data types are not byte-addressable, so a nonzero offset is unsupported.
+    if((batch_offset_a != 0 && hip_datatype_is_mxtype(type_a))
+       || (batch_offset_b != 0 && hip_datatype_is_mxtype(type_b))
+       || (batch_offset_c != 0 && hip_datatype_is_mxtype(type_c))
+       || (batch_offset_d != 0 && hip_datatype_is_mxtype(type_d)))
+    {
+#ifndef CODE_COVERAGE
+        std::cerr << "matrix offset is not supported for sub-byte (MX) data types" << std::endl;
+#endif
+        return rocblaslt_status_not_implemented;
+    }
+
     // number of batches of matrics A,B,C,D must be the same and positive
     if(num_batches_a != num_batches_b || num_batches_a != num_batches_c
        || num_batches_a != num_batches_d || num_batches_a < 1)

projects/hipblaslt/library/src/amd_detail/rocblaslt/src/tensile_host.cpp

@@ -2394,11 +2394,22 @@ namespace
         inputs.batchC = reinterpret_cast<void const* const*>(prob.batch_C);
         inputs.batchD = reinterpret_cast<void* const*>(prob.batch_D);
 
-        // set the batched offset of A, B, C, and D.
-        inputs.batchOffsetA = prob.batch_offset_a;
-        inputs.batchOffsetB = prob.batch_offset_b;
-        inputs.batchOffsetC = prob.batch_offset_c;
-        inputs.batchOffsetD = prob.batch_offset_d;
+        // The batch offsets are specified by the user in elements; convert them to
+        // bytes here so the kernel/assembly can add them straight to byte addresses.
+        // Only data types whose element size is at least one byte are supported
+        // (sub-byte types such as fp4/fp6 are rejected during argument validation).
+        inputs.batchOffsetA
+            = prob.batch_offset_a
+              * size_t(TensileLite::DataTypeInfo::Get(hip2TensileType(prob.a_type)).elementSize);
+        inputs.batchOffsetB
+            = prob.batch_offset_b
+              * size_t(TensileLite::DataTypeInfo::Get(hip2TensileType(prob.b_type)).elementSize);
+        inputs.batchOffsetC
+            = prob.batch_offset_c
+              * size_t(TensileLite::DataTypeInfo::Get(hip2TensileType(prob.c_type)).elementSize);
+        inputs.batchOffsetD
+            = prob.batch_offset_d
+              * size_t(TensileLite::DataTypeInfo::Get(hip2TensileType(prob.d_type)).elementSize);
 
         // Set the GSU workspace
         inputs.ws            = prob.workspace;