From 0c1950df1a9b8792341a5d361f4f0dbefc5e9f5d Mon Sep 17 00:00:00 2001
From: vraspar <v2parikh@uwaterloo.ca>
Date: Fri, 31 Jan 2025 15:58:19 -0800
Subject: [PATCH] add matmul kernal class

---
 .../core/providers/webgpu/math/matmul.cc      | 97 +++++++++++++++++++
 .../core/providers/webgpu/math/matmul.h       | 52 ++++++++++
 2 files changed, 149 insertions(+)
 create mode 100644 onnxruntime/core/providers/webgpu/math/matmul.cc
 create mode 100644 onnxruntime/core/providers/webgpu/math/matmul.h

diff --git a/onnxruntime/core/providers/webgpu/math/matmul.cc b/onnxruntime/core/providers/webgpu/math/matmul.cc
new file mode 100644
index 0000000000..231230bc55
--- /dev/null
+++ b/onnxruntime/core/providers/webgpu/math/matmul.cc
@@ -0,0 +1,97 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/common/inlined_containers.h"
+#include "core/providers/webgpu/tensor/matmul.h"
+#include "core/providers/cpu/tensor/utils.h"
+#include "core/providers/webgpu/shader_helper.h"
+#include "core/providers/webgpu/webgpu_supported_types.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+
+ONNX_OPERATOR_VERSIONED_KERNEL_EX(
+    MatMul,
+    kOnnxDomain,
+    1, 12,
+    kWebGpuExecutionProvider,
+    (*KernelDefBuilder::Create())
+        .TypeConstraint("T", WebGpuSupportedNumberTypes()),
+    MatMul);
+
+ONNX_OPERATOR_KERNEL_EX(
+    MatMul,
+    kOnnxDomain,
+    13,
+    kWebGpuExecutionProvider,
+    (*KernelDefBuilder::Create())
+        .TypeConstraint("T", WebGpuSupportedNumberTypes()),
+    MatMul);
+
+
+
+Status MatMulNativeProgram::GenerateShaderCode(ShaderHelper& sh) const {
+
+    return Status::OK();
+}
+
+
+Status MatMulProgram::GenerateShaderCode(ShaderHelper& sh) const {
+    return Status::OK();
+}
+
+Status MatMul::ComputeInternal(ComputeContext& context) const {
+    // calculate output shape
+    MatMulComputeHelper helper;
+    const auto* a = context.Input(0);
+    const auto* b = context.Input(1);
+    ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b->Shape()));
+
+    if (helper.N() < 8 && helper.K() < 8) {
+        // call MatMulNativeProgram
+        MatMulNativeProgram program{helper.OutputShape()};
+
+
+    } else {
+
+        // const batchA = ShapeUtil.size(context.inputs[0].dims.slice(0, -2));
+        // const batchB = ShapeUtil.size(context.inputs[1].dims.slice(0, -2));
+        int64_t batchA = a->Shape().SizeToDimension(a->Shape().NumDimensions() - 2);
+        int64_t batchB = b->Shape().SizeToDimension(b->Shape().NumDimensions() - 2);
+
+        // check if A is  batch of vector (bach is not 1, M is 1) and B is a matrix (batch is 1)
+        if (batchA != 1 && m == 1 && batchB == 1) {
+            // optimization for batched vector matrix multiplication
+            // const reshapedA = context.inputs[0].reshape([1, batchA, K]);
+            // const reshapedB = context.inputs[1].reshape([1, K, N]);
+            // const matmulOutputShape = [1, batchA, N];
+            // const matmulInputs = [reshapedA, reshapedB];
+
+            // dimensions of A: [1,`batchA`,K]
+            const gsl::span<const int64_t>& dims_a = {1, batchA, helper.K()};
+            // dimensions of B: [1,K,N]
+            const gsl::span<const int64_t>& dims_b = {1, helper.K(), helper.N()};
+
+            a.Reshape(TensorShape(dims_a));
+            b.Reshape(TensorShape(dims_b));
+
+            TensorShape output_shape = {1, batchA, helper.N()};
+
+
+            MatMulProgram program;
+        } else {
+            // call MatMulProgram
+            MatMulProgram program;
+        }
+
+
+    }
+
+    return Status::OK();
+
+}
+
+
+}  // namespace webgpu
+}  // namespace onnxruntime
diff --git a/onnxruntime/core/providers/webgpu/math/matmul.h b/onnxruntime/core/providers/webgpu/math/matmul.h
new file mode 100644
index 0000000000..7cdf26d0d0
--- /dev/null
+++ b/onnxruntime/core/providers/webgpu/math/matmul.h
@@ -0,0 +1,52 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/providers/webgpu/webgpu_kernel.h"
+#include "core/providers/webgpu/program.h"
+#include "core/providers/cpu/math/matmul_helper.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+class MatMul final : public WebGpuKernel {
+ public:
+  MatMul(const OpKernelInfo& info) : WebGpuKernel{info} {}
+
+  Status ComputeInternal(ComputeContext& context) const override;
+};
+
+class MatMulProgram final : public Program<MatMulProgram> {
+ public:
+  MatMulProgram() : Program{"MatMul"} {}
+
+  Status GenerateShaderCode(ShaderHelper& sh) const override;
+
+  // uniform variables
+
+
+};
+
+class MatMulNativeProgram final: public Program<MatMulNativeProgram> {
+ public:
+  MatMulNativeProgram(const int64_t output_size, const gsl::span<const int64_t>& outer_dims)
+      : Program{"MatMulNative"}, output_size_(output_Size), outer_dims_(outer_dims.begin(), outer_dims.end()) {
+  }
+
+  Status GenerateShaderCode(ShaderHelper& sh) const override;
+
+  // uniform variables output_size, M,N, K
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"output_size", ProgramUniformVariableDataType::Uint32},
+                                         {"M", ProgramUniformVariableDataType::Uint32},
+                                         {"N", ProgramUniformVariableDataType::Uint32},
+                                         {"K", ProgramUniformVariableDataType::Uint32});
+
+
+  private:
+    const int64_t output_size_;
+    const TensorShapeVector outer_dims_;
+};
+
+}  // namespace webgpu
+} // namespace onnxruntime