Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[CUDA] Special case for K==0 in CUDA MatMul #21525

Merged
merged 5 commits into from
Aug 13, 2024
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
Special case for K==0
yuslepukhin committed Jul 26, 2024
commit 9c6b124646bc80406fa531094e720fc7d25a6d63
11 changes: 10 additions & 1 deletion onnxruntime/core/providers/cuda/math/matmul.cc
Original file line number Diff line number Diff line change
@@ -110,7 +110,16 @@ Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {

Tensor* Y = ctx->Output(0, helper.OutputShape());
// Bail out early if the output is going to be empty
if (Y->Shape().Size() == 0) return Status::OK();
const auto output_size = Y->Shape().Size();
if (output_size == 0) return Status::OK();

if (helper.K() == 0) {
// When we have (M, 0, N) then the inputs are empty, but the output is not
// fill out with zeros consistent with CPU Eigen behavior.
using CudaT = typename ToCudaType<T>::MappedType;
Fill<CudaT>(Stream(ctx), reinterpret_cast<CudaT*>(Y->MutableData<T>()), CudaT(0.f), narrow<int64_t>(output_size));
return Status::OK();
}

if (GetTuningContext()->IsTunableOpEnabled()) {
return tunable::TunableMatMul<T>(alpha_, trans_a, trans_b, trans_batch_a_, trans_batch_b_, helper, this, ctx);
10 changes: 10 additions & 0 deletions onnxruntime/test/providers/cpu/math/matmul_test.cc
Original file line number Diff line number Diff line change
@@ -37,6 +37,16 @@ template <typename T>
std::vector<MatMulTestData<T>> GenerateTestCases() {
std::vector<MatMulTestData<T>> test_cases;

test_cases.push_back(
{"test with empty inputs and zero filled output",
{4, 0},
{0, 4},
{4, 4},
{0, 0, 0, 0,
0, 0, 0, 0,
0, 0, 0, 0,
0, 0, 0, 0}});

test_cases.push_back(
{"test padding and broadcast A > B",
{3, 1, 1, 2},