diff --git a/projects/hipblaslt/clients/common/include/hipblaslt_arguments.hpp b/projects/hipblaslt/clients/common/include/hipblaslt_arguments.hpp index d3d2040686ad..ac50c9c48a21 100644 --- a/projects/hipblaslt/clients/common/include/hipblaslt_arguments.hpp +++ b/projects/hipblaslt/clients/common/include/hipblaslt_arguments.hpp @@ -112,6 +112,12 @@ struct Arguments int32_t batch_count; int32_t batch_mode; + // Batch offset support for general batched GEMM + int64_t batch_offset_a; + int64_t batch_offset_b; + int64_t batch_offset_c; + int64_t batch_offset_d; + int32_t iters; int32_t cold_iters; @@ -251,6 +257,10 @@ struct Arguments OPER(lde) SEP \ OPER(batch_count) SEP \ OPER(batch_mode) SEP \ + OPER(batch_offset_a) SEP \ + OPER(batch_offset_b) SEP \ + OPER(batch_offset_c) SEP \ + OPER(batch_offset_d) SEP \ OPER(iters) SEP \ OPER(cold_iters) SEP \ OPER(warmup_time) SEP \ diff --git a/projects/hipblaslt/clients/common/include/testing_matmul_batch_offset.hpp b/projects/hipblaslt/clients/common/include/testing_matmul_batch_offset.hpp new file mode 100644 index 000000000000..4ca921f119b7 --- /dev/null +++ b/projects/hipblaslt/clients/common/include/testing_matmul_batch_offset.hpp @@ -0,0 +1,458 @@ +/******************************************************************************* + * + * MIT License + * + * Copyright (C) 2022-2026 Advanced Micro Devices, Inc. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + *******************************************************************************/ + +#pragma once + +#include "cblas_interface.hpp" +#include "datatype_interface.hpp" +#include "flops.hpp" +#include "hipblaslt_datatype2string.hpp" +#include "hipblaslt_init.hpp" +#include "hipblaslt_math.hpp" +#include "hipblaslt_random.hpp" +#include "hipblaslt_test.hpp" +#include "hipblaslt_vector.hpp" +#include "near.hpp" +#include "norm.hpp" +#include "unit.hpp" +#include "utility.hpp" +#include +#include + +/* ============================================================================================ */ +/*! \brief Test for 64-bit batch offset support in general batched GEMM */ + +template +void testing_matmul_batch_offset_impl(const Arguments& arg) +{ + hipblasOperation_t transA = char_to_hipblas_operation(arg.transA); + hipblasOperation_t transB = char_to_hipblas_operation(arg.transB); + + // Use first element from arrays (grouped GEMM uses arrays, we use first element) + int64_t M = arg.M[0]; + int64_t N = arg.N[0]; + int64_t K = arg.K[0]; + + int64_t lda = arg.lda[0]; + int64_t ldb = arg.ldb[0]; + int64_t ldc = arg.ldc[0]; + int64_t ldd = arg.ldd[0]; + + int32_t batch_count = arg.batch_count; + + // 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; + int64_t offset_d = arg.batch_offset_d; + + // Only test general batched mode (pointer array) + if(arg.batch_mode != 1) // 1 = pointer array + { + GTEST_SKIP() << "Batch offset only supported for general batched mode (batch_mode=1)"; + } + + // Calculate matrix sizes + int64_t A_row = transA == HIPBLAS_OP_N ? M : K; + int64_t A_col = transA == HIPBLAS_OP_N ? K : M; + int64_t B_row = transB == HIPBLAS_OP_N ? K : N; + int64_t B_col = transB == HIPBLAS_OP_N ? N : K; + + // Sub-matrix sizes (actual GEMM size) in elements + size_t size_A_sub = size_t(lda) * size_t(A_col); + size_t size_B_sub = size_t(ldb) * size_t(B_col); + size_t size_C_sub = size_t(ldc) * size_t(N); + size_t size_D_sub = size_t(ldd) * size_t(N); + + // Calculate padding needed for negative offsets (in elements) + // If offset is negative, we need padding BEFORE the base pointer + // If offset is positive, padding is at the beginning (offset space) + size_t padding_a = (offset_a < 0) ? size_t(-offset_a) : 0; + size_t padding_b = (offset_b < 0) ? size_t(-offset_b) : 0; + size_t padding_c = (offset_c < 0) ? size_t(-offset_c) : 0; + size_t padding_d = (offset_d < 0) ? size_t(-offset_d) : 0; + + // Full buffer sizes: [padding for negative offsets] + [matrix data] + [positive offset space] + // Layout for negative offset: [padding|matrix_data] base points to start of matrix_data + // Layout for positive offset: [matrix_data|padding] base points to start of buffer + size_t size_A_full = padding_a + size_A_sub + (offset_a > 0 ? size_t(offset_a) : 0); + size_t size_B_full = padding_b + size_B_sub + (offset_b > 0 ? size_t(offset_b) : 0); + size_t size_C_full = padding_c + size_C_sub + (offset_c > 0 ? size_t(offset_c) : 0); + size_t size_D_full = padding_d + size_D_sub + (offset_d > 0 ? size_t(offset_d) : 0); + + // Allocate host memory for full buffers + host_vector h_A_full(size_A_full * batch_count); + host_vector h_B_full(size_B_full * batch_count); + host_vector h_C_full(size_C_full * batch_count); + host_vector h_D_full(size_D_full * batch_count); // GPU result with offset API + host_vector h_D_gold(size_D_sub * batch_count); // CPU reference + + // Initialize matrices with known pattern + // For negative offsets: base points to (padding + 0), data at (base + negative_offset) = padding region + // For positive offsets: base points to 0, data at (base + positive_offset) = offset region + for(int b = 0; b < batch_count; b++) + { + // Calculate base pointer for this batch (after padding for negative offsets) + Ti* A_base = h_A_full.data() + b * size_A_full + padding_a; + Ti* B_base = h_B_full.data() + b * size_B_full + padding_b; + To* C_base = h_C_full.data() + b * size_C_full + padding_c; + + // Data location is at (base + offset) + // For negative offset: base + (-N) points backward into padding region + // For positive offset: base + (+N) points forward into buffer + Ti* A_batch = A_base + offset_a; + Ti* B_batch = B_base + offset_b; + To* C_batch = C_base + offset_c; + + // Simple initialization: A and B with small integers + for(int64_t j = 0; j < A_col; j++) + for(int64_t i = 0; i < A_row; i++) + A_batch[i + j * lda] = Ti((i + j + b) % 7 + 1); + + for(int64_t j = 0; j < B_col; j++) + for(int64_t i = 0; i < B_row; i++) + B_batch[i + j * ldb] = Ti((i - j + b) % 5 + 1); + + for(int64_t j = 0; j < N; j++) + for(int64_t i = 0; i < M; i++) + C_batch[i + j * ldc] = To((i + j) % 3); + } + + // Allocate device memory + device_vector d_A_full(size_A_full * batch_count); + device_vector d_B_full(size_B_full * batch_count); + device_vector d_C_full(size_C_full * batch_count); + device_vector d_D_full(size_D_full * batch_count); + + // Copy to device + CHECK_HIP_ERROR( + hipMemcpy(d_A_full, h_A_full.data(), sizeof(Ti) * size_A_full * batch_count, hipMemcpyHostToDevice)); + CHECK_HIP_ERROR( + hipMemcpy(d_B_full, h_B_full.data(), sizeof(Ti) * size_B_full * batch_count, hipMemcpyHostToDevice)); + CHECK_HIP_ERROR( + hipMemcpy(d_C_full, h_C_full.data(), sizeof(To) * size_C_full * batch_count, hipMemcpyHostToDevice)); + + // Setup pointer arrays for base addresses + // Base pointers must point AFTER the padding (for negative offset support) + std::vector h_batch_A(batch_count); + std::vector h_batch_B(batch_count); + std::vector h_batch_C(batch_count); + std::vector h_batch_D(batch_count); + + for(int b = 0; b < batch_count; b++) + { + h_batch_A[b] = d_A_full + b * size_A_full + padding_a; + h_batch_B[b] = d_B_full + b * size_B_full + padding_b; + h_batch_C[b] = d_C_full + b * size_C_full + padding_c; + h_batch_D[b] = d_D_full + b * size_D_full + padding_d; + } + + // Allocate device memory for pointer arrays + Ti** d_batch_A; + Ti** d_batch_B; + To** d_batch_C; + To** d_batch_D; + + CHECK_HIP_ERROR(hipMalloc(&d_batch_A, sizeof(Ti*) * batch_count)); + CHECK_HIP_ERROR(hipMalloc(&d_batch_B, sizeof(Ti*) * batch_count)); + CHECK_HIP_ERROR(hipMalloc(&d_batch_C, sizeof(To*) * batch_count)); + CHECK_HIP_ERROR(hipMalloc(&d_batch_D, sizeof(To*) * batch_count)); + + CHECK_HIP_ERROR( + hipMemcpy(d_batch_A, h_batch_A.data(), sizeof(Ti*) * batch_count, hipMemcpyHostToDevice)); + CHECK_HIP_ERROR( + hipMemcpy(d_batch_B, h_batch_B.data(), sizeof(Ti*) * batch_count, hipMemcpyHostToDevice)); + CHECK_HIP_ERROR( + hipMemcpy(d_batch_C, h_batch_C.data(), sizeof(To*) * batch_count, hipMemcpyHostToDevice)); + CHECK_HIP_ERROR( + hipMemcpy(d_batch_D, h_batch_D.data(), sizeof(To*) * batch_count, hipMemcpyHostToDevice)); + + // Alpha and beta + Tc h_alpha = arg.get_alpha(); + Tc h_beta = arg.get_beta(); + + // Setup hipBLASLt + hipblasLtHandle_t handle; + CHECK_HIPBLASLT_ERROR(hipblasLtCreate(&handle)); + + hipblasLtMatmulDesc_t matmul_desc; + CHECK_HIPBLASLT_ERROR( + hipblasLtMatmulDescCreate(&matmul_desc, arg.compute_type, arg.scale_type)); + CHECK_HIPBLASLT_ERROR( + hipblasLtMatmulDescSetAttribute( + matmul_desc, HIPBLASLT_MATMUL_DESC_TRANSA, &transA, sizeof(transA))); + CHECK_HIPBLASLT_ERROR( + hipblasLtMatmulDescSetAttribute( + matmul_desc, HIPBLASLT_MATMUL_DESC_TRANSB, &transB, sizeof(transB))); + + // Create matrix layouts + hipblasLtMatrixLayout_t matA, matB, matC, matD; + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutCreate(&matA, arg.a_type, A_row, A_col, lda)); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutCreate(&matB, arg.b_type, B_row, B_col, ldb)); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutCreate(&matC, arg.c_type, M, N, ldc)); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutCreate(&matD, arg.d_type, M, N, ldd)); + + // Set batch count and mode + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutSetAttribute( + matA, HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count))); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutSetAttribute( + matB, HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count))); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutSetAttribute( + matC, HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count))); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutSetAttribute( + matD, HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count))); + + int32_t batch_mode = 1; // Pointer array + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutSetAttribute( + matA, HIPBLASLT_MATRIX_LAYOUT_BATCH_MODE, &batch_mode, sizeof(batch_mode))); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutSetAttribute( + matB, HIPBLASLT_MATRIX_LAYOUT_BATCH_MODE, &batch_mode, sizeof(batch_mode))); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutSetAttribute( + matC, HIPBLASLT_MATRIX_LAYOUT_BATCH_MODE, &batch_mode, sizeof(batch_mode))); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutSetAttribute( + matD, HIPBLASLT_MATRIX_LAYOUT_BATCH_MODE, &batch_mode, sizeof(batch_mode))); + + // ======================================== + // GPU GEMM with offset API + // ======================================== + + // Set offsets for all matrices + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutSetAttribute( + matA, HIPBLASLT_MATRIX_LAYOUT_OFFSET, &offset_a, sizeof(offset_a))); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutSetAttribute( + matB, HIPBLASLT_MATRIX_LAYOUT_OFFSET, &offset_b, sizeof(offset_b))); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutSetAttribute( + matC, HIPBLASLT_MATRIX_LAYOUT_OFFSET, &offset_c, sizeof(offset_c))); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutSetAttribute( + matD, HIPBLASLT_MATRIX_LAYOUT_OFFSET, &offset_d, sizeof(offset_d))); + + // Find algorithm + hipblasLtMatmulPreference_t pref; + CHECK_HIPBLASLT_ERROR(hipblasLtMatmulPreferenceCreate(&pref)); + size_t prefMaxWorkspaceSize = 128 * 1024 * 1024; // 128 MB + CHECK_HIPBLASLT_ERROR(hipblasLtMatmulPreferenceSetAttribute( + pref, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &prefMaxWorkspaceSize, sizeof(prefMaxWorkspaceSize))); + + // Support testing multiple solutions via requested_solution_num parameter + // Default to 1 (test only best solution) for backward compatibility + // Use HIPBLASLT_MAX_REQUESTED_SOLUTION_NUM when -1 to get all available solutions + int32_t requestedAlgos = (arg.requested_solution_num < 0) ? HIPBLASLT_MAX_REQUESTED_SOLUTION_NUM + : (arg.requested_solution_num == 0) ? 1 + : arg.requested_solution_num; + + std::vector heuristicResult(requestedAlgos); + int numAlgos = 0; + CHECK_HIPBLASLT_ERROR(hipblasLtMatmulAlgoGetHeuristic( + handle, matmul_desc, matA, matB, matC, matD, pref, requestedAlgos, heuristicResult.data(), &numAlgos)); + + if(numAlgos == 0) + { + CHECK_HIPBLASLT_ERROR(hipblasLtMatmulPreferenceDestroy(pref)); + GTEST_SKIP() << "No algorithm found for this configuration"; + } + + // Find maximum workspace size across all solutions + size_t maxWorkspaceSize = 0; + for(int i = 0; i < numAlgos; i++) + { + maxWorkspaceSize = std::max(maxWorkspaceSize, heuristicResult[i].workspaceSize); + } + + // Allocate workspace + void* d_workspace = nullptr; + if(maxWorkspaceSize > 0) + { + CHECK_HIP_ERROR(hipMalloc(&d_workspace, maxWorkspaceSize)); + } + + // ======================================== + // CPU Reference: D = alpha * A * B + beta * C + // Computed once before testing any solutions + // ======================================== + for(int b = 0; b < batch_count; b++) + { + // Get pointers to sub-matrices (with element offset applied) + // Base is at padding, then add offset (which may be negative) + Ti* A_sub = h_A_full.data() + b * size_A_full + padding_a + offset_a; + Ti* B_sub = h_B_full.data() + b * size_B_full + padding_b + offset_b; + To* C_sub = h_C_full.data() + b * size_C_full + padding_c + offset_c; + To* D_sub = h_D_gold.data() + b * size_D_sub; + + // Simple GEMM: D = alpha * A * B + beta * C + for(int64_t i = 0; i < M; i++) + { + for(int64_t j = 0; j < N; j++) + { + Tc sum = 0; + for(int64_t k = 0; k < K; k++) + { + // A is A_row x A_col, B is B_row x B_col + // For transA=N: A(i,k) = A[i + k*lda] + // For transA=T: A(i,k) = A[k + i*lda] + // For transB=N: B(k,j) = B[k + j*ldb] + // For transB=T: B(k,j) = B[j + k*ldb] + Tc a_val = (transA == HIPBLAS_OP_N) + ? Tc(A_sub[i + k * lda]) + : Tc(A_sub[k + i * lda]); + Tc b_val = (transB == HIPBLAS_OP_N) + ? Tc(B_sub[k + j * ldb]) + : Tc(B_sub[j + k * ldb]); + sum += a_val * b_val; + } + D_sub[i + j * ldd] = To(h_alpha * sum + h_beta * Tc(C_sub[i + j * ldc])); + } + } + } + + // Tolerance: epsilon * factor * K (accumulation over K elements) + double tol = std::numeric_limits::epsilon() * 100 * K; + + // Track validation results across all solutions + int numPassed = 0; + int numFailed = 0; + + // Test each solution + for(int algoIdx = 0; algoIdx < numAlgos; algoIdx++) + { + // Reset only D output buffer to sentinel values before testing each solution + // A, B, C are inputs and don't change, so we can reuse them + CHECK_HIP_ERROR(hipMemcpy(d_D_full, + h_D_full.data(), + sizeof(To) * size_D_full * batch_count, + hipMemcpyHostToDevice)); + + CHECK_HIPBLASLT_ERROR(hipblasLtMatmul(handle, + matmul_desc, + &h_alpha, + d_batch_A, + matA, + d_batch_B, + matB, + &h_beta, + d_batch_C, + matC, + d_batch_D, + matD, + &heuristicResult[algoIdx].algo, + d_workspace, + heuristicResult[algoIdx].workspaceSize, + 0)); + + // Ensure kernel completes before reading result + CHECK_HIP_ERROR(hipDeviceSynchronize()); + + // Copy GPU result back for verification + CHECK_HIP_ERROR(hipMemcpy(h_D_full.data(), + d_D_full, + sizeof(To) * size_D_full * batch_count, + hipMemcpyDeviceToHost)); + + // ======================================== + // VALIDATION: Compare GPU vs CPU + // ======================================== + double max_error = 0.0; + for(int b = 0; b < batch_count; b++) + { + // GPU result is at (base + offset) within each batch's buffer + // Base is at padding_d, then add offset_d (which may be negative) + To* result_gpu = h_D_full.data() + b * size_D_full + padding_d + offset_d; + To* result_cpu = h_D_gold.data() + b * size_D_sub; + + for(size_t i = 0; i < size_D_sub; i++) + { + double diff = std::abs(double(result_gpu[i]) - double(result_cpu[i])); + max_error = std::max(max_error, diff); + } + } + + // Check and count per-solution results + if(arg.unit_check) + { + if(max_error >= tol) + { + numFailed++; + EXPECT_LT(max_error, tol) + << "Solution " << algoIdx << "/" << numAlgos + << " FAILED (error: " << max_error << ", tol: " << tol << ")"; + } + else + { + numPassed++; + } + } + } + + // Report summary when testing multiple solutions + if(numAlgos > 1 && arg.unit_check) + { + hipblaslt_cout << "Tested " << numAlgos << " solutions: " + << numPassed << " passed, " << numFailed << " failed" << std::endl; + } + + // Cleanup + if(d_workspace) + { + CHECK_HIP_ERROR(hipFree(d_workspace)); + } + CHECK_HIP_ERROR(hipFree(d_batch_A)); + CHECK_HIP_ERROR(hipFree(d_batch_B)); + CHECK_HIP_ERROR(hipFree(d_batch_C)); + CHECK_HIP_ERROR(hipFree(d_batch_D)); + + CHECK_HIPBLASLT_ERROR(hipblasLtMatmulPreferenceDestroy(pref)); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutDestroy(matA)); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutDestroy(matB)); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutDestroy(matC)); + CHECK_HIPBLASLT_ERROR(hipblasLtMatrixLayoutDestroy(matD)); + CHECK_HIPBLASLT_ERROR(hipblasLtMatmulDescDestroy(matmul_desc)); + CHECK_HIPBLASLT_ERROR(hipblasLtDestroy(handle)); +} + +// Type dispatcher based on Arguments +void testing_matmul_batch_offset(const Arguments& arg) +{ + // Dispatch based on data types in Arguments + // For now, support only f32, f16, bf16 with matching input/output types + if(arg.a_type == HIP_R_32F && arg.b_type == HIP_R_32F && + arg.c_type == HIP_R_32F && arg.d_type == HIP_R_32F) + { + testing_matmul_batch_offset_impl(arg); + } + else if(arg.a_type == HIP_R_16F && arg.b_type == HIP_R_16F && + arg.c_type == HIP_R_16F && arg.d_type == HIP_R_16F) + { + testing_matmul_batch_offset_impl(arg); + } + else if(arg.a_type == HIP_R_16BF && arg.b_type == HIP_R_16BF && + arg.c_type == HIP_R_16BF && arg.d_type == HIP_R_16BF) + { + testing_matmul_batch_offset_impl(arg); + } + else + { + GTEST_SKIP() << "Unsupported type combination for batch_offset test"; + } +} diff --git a/projects/hipblaslt/clients/tests/data/hipblaslt_common.yaml b/projects/hipblaslt/clients/tests/data/hipblaslt_common.yaml index ab4e28b116be..e4d4b63fdb66 100644 --- a/projects/hipblaslt/clients/tests/data/hipblaslt_common.yaml +++ b/projects/hipblaslt/clients/tests/data/hipblaslt_common.yaml @@ -597,6 +597,10 @@ Arguments: - lde: c_int64*32 - batch_count: c_int32 - batch_mode: c_int32 + - batch_offset_a: c_int64 + - batch_offset_b: c_int64 + - batch_offset_c: c_int64 + - batch_offset_d: c_int64 - iters: c_int32 - cold_iters: c_int32 - warmup_time: c_float @@ -723,6 +727,10 @@ Defaults: transB: '*' batch_count: 1 batch_mode: 0 + batch_offset_a: 0 + batch_offset_b: 0 + batch_offset_c: 0 + batch_offset_d: 0 HMM: false pad: 4096 threads: 0 diff --git a/projects/hipblaslt/clients/tests/data/matmul_gtest.yaml b/projects/hipblaslt/clients/tests/data/matmul_gtest.yaml index 915507cbb655..5f0f5815dac6 100755 --- a/projects/hipblaslt/clients/tests/data/matmul_gtest.yaml +++ b/projects/hipblaslt/clients/tests/data/matmul_gtest.yaml @@ -3211,4 +3211,204 @@ Tests: requested_solution_num: 10 gpu_arch: '950' +# ============================================================================== +# Batch Offset Tests - 64-bit offset support for general batched GEMM +# ============================================================================== + +# Quick validation test - verifies basic offset functionality +- name: matmul_batch_offset_quick + category: quick + function: matmul_batch_offset + precision: *real_precisions + transA: N + transB: N + M: 256 + N: 128 + K: 64 + lda: 256 + ldb: 64 + ldc: 256 + ldd: 256 + batch_mode: 1 # Pointer array mode + batch_count: 2 + batch_offset_a: 0 + batch_offset_b: 64 + batch_offset_c: 128 + batch_offset_d: 256 + alpha: 1.0 + beta: 0.0 + unit_check: 1 + norm_check: 1 + +# Offset variation test - various offset values +# Note: Uses M=256 to avoid known General Batched GEMM issue with larger sizes +- name: matmul_batch_offset_values + category: pre_checkin + function: matmul_batch_offset + precision: *real_precisions + transA: N + transB: N + M: 256 + N: 128 + K: 128 + lda: 256 + ldb: 128 + ldc: 256 + ldd: 256 + batch_mode: 1 + batch_count: 3 + batch_offset_a: [0, 64, 256, 512] + batch_offset_b: [0, 64, 256, 512] + batch_offset_c: [0, 64, 256, 512] + batch_offset_d: [0, 64, 256, 512] + alpha: 1.0 + beta: [0.0, 1.0] + unit_check: 1 + norm_check: 1 + +# Transpose with offset +# Note: lda/ldb must be valid for all transpose combinations +# For transA=N: lda >= M, for transA=T: lda >= K +# For transB=N: ldb >= K, for transB=T: ldb >= N +# Using max(M,K)=256 for lda and max(K,N)=256 for ldb to cover all cases +- name: matmul_batch_offset_transpose + category: pre_checkin + function: matmul_batch_offset + precision: *real_precisions + transA: [N, T] + transB: [N, T] + M: 256 + N: 256 + K: 128 + lda: 256 + ldb: 256 + ldc: 256 + ldd: 256 + batch_mode: 1 + batch_count: 3 + batch_offset_a: 128 + batch_offset_b: 128 + batch_offset_c: 128 + batch_offset_d: 128 + alpha: 1.0 + beta: 0.5 + unit_check: 1 + norm_check: 1 + +# Alpha/Beta edge cases with offset +- name: matmul_batch_offset_alpha_beta + category: pre_checkin + function: matmul_batch_offset + precision: *real_precisions + transA: N + transB: N + M: 256 + N: 128 + K: 64 + batch_mode: 1 + batch_count: 4 + batch_offset_a: 64 + batch_offset_b: 64 + batch_offset_c: 64 + batch_offset_d: 64 + alpha_beta: *alpha_beta_range + unit_check: 1 + norm_check: 1 + +# Large offset test +- name: matmul_batch_offset_large + category: nightly + function: matmul_batch_offset + precision: [*hpa_half_precision, *hpa_bf16_precision] + transA: N + transB: N + M: 1024 + N: 1024 + K: 256 + lda: 1024 + ldb: 256 + ldc: 1024 + ldd: 1024 + batch_mode: 1 + batch_count: 2 + batch_offset_a: 4294967296 + batch_offset_b: 4294967297 + batch_offset_c: 4294967298 + batch_offset_d: 4097 + alpha: 1.0 + beta: 1.0 + unit_check: 1 + norm_check: 1 + gpu_arch: '9(42|50)' + +# Test all solutions with batch offsets and all transpose combinations +- name: matmul_batch_offset_all_solutions + category: nightly + function: matmul_batch_offset + precision: *real_precisions + transA_transB: *transA_transB_range + M: 256 + N: 256 + K: 128 + lda: 256 + ldb: 256 + ldc: 256 + ldd: 256 + batch_mode: 1 + batch_count: 2 + batch_offset_a: 256 + batch_offset_b: 128 + batch_offset_c: 512 + batch_offset_d: 256 + alpha: 1.0 + beta: 1.0 + requested_solution_num: -1 + unit_check: 1 + +# Test with negative batch offsets to verify proper memory layout handling +- name: matmul_batch_offset_negative + category: nightly + function: matmul_batch_offset + precision: *real_precisions + transA_transB: *transA_transB_range + M: 256 + N: 256 + K: 128 + lda: 256 + ldb: 256 + ldc: 256 + ldd: 256 + batch_mode: 1 + batch_count: 2 + batch_offset_a: -128 + batch_offset_b: -64 + batch_offset_c: -256 + batch_offset_d: -128 + alpha: 1.0 + beta: 1.0 + unit_check: 1 + +# Test with mixed positive and negative offsets +- name: matmul_batch_offset_mixed + category: nightly + function: matmul_batch_offset + precision: *real_precisions + transA_transB: *transA_transB_range + M: 256 + N: 256 + K: 128 + lda: 256 + ldb: 256 + ldc: 256 + ldd: 256 + batch_mode: 1 + batch_count: 2 + batch_offset_a: -64 + batch_offset_b: 128 + batch_offset_c: -128 + batch_offset_d: 256 + alpha: 1.0 + beta: 1.0 + unit_check: 1 + ... diff --git a/projects/hipblaslt/clients/tests/src/matmul_gtest.cpp b/projects/hipblaslt/clients/tests/src/matmul_gtest.cpp index 5d278d315add..86f234388aef 100644 --- a/projects/hipblaslt/clients/tests/src/matmul_gtest.cpp +++ b/projects/hipblaslt/clients/tests/src/matmul_gtest.cpp @@ -27,6 +27,7 @@ #include "hipblaslt_datatype2string.hpp" #include "hipblaslt_test.hpp" #include "testing_matmul.hpp" +#include "testing_matmul_batch_offset.hpp" #include #include #include @@ -48,6 +49,8 @@ namespace testing_matmul(arg); else if(!strcmp(arg.function, "matmul_bad_arg")) testing_matmul_bad_arg(arg); + else if(!strcmp(arg.function, "matmul_batch_offset")) + testing_matmul_batch_offset(arg); else FAIL() << "Internal error: Test called with unknown function: " << arg.function; } @@ -64,7 +67,8 @@ namespace // Filter for which functions apply to this suite static bool function_filter(const Arguments& arg) { - return !strcmp(arg.function, "matmul") || !strcmp(arg.function, "matmul_bad_arg"); + return !strcmp(arg.function, "matmul") || !strcmp(arg.function, "matmul_bad_arg") + || !strcmp(arg.function, "matmul_batch_offset"); } // Google Test name suffix based on parameters diff --git a/projects/hipblaslt/library/include/hipblaslt/hipblaslt.h b/projects/hipblaslt/library/include/hipblaslt/hipblaslt.h index dd391231c0f0..795e58661493 100644 --- a/projects/hipblaslt/library/include/hipblaslt/hipblaslt.h +++ b/projects/hipblaslt/library/include/hipblaslt/hipblaslt.h @@ -163,15 +163,25 @@ typedef enum { * ``int64_t`` */ HIPBLASLT_MATRIX_LAYOUT_LD = 6, + /** Matrix Batch Mode. * Batched GEMM can be either: * 1. Strided Batch: Single contiguous memory allocation and stride between matrices in * the batch is specified in terms of number of elements. - * 2. General Batched: This uses pointer array with each pointer storing the base address + * 2. General Batched: This uses pointer array with each pointer storing the base address * of the matrices in the batch. * See hipblasLtBatchMode_t */ - HIPBLASLT_MATRIX_LAYOUT_BATCH_MODE = 7, + HIPBLASLT_MATRIX_LAYOUT_BATCH_MODE = 7, + + /** 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 (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. + */ + HIPBLASLT_MATRIX_LAYOUT_OFFSET = 8 } hipblasLtMatrixLayoutAttribute_t; /*! \ingroup types_module diff --git a/projects/hipblaslt/library/src/amd_detail/hipblaslt.cpp b/projects/hipblaslt/library/src/amd_detail/hipblaslt.cpp index 7191a2a2bd13..eed7a873d66a 100644 --- a/projects/hipblaslt/library/src/amd_detail/hipblaslt.cpp +++ b/projects/hipblaslt/library/src/amd_detail/hipblaslt.cpp @@ -262,7 +262,7 @@ try { rocblaslt::Debug::Instance().markerStart("hipblasLtMatrixLayoutDestroy"); auto status = RocBlasLtStatusToHIPStatus( - rocblaslt_matrix_layout_destory((const rocblaslt_matrix_layout)descr)); + rocblaslt_matrix_layout_destroy((const rocblaslt_matrix_layout)descr)); rocblaslt::Debug::Instance().markerStop(); return status; } diff --git a/projects/hipblaslt/library/src/amd_detail/include/auxiliary.hpp b/projects/hipblaslt/library/src/amd_detail/include/auxiliary.hpp index c5dd621ad5f3..320de121e4ca 100644 --- a/projects/hipblaslt/library/src/amd_detail/include/auxiliary.hpp +++ b/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) diff --git a/projects/hipblaslt/library/src/amd_detail/rocblaslt/include/rocblaslt-auxiliary.h b/projects/hipblaslt/library/src/amd_detail/rocblaslt/include/rocblaslt-auxiliary.h index ffa3191ef4c8..796fc86a2c6d 100644 --- a/projects/hipblaslt/library/src/amd_detail/rocblaslt/include/rocblaslt-auxiliary.h +++ b/projects/hipblaslt/library/src/amd_detail/rocblaslt/include/rocblaslt-auxiliary.h @@ -117,7 +117,7 @@ rocblaslt_status rocblaslt_get_sm_count_target(rocblaslt_handle handle, * \brief Create a descriptor for matrix * \details * \p rocblaslt_matrix_layout_create creates a matrix descriptor It initializes - * It should be destroyed at the end using rocblaslt_matrix_layout_destory(). + * It should be destroyed at the end using rocblaslt_matrix_layout_destroy(). * * @param[out] * matDescr the pointer to the matrix descriptor @@ -136,7 +136,7 @@ rocblaslt_status rocblaslt_matrix_layout_create(rocblaslt_matrix_layout* matDesc * \brief Destroy a matrix descriptor * * \details - * \p rocblaslt_matrix_layout_destory destroys a matrix descriptor and releases + * \p rocblaslt_matrix_layout_destroy destroys a matrix descriptor and releases * all resources used by the descriptor * * @param[in] @@ -145,7 +145,7 @@ rocblaslt_status rocblaslt_matrix_layout_create(rocblaslt_matrix_layout* matDesc * \retval rocblaslt_status_success the operation completed successfully. * \retval rocblaslt_status_invalid_pointer \p descr is invalid. */ -rocblaslt_status rocblaslt_matrix_layout_destory(const rocblaslt_matrix_layout descr); +rocblaslt_status rocblaslt_matrix_layout_destroy(const rocblaslt_matrix_layout descr); rocblaslt_status rocblaslt_matrix_layout_set_attribute(rocblaslt_matrix_layout matLayout, rocblaslt_matrix_layout_attribute attr, @@ -187,7 +187,7 @@ rocblaslt_status rocblaslt_matmul_desc_create(rocblaslt_matmul_desc* matmulDesc, * \brief Destroy a matrix multiplication descriptor * * \details - * \p rocblaslt_matrix_layout_destory destroys a multiplication matrix descr. + * \p rocblaslt_matrix_layout_destroy destroys a multiplication matrix descr. * * @param[in] * descr the matrix multiplication descriptor diff --git a/projects/hipblaslt/library/src/amd_detail/rocblaslt/include/rocblaslt-types.h b/projects/hipblaslt/library/src/amd_detail/rocblaslt/include/rocblaslt-types.h index 6c1f9cf7b727..7a983557332d 100644 --- a/projects/hipblaslt/library/src/amd_detail/rocblaslt/include/rocblaslt-types.h +++ b/projects/hipblaslt/library/src/amd_detail/rocblaslt/include/rocblaslt-types.h @@ -337,13 +337,14 @@ typedef enum rocblaslt_matrix_layout_attribute_ ROCBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET = 1, /**< stride between consecutive matrices in a batch expressed in terms of matrix elements. */ - ROCBLASLT_MATRIX_LAYOUT_TYPE = 2, - ROCBLASLT_MATRIX_LAYOUT_ORDER = 3, - ROCBLASLT_MATRIX_LAYOUT_ROWS = 4, - ROCBLASLT_MATRIX_LAYOUT_COLS = 5, - ROCBLASLT_MATRIX_LAYOUT_LD = 6, + ROCBLASLT_MATRIX_LAYOUT_TYPE = 2, + ROCBLASLT_MATRIX_LAYOUT_ORDER = 3, + ROCBLASLT_MATRIX_LAYOUT_ROWS = 4, + ROCBLASLT_MATRIX_LAYOUT_COLS = 5, + ROCBLASLT_MATRIX_LAYOUT_LD = 6, ROCBLASLT_MATRIX_LAYOUT_BATCH_MODE = 7, - ROCBLASLT_MATRIX_LAYOUT_MAX = 8 + ROCBLASLT_MATRIX_LAYOUT_OFFSET = 8, + ROCBLASLT_MATRIX_LAYOUT_MAX = 9 } rocblaslt_matrix_layout_attribute; typedef enum @@ -523,6 +524,7 @@ struct RocblasltContractionProblem size_t row_stride_a; size_t col_stride_a; size_t batch_stride_a; + int64_t batch_offset_a; hipDataType b_type; const void* B; @@ -530,6 +532,7 @@ struct RocblasltContractionProblem size_t row_stride_b; size_t col_stride_b; size_t batch_stride_b; + int64_t batch_offset_b; const void* beta; @@ -539,6 +542,7 @@ struct RocblasltContractionProblem size_t row_stride_c; size_t col_stride_c; size_t batch_stride_c; + int64_t batch_offset_c; hipDataType d_type; void* D; @@ -546,6 +550,7 @@ struct RocblasltContractionProblem size_t row_stride_d; size_t col_stride_d; size_t batch_stride_d; + int64_t batch_offset_d; void* E; void* const* batch_E; @@ -612,22 +617,26 @@ struct RocblasltContractionProblem const void* const* batch_A, int64_t ld_a, int64_t batch_stride_a, + int64_t batch_offset_a, hipDataType b_type, const void* B, const void* const* batch_B, int64_t ld_b, int64_t batch_stride_b, + int64_t batch_offset_b, const void* beta, hipDataType c_type, const void* C, const void* const* batch_C, int64_t ld_c, int64_t batch_stride_c, + int64_t batch_offset_c, hipDataType d_type, void* D, void* const* batch_D, int64_t ld_d, int64_t batch_stride_d, + int64_t batch_offset_d, void* E, void* const* batch_E, int64_t ld_e, diff --git a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/include/handle.h b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/include/handle.h index 32181747a0ba..130980ee9740 100644 --- a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/include/handle.h +++ b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/include/handle.h @@ -133,7 +133,7 @@ struct _rocblaslt_handle * content. It must be initialized using rocblaslt_matrix_layout_create() * and the retured handle must be passed * to all subsequent library function calls that involve the matrix. - * It should be destroyed at the end using rocblaslt_matrix_layout_destory(). + * It should be destroyed at the end using rocblaslt_matrix_layout_destroy(). *******************************************************************************/ struct _rocblaslt_matrix_layout { @@ -152,6 +152,7 @@ struct _rocblaslt_matrix_layout hipDataType type; int32_t batch_count = 1; int64_t batch_stride = 0; + int64_t batch_offset = 0; hipblasLtOrder_t order = HIPBLASLT_ORDER_COL; // Batch Mode hipblasLtBatchMode_t batch_mode = HIPBLASLT_BATCH_MODE_STRIDED; diff --git a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/include/rocblaslt_mat_utils.hpp b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/include/rocblaslt_mat_utils.hpp index 6eca0b491472..90db7c08ac0f 100644 --- a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/include/rocblaslt_mat_utils.hpp +++ b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/include/rocblaslt_mat_utils.hpp @@ -151,6 +151,10 @@ inline rocblaslt_status validateMatmulArgs(int64_t m, int64_t batch_stride_b = 0, int64_t batch_stride_c = 0, int64_t batch_stride_d = 0, + int64_t batch_offset_a = 0, + int64_t batch_offset_b = 0, + int64_t batch_offset_c = 0, + int64_t batch_offset_d = 0, const rocblaslt_pointer_mode& pointermode = rocblaslt_pointer_mode_host) { @@ -211,17 +215,30 @@ inline rocblaslt_status validateMatmulArgs(int64_t m, if(batch_stride_a < 0 || batch_stride_b < 0 || batch_stride_c < 0 || batch_stride_d < 0) { #ifndef CODE_COVERAGE - std::cerr << "matrix and stride size must be positive" << std::endl; + std::cerr << "matrix and stride size must be zero or positive" << std::endl; #endif return rocblaslt_status_invalid_size; } - // number of batches of matrics A,B,C,D must be the same and negative + // 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) { #ifndef CODE_COVERAGE - std::cerr << " number of batches of matrics A,B,C,D must be the same and negative" + std::cerr << " number of batches of matrics A,B,C,D must be the same and positive" << std::endl; #endif return rocblaslt_status_invalid_size; @@ -329,15 +346,19 @@ inline rocblaslt_status rocblaslt_matmul_valid_args(const rocblaslt_matmul_desc hipDataType& a_type, int64_t& lda, int64_t& batch_stride_a, + int64_t& batch_offset_a, hipDataType& b_type, int64_t& ldb, int64_t& batch_stride_b, + int64_t& batch_offset_b, hipDataType& c_type, int64_t& ldc, int64_t& batch_stride_c, + int64_t& batch_offset_c, hipDataType& d_type, int64_t& ldd, int64_t& batch_stride_d, + int64_t& batch_offset_d, int64_t& lde, int64_t& batch_stride_e, void*& bias, @@ -367,18 +388,21 @@ inline rocblaslt_status rocblaslt_matmul_valid_args(const rocblaslt_matmul_desc a_type = matA->type; lda = matA->ld; batch_stride_a = matA->batch_stride; + batch_offset_a = matA->batch_offset; // matrix B int num_batches_b = matB->batch_count; b_type = matB->type; ldb = matB->ld; batch_stride_b = matB->batch_stride; + batch_offset_b = matB->batch_offset; // matrix C int num_batches_c = matC->batch_count; c_type = matC->type; ldc = matC->ld; batch_stride_c = matC->batch_stride; + batch_offset_c = matC->batch_offset; // matrix D int64_t num_rows_d = matD->m; @@ -387,6 +411,7 @@ inline rocblaslt_status rocblaslt_matmul_valid_args(const rocblaslt_matmul_desc d_type = matD->type; ldd = matD->ld; batch_stride_d = matD->batch_stride; + batch_offset_d = matD->batch_offset; compute_type = matmul_descr->compute_type; @@ -394,6 +419,32 @@ inline rocblaslt_status rocblaslt_matmul_valid_args(const rocblaslt_matmul_desc n = num_cols_d; k = (opA == HIPBLAS_OP_N) ? num_cols_a : num_rows_a; + // Validate: batch offsets are only valid with POINTER_ARRAY mode (general batched) + bool hasNonZeroOffset = (batch_offset_a != 0) || (batch_offset_b != 0) + || (batch_offset_c != 0) || (batch_offset_d != 0); + + if(hasNonZeroOffset) + { + // Check that all matrices are using POINTER_ARRAY mode + if(matA->batch_mode != HIPBLASLT_BATCH_MODE_POINTER_ARRAY + || matB->batch_mode != HIPBLASLT_BATCH_MODE_POINTER_ARRAY + || matC->batch_mode != HIPBLASLT_BATCH_MODE_POINTER_ARRAY + || matD->batch_mode != HIPBLASLT_BATCH_MODE_POINTER_ARRAY) + { + log_error(__func__, + "Batch offsets require all matrices to use batch_mode=POINTER_ARRAY. ", + "Current modes: A=", matA->batch_mode, + ", B=", matB->batch_mode, + ", C=", matC->batch_mode, + ", D=", matD->batch_mode, + ". Offsets: A=", batch_offset_a, + ", B=", batch_offset_b, + ", C=", batch_offset_c, + ", D=", batch_offset_d); + return rocblaslt_status_invalid_value; + } + } + auto matmul_status = validateMatmulArgs(m, n, k, @@ -418,6 +469,10 @@ inline rocblaslt_status rocblaslt_matmul_valid_args(const rocblaslt_matmul_desc batch_stride_b, batch_stride_c, batch_stride_d, + batch_offset_a, + batch_offset_b, + batch_offset_c, + batch_offset_d, matmul_descr->pointermode); const void* alphaVecPtr = matmul_descr->pointermode ? alpha : nullptr; diff --git a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_auxiliary.cpp b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_auxiliary.cpp index 79dc03366528..37a6a1d2debc 100644 --- a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_auxiliary.cpp +++ b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_auxiliary.cpp @@ -367,8 +367,9 @@ RocblasltContractionProblem construct_rocblaslt_problem(rocblaslt_handle { int8_t dummy; const void* dummy_ptr = &dummy; - int64_t m, n, k, lda, ldb, ldc, ldd, lde, batch_stride_a, batch_stride_b, batch_stride_c, - batch_stride_d, batch_stride_e; + int64_t m, n, k, lda, ldb, ldc, ldd, lde; + int64_t batch_stride_a, batch_stride_b, batch_stride_c, batch_stride_d, batch_stride_e; + int64_t batch_offset_a, batch_offset_b, batch_offset_c, batch_offset_d; int32_t bias_stride = matmul_descr->bias_stride; hipDataType bias_type; hipDataType aux_type; @@ -396,15 +397,19 @@ RocblasltContractionProblem construct_rocblaslt_problem(rocblaslt_handle a_type, lda, batch_stride_a, + batch_offset_a, b_type, ldb, batch_stride_b, + batch_offset_b, c_type, ldc, batch_stride_c, + batch_offset_c, d_type, ldd, batch_stride_d, + batch_offset_d, lde, batch_stride_e, bias, @@ -460,22 +465,26 @@ RocblasltContractionProblem construct_rocblaslt_problem(rocblaslt_handle nullptr, lda, batch_stride_a, + batch_offset_a, b_type, nullptr, nullptr, ldb, batch_stride_b, + batch_offset_b, beta, c_type, nullptr, nullptr, ldc, batch_stride_c, + batch_offset_c, d_type, nullptr, nullptr, ldd, batch_stride_d, + batch_offset_d, e, nullptr, lde, @@ -634,7 +643,7 @@ rocblaslt_status rocblaslt_get_sm_count_target(rocblaslt_handle handle, * content. It must be initialized using rocblaslt_matrix_layout_create() * and the retured handle must be passed * to all subsequent library function calls that involve the matrix. - * It should be destroyed at the end using rocblaslt_matrix_layout_destory(). + * It should be destroyed at the end using rocblaslt_matrix_layout_destroy(). *******************************************************************************/ rocblaslt_status rocblaslt_matrix_layout_create(rocblaslt_matrix_layout* matDescr, hipDataType valueType, @@ -682,7 +691,7 @@ rocblaslt_status rocblaslt_matrix_layout_create(rocblaslt_matrix_layout* matDesc /******************************************************************************** * \brief destroy matrix descriptor *******************************************************************************/ -rocblaslt_status rocblaslt_matrix_layout_destory(const rocblaslt_matrix_layout matDescr) +rocblaslt_status rocblaslt_matrix_layout_destroy(const rocblaslt_matrix_layout matDescr) { if(matDescr == nullptr) { @@ -804,8 +813,17 @@ rocblaslt_status rocblaslt_matrix_layout_set_attribute(rocblaslt_matrix_layout { log_error(__func__, "invalid buf size", sizeInBytes); return rocblaslt_status_invalid_value; - } - break; + } + break; + case ROCBLASLT_MATRIX_LAYOUT_OFFSET: + if(sizeof(int64_t) <= sizeInBytes) + memcpy(&matLayout->batch_offset, buf, sizeof(int64_t)); + else + { + log_error(__func__, "invalid buf size", sizeInBytes); + return rocblaslt_status_invalid_value; + } + break; default: log_error(__func__, "invalid attribute", attr); return rocblaslt_status_invalid_value; @@ -892,6 +910,16 @@ rocblaslt_status rocblaslt_matrix_layout_get_attribute(rocblaslt_matrix_layout } memcpy(buf, &matLayout->batch_mode, sizeof(int32_t)); break; + case ROCBLASLT_MATRIX_LAYOUT_OFFSET: + if(sizeWritten) + *sizeWritten = sizeof(int64_t); + if(sizeInBytes < sizeof(int64_t)) + { + log_error(__func__, "invalid buf size", sizeInBytes); + return rocblaslt_status_invalid_value; + } + memcpy(buf, &matLayout->batch_offset, sizeof(int64_t)); + break; default: log_error(__func__, "invalid attribute", attr); return rocblaslt_status_invalid_value; diff --git a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_mat.cpp b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_mat.cpp index 5f001b8ef88d..679b19eb8779 100644 --- a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_mat.cpp +++ b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_mat.cpp @@ -58,8 +58,10 @@ rocblaslt_status rocblaslt_matmul_impl(const rocblaslt_handle handle, size_t workspaceSizeInBytes, hipStream_t stream) { - int64_t m, n, k, lda, ldb, ldc, ldd, lde, batch_stride_a, batch_stride_b, batch_stride_c, - batch_stride_d, batch_stride_e; + int64_t m, n, k, lda, ldb, ldc, ldd, lde; + int64_t batch_stride_a, batch_stride_b, batch_stride_c, batch_stride_d, batch_stride_e; + int64_t batch_offset_a, batch_offset_b, batch_offset_c, batch_offset_d; + hipDataType bias_type; hipDataType aux_type; hipDataType type_a, type_b, type_c, type_d; @@ -87,15 +89,19 @@ rocblaslt_status rocblaslt_matmul_impl(const rocblaslt_handle handle, type_a, lda, batch_stride_a, + batch_offset_a, type_b, ldb, batch_stride_b, + batch_offset_b, type_c, ldc, batch_stride_c, + batch_offset_c, type_d, ldd, batch_stride_d, + batch_offset_d, lde, batch_stride_e, bias, @@ -125,6 +131,8 @@ rocblaslt_status rocblaslt_matmul_impl(const rocblaslt_handle handle, hipDataType scale_type = matmul_descr->scale_type; // Others + // Use strided_batch=true for kernel selection (StridedBatched=true kernels with SupportUserArgs) + // The actual batch mode is tracked via problem.batchMode() for argument passing bool strided_batch = true; bool grouped_gemm = false; @@ -174,22 +182,26 @@ rocblaslt_status rocblaslt_matmul_impl(const rocblaslt_handle handle, nullptr, lda, batch_stride_a, + batch_offset_a, type_b, B, nullptr, ldb, batch_stride_b, + batch_offset_b, beta, type_c, C, nullptr, ldc, batch_stride_c, + batch_offset_c, type_d, D, nullptr, ldd, batch_stride_d, + batch_offset_d, E, nullptr, lde, @@ -265,8 +277,10 @@ rocblaslt_status rocblaslt_gemm_create_cpp_impl(const rocblaslt_handle std::shared_ptr& gemmData, size_t& gemmCount) { - int64_t m, n, k, lda, ldb, ldc, ldd, lde, batch_stride_a, batch_stride_b, batch_stride_c, - batch_stride_d, batch_stride_e; + int64_t m, n, k, lda, ldb, ldc, ldd, lde; + int64_t batch_stride_a, batch_stride_b, batch_stride_c, batch_stride_d, batch_stride_e; + int64_t batch_offset_a, batch_offset_b, batch_offset_c, batch_offset_d; + hipDataType bias_type; hipDataType aux_type; hipDataType type_a, type_b, type_c, type_d; @@ -293,15 +307,19 @@ rocblaslt_status rocblaslt_gemm_create_cpp_impl(const rocblaslt_handle type_a, lda, batch_stride_a, + batch_offset_a, type_b, ldb, batch_stride_b, + batch_offset_b, type_c, ldc, batch_stride_c, + batch_offset_c, type_d, ldd, batch_stride_d, + batch_offset_d, lde, batch_stride_e, bias, @@ -330,6 +348,8 @@ rocblaslt_status rocblaslt_gemm_create_cpp_impl(const rocblaslt_handle void* amaxD = matmul_descr->amaxD; // Others + // Use strided_batch=true for kernel selection (StridedBatched=true kernels with SupportUserArgs) + // The actual batch mode is tracked via problem.batchMode() for argument passing bool strided_batch = true; bool grouped_gemm = false; @@ -358,22 +378,26 @@ rocblaslt_status rocblaslt_gemm_create_cpp_impl(const rocblaslt_handle nullptr, lda, batch_stride_a, + batch_offset_a, type_b, B, nullptr, ldb, batch_stride_b, + batch_offset_b, beta, type_c, C, nullptr, ldc, batch_stride_c, + batch_offset_c, type_d, D, nullptr, ldd, batch_stride_d, + batch_offset_d, E, nullptr, lde, @@ -652,22 +676,26 @@ rocblaslt_status nullptr, lda_vec[i], batch_stride_a_vec[i], + 0, // batch_offset_a type_b, B_vec[i], nullptr, ldb_vec[i], batch_stride_b_vec[i], + 0, // batch_offset_b beta_vec[i], type_c, C_vec[i], nullptr, ldc_vec[i], batch_stride_c_vec[i], + 0, // batch_offset_c type_d, D_vec[i], nullptr, ldd_vec[i], batch_stride_d_vec[i], + 0, // batch_offset_d E_vec[i], nullptr, lde_vec[i], @@ -997,22 +1025,26 @@ rocblaslt_status rocblaslt_gemm_create_cpp_impl_2(const rocblaslt_handle handle, nullptr, lda, batch_stride_a, + 0, // batch_offset_a, type_b, B, nullptr, ldb, batch_stride_b, + 0, // batch_offset_b, beta, type_c, C, nullptr, ldc, batch_stride_c, + 0, // batch_offset_c, type_d, D, nullptr, ldd, batch_stride_d, + 0, // batch_offset_d, E, nullptr, lde, @@ -1317,22 +1349,26 @@ rocblaslt_status rocblaslt_groupedgemm_create_cpp_impl_2(const rocblaslt_handle nullptr, lda[i], strideA[i], + 0, // batch_offset_a type_b, B_vec[i], nullptr, ldb[i], strideB[i], + 0, // batch_offset_b beta_vec[i], type_c, C_vec[i], nullptr, ldc[i], strideC[i], + 0, // batch_offset_c type_d, D_vec[i], nullptr, ldd[i], strideD[i], + 0, // batch_offset_d E_vec[i], nullptr, lde_vec[i], diff --git a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/tensile_host.cpp b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/tensile_host.cpp index 323e802910b2..da394dcd3b5f 100644 --- a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/tensile_host.cpp +++ b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/tensile_host.cpp @@ -89,22 +89,26 @@ RocblasltContractionProblem::RocblasltContractionProblem(hipblasOperation_t const void* const* batch_A, int64_t ld_a, int64_t batch_stride_a, + int64_t batch_offset_a, hipDataType b_type, const void* B, const void* const* batch_B, int64_t ld_b, int64_t batch_stride_b, + int64_t batch_offset_b, const void* beta, hipDataType c_type, const void* C, const void* const* batch_C, int64_t ld_c, int64_t batch_stride_c, + int64_t batch_offset_c, hipDataType d_type, void* D, void* const* batch_D, int64_t ld_d, int64_t batch_stride_d, + int64_t batch_offset_d, void* E, void* const* batch_E, int64_t ld_e, @@ -152,12 +156,14 @@ RocblasltContractionProblem::RocblasltContractionProblem(hipblasOperation_t , row_stride_a(1) , col_stride_a(ld_a) , batch_stride_a(batch_stride_a) + , batch_offset_a(batch_offset_a) , b_type(b_type) , B(B) , batch_B(batch_B) , row_stride_b(1) , col_stride_b(ld_b) , batch_stride_b(batch_stride_b) + , batch_offset_b(batch_offset_b) , beta(beta) , c_type(c_type) , C(C) @@ -165,12 +171,14 @@ RocblasltContractionProblem::RocblasltContractionProblem(hipblasOperation_t , row_stride_c(1) , col_stride_c(ld_c) , batch_stride_c(batch_stride_c) + , batch_offset_c(batch_offset_c) , d_type(d_type) , D(D) , batch_D(batch_D) , row_stride_d(1) , col_stride_d(ld_d) , batch_stride_d(batch_stride_d) + , batch_offset_d(batch_offset_d) , E(E) , batch_E(batch_E) , row_stride_e(1) @@ -2416,6 +2424,23 @@ namespace inputs.batchC = reinterpret_cast(prob.batch_C); inputs.batchD = reinterpret_cast(prob.batch_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; inputs.workspaceSize = prob.workspaceSize; @@ -4456,6 +4481,24 @@ rocblaslt_status getAllSolutions(MyProblem& int duplicated_counts = 0; for(auto solution : solutions) { + // Custom kernels don't support general batched mode (pointer arrays) + // Only check for ContractionProblemGemm (grouped gemm doesn't use batchMode) + if constexpr(std::is_same::value) + { + if(prob.batchMode() == TensileLite::ContractionProblemGemm::BATCHMODE::POINTER_ARRAY + && !solution->sizeMapping.customKernelName.empty()) + { + if(get_logger_layer_mode() & rocblaslt_layer_mode_log_info) + { + std::ostringstream msg; + msg << "Skipping custom kernel " << solution->sizeMapping.customKernelName + << " - does not support batch_mode=POINTER_ARRAY" << std::endl; + log_info(__func__, msg.str()); + } + continue; + } + } + //workaround: findAllSolutions should get all solutions without duplications bool duplicated_sol = false; for(int j = 0; j < i; j++) diff --git a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/utility.cpp b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/utility.cpp index 57cfc325a651..d56ab2f29896 100644 --- a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/utility.cpp +++ b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/utility.cpp @@ -202,9 +202,9 @@ const char* rocblaslt_matrix_layout_attributes_to_string(rocblaslt_matrix_layout switch(type) { case ROCBLASLT_MATRIX_LAYOUT_BATCH_COUNT: - return "MATRIX_LAYOUT_BATCH_COUNT"; + return "ROCBLASLT_MATRIX_LAYOUT_BATCH_COUNT"; case ROCBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET: - return "MATRIX_LAYOUT_STRIDED_BATCH_OFFSET"; + return "ROCBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET"; case ROCBLASLT_MATRIX_LAYOUT_TYPE: return "ROCBLASLT_MATRIX_LAYOUT_TYPE"; case ROCBLASLT_MATRIX_LAYOUT_ORDER: @@ -217,6 +217,8 @@ const char* rocblaslt_matrix_layout_attributes_to_string(rocblaslt_matrix_layout return "ROCBLASLT_MATRIX_LAYOUT_LD"; case ROCBLASLT_MATRIX_LAYOUT_BATCH_MODE: return "ROCBLASLT_MATRIX_LAYOUT_BATCH_MODE"; + case ROCBLASLT_MATRIX_LAYOUT_OFFSET: + return "ROCBLASLT_MATRIX_LAYOUT_OFFSET"; case ROCBLASLT_MATRIX_LAYOUT_MAX: return "ROCBLASLT_MATRIX_LAYOUT_MAX"; default: diff --git a/projects/hipblaslt/tensilelite/Tensile/Components/Signature.py b/projects/hipblaslt/tensilelite/Tensile/Components/Signature.py index 7d101e4ea4f5..847a17e510e7 100644 --- a/projects/hipblaslt/tensilelite/Tensile/Components/Signature.py +++ b/projects/hipblaslt/tensilelite/Tensile/Components/Signature.py @@ -311,6 +311,28 @@ def __call__(self, writer) -> SignatureBase: signature.addArg( "Synchronizer", SVK.SIG_GLOBALBUFFER, cptValueType, "generic") signature.addArg( "GSUSync", SVK.SIG_VALUE, "u32") + # Batch offset support for general batched mode (pointer array). + # Placed at the tail of the kernarg buffer (after the dstD/Synchronizer block) + # so no later arg is shifted; the host appends them in the same position, + # after the dstD/Synchronizer/seed block. Record each arg's kernarg byte + # offset so the assembly loads them from the accurate position rather + # than re-deriving it. + # + # signature.offset counts from the very first arg including the common header. + # The assembly loads these args with KernArgAddress already advanced past + # that header by commonArgsSize, so subtract it. + if not kernel["ProblemType"]["GroupedGemm"]: + commonArgsSize = userArgumentsInfo.commonArgsSize + writer.states.batchOffsetDKernArgOffset = signature.offset - commonArgsSize + signature.addArg("batchOffsetD", SVK.SIG_VALUE, "u64") + writer.states.batchOffsetCKernArgOffset = signature.offset - commonArgsSize + signature.addArg("batchOffsetC", SVK.SIG_VALUE, "u64") + writer.states.batchOffsetAKernArgOffset = signature.offset - commonArgsSize + signature.addArg("batchOffsetA", SVK.SIG_VALUE, "u64") + writer.states.batchOffsetBKernArgOffset = signature.offset - commonArgsSize + signature.addArg("batchOffsetB", SVK.SIG_VALUE, "u64") + userArgumentsInfo.gemmArgumentSize += 32 # 4 offsets * 8 bytes each + activationType = ActivationType("all") for name in activationType.getAdditionalArgStringList(): userArgumentsInfo.activationSize += userArgumentsInfo.actMaxSize diff --git a/projects/hipblaslt/tensilelite/Tensile/KernelWriter.py b/projects/hipblaslt/tensilelite/Tensile/KernelWriter.py index a63a57da1e7c..d6926f0d24df 100644 --- a/projects/hipblaslt/tensilelite/Tensile/KernelWriter.py +++ b/projects/hipblaslt/tensilelite/Tensile/KernelWriter.py @@ -310,6 +310,12 @@ class StateValues: numSgprToLoad: int = 0 # For kernel args preloadGuard: List[int] = field(init=False) # For preload kernel args guard numSgprPreload: int = 0 # For kernel args + # Kernarg byte offsets of the general-batched offsets, set when the signature + # appends them at the tail (0 for grouped-gemm kernels that omit them) + batchOffsetDKernArgOffset: int = 0 + batchOffsetCKernArgOffset: int = 0 + batchOffsetAKernArgOffset: int = 0 + batchOffsetBKernArgOffset: int = 0 numSgprAlpha: int = 0 # For user arguments numSgprBeta: int = 0 # For user arguments numStoreSgprNames: List[str] = field(init=False) # For post-loop kernel args @@ -9233,6 +9239,10 @@ def vgprAllocationImplSubtile(): if kernel["ProblemType"]["Sparse"]: self.defineSgpr("StridesMetadata", self.states.m.numSgprStrides) + # Batch offset support for general batched GEMM (pointer array mode) + # Offsets are loaded on-demand from kernel arguments to avoid using persistent SGPRs + # No SGPR allocation here - offsets loaded directly when applying to addresses + # for packed batches without stride restrictions need to do something different here assert sorted(kernel["PackedC0IdxChars"]+kernel["PackedC1IdxChars"]) == \ sorted(set(kernel["PackedC0IdxChars"]+kernel["PackedC1IdxChars"])) diff --git a/projects/hipblaslt/tensilelite/Tensile/KernelWriterAssembly.py b/projects/hipblaslt/tensilelite/Tensile/KernelWriterAssembly.py index eb25ac4b68e4..62e1d7d18a11 100644 --- a/projects/hipblaslt/tensilelite/Tensile/KernelWriterAssembly.py +++ b/projects/hipblaslt/tensilelite/Tensile/KernelWriterAssembly.py @@ -2184,6 +2184,12 @@ def getKernelArgLoadModule(self, kernel, sgprStartIdx, numsOfLoad, preloadNum): kernelArgs.add(self.argLoader.loadKernArg(self.states.esmRuntimeFlagSgpr, "KernArgAddress", sgprOffset=hex(sgprOffset), dword=1)) sgprOffset += 4 self.argLoader.setOffset(sgprOffset) + + # Batch offset arguments for general batched mode are loaded on-demand from + # their kernarg byte offsets. Those offsets are recorded accurately by + # the signature builder (see Signature.py) into self.states.batchOffset*KernArgOffset, + # so nothing is computed here. + return kernelArgs def localReadAddresses(self, kernel, tPA, tPB, tPM): @@ -4638,7 +4644,16 @@ def computeLoadSrd(self, kernel, tP, tc, indices, bpe): moduleLoadGeneralBatch.add(SAddU32(dst=sgpr(stmp+0), src0=sgpr(stmp+0), src1=sgpr("Address%s+0"%tc), comment="Offsetting to the location [Lower half of address]")) moduleLoadGeneralBatch.add(SAddCU32(dst=sgpr(stmp+1), src0=sgpr("Address%s+1"%tc), src1=0, comment="Offsetting to the location [Higher half of address]")) moduleLoadGeneralBatch.add(SLoadB64(dst=sgpr("Srd%s"%tc, 2), base=sgpr(stmp, 2), soffset=0, comment="Load the Matrix Address in the Pointer Array")) - moduleLoadGeneralBatch.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address Load from the Pointer Array")) + # Load and apply batch offset for General Batched GEMM + if not kernel["ProblemType"]["GroupedGemm"]: + batchOffsetKernArgOffset = self.states.batchOffsetAKernArgOffset if tc == "A" else self.states.batchOffsetBKernArgOffset + moduleLoadGeneralBatch.add(SLoadB64(dst=sgpr(stmp, 2), base=sgpr("KernArgAddress", 2), soffset=hex(batchOffsetKernArgOffset), comment="Load batchOffset%s from kernel args"%tc)) + moduleLoadGeneralBatch.add(SWaitCnt(kmcnt=0, comment="Wait for Matrix Address and Batch Offset Loads")) + moduleLoadGeneralBatch.add(SAddU32(dst=sgpr("Srd%s+0"%tc), src0=sgpr("Srd%s+0"%tc), src1=sgpr(stmp+0), comment="Add batch offset to %s address (low)"%tc)) + moduleLoadGeneralBatch.add(SAddCU32(dst=sgpr("Srd%s+1"%tc), src0=sgpr("Srd%s+1"%tc), src1=sgpr(stmp+1), comment="Add batch offset to %s address (high)"%tc)) + else: + moduleLoadGeneralBatch.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address Load from the Pointer Array")) + if self.states.groOffsetInMacroTile and ((tc == "A" and not kernel["enableTDMA"]) or (tc == "B" and not kernel["enableTDMB"])): prePad1 = int(self.states.srdShiftLeft[tc] * tP["bpeGR"]) # leave room in case we have to pointer shift moduleLoadGeneralBatch.add(SSubU32(dst=sgpr("Srd%s+0"%tc), src0=sgpr("Srd%s+0"%tc), src1=prePad1, comment="pre-pad to make room for possible pointer shift")) @@ -13271,7 +13286,21 @@ def computeStoreSrdStart(self, kernel, srdTcList: list, sgprBpeList = [], useSiz module.add(SLoadB64(dst=sgpr(tmpS0, 2), base=sgpr(tmpS0, 2), soffset=0, comment="Load the Matrix Address in the Pointer Array")) module.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address Load from the Pointer Array")) module.add(SAddU32(dst=sgpr("Srd%s+0"%mat), src0=sgpr("Srd%s+0"%mat), src1=sgpr(tmpS0), comment="Offsetting within the Batch Matrix [Lower half of address]")) - module.add(SAddCU32(dst=sgpr("Srd%s+1"%mat), src0=sgpr("Srd%s+1"%mat), src1=sgpr(tmpS1), comment="Offsetting within the Batch Matrix [Higher half of address]")) + module.add(SAddCU32(dst=sgpr("Srd%s+1"%mat), src0=sgpr("Srd%s+1"%mat), src1=sgpr(tmpS1), comment="Offsetting within the Batch Matrix [Higher half of address]")) + # Now, we have starting matrix address of a specific batch in the corresponding Srd. + # Load and apply batch offset for General Batched GEMM (C or D matrix) as necessary. + # This block sits inside the generalBatchedGemmLoad label, which the GSU routing + # only reaches at runtime GSU==1 -- where SrdC/SrdD are the real user pointer arrays + # (not the GSU workspace), so the offset is correctly applied here. When GSU>1 the + # routing branches to the strided/workspace path and skips this block; the PostGSU + # conversion kernel applies the offset when it writes the final result to C/D. + if not kernel["ProblemType"]["GroupedGemm"]: + batchOffsetKernArgOffset = self.states.batchOffsetCKernArgOffset if mat == "C" else self.states.batchOffsetDKernArgOffset + module.add(SLoadB64(dst=sgpr(tmpS0, 2), base=sgpr("KernArgAddress", 2), soffset=hex(batchOffsetKernArgOffset), comment="Load batchOffset%s from kernel args"%mat)) + module.add(SWaitCnt(kmcnt=0, comment="Wait for Matrix Address and Batch Offset Loads")) + # Add loaded matrix address to SRD + module.add(SAddU32(dst=sgpr("Srd%s+0"%mat), src0=sgpr("Srd%s+0"%mat), src1=sgpr(tmpS0), comment="Add matrix address to SRD (low)")) + module.add(SAddCU32(dst=sgpr("Srd%s+1"%mat), src0=sgpr("Srd%s+1"%mat), src1=sgpr(tmpS1), comment="Add matrix address to SRD (high)")) module.add(generalBatchedGemmLoad_End) module.addSpaceLine() @@ -13364,6 +13393,12 @@ def SrdTDInit(self, kernel): module.add(SAddCU32(dst=sgpr(tmpspgr2+1), src0=sgpr("AddressTD+1"), src1=0, comment="Offsetting to the location [Higher half of address]")) module.add(SLoadB64(dst=sgpr("SrdTD", 2), base=sgpr(tmpspgr2, 2), soffset=0, comment="Load the Matrix Address in the Pointer Array")) module.add(SWaitCnt(kmcnt=0, comment="Wait for the Matrix Address Load from the Pointer Array")) + # Load and apply batch offset for General Batched GEMM (dstD) as necessary. + if not kernel["ProblemType"]["GroupedGemm"]: + module.add(SLoadB64(dst=sgpr(tmpspgr2, 2), base=sgpr("KernArgAddress", 2), soffset=hex(self.states.batchOffsetDKernArgOffset), comment="Load batchOffsetD from kernel args")) + module.add(SWaitCnt(kmcnt=0, comment="Wait for batchOffsetD Load")) + module.add(SAddU32(dst=sgpr("SrdTD+0"), src0=sgpr("SrdTD+0"), src1=sgpr(tmpspgr2+0), comment="Apply batchOffsetD to SrdTD (low)")) + module.add(SAddCU32(dst=sgpr("SrdTD+1"), src0=sgpr("SrdTD+1"), src1=sgpr(tmpspgr2+1), comment="Apply batchOffsetD to SrdTD (high)")) module.add(SrdTDGeneralBatched_End) module.add(SAddU32(dst=sgpr("SrdTD+0"), src0=sgpr("SrdTD+0"), src1=sgpr(tmpspgr1+0), comment="add lo to SRTD" )) module.add(SAddCU32(dst=sgpr("SrdTD+1"), src0=sgpr("SrdTD+1"), src1=sgpr(tmpspgr1+1), comment="add hi to SRTD" )) diff --git a/projects/hipblaslt/tensilelite/Tensile/KernelWriterConversion.py b/projects/hipblaslt/tensilelite/Tensile/KernelWriterConversion.py index fa7fa8145f03..d685abe2a73e 100644 --- a/projects/hipblaslt/tensilelite/Tensile/KernelWriterConversion.py +++ b/projects/hipblaslt/tensilelite/Tensile/KernelWriterConversion.py @@ -207,7 +207,8 @@ def functionSignature(self): # Additional argument batch_mode is added to distinguish between Strided Batch and General Batched GEMM # batch_mode will dictate how the GLOBAL_C and GLOBAL_D macros are defined and used in the kernel body # since the index calculation for Strided Batch and General Batch GEMM are different. - kStr += " argument_%s arg, uint32_t batch_mode, uint32_t additionalPaddingPerBatch)" % ( self.kernelName ) + self.endLine + # Also, batchOffsetD and batchOffsetC arguments added at the end. + kStr += " argument_%s arg, uint32_t batch_mode, uint32_t additionalPaddingPerBatch, int64_t batchOffsetD, int64_t batchOffsetC)" % ( self.kernelName ) + self.endLine return kStr @@ -349,9 +350,13 @@ def kernelBody(self): ######################################## # kernel start kStr += self.endLine + # Declare batchIdx for indexing pointer arrays in general batched mode + if not self.state["ProblemType"]["GroupedGemm"]: + kStr += " uint64_t batchIdx = 0;%s" % self.endLine + if not self.state["ProblemType"]["GroupedGemm"]: kStr += " if(batch_mode == 0)" + self.endLine - kStr += " {" + self.endLine + kStr += " {" + self.endLine kStr += " if (id*NUM_ELEMENT_LOAD >= (arg.size%s" % self.indexChars[0] for i in range(1, problemType["NumIndicesC"]): kStr += " * arg.size%s" % self.indexChars[i] @@ -361,14 +366,14 @@ def kernelBody(self): kStr += " }" + self.endLine kStr += " else" + self.endLine kStr += " {" + self.endLine - kStr += " uint64_t index2 = ((id*NUM_ELEMENT_LOAD) / (arg.size%s" % self.indexChars[0] + kStr += " batchIdx = ((id*NUM_ELEMENT_LOAD) / (arg.size%s" % self.indexChars[0] for i in range(1, problemType["NumIndicesC"]-1): kStr += " * arg.size%s" % self.indexChars[i] kStr += " + additionalPaddingPerBatch));%s" % self.endLine - kStr += " if (id*NUM_ELEMENT_LOAD >= ((index2+1) * (arg.size%s * arg.size%s)) + index2 * additionalPaddingPerBatch)%s" % (self.indexChars[0], self.indexChars[1], self.endLine) + kStr += " if (id*NUM_ELEMENT_LOAD >= ((batchIdx+1) * (arg.size%s * arg.size%s)) + batchIdx * additionalPaddingPerBatch)%s" % (self.indexChars[0], self.indexChars[1], self.endLine) kStr += " return;%s" % self.endLine - kStr += " if(index2 > 0)%s" % self.endLine - kStr += " id = id - (index2 * additionalPaddingPerBatch) / NUM_ELEMENT_LOAD;%s" % self.endLine + kStr += " if(batchIdx > 0)%s" % self.endLine + kStr += " id = id - (batchIdx * additionalPaddingPerBatch) / NUM_ELEMENT_LOAD;%s" % self.endLine kStr += " }" + self.endLine kStr += self.endLine @@ -385,6 +390,10 @@ def kernelBody(self): kStr += " id%d = id %% arg.size%s;%s" % (i, self.indexChars[i], self.endLine) kStr += " id = id / arg.size%s;%s" % (self.indexChars[i], self.endLine) + # Set batchIdx = id2 for strided batched mode (batch_mode == 0) + if not self.state["ProblemType"]["GroupedGemm"]: + kStr += " if(batch_mode == 0) batchIdx = id2;%s" % self.endLine + nonTileFreeIndices = [] ######################################## @@ -726,7 +735,8 @@ def kernelBody(self): kStr += " }" + self.endLine kStr += " else" + self.endLine kStr += " {" + self.endLine - kStr += " %s *ptr = *(reinterpret_cast<%s **>(((char *)arg.C) + (8*id2)));" % (destTypeStr, destTypeStr) + self.endLine + # Dereference C pointer array and apply the batch offset (in bytes). + kStr += " %s *ptr = *(reinterpret_cast<%s **>(((char *)arg.C) + (8*batchIdx))) + batchOffsetC/sizeof(%s);" % (destTypeStr, destTypeStr, destTypeStr) + self.endLine for vIdx in range(self.num_dword_load): kStr += " %s[%d] += arg.beta * (%s)ptr[idxC+%d];%s" % (accumStr, vIdx, intermediateDataType, vIdx, self.endLine) kStr += " }" + self.endLine @@ -822,7 +832,8 @@ def kernelBody(self): kStr += " if(batch_mode == 0) {" + self.endLine kStr += " buffer_store<%s, sizeof(%s), CacheOperation::Kind::Always>(*(%s *)%s, arg.D, idxD * sizeof(%s), 0);%s" % (storeTypeStr, storeTypeStr, storeTypeStr, resultStr, destTypeStr, self.endLine) kStr += " } else {" + self.endLine - kStr += " %s *ptr = *(reinterpret_cast<%s **>(((char *)arg.D) + (8*id2)));" % (destTypeStr, destTypeStr) + self.endLine + # Dereference D pointer array and apply the batch offset (in bytes). + kStr += " %s *ptr = *(reinterpret_cast<%s **>(((char *)arg.D) + (8*batchIdx))) + batchOffsetD/sizeof(%s);" % (destTypeStr, destTypeStr, destTypeStr) + self.endLine kStr += " buffer_store<%s, sizeof(%s), CacheOperation::Kind::Always>(*(%s *)%s, ptr, idxD * sizeof(%s), 0);%s" % (storeTypeStr, storeTypeStr, storeTypeStr, resultStr, destTypeStr, self.endLine) kStr += " }" + self.endLine else: diff --git a/projects/hipblaslt/tensilelite/include/Tensile/ContractionProblem.hpp b/projects/hipblaslt/tensilelite/include/Tensile/ContractionProblem.hpp index 0e91852bb1be..051ce6acdded 100644 --- a/projects/hipblaslt/tensilelite/include/Tensile/ContractionProblem.hpp +++ b/projects/hipblaslt/tensilelite/include/Tensile/ContractionProblem.hpp @@ -1587,6 +1587,11 @@ namespace TensileLite unsigned char const* metadata = nullptr; void const* compressed = nullptr; + int64_t batchOffsetA = 0; + int64_t batchOffsetB = 0; + int64_t batchOffsetC = 0; + int64_t batchOffsetD = 0; + // Constants ConstantVariant alpha = static_cast(0); ConstantVariant beta = static_cast(0); diff --git a/projects/hipblaslt/tensilelite/include/Tensile/ContractionSolution.hpp b/projects/hipblaslt/tensilelite/include/Tensile/ContractionSolution.hpp index f2fc1d8fa5c3..0c68cca2d1fd 100644 --- a/projects/hipblaslt/tensilelite/include/Tensile/ContractionSolution.hpp +++ b/projects/hipblaslt/tensilelite/include/Tensile/ContractionSolution.hpp @@ -529,7 +529,7 @@ namespace TensileLite KA& args, StreamKSettings const& sk, uint32_t autoGsuVal, - uint32_t additionalPaddingPerBatchGeneralBatch=0) const; + uint32_t additionalPaddingPerBatchGeneralBatch=0) const; template inline void calculateConversionCallWorkGroupItems( diff --git a/projects/hipblaslt/tensilelite/include/Tensile/KernelArguments.hpp b/projects/hipblaslt/tensilelite/include/Tensile/KernelArguments.hpp index c4a1518bf2dc..210ee40563cb 100644 --- a/projects/hipblaslt/tensilelite/include/Tensile/KernelArguments.hpp +++ b/projects/hipblaslt/tensilelite/include/Tensile/KernelArguments.hpp @@ -738,6 +738,20 @@ namespace TensileLite counter += sizeof(value); } + void alignTo(size_t alignment) + { + size_t extraElements = counter % alignment; + size_t padding = (alignment - extraElements) % alignment; + counter += padding; + } + + template + inline void appendAligned(std::string const& name, T value) + { + alignTo(alignof(T)); + append(name, value); + } + template inline void appendUnbound(std::string const& name) { diff --git a/projects/hipblaslt/tensilelite/rocisa/rocisa/src/code.cpp b/projects/hipblaslt/tensilelite/rocisa/rocisa/src/code.cpp index 9fdff89722ca..a7eafed0e3b3 100644 --- a/projects/hipblaslt/tensilelite/rocisa/rocisa/src/code.cpp +++ b/projects/hipblaslt/tensilelite/rocisa/rocisa/src/code.cpp @@ -367,6 +367,7 @@ void init_code(nb::module_ m) nb::arg("totalVgprs") = 0, nb::arg("totalSgprs") = 0) .def("setGprs", &rocisa::SignatureCodeMeta::setGprs) + .def_ro("offset", &rocisa::SignatureCodeMeta::offset) .def("addArg", &rocisa::SignatureCodeMeta::addArg, nb::arg("name"), @@ -406,6 +407,8 @@ void init_code(nb::module_ m) nb::arg("totalSgprs") = 0, nb::arg("numSgprPreload") = 0) .def("setGprs", &rocisa::SignatureBase::setGprs) + .def_prop_ro("offset", + [](const rocisa::SignatureBase& self) { return self.codeMeta.offset; }) .def("addArg", &rocisa::SignatureBase::addArg, nb::arg("name"), diff --git a/projects/hipblaslt/tensilelite/src/ContractionSolution.cpp b/projects/hipblaslt/tensilelite/src/ContractionSolution.cpp index 9a60778d3580..02689c9b6b5f 100644 --- a/projects/hipblaslt/tensilelite/src/ContractionSolution.cpp +++ b/projects/hipblaslt/tensilelite/src/ContractionSolution.cpp @@ -1835,12 +1835,23 @@ namespace TensileLite rv.args.append("dstD", inputs.d); // MBSK: synchronizer address, MB: null address rv.args.append("Synchronizer", - gsuSettings.globalAccumulation == 3 - ? inputs.Synchronizer + gsuSettings.globalAccumulation == 3 + ? inputs.Synchronizer : NULL); rv.args.append("GSUSync", 0); } + // Batch offset support for General Batched GEMM (SupportUserArgs kernels). + // Appended at the tail, after the dstD/Synchronizer block, to match the + // kernel signature order (see Signature.py). + if(!problemType.groupedGemm && sizeMapping.customKernelName.empty()) + { + rv.args.append("batchOffsetD", inputs.batchOffsetD); + rv.args.append("batchOffsetC", inputs.batchOffsetC); + rv.args.append("batchOffsetA", inputs.batchOffsetA); + rv.args.append("batchOffsetB", inputs.batchOffsetB); + } + if(problemType.stochasticRounding) { // generate seed from random generator @@ -2263,7 +2274,7 @@ namespace TensileLite KA& args, StreamKSettings const& sk, uint32_t autoGsuVal, - uint32_t additionalPaddingPerBatchGeneralBatch) const + uint32_t additionalPaddingPerBatchGeneralBatch) const { TensorDescriptor const& c = problem.c(); TensorDescriptor const& d = problem.d(); @@ -2438,11 +2449,20 @@ namespace TensileLite } // Adding the batchmode kernel argument for post GSU kernel to determine // how to index the batch dimension in Strided Batch versus General Batched. - if(problemType.groupedGemm == false) + if(problemType.groupedGemm == false && sizeMapping.customKernelName.empty()) { ContractionProblemGemm::BATCHMODE batchMode = problem.batchMode(); args.template append("batchMode", static_cast(batchMode)); - args.template append("additionalPaddingPerBatch", additionalPaddingPerBatchGeneralBatch); + args.template append("additionalPaddingPerBatch", additionalPaddingPerBatchGeneralBatch); + + // The HIP-compiled conversion kernel lays out these int64_t params on + // 8-byte-aligned kernarg slots. Match that alignment on the host so the + // bytes line up; a bare append() leaves them 4-byte-shifted when the + // preceding args don't end on an 8-byte boundary (e.g. the non-HAS + // variant), causing the kernel to read the neighboring offset into the + // high dword of the address and fault. + args.template appendAligned("batchOffsetD", inputs.batchOffsetD); + args.template appendAligned("batchOffsetC", inputs.batchOffsetC); } }