Skip to content

Commit aa34cc0

Browse files
Fix for triangular_solve BEF executable unit test failure on ROCm.
1 parent aae8198 commit aa34cc0

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

backends/gpu/lib/kernels/blas_kernels.cc

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -222,7 +222,10 @@ static Error BlasTrsmBatch(
222222
const void** a_array = const_cast<const void**>(b_array + batchCount);
223223

224224
auto side_mode = wrapper::BlasSideMode::FromOpaqueValue(*sideMode);
225-
int32_t a_num_elements = side_mode == CUBLAS_SIDE_LEFT ? m * m : n * n;
225+
int32_t a_num_elements = n * n;
226+
if ((platform == wrapper::Platform::CUDA && side_mode == CUBLAS_SIDE_LEFT) ||
227+
(platform == wrapper::Platform::ROCm && side_mode == rocblas_side_left))
228+
a_num_elements = m * m;
226229
ptrdiff_t a_batch_stride_bytes = *data_type_size_bytes * a_num_elements;
227230
ptrdiff_t b_batch_stride_bytes = *data_type_size_bytes * m * n;
228231
const char* a_ptr = static_cast<const char*>(A.pointer().raw(platform));

0 commit comments

Comments
 (0)