Initial draft CMSIS-NN integration (WIP)

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Author: sidart

backends/cortex_m/CMakeLists.txt

@@ -25,9 +25,11 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Codegen.cmake)
 
 # Cortex-M ops kernel sources
 set(_cortex_m_kernels__srcs
+    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_add.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_aten_add_tensor.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_softmax.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantize_per_tensor.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_dequantize_per_tensor.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_add.cpp
 )
 
 # Generate C++ bindings to register kernels into Executorch (for runtime).

backends/cortex_m/ops/op_add.cpp

@@ -1,5 +1,6 @@
 #include <executorch/runtime/kernel/kernel_includes.h>
-#include <cinttypes>
+#include <iostream>
+
 namespace cortex_m {
 namespace native {
 
@@ -13,7 +14,9 @@ Tensor& add_out(
     const Tensor& input2,
     const ScalarType dtype,
     Tensor& out) {
-  
+  std::cout << "add_out kernel called" << std::endl;
+  ET_LOG(Info, "xxxxxxxxxx add_out kernel called");
+
   // Ensure input is char type
   ET_CHECK_MSG(
       input1.scalar_type() == ScalarType::Char,
@@ -37,6 +40,7 @@ Tensor& add_out(
       "dtype %" PRId8 " is not int8 (Char)",
       static_cast<int8_t>(dtype));
 
+  assert(false);
 
   return out;
 }

backends/cortex_m/ops/op_aten_add_tensor.cpp

@@ -0,0 +1,49 @@
+#include <executorch/runtime/kernel/kernel_includes.h>
+#include <iostream>
+
+namespace cortex_m {
+namespace native {
+
+using Tensor = executorch::aten::Tensor;
+using ScalarType = executorch::aten::ScalarType;
+using KernelRuntimeContext = torch::executor::KernelRuntimeContext;
+
+Tensor& aten_add_tensor(
+    KernelRuntimeContext& ctx,
+    const Tensor& self,
+    const Tensor& other,
+    const ScalarType dtype,
+    Tensor& out) {
+  ET_LOG(Info, "xxxxxxxxxx aten_add_tensor kernel called");
+
+  // Ensure input is char type
+  ET_CHECK_MSG(
+      self.scalar_type() == ScalarType::Char,
+      "self.scalar_type() %" PRId8 " is not char type",
+      static_cast<int8_t>(self.scalar_type()));
+
+  ET_CHECK_MSG(
+      other.scalar_type() == ScalarType::Char,
+      "other.scalar_type() %" PRId8 " is not char type",
+      static_cast<int8_t>(other.scalar_type()));
+
+  // Check dtype is int8 (Char)
+  ET_CHECK_MSG(
+      dtype == ScalarType::Char,
+      "dtype %" PRId8 " is not int8 (Char)",
+      static_cast<int8_t>(dtype));
+  
+  // Example: element-wise add self and other into out
+  // (Assuming Tensor has data() and size() methods)
+ const int8_t* self_data = self.const_data_ptr<int8_t>();
+  const int8_t* other_data = other.const_data_ptr<int8_t>();
+  int8_t* out_data = out.mutable_data_ptr<int8_t>();
+  size_t numel = self.numel(); // or self.size() if that's the API
+  for (size_t i = 0; i < numel; ++i) {
+    out_data[i] = self_data[i] + other_data[i];
+  }
+  return out;
+}
+
+} // namespace native
+} // namespace cortex_m

backends/cortex_m/ops/op_softmax.cpp

@@ -0,0 +1,38 @@
+#include <executorch/runtime/kernel/kernel_includes.h>
+#include <iostream>
+
+namespace cortex_m {
+namespace native {
+
+using Tensor = executorch::aten::Tensor;
+using ScalarType = executorch::aten::ScalarType;
+using KernelRuntimeContext = torch::executor::KernelRuntimeContext;
+
+Tensor& _softmax_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& self,
+    int64_t dim,
+    bool half_to_float,
+    Tensor& out) {
+    // Your optimized implementation here
+    // Fill 'out' with the result and return it
+    std::cout << "xxxxxxxxxx softmax_out kernel called" << std::endl;
+    std::cout.flush();
+    ET_LOG(Error, "xxxxxxxxxx softmax_out kernel called");
+
+  return out;
+}
+
+Tensor _softmax(
+    KernelRuntimeContext& ctx,
+    const Tensor& self,
+    int64_t dim,
+    bool half_to_float) {
+    std::cout << "xxxxxxxxxx softmax_default kernel called" << std::endl;
+    std::cout.flush();
+    ET_LOG(Error, "xxxxxxxxxx softmax_default kernel called");
+    return self;
+}
+
+} // namespace native
+} // namespace cortex_m

backends/cortex_m/ops/operators.py

@@ -13,6 +13,57 @@
 # New operator library with a custom namespace to allow fusion etc.
 lib = Library("cortex_m", "DEF")
 
+###
+# add.Tensor
+###
+
+lib.define(
+    "add.Tensor(Tensor self, Tensor other, ScalarType dtype) -> (Tensor Z)"
+)
+
+lib.define(
+    "add_Tensor.out(Tensor self, Tensor other, ScalarType dtype, Tensor(a!) out) -> Tensor(a!)"
+)
+
+@impl(lib, "add.Tensor", "CompositeExplicitAutograd")
+def aten_add_tensor_impl(
+    input1: torch.Tensor,
+    input2: torch.Tensor,
+    dtype: torch.dtype,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    """
+    The implementation of aten add.Tensor.
+    """
+    return exir_ops.edge.cortex_m.add.Tensor(input1, input2, dtype)
+
+###
+# add.out
+###
+
+lib.define(
+    "add(Tensor input1, Tensor input2, ScalarType dtype) -> (Tensor Z)"
+)
+
+lib.define(
+    "add.out(Tensor input1, Tensor input2, ScalarType dtype, Tensor(a!) out) -> Tensor(a!)"
+)
+
+@impl(lib, "add.out", "CompositeExplicitAutograd")
+def add_out_impl(
+    input1: torch.Tensor,
+    input2: torch.Tensor,
+    dtype: torch.dtype,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    """
+    The implementation of cmsis-nn add.out.
+    """
+
+    return exir_ops.edge.cortex_m.add.default(
+        input1, input2, dtype, dtype
+    )
+
 ###
 # dequantize_per_tensor
 ###
@@ -25,7 +76,6 @@
     "quantize_per_tensor.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
 )
 
-
 @register_fake("cortex_m::quantize_per_tensor")
 def quantize_per_tensor_meta(
     input: torch.Tensor,
@@ -37,7 +87,6 @@ def quantize_per_tensor_meta(
 ) -> torch.Tensor:
     return torch.empty_like(input, dtype=dtype)
 
-
 @impl(lib, "quantize_per_tensor", "CompositeExplicitAutograd")
 def quantize_per_tensor_impl(
     input: torch.Tensor,
@@ -96,3 +145,17 @@ def dequantize_per_tensor_impl(
     return exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default(
         input, scale, zero_point, quant_min, quant_max, dtype
     )
+
+lib.define(
+    "softmax(Tensor self, int dim, bool half_to_float) -> Tensor"
+)
+lib.define(
+    "softmax.out(Tensor self, int dim, bool half_to_float, *, Tensor(a!) out) -> Tensor(a!)"
+)
+@impl(lib, "softmax", "CompositeExplicitAutograd")
+def softmax_impl(self: torch.Tensor, dim: int, half_to_float: bool) -> torch.Tensor:
+    # Call your custom edge op or fallback
+    return exir_ops.edge.cortex_m._softmax(self, dim, half_to_float)
+@impl(lib, "softmax.out", "CompositeExplicitAutograd")
+def softmax_out_impl(self: torch.Tensor, dim: int, half_to_float: bool, out: torch.Tensor) -> torch.Tensor:
+    return exir_ops.edge.cortex_m._softmax_out(self, dim, half_to_float, out)

backends/cortex_m/ops/operators.yaml

@@ -16,8 +16,26 @@
     - arg_meta: null
       kernel_name: cortex_m::dequantize_per_tensor_out
 
-- func: cortex_m::add.out(Tensor a, Tensor b, Scalar alpha, *, Tensor(a!) out) -> Tensor(a!)
+- func: cortex_m::add.out(Tensor a, Tensor b, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
   variants: function
   kernels:
     - arg_meta: null
       kernel_name: cortex_m::add_out
+
+- func: cortex_m::add.Tensor(Tensor self, Tensor other, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
+  variants: function
+  kernels:
+    - arg_meta: null
+      kernel_name: cortex_m::aten_add_tensor
+
+- func: cortex_m::_softmax(Tensor self, int dim, bool half_to_float) -> Tensor
+  variants: function
+  kernels:
+    - arg_meta: null
+      kernel_name: cortex_m::_softmax
+
+- func: cortex_m::_softmax.out(Tensor self, int dim, bool half_to_float, *, Tensor(a!) out) -> Tensor(a!)
+  variants: function
+  kernels:
+    - arg_meta: null
+      kernel_name: cortex_m::_softmax_out

backends/cortex_m/passes/replace_quant_nodes_pass.py

@@ -31,6 +31,18 @@ def _is_qualified_int8_node(args) -> bool:
     def __init__(self):
         super().__init__()
         self.op_replacements = {
+            exir_ops.edge.add: {
+                "new_target": exir_ops.edge.cortex_m.add,
+                "qualifier": lambda args: True,
+            },
+            exir_ops.edge.aten.add.Tensor: {
+                "new_target": exir_ops.edge.cortex_m.add.Tensor,
+                "qualifier": lambda args: True,
+            },
+            exir_ops.edge.aten._softmax.out: {
+                "new_target": exir_ops.edge.cortex_m.softmax.out,
+                "qualifier": lambda args: True,
+            },
             exir_ops.edge.quantized_decomposed.quantize_per_tensor.default: {
                 "new_target": exir_ops.edge.cortex_m.quantize_per_tensor.default,
                 "qualifier": self._is_qualified_int8_node,
@@ -51,12 +63,14 @@ def call_operator(
         assert isinstance(
             op, EdgeOpOverload
         ), "Op must be an EdgeOpOverload. Run this pass after to_edge()."
+        print(f"[ReplaceQuantNodesPass] Operator called: {op}, Args: {args}")
 
-        if op in self.op_replacements and self.op_replacements[op]["qualifier"](args):
+        if op in self.op_replacements and self.op_replacements[op]["qualifier"](args):            
+            print(f"[ReplaceQuantNodesPass] Replacing {op} with {self.op_replacements[op]['new_target']}")       
             return super().call_operator(
-                self.op_replacements[op]["new_target"],
-                args,
-                kwargs,
-                meta,
-            )
+                    self.op_replacements[op]["new_target"],
+                    args,
+                    kwargs,
+                    meta,
+                )
         return super().call_operator(op, args, kwargs, meta)