issue/497 - add add interface to assist test

wooway777 · wooway777 · commit 7b7983769d5d · 2025-10-13T18:01:05.000+08:00
diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp
@@ -1,5 +1,6 @@
 #pragma once
 
-#include "op/matmul.hpp"
-#include "op/ones.hpp"
-#include "op/rearrange.hpp"
+#include "ops/add.hpp"
+#include "ops/matmul.hpp"
+#include "ops/ones.hpp"
+#include "ops/rearrange.hpp"
diff --git a/include/infinicore/ops/add.hpp b/include/infinicore/ops/add.hpp
@@ -0,0 +1,17 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+class Add {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor);
+    static void execute(Tensor c, Tensor a, Tensor b);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor add(Tensor a, Tensor b);
+void add_(Tensor c, Tensor a, Tensor b);
+Tensor operator+(Tensor a, Tensor b);
+} // namespace infinicore::op
diff --git a/include/infinicore/ops/common/cache.hpp b/include/infinicore/ops/common/cache.hpp
diff --git a/include/infinicore/ops/common/dispatcher.hpp b/include/infinicore/ops/common/dispatcher.hpp
@@ -8,19 +8,19 @@ namespace infinicore::op::common {
 template <typename Fn>
 class OpDispatcher {
 public:
-    void registerDevice(Device::Type device_type, Fn fn, bool override_existing=true) {
-        if (table_[(size_t)device_type] == nullptr || override_existing){
+    void registerDevice(Device::Type device_type, Fn fn, bool override_existing = true) {
+        if (table_[(size_t)device_type] == nullptr || override_existing) {
             table_[(size_t)device_type] = fn;
         }
     }
 
-    void registerDevice(std::initializer_list<Device::Type> device_types, Fn fn, bool override_existing=true) {
+    void registerDevice(std::initializer_list<Device::Type> device_types, Fn fn, bool override_existing = true) {
         for (auto device_type : device_types) {
             registerDevice(device_type, fn, override_existing);
         }
     }
 
-    void registerAll(Fn fn, bool override_existing=true) {
+    void registerAll(Fn fn, bool override_existing = true) {
         for (size_t device_type = 0; device_type < static_cast<size_t>(Device::Type::COUNT); ++device_type) {
             registerDevice((Device::Type)device_type, fn, override_existing);
         }
diff --git a/include/infinicore/ops/common/op.hpp b/include/infinicore/ops/common/op.hpp
diff --git a/include/infinicore/ops/matmul.hpp b/include/infinicore/ops/matmul.hpp
diff --git a/include/infinicore/ops/ones.hpp b/include/infinicore/ops/ones.hpp
diff --git a/include/infinicore/ops/rearrange.hpp b/include/infinicore/ops/rearrange.hpp
diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py
@@ -24,6 +24,7 @@
     short,
     uint8,
 )
+from infinicore.ops.add import add
 from infinicore.ops.matmul import matmul
 from infinicore.ops.rearrange import rearrange
 from infinicore.tensor import (
@@ -63,6 +64,7 @@
     "short",
     "uint8",
     # Operations.
+    "add",
     "matmul",
     "rearrange",
     "empty",
diff --git a/python/infinicore/ops/add.py b/python/infinicore/ops/add.py
@@ -0,0 +1,9 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def add(input, other, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.add(input._underlying, other._underlying))
+
+    _infinicore.add_(out._underlying, input._underlying, other._underlying)
diff --git a/src/infinicore/op/add/add.cc b/src/infinicore/op/add/add.cc
@@ -0,0 +1,24 @@
+#include "infinicore/ops/add.hpp"
+
+namespace infinicore::op {
+
+common::OpDispatcher<Add::schema> &Add::dispatcher() {
+    static common::OpDispatcher<Add::schema> dispatcher_;
+    return dispatcher_;
+};
+
+void Add::execute(Tensor c, Tensor a, Tensor b) {
+    dispatcher().lookup(context::getDevice().getType())(c, a, b);
+}
+
+Tensor add(Tensor a, Tensor b) {
+    auto c = Tensor::empty(a->shape(), a->dtype(), a->device());
+    add_(c, a, b);
+    return c;
+}
+
+void add_(Tensor c, Tensor a, Tensor b) {
+    Add::execute(c, a, b);
+}
+
+} // namespace infinicore::op
diff --git a/src/infinicore/op/add/add_infiniop.cc b/src/infinicore/op/add/add_infiniop.cc
@@ -0,0 +1,52 @@
+#include "../../utils.hpp"
+#include "infinicore/common/hash.hpp"
+#include "infinicore/ops/add.hpp"
+#include "infinicore/ops/common/cache.hpp"
+#include <infiniop.h>
+
+namespace infinicore::op::add_impl::infiniop {
+
+thread_local common::OpCache<size_t, infiniopAddDescriptor_t> caches(
+    100, // capacity
+    [](infiniopAddDescriptor_t &desc) {
+        if (desc != nullptr) {
+            INFINICORE_CHECK_ERROR(infiniopDestroyAddDescriptor(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);
+    infiniopAddDescriptor_t desc = nullptr;
+
+    if (!desc_opt) {
+        INFINICORE_CHECK_ERROR(infiniopCreateAddDescriptor(
+            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(infiniopGetAddWorkspaceSize(desc, &workspace_size));
+    std::shared_ptr<Memory> workspace = context::allocateMemory(workspace_size);
+
+    INFINICORE_CHECK_ERROR(infiniopAdd(
+        desc, workspace->data(), workspace_size,
+        c->data(), a->data(), b->data(), context::getStream()));
+}
+
+static bool registered = []() {
+    Add::dispatcher().registerAll(&calculate, false);
+    return true;
+}();
+
+} // namespace infinicore::op::add_impl::infiniop
diff --git a/src/infinicore/op/matmul/matmul.cc b/src/infinicore/op/matmul/matmul.cc
@@ -1,4 +1,4 @@
-#include "infinicore/op/matmul.hpp"
+#include "infinicore/ops/matmul.hpp"
 
 namespace infinicore::op {
 
diff --git a/src/infinicore/op/matmul/matmul_infiniop.cc b/src/infinicore/op/matmul/matmul_infiniop.cc
@@ -1,7 +1,7 @@
 #include "../../utils.hpp"
 #include "infinicore/common/hash.hpp"
-#include "infinicore/op/common/cache.hpp"
-#include "infinicore/op/matmul.hpp"
+#include "infinicore/ops/common/cache.hpp"
+#include "infinicore/ops/matmul.hpp"
 #include <infiniop.h>
 
 namespace infinicore::op::matmul_impl::infiniop {
@@ -27,7 +27,9 @@ void calculate(Tensor c, Tensor a, Tensor b) {
     infiniopGemmDescriptor_t desc = nullptr;
 
     if (!desc_opt) {
-        INFINICORE_CHECK_ERROR(infiniopCreateGemmDescriptor(context::getInfiniopHandle(), &desc, c->desc(), a->desc(), b->desc()));
+        INFINICORE_CHECK_ERROR(infiniopCreateGemmDescriptor(
+            context::getInfiniopHandle(), &desc,
+            c->desc(), a->desc(), b->desc()));
         cache.put(seed, desc);
     } else {
         desc = *desc_opt;
diff --git a/src/infinicore/op/ones/ones.cc b/src/infinicore/op/ones/ones.cc
@@ -1,4 +1,4 @@
-#include "infinicore/op/ones.hpp"
+#include "infinicore/ops/ones.hpp"
 
 namespace infinicore::op {
 
diff --git a/src/infinicore/op/rearrange/rearrange.cc b/src/infinicore/op/rearrange/rearrange.cc
@@ -1,4 +1,4 @@
-#include "infinicore/op/rearrange.hpp"
+#include "infinicore/ops/rearrange.hpp"
 
 namespace infinicore::op {
 
diff --git a/src/infinicore/op/rearrange/rearrange_infiniop.cc b/src/infinicore/op/rearrange/rearrange_infiniop.cc
@@ -1,7 +1,7 @@
 #include "../../utils.hpp"
 #include "infinicore/common/hash.hpp"
-#include "infinicore/op/common/cache.hpp"
-#include "infinicore/op/rearrange.hpp"
+#include "infinicore/ops/common/cache.hpp"
+#include "infinicore/ops/rearrange.hpp"
 #include <infiniop.h>
 
 namespace infinicore::op::rearrange_impl::infiniop {
diff --git a/src/infinicore/pybind11/infinicore.cc b/src/infinicore/pybind11/infinicore.cc
@@ -4,7 +4,7 @@
 #include "context.hpp"
 #include "device.hpp"
 #include "dtype.hpp"
-#include "op.hpp"
+#include "ops.hpp"
 #include "tensor.hpp"
 
 namespace infinicore {
@@ -13,7 +13,7 @@ PYBIND11_MODULE(_infinicore, m) {
     context::bind(m);
     device::bind(m);
     dtype::bind(m);
-    op::bind(m);
+    ops::bind(m);
     tensor::bind(m);
 }
 
diff --git a/src/infinicore/pybind11/ops.hpp b/src/infinicore/pybind11/ops.hpp
@@ -2,16 +2,18 @@
 
 #include <pybind11/pybind11.h>
 
-#include "op/matmul.hpp"
-#include "op/rearrange.hpp"
+#include "ops/add.hpp"
+#include "ops/matmul.hpp"
+#include "ops/rearrange.hpp"
 
 namespace py = pybind11;
 
-namespace infinicore::op {
+namespace infinicore::ops {
 
 inline void bind(py::module &m) {
+    bind_add(m);
     bind_matmul(m);
     bind_rearrange(m);
 }
 
-} // namespace infinicore::op
+} // namespace infinicore::ops
diff --git a/src/infinicore/pybind11/ops/add.hpp b/src/infinicore/pybind11/ops/add.hpp
@@ -0,0 +1,26 @@
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+#include "infinicore/ops/add.hpp"
+
+namespace py = pybind11;
+
+namespace infinicore::ops {
+
+inline void bind_add(py::module &m) {
+    m.def("add",
+          &op::add,
+          py::arg("a"),
+          py::arg("b"),
+          R"doc(Addition of two tensors.)doc");
+
+    m.def("add_",
+          &op::add_,
+          py::arg("c"),
+          py::arg("a"),
+          py::arg("b"),
+          R"doc(In-place tensor addition.)doc");
+}
+
+} // namespace infinicore::ops
diff --git a/src/infinicore/pybind11/ops/matmul.hpp b/src/infinicore/pybind11/ops/matmul.hpp
@@ -2,11 +2,11 @@
 
 #include <pybind11/pybind11.h>
 
-#include "infinicore/op/matmul.hpp"
+#include "infinicore/ops/matmul.hpp"
 
 namespace py = pybind11;
 
-namespace infinicore::op {
+namespace infinicore::ops {
 
 inline void bind_matmul(py::module &m) {
     m.def("matmul",
@@ -23,4 +23,4 @@ inline void bind_matmul(py::module &m) {
           R"doc(In-place matrix multiplication.)doc");
 }
 
-} // namespace infinicore::op
+} // namespace infinicore::ops
diff --git a/src/infinicore/pybind11/ops/rearrange.hpp b/src/infinicore/pybind11/ops/rearrange.hpp
@@ -2,11 +2,11 @@
 
 #include <pybind11/pybind11.h>
 
-#include "infinicore/op/rearrange.hpp"
+#include "infinicore/ops/rearrange.hpp"
 
 namespace py = pybind11;
 
-namespace infinicore::op {
+namespace infinicore::ops {
 
 inline void bind_rearrange(py::module &m) {
     m.def("rearrange",
@@ -21,4 +21,4 @@ inline void bind_rearrange(py::module &m) {
           R"doc(In-place tensor rearrangement.)doc");
 }
 
-} // namespace infinicore::op
+} // namespace infinicore::ops

Original file line number	Diff line number	Diff line change
`@@ -8,19 +8,19 @@ namespace infinicore::op::common {`
`8`	`8`	`template <typename Fn>`
`9`	`9`	`class OpDispatcher {`
`10`	`10`	`public:`
`11`		`- void registerDevice(Device::Type device_type, Fn fn, bool override_existing=true) {`
`12`		`- if (table_[(size_t)device_type] == nullptr \|\| override_existing){`
	`11`	`+ void registerDevice(Device::Type device_type, Fn fn, bool override_existing = true) {`
	`12`	`+ if (table_[(size_t)device_type] == nullptr \|\| override_existing) {`
`13`	`13`	`table_[(size_t)device_type] = fn;`
`14`	`14`	`}`
`15`	`15`	`}`
`16`	`16`
`17`		`- void registerDevice(std::initializer_list<Device::Type> device_types, Fn fn, bool override_existing=true) {`
	`17`	`+ void registerDevice(std::initializer_list<Device::Type> device_types, Fn fn, bool override_existing = true) {`
`18`	`18`	`for (auto device_type : device_types) {`
`19`	`19`	`registerDevice(device_type, fn, override_existing);`
`20`	`20`	`}`
`21`	`21`	`}`
`22`	`22`
`23`		`- void registerAll(Fn fn, bool override_existing=true) {`
	`23`	`+ void registerAll(Fn fn, bool override_existing = true) {`
`24`	`24`	`for (size_t device_type = 0; device_type < static_cast<size_t>(Device::Type::COUNT); ++device_type) {`
`25`	`25`	`registerDevice((Device::Type)device_type, fn, override_existing);`
`26`	`26`	`}`