feat(cambricon): extend native model operator coverage

include/infini/convolution_layout.h

@@ -0,0 +1,13 @@
+#ifndef INFINI_OPS_CONVOLUTION_LAYOUT_H_
+#define INFINI_OPS_CONVOLUTION_LAYOUT_H_
+
+#include <infini/rt.h>
+
+namespace infini::ops {
+
+// Logical NCHW tensors may use packed channels-last strides.
+bool SupportsConvolutionChannelsLast(const infini::rt::Device& device);
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_CONVOLUTION_LAYOUT_H_

src/CMakeLists.txt

@@ -434,9 +434,23 @@ if(WITH_CAMBRICON)
     if(CNCC_COMPILER)
         message(STATUS "Found cncc: ${CNCC_COMPILER}")
         separate_arguments(INFINI_MLU_HOST_ABI_FLAGS NATIVE_COMMAND "${CMAKE_CXX_FLAGS}")
+        set(INFINIOPS_CAMBRICON_MLU_ARCH "mtp_592" CACHE STRING "BANG target architecture")
+        option(INFINIOPS_CAMBRICON_BANG_BF16 "Enable BF16 in native BANG kernels" ON)
+        target_compile_definitions(infiniops PRIVATE
+            INFINIOPS_CAMBRICON_BANG_BF16=$<BOOL:${INFINIOPS_CAMBRICON_BANG_BF16}>)
+        set(_cncc_host_includes "")
+        foreach(_dir IN LISTS CMAKE_CXX_IMPLICIT_INCLUDE_DIRECTORIES)
+            if(_dir MATCHES "/c[+][+]/")
+                list(APPEND _cncc_host_includes -isystem "${_dir}")
+            endif()
+        endforeach()
         set(MLU_COMPILE_OPTS
-            -c --bang-mlu-arch=mtp_592 -O3 -fPIC -Wall -Werror -std=c++17 -pthread
+            --bang-mlu-arch=${INFINIOPS_CAMBRICON_MLU_ARCH}
+            -DINFINIOPS_CAMBRICON_NATIVE=1
+            -DINFINIOPS_CAMBRICON_BANG_BF16=$<BOOL:${INFINIOPS_CAMBRICON_BANG_BF16}>
+            -O3 -fPIC -Wall -Werror -Wno-error=return-type -std=c++17 -pthread
             -I${CMAKE_CURRENT_SOURCE_DIR} -I${NEUWARE_HOME}/include
+            ${_cncc_host_includes}
             ${INFINI_RT_INCLUDE_FLAGS}
             ${INFINI_MLU_HOST_ABI_FLAGS}
             -idirafter /usr/local/neuware/lib/clang/11.1.0/include
@@ -1490,6 +1504,9 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS)
             set_source_files_properties(${ILUVATAR_OPERATOR_DISPATCH_OBJECTS}
                 PROPERTIES EXTERNAL_OBJECT TRUE GENERATED TRUE)
             target_sources(infiniops PRIVATE ${ILUVATAR_OPERATOR_DISPATCH_OBJECTS})
+        elseif(WITH_CAMBRICON AND NOT GENERATE_PYTHON_BINDINGS)
+            # C++ consumers need dispatch entry points without the Python extension.
+            target_sources(infiniops PRIVATE ${OPERATOR_DISPATCH_SOURCES})
         endif()
     endif()
 endif()

src/base/gather_axis.h

@@ -0,0 +1,20 @@
+#ifndef INFINI_OPS_BASE_GATHER_AXIS_H_
+#define INFINI_OPS_BASE_GATHER_AXIS_H_
+
+#include "operator.h"
+
+namespace infini::ops {
+
+// Select a one-dimensional list of indices, wrapping negatives on the axis.
+class GatherAxis : public Operator<GatherAxis> {
+ public:
+  GatherAxis(const Tensor input, const Tensor index, const int64_t dim,
+             Tensor out) {}
+
+  virtual void operator()(const Tensor input, const Tensor index,
+                          const int64_t dim, Tensor out) const = 0;
+};
+
+}  // namespace infini::ops
+
+#endif

src/base/grid_sample.h

@@ -0,0 +1,19 @@
+#ifndef INFINI_OPS_BASE_GRID_SAMPLE_H_
+#define INFINI_OPS_BASE_GRID_SAMPLE_H_
+
+#include "operator.h"
+
+namespace infini::ops {
+
+// Spatial sampling in NHWC layout; the last grid dimension stores (x, y).
+class GridSample : public Operator<GridSample> {
+ public:
+  GridSample(const Tensor input, const Tensor grid, bool nearest,
+             bool reflection, bool align_corners, Tensor out) {}
+  virtual void operator()(const Tensor input, const Tensor grid, bool nearest,
+                          bool reflection, bool align_corners, Tensor out) const = 0;
+};
+
+}  // namespace infini::ops
+
+#endif

src/base/instance_norm.h

@@ -0,0 +1,18 @@
+#ifndef INFINI_OPS_BASE_INSTANCE_NORM_H_
+#define INFINI_OPS_BASE_INSTANCE_NORM_H_
+
+#include "operator.h"
+
+namespace infini::ops {
+
+class InstanceNorm : public Operator<InstanceNorm> {
+ public:
+  InstanceNorm(const Tensor input, const Tensor scale, const Tensor bias,
+               float epsilon, Tensor out) {}
+  virtual void operator()(const Tensor input, const Tensor scale, const Tensor bias,
+                          float epsilon, Tensor out) const = 0;
+};
+
+}  // namespace infini::ops
+
+#endif

src/base/masked_scatter.h

@@ -0,0 +1,19 @@
+#ifndef INFINI_OPS_BASE_MASKED_SCATTER_H_
+#define INFINI_OPS_BASE_MASKED_SCATTER_H_
+
+#include "operator.h"
+
+namespace infini::ops {
+
+// Replace true positions in row-major order with consecutive source values.
+class MaskedScatter : public Operator<MaskedScatter> {
+ public:
+  MaskedScatter(const Tensor input, const Tensor mask, const Tensor source,
+                Tensor out) {}
+  virtual void operator()(const Tensor input, const Tensor mask,
+                          const Tensor source, Tensor out) const = 0;
+};
+
+}  // namespace infini::ops
+
+#endif

src/base/resize2d.h

@@ -0,0 +1,23 @@
+#ifndef INFINI_OPS_BASE_RESIZE2D_H_
+#define INFINI_OPS_BASE_RESIZE2D_H_
+
+#include "operator.h"
+
+namespace infini::ops {
+
+class Resize2d : public Operator<Resize2d> {
+ public:
+  enum Mode : int64_t { kNearest, kLinear };
+  enum Coordinates : int64_t { kHalfPixel, kPytorchHalfPixel, kAlignCorners, kAsymmetric };
+  enum Rounding : int64_t { kPreferFloor, kPreferCeil, kFloor, kCeil };
+
+  Resize2d(const Tensor input, int64_t mode, int64_t coordinates,
+           int64_t rounding, float height_scale, float width_scale, Tensor out) {}
+  virtual void operator()(const Tensor input, int64_t mode, int64_t coordinates,
+                          int64_t rounding, float height_scale, float width_scale,
+                          Tensor out) const = 0;
+};
+
+}  // namespace infini::ops
+
+#endif

src/base/roi_align.h

@@ -0,0 +1,20 @@
+#ifndef INFINI_OPS_BASE_ROI_ALIGN_H_
+#define INFINI_OPS_BASE_ROI_ALIGN_H_
+
+#include "operator.h"
+
+namespace infini::ops {
+
+// Average pooling over NHWC feature maps. Boxes are [batch, x1, y1, x2, y2].
+class RoiAlign : public Operator<RoiAlign> {
+ public:
+  RoiAlign(const Tensor input, const Tensor boxes, int64_t sampling_ratio,
+           float spatial_scale, bool aligned, Tensor out) {}
+  virtual void operator()(const Tensor input, const Tensor boxes,
+                          int64_t sampling_ratio, float spatial_scale,
+                          bool aligned, Tensor out) const = 0;
+};
+
+}  // namespace infini::ops
+
+#endif

src/base/scatter_nd.h

@@ -0,0 +1,18 @@
+#ifndef INFINI_OPS_BASE_SCATTER_ND_H_
+#define INFINI_OPS_BASE_SCATTER_ND_H_
+
+#include "operator.h"
+
+namespace infini::ops {
+
+class ScatterNd : public Operator<ScatterNd> {
+ public:
+  ScatterNd(const Tensor input, const Tensor indices, const Tensor updates,
+            bool add, Tensor out) {}
+  virtual void operator()(const Tensor input, const Tensor indices,
+                          const Tensor updates, bool add, Tensor out) const = 0;
+};
+
+}  // namespace infini::ops
+
+#endif

src/base/tile.h

@@ -0,0 +1,16 @@
+#ifndef INFINI_OPS_BASE_TILE_H_
+#define INFINI_OPS_BASE_TILE_H_
+
+#include "operator.h"
+
+namespace infini::ops {
+
+class Tile : public Operator<Tile> {
+ public:
+  Tile(const Tensor input, Tensor out) {}
+  virtual void operator()(const Tensor input, Tensor out) const = 0;
+};
+
+}  // namespace infini::ops
+
+#endif

src/convolution_layout.cc

@@ -0,0 +1,14 @@
+#include <infini/convolution_layout.h>
+
+namespace infini::ops {
+
+bool SupportsConvolutionChannelsLast(const infini::rt::Device& device) {
+#ifdef WITH_CAMBRICON
+  return device.type() == infini::rt::Device::Type::kCambricon;
+#else
+  (void)device;
+  return false;
+#endif
+}
+
+}  // namespace infini::ops

src/native/cambricon/activation_cnnl.h

@@ -0,0 +1,50 @@
+#ifndef INFINI_OPS_CAMBRICON_ACTIVATION_CNNL_H_
+#define INFINI_OPS_CAMBRICON_ACTIVATION_CNNL_H_
+
+#include <stdexcept>
+
+#include "native/cambricon/cnnl_utils.h"
+
+namespace infini::ops::cnnl_utils {
+
+template <typename Base, cnnlActivationMode_t Mode>
+class Activation : public Base {
+ public:
+  Activation(const Tensor input, Tensor out) : Base{input, out} {
+    if (input.shape() != out.shape() || input.dtype() != out.dtype() ||
+        input.device() != out.device() || !input.IsContiguous() ||
+        !out.IsContiguous() ||
+        (out.dtype() != DataType::kFloat32 && out.dtype() != DataType::kFloat16))
+      throw std::invalid_argument("Cambricon activation requires matching FP32 or FP16 tensors");
+    if (out.numel() == 0) return;
+    handle_ = cnnl_utils::CreateHandle();
+    descriptor_ = cnnl_utils::MakeTensorDescriptor(
+        input.dtype(), Tensor::Shape(input.shape()), Tensor::Strides(input.strides()));
+    cnnlActivationDescriptor_t value = nullptr;
+    INFINI_OPS_CNNL_CHECK(cnnlCreateActivationDescriptor(&value));
+    operation_.reset(value);
+    INFINI_OPS_CNNL_CHECK(cnnlSetActivationDescriptor_v5(
+        value, Mode, CNNL_ACTIVATION_HIGH_PRECISION,
+        CNNL_PROPAGATE_NAN, 0.0f, 0.0f, 0.0f, 0.0f, true));
+  }
+
+  void operator()(const Tensor input, Tensor out) const override {
+    if (!descriptor_) return;
+    INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast<cnrtQueue_t>(this->stream_)));
+    const float one = 1.0f, zero = 0.0f;
+    INFINI_OPS_CNNL_CHECK(cnnlActivationForward(
+        handle_.get(), operation_.get(), &one, descriptor_.get(), input.data(),
+        &zero, descriptor_.get(), out.data()));
+  }
+
+ private:
+  cnnl_utils::Handle handle_;
+  cnnl_utils::TensorDescriptor descriptor_;
+  std::unique_ptr<std::remove_pointer_t<cnnlActivationDescriptor_t>,
+                  decltype(&cnnlDestroyActivationDescriptor)>
+      operation_{nullptr, cnnlDestroyActivationDescriptor};
+};
+
+}  // namespace infini::ops::cnnl_utils
+
+#endif

src/native/cambricon/cnnl_utils.h

@@ -4,6 +4,8 @@
 #include <cassert>
 #include <cstdint>
 #include <memory>
+#include <stdexcept>
+#include <string>
 #include <vector>
 
 #include "native/cambricon/common.h"
@@ -12,8 +14,9 @@
 #define INFINI_OPS_CNNL_CHECK(call)                                      \
   do {                                                                   \
     const auto cnnl_status = (call);                                     \
-    assert(cnnl_status == CNNL_STATUS_SUCCESS && "`" #call "` failed."); \
-    (void)cnnl_status;                                                   \
+    if (cnnl_status != CNNL_STATUS_SUCCESS)                               \
+      throw std::runtime_error(std::string(#call ": ") +                  \
+                               cnnlGetErrorString(cnnl_status));         \
   } while (false)
 
 namespace infini::ops::cnnl_utils {
@@ -92,6 +95,11 @@ inline TensorDescriptor MakeTensorDescriptor(DataType dtype,
   return desc;
 }
 
+inline TensorDescriptor MakeTensorDescriptor(const Tensor& tensor) {
+  return MakeTensorDescriptor(tensor.dtype(), Tensor::Shape(tensor.shape()),
+                              Tensor::Strides(tensor.strides()));
+}
+
 }  // namespace infini::ops::cnnl_utils
 
 #endif

src/native/cambricon/common.h

@@ -38,10 +38,16 @@ namespace infini::ops::cnnl_utils {
 
 inline cnnlDataType_t GetDataType(DataType dtype) {
   switch (dtype) {
+    case DataType::kBool:
+      return CNNL_DTYPE_BOOL;
     case DataType::kInt8:
       return CNNL_DTYPE_INT8;
     case DataType::kUInt8:
       return CNNL_DTYPE_UINT8;
+    case DataType::kInt16:
+      return CNNL_DTYPE_INT16;
+    case DataType::kUInt32:
+      return CNNL_DTYPE_UINT32;
     case DataType::kInt32:
       return CNNL_DTYPE_INT32;
     case DataType::kInt64:

src/native/cambricon/compare_cnnl.h

@@ -0,0 +1,64 @@
+#ifndef INFINI_OPS_CAMBRICON_COMPARE_CNNL_H_
+#define INFINI_OPS_CAMBRICON_COMPARE_CNNL_H_
+
+#include "native/cambricon/cnnl_utils.h"
+#include "native/cambricon/cnrt_utils.h"
+
+namespace infini::ops::cnnl_utils {
+
+template <typename Base, cnnlLogicOp_t Kind>
+class Compare : public Base {
+ public:
+  Compare(const Tensor input, const Tensor other, Tensor out)
+      : Base{input, other, out} {
+    if (input.dtype() != other.dtype() || out.dtype() != DataType::kBool ||
+        input.device() != out.device() || other.device() != out.device() ||
+        !input.IsContiguous() || !other.IsContiguous() || !out.IsContiguous())
+      throw std::invalid_argument("Cambricon comparison requires matching inputs and BOOL output");
+    if (out.numel() == 0) return;
+    handle_ = CreateHandle();
+    lhs_ = MakeTensorDescriptor(input);
+    rhs_ = MakeTensorDescriptor(other);
+    out_ = MakeTensorDescriptor(out);
+    INFINI_OPS_CNNL_CHECK(cnnlGetLogicOpWorkspaceSize(
+        handle_.get(), lhs_.get(), rhs_.get(), out_.get(), &workspace_size_));
+    scratch_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_);
+    if (workspace_size_ && !scratch_)
+      throw std::runtime_error("Cambricon comparison workspace allocation failed");
+  }
+  Compare(const Tensor input, const double other, Tensor out)
+      : Base{input, other, out} { Unsupported(); }
+  Compare(Tensor input, const Tensor other)
+      : Base{input, other, input} { Unsupported(); }
+  Compare(Tensor input, const double other)
+      : Base{input, other, input} { Unsupported(); }
+
+  std::size_t workspace_size_in_bytes() const override { return workspace_size_; }
+
+  void operator()(const Tensor input, const Tensor other, Tensor out) const override {
+    if (!out_) return;
+    INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast<cnrtQueue_t>(this->stream_)));
+    auto* scratch = this->workspace_ && this->workspace_size_in_bytes_ >= workspace_size_
+                        ? this->workspace_ : scratch_.get();
+    INFINI_OPS_CNNL_CHECK(cnnlLogicOp(
+        handle_.get(), Kind, lhs_.get(), input.data(), rhs_.get(), other.data(),
+        scratch, workspace_size_, out_.get(), out.data()));
+  }
+  void operator()(const Tensor, const double, Tensor) const override { Unsupported(); }
+  void operator()(Tensor, const Tensor) const override { Unsupported(); }
+  void operator()(Tensor, const double) const override { Unsupported(); }
+
+ private:
+  [[noreturn]] static void Unsupported() {
+    throw std::invalid_argument("Cambricon comparison requires two tensors and a separate BOOL output");
+  }
+
+  Handle handle_;
+  TensorDescriptor lhs_, rhs_, out_;
+  cnrt_utils::DeviceBuffer scratch_;
+  std::size_t workspace_size_ = 0;
+};
+
+}  // namespace infini::ops::cnnl_utils
+
+#endif