Add mul python interface and tests.

Author: zhuyue

include/infinicore/ops/mul.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+class Mul {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor);
+    static void execute(Tensor c, Tensor a, Tensor b);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor mul(Tensor a, Tensor b);
+void mul_(Tensor c, Tensor a, Tensor b);
+} // namespace infinicore::op

python/infinicore/__init__.py

@@ -29,6 +29,7 @@
 from infinicore.ops.attention import attention
 from infinicore.ops.causal_softmax import causal_softmax
 from infinicore.ops.matmul import matmul
+from infinicore.ops.mul import mul
 from infinicore.ops.rearrange import rearrange
 from infinicore.ops.rms_norm import rms_norm
 from infinicore.ops.silu import silu
@@ -76,6 +77,7 @@
     "attention",
     "causal_softmax",
     "matmul",
+    "mul",
     "rearrange",
     "rms_norm",
     "silu",

python/infinicore/ops/mul.py

@@ -0,0 +1,9 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def mul(input, other, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.mul(input._underlying, other._underlying))
+
+    _infinicore.mul_(out._underlying, input._underlying, other._underlying)

src/infinicore/ops/mul/mul.cc

@@ -0,0 +1,24 @@
+#include "infinicore/ops/mul.hpp"
+
+namespace infinicore::op {
+
+common::OpDispatcher<Mul::schema> &Mul::dispatcher() {
+    static common::OpDispatcher<Mul::schema> dispatcher_;
+    return dispatcher_;
+};
+
+void Mul::execute(Tensor c, Tensor a, Tensor b) {
+    dispatcher().lookup(context::getDevice().getType())(c, a, b);
+}
+
+Tensor mul(Tensor a, Tensor b) {
+    auto c = Tensor::empty(a->shape(), a->dtype(), a->device());
+    mul_(c, a, b);
+    return c;
+}
+
+void mul_(Tensor c, Tensor a, Tensor b) {
+    Mul::execute(c, a, b);
+}
+
+} // namespace infinicore::op

src/infinicore/ops/mul/mul_infiniop.cc

@@ -0,0 +1,52 @@
+#include "../../utils.hpp"
+#include "infinicore/common/hash.hpp"
+#include "infinicore/ops/common/cache.hpp"
+#include "infinicore/ops/mul.hpp"
+#include <infiniop.h>
+
+namespace infinicore::op::mul_impl::infiniop {
+
+thread_local common::OpCache<size_t, infiniopMulDescriptor_t> caches(
+    100, // capacity
+    [](infiniopMulDescriptor_t &desc) {
+        if (desc != nullptr) {
+            INFINICORE_CHECK_ERROR(infiniopDestroyMulDescriptor(desc));
+            desc = nullptr;
+        }
+    });
+
+void calculate(Tensor c, Tensor a, Tensor b) {
+    size_t seed = hash_combine(c, b, a);
+
+    auto device_type = context::getDevice().getType();
+    auto device_index = context::getDevice().getIndex();
+
+    auto &cache = caches.getCache(device_type, device_index);
+
+    auto desc_opt = cache.get(seed);
+    infiniopMulDescriptor_t desc = nullptr;
+
+    if (!desc_opt) {
+        INFINICORE_CHECK_ERROR(infiniopCreateMulDescriptor(
+            context::getInfiniopHandle(), &desc,
+            c->desc(), a->desc(), b->desc()));
+        cache.put(seed, desc);
+    } else {
+        desc = *desc_opt;
+    }
+
+    size_t workspace_size = 0;
+    INFINICORE_CHECK_ERROR(infiniopGetMulWorkspaceSize(desc, &workspace_size));
+    std::shared_ptr<Memory> workspace = context::allocateMemory(workspace_size);
+
+    INFINICORE_CHECK_ERROR(infiniopMul(
+        desc, workspace->data(), workspace_size,
+        c->data(), a->data(), b->data(), context::getStream()));
+}
+
+static bool registered = []() {
+    Mul::dispatcher().registerAll(&calculate, false);
+    return true;
+}();
+
+} // namespace infinicore::op::mul_impl::infiniop

src/infinicore/pybind11/ops.hpp

@@ -6,6 +6,7 @@
 #include "ops/attention.hpp"
 #include "ops/causal_softmax.hpp"
 #include "ops/matmul.hpp"
+#include "ops/mul.hpp"
 #include "ops/rearrange.hpp"
 #include "ops/rms_norm.hpp"
 #include "ops/silu.hpp"
@@ -20,6 +21,7 @@ inline void bind(py::module &m) {
     bind_attention(m);
     bind_causal_softmax(m);
     bind_matmul(m);
+    bind_mul(m);
     bind_rearrange(m);
     bind_rms_norm(m);
     bind_silu(m);

src/infinicore/pybind11/ops/mul.hpp

@@ -0,0 +1,26 @@
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+#include "infinicore/ops/mul.hpp"
+
+namespace py = pybind11;
+
+namespace infinicore::ops {
+
+inline void bind_mul(py::module &m) {
+    m.def("mul",
+          &op::mul,
+          py::arg("a"),
+          py::arg("b"),
+          R"doc(Element-wise multiplication of two tensors.)doc");
+
+    m.def("mul_",
+          &op::mul_,
+          py::arg("c"),
+          py::arg("a"),
+          py::arg("b"),
+          R"doc(In-place element-wise tensor multiplication.)doc");
+}
+
+} // namespace infinicore::ops

test/infinicore/ops/mul.py

@@ -0,0 +1,106 @@
+import sys
+import os
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+import torch
+import infinicore
+from framework.base import BaseOperatorTest, TensorSpec, TestCase
+from framework.runner import GenericTestRunner
+
+# ==============================================================================
+# Operator-specific configuration
+# ==============================================================================
+
+# Test cases format: (operation_mode, shape, a_strides, b_strides, c_strides)
+_TEST_CASES_DATA = [
+    (TestCase.BOTH, (13, 4), None, None, None),
+    (TestCase.BOTH, (13, 4), (10, 1), (10, 1), (10, 1)),
+    (TestCase.BOTH, (13, 4), (0, 1), None, None),
+    (TestCase.BOTH, (13, 4, 4), None, None, None),
+    (TestCase.BOTH, (13, 4, 4), (20, 4, 1), (20, 4, 1), (20, 4, 1)),
+    (TestCase.BOTH, (13, 4, 4), (4, 0, 1), (0, 4, 1), None),
+    (TestCase.BOTH, (16, 5632), None, None, None),
+    (TestCase.BOTH, (16, 5632), (13312, 1), (13312, 1), (13312, 1)),
+]
+
+
+def parse_test_cases(data):
+    """
+    Parse mul test case data according to format:
+    (operation_mode, shape, a_strides, b_strides, c_strides)
+    """
+    operation_mode = data[0]
+    shape = data[1]
+    a_strides = data[2] if len(data) > 2 else None
+    b_strides = data[3] if len(data) > 3 else None
+    c_strides = data[4] if len(data) > 4 else None
+
+    # Create input specifications
+    inputs = []
+
+    # Input tensor a
+    if a_strides is not None:
+        inputs.append(TensorSpec.from_strided_tensor(shape, a_strides))
+    else:
+        inputs.append(TensorSpec.from_tensor(shape))
+
+    # Input tensor b (same shape as a)
+    if b_strides is not None:
+        inputs.append(TensorSpec.from_strided_tensor(shape, b_strides))
+    else:
+        inputs.append(TensorSpec.from_tensor(shape))
+
+    # Output tensor
+    if c_strides is not None:
+        output = TensorSpec.from_strided_tensor(shape, c_strides)
+    else:
+        output = TensorSpec.from_tensor(shape)
+
+    return TestCase(operation_mode, inputs, output)
+
+
+# Parse test cases
+_TEST_CASES = [parse_test_cases(data) for data in _TEST_CASES_DATA]
+
+# Data types
+_TENSOR_DTYPES = [infinicore.float16, infinicore.bfloat16, infinicore.float32]
+
+# Tolerance
+_TOLERANCE_MAP = {
+    infinicore.float16: {"atol": 0, "rtol": 1e-2},
+    infinicore.float32: {"atol": 0, "rtol": 1e-3},
+    infinicore.bfloat16: {"atol": 0, "rtol": 5e-2},
+}
+
+
+class OpTest(BaseOperatorTest):
+    """Mul test with simplified test case parsing"""
+
+    def __init__(self):
+        super().__init__("Mul")
+
+    def get_test_cases(self):
+        return _TEST_CASES
+
+    def get_tensor_dtypes(self):
+        return _TENSOR_DTYPES
+
+    def get_tolerance_map(self):
+        return _TOLERANCE_MAP
+
+    def torch_operator(self, a, b, out=None, **kwargs):
+        return torch.mul(a, b, out=out)
+
+    def infinicore_operator(self, a, b, out=None, **kwargs):
+        return infinicore.mul(a, b, out=out)
+
+
+def main():
+    """Main entry point"""
+    runner = GenericTestRunner(OpTest)
+    runner.run_and_exit()
+
+
+if __name__ == "__main__":
+    main()