gpu:cuda: Fix matmul parameters for inner_product usages

oneapi-src · Dec 19, 2024 · 8eddd9b · 8eddd9b
1 parent c3b5d23
commit 8eddd9b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/src/gpu/nvidia/cudnn_matmul_executor.hpp b/src/gpu/nvidia/cudnn_matmul_executor.hpp
@@ -392,12 +392,12 @@ struct cudnn_matmul_lt_exec_t final : public cudnn_matmul_lt_base_exec_t {
                             memory_tracking::names::key_matmul_dst_in_acc_dt)
                     : xpu::sycl::interop_memory_arg_t<
                             ::sycl::access::mode::read_write>();
-            auto arg_block_a_scratch = params->source_size_ != 0
+            auto arg_block_a_scratch = params->weight_size_ != 0
                     ? CTX_SCRATCH_SYCL_MEMORY(
                             memory_tracking::names::key_gemm_blocked_a)
                     : xpu::sycl::interop_memory_arg_t<
                             ::sycl::access::mode::read_write>();
-            auto arg_block_b_scratch = params->weight_size_ != 0
+            auto arg_block_b_scratch = params->source_size_ != 0
                     ? CTX_SCRATCH_SYCL_MEMORY(
                             memory_tracking::names::key_gemm_blocked_b)
                     : xpu::sycl::interop_memory_arg_t<
@@ -457,10 +457,10 @@ struct cudnn_matmul_lt_runtime_args_exec_t final
                 matmul_params->reorder_scratch_size_, cuda_stream->queue());
 
         uint8_t *block_a_scratch_ptr
-                = alloc_ptr(matmul_params->source_size_, cuda_stream->queue());
+                = alloc_ptr(matmul_params->weight_size_, cuda_stream->queue());
 
         uint8_t *block_b_scratch_ptr
-                = alloc_ptr(matmul_params->weight_size_, cuda_stream->queue());
+                = alloc_ptr(matmul_params->source_size_, cuda_stream->queue());
 
         uint8_t *block_c_scratch_ptr
                 = alloc_ptr(matmul_params->dest_size_, cuda_stream->queue());

diff --git a/src/gpu/nvidia/cudnn_matmul_lt_impl.hpp b/src/gpu/nvidia/cudnn_matmul_lt_impl.hpp
@@ -717,7 +717,7 @@ struct cudnn_matmul_lt_impl_t {
             }
             if (!params->w_blocked_) {
                 transform_matrix(lt_handle, params, a_layout, a,
-                        blocked_a_layout, block_a_scratch, !params->trans_a_,
+                        blocked_a_layout, block_a_scratch, params->trans_a_,
                         streamId);
                 a = block_a_scratch;
             }