saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-26 19:52:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add K=0 check to MatMul<float>::Compute() specialization. (#21803)

Browse source 
Add K=0 check to `MatMul<float>::Compute()` specialization.
Add unit test to cover both primary template and float specialization.
This commit is contained in:
Edward Chen 2024-08-21 09:15:58 -07:00 committed by GitHub
parent 0e827c27fb
commit fb9ce18e88
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 24 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
onnxruntime/core/providers/cpu/math/matmul.cc
View file

@ -238,6 +238,14 @@ Status MatMul<float>::Compute(OpKernelContext* ctx) const {
if (y->Shape().Size() == 0)
return Status::OK();
if (helper.K() == 0) {
// When we have (M, 0, N) then the inputs are empty, but the output should
// be filled out with zeros.
auto output_span = y->MutableDataAsSpan<float>();
std::fill(output_span.begin(), output_span.end(), float{});
return Status::OK();
}
const auto* a_data = a->Data<float>();
const auto* b_data = b ? b->Data<float>() : nullptr;
auto* y_data = y->MutableData<float>();

26
onnxruntime/test/providers/cpu/math/matmul_test.cc
View file

@ -219,18 +219,16 @@ TEST(MathOpTest, MatMulUint64Type) {
RunMatMulTest<uint64_t>(9);
}
TEST(MathOpTest, MatMul_ZeroK) {
template <typename T>
void RunMatMulZeroKTest() {
// test with empty inputs and zero filled output
constexpr const std::array<float, 0> empty_input{};
const std::vector<float> expected_output{0, 0, 0, 0,
0, 0, 0, 0,
0, 0, 0, 0,
0, 0, 0, 0};
OpTester test("MatMul", 14);
constexpr const std::array<T, 0> empty_input{};
const std::vector<T> expected_output(4 * 4, T{});
OpTester test("MatMul", 13);
test.AddInput<float>("A", {4, 0}, empty_input);
test.AddInput<float>("B", {0, 4}, empty_input);
test.AddOutput<float>("Y", {4, 4}, expected_output);
test.AddInput<T>("A", {4, 0}, empty_input);
test.AddInput<T>("B", {0, 4}, empty_input);
test.AddOutput<T>("Y", {4, 4}, expected_output);
// No special case is implemented.
test.ConfigExcludeEps({kCoreMLExecutionProvider, kNnapiExecutionProvider,
@ -240,6 +238,14 @@ TEST(MathOpTest, MatMul_ZeroK) {
.RunWithConfig();
}
TEST(MathOpTest, MatMulZeroKFloatType) {
RunMatMulZeroKTest<float>();
}
TEST(MathOpTest, MatMulZeroKInt32Type) {
RunMatMulZeroKTest<int32_t>();
}
#if defined(USE_CUDA) || defined(USE_ROCM)
TEST(MathOpTest, MatMul_Float16) {
#ifdef USE_CUDA