feat(cambricon): add mul provider

src/native/cambricon/kernel_utils.h

@@ -0,0 +1,40 @@
+#ifndef INFINI_OPS_CAMBRICON_KERNEL_UTILS_H_
+#define INFINI_OPS_CAMBRICON_KERNEL_UTILS_H_
+
+#include <cstddef>
+
+#ifdef __BANG__
+
+namespace infini::ops::cambricon::kernel_utils {
+
+struct TaskRange {
+  std::size_t begin;
+  std::size_t end;
+};
+
+__mlu_device__ inline TaskRange GetTaskRange(std::size_t size) {
+  const std::size_t elements_per_task = (size + taskDim - 1) / taskDim;
+  const std::size_t begin = taskId * elements_per_task;
+  const std::size_t end =
+      begin + elements_per_task < size ? begin + elements_per_task : size;
+  return {begin, end};
+}
+
+__mlu_device__ inline ptrdiff_t LogicalToOffset(std::size_t logical_index,
+                                                int ndim,
+                                                const std::size_t* shape,
+                                                const ptrdiff_t* strides) {
+  ptrdiff_t offset = 0;
+  for (int dim = ndim - 1; dim >= 0; --dim) {
+    const std::size_t coordinate = logical_index % shape[dim];
+    logical_index /= shape[dim];
+    offset += static_cast<ptrdiff_t>(coordinate) * strides[dim];
+  }
+  return offset;
+}
+
+}  // namespace infini::ops::cambricon::kernel_utils
+
+#endif  // __BANG__
+
+#endif  // INFINI_OPS_CAMBRICON_KERNEL_UTILS_H_

src/native/cambricon/ops/mul/kernel.h

@@ -0,0 +1,79 @@
+#ifndef INFINI_OPS_CAMBRICON_MUL_KERNEL_H_
+#define INFINI_OPS_CAMBRICON_MUL_KERNEL_H_
+
+#include <cstddef>
+
+#include "base/mul.h"
+#include "data_type.h"
+#include "dispatcher.h"
+#include "native/cambricon/cnrt_utils.h"
+#include "native/cambricon/common.h"
+#include "native/cambricon/data_type_.h"
+
+namespace infini::ops {
+
+template <typename T>
+void MulUnion(void* workspace, cnrtQueue_t queue, int core_per_cluster,
+              int cluster_count, const T* input, const T* other, T* out,
+              const std::size_t* out_shape, const ptrdiff_t* input_strides,
+              const ptrdiff_t* other_strides, const ptrdiff_t* out_strides,
+              std::size_t output_size, int ndim, bool fast_path,
+              bool out_contiguous);
+
+template <>
+class Operator<Mul, Device::Type::kCambricon> : public Mul {
+ public:
+  Operator(const Tensor input, const Tensor other, Tensor out)
+      : Mul{input, other, out},
+        default_workspace_{
+            cnrt_utils::AllocateDeviceBuffer(workspace_size_in_bytes())} {
+    cnrt_utils::GetLaunchConfig(input.device(), &core_per_cluster_,
+                                &cluster_count_);
+  }
+
+  void operator()(const Tensor input, const Tensor other,
+                  Tensor out) const override {
+    if (output_size_ == 0) {
+      return;
+    }
+
+    const bool fast_path = is_input_contiguous_ && is_other_contiguous_ &&
+                           is_out_contiguous_ && input.shape() == out.shape() &&
+                           other.shape() == out.shape();
+    auto queue = static_cast<cnrtQueue_t>(stream_ ? stream_ : 0);
+    void* workspace = workspace_ ? workspace_ : default_workspace_.get();
+
+    using SupportedTypes =
+        List<DataType::kFloat16, DataType::kBFloat16, DataType::kFloat32,
+             DataType::kInt16, DataType::kInt32, DataType::kInt64,
+             DataType::kUInt16, DataType::kUInt32, DataType::kUInt64>;
+    DispatchFunc<Device::Type::kCambricon, SupportedTypes>(
+        out_type_,
+        [&](auto tag) {
+          using T = typename decltype(tag)::type;
+          MulUnion<T>(workspace, queue, core_per_cluster_, cluster_count_,
+                      static_cast<const T*>(input.data()),
+                      static_cast<const T*>(other.data()),
+                      static_cast<T*>(out.data()), out_shape_.data(),
+                      input_strides_.data(), other_strides_.data(),
+                      out_strides_.data(), output_size_,
+                      static_cast<int>(ndim_), fast_path, is_out_contiguous_);
+        },
+        "CambriconMul::operator()");
+  }
+
+  std::size_t workspace_size_in_bytes() const override {
+    return ndim_ * (sizeof(std::size_t) + 3 * sizeof(ptrdiff_t));
+  }
+
+ private:
+  cnrt_utils::DeviceBuffer default_workspace_{};
+
+  int core_per_cluster_{0};
+
+  int cluster_count_{0};
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_CAMBRICON_MUL_KERNEL_H_

src/native/cambricon/ops/mul/kernel.mlu

@@ -0,0 +1,144 @@
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <type_traits>
+
+#include "kernel.h"
+#include "native/cambricon/kernel_utils.h"
+
+namespace infini::ops {
+namespace {
+
+__nram__ char mul_nram_buffer[NRAM_MAX_SIZE] __attribute__((aligned(128)));
+
+template <typename T>
+__mlu_device__ void Multiply(const T* input, const T* other, T* out,
+                             std::size_t count) {
+  if constexpr (std::is_same_v<T, __half>) {
+    __bang_mul(reinterpret_cast<half*>(out),
+               reinterpret_cast<const half*>(input),
+               reinterpret_cast<const half*>(other), count);
+  } else if constexpr (std::is_same_v<T, __bang_bfloat16> ||
+                       std::is_same_v<T, float>) {
+    __bang_mul(out, input, other, count);
+  } else {
+    for (std::size_t i = 0; i < count; ++i) {
+      out[i] = static_cast<T>(input[i] * other[i]);
+    }
+  }
+}
+
+template <typename T>
+__mlu_global__ void MulKernel(const T* input, const T* other, T* out,
+                              const std::size_t* out_shape,
+                              const ptrdiff_t* input_strides,
+                              const ptrdiff_t* other_strides,
+                              const ptrdiff_t* out_strides,
+                              std::size_t output_size, int ndim, bool fast_path,
+                              bool out_contiguous) {
+  const auto range = cambricon::kernel_utils::GetTaskRange(output_size);
+  if (range.begin >= range.end) {
+    return;
+  }
+
+  std::size_t block_size = NRAM_MAX_SIZE / (3 * sizeof(T));
+  if (block_size >= 64) {
+    block_size = block_size / 64 * 64;
+  }
+
+  auto* input_buffer = reinterpret_cast<T*>(mul_nram_buffer);
+  auto* other_buffer = input_buffer + block_size;
+  auto* output_buffer = other_buffer + block_size;
+
+  for (std::size_t processed = range.begin; processed < range.end;) {
+    const std::size_t current = std::min(block_size, range.end - processed);
+
+    if (fast_path) {
+      __memcpy(input_buffer, input + processed, current * sizeof(T),
+               GDRAM2NRAM);
+      __memcpy(other_buffer, other + processed, current * sizeof(T),
+               GDRAM2NRAM);
+    } else {
+      for (std::size_t i = 0; i < current; ++i) {
+        const std::size_t logical_index = processed + i;
+        input_buffer[i] = input[cambricon::kernel_utils::LogicalToOffset(
+            logical_index, ndim, out_shape, input_strides)];
+        other_buffer[i] = other[cambricon::kernel_utils::LogicalToOffset(
+            logical_index, ndim, out_shape, other_strides)];
+      }
+    }
+
+    Multiply(input_buffer, other_buffer, output_buffer, current);
+
+    if (out_contiguous) {
+      __memcpy(out + processed, output_buffer, current * sizeof(T), NRAM2GDRAM);
+    } else {
+      for (std::size_t i = 0; i < current; ++i) {
+        out[cambricon::kernel_utils::LogicalToOffset(
+            processed + i, ndim, out_shape, out_strides)] = output_buffer[i];
+      }
+    }
+
+    processed += current;
+  }
+}
+
+}  // namespace
+
+template <typename T>
+void MulUnion(void* workspace, cnrtQueue_t queue, int core_per_cluster,
+              int cluster_count, const T* input, const T* other, T* out,
+              const std::size_t* out_shape, const ptrdiff_t* input_strides,
+              const ptrdiff_t* other_strides, const ptrdiff_t* out_strides,
+              std::size_t output_size, int ndim, bool fast_path,
+              bool out_contiguous) {
+  auto* bytes = static_cast<char*>(workspace);
+  auto* device_shape = reinterpret_cast<std::size_t*>(bytes);
+  auto* device_input_strides =
+      reinterpret_cast<ptrdiff_t*>(device_shape + ndim);
+  auto* device_other_strides = device_input_strides + ndim;
+  auto* device_out_strides = device_other_strides + ndim;
+
+  if (ndim != 0) {
+    CNRT_CHECK(cnrtMemcpyAsync(
+        device_shape, const_cast<std::size_t*>(out_shape),
+        ndim * sizeof(std::size_t), queue, cnrtMemcpyHostToDev));
+    CNRT_CHECK(cnrtMemcpyAsync(
+        device_input_strides, const_cast<ptrdiff_t*>(input_strides),
+        ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev));
+    CNRT_CHECK(cnrtMemcpyAsync(
+        device_other_strides, const_cast<ptrdiff_t*>(other_strides),
+        ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev));
+    CNRT_CHECK(
+        cnrtMemcpyAsync(device_out_strides, const_cast<ptrdiff_t*>(out_strides),
+                        ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev));
+  }
+
+  const cnrtDim3_t kernel_dim = {static_cast<unsigned int>(core_per_cluster),
+                                 static_cast<unsigned int>(cluster_count), 1};
+  MulKernel<T><<<kernel_dim, cnrtFuncTypeUnion1, queue>>>(
+      input, other, out, device_shape, device_input_strides,
+      device_other_strides, device_out_strides, output_size, ndim, fast_path,
+      out_contiguous);
+  CNRT_CHECK(cnrtGetLastError());
+}
+
+#define INSTANTIATE_MUL(T)                                                    \
+  template void MulUnion<T>(void*, cnrtQueue_t, int, int, const T*, const T*, \
+                            T*, const std::size_t*, const ptrdiff_t*,         \
+                            const ptrdiff_t*, const ptrdiff_t*, std::size_t,  \
+                            int, bool, bool)
+
+INSTANTIATE_MUL(__half);
+INSTANTIATE_MUL(__bang_bfloat16);
+INSTANTIATE_MUL(float);
+INSTANTIATE_MUL(int16_t);
+INSTANTIATE_MUL(int32_t);
+INSTANTIATE_MUL(int64_t);
+INSTANTIATE_MUL(uint16_t);
+INSTANTIATE_MUL(uint32_t);
+INSTANTIATE_MUL(uint64_t);
+
+#undef INSTANTIATE_MUL
+
+}  // namespace infini::ops

tests/test_mul.py

@@ -49,9 +49,9 @@
 def test_mul(
     shape, input_strides, other_strides, out_strides, dtype, device, rtol, atol
 ):
-    if device == "musa" and dtype in _UINT_DTYPES:
+    if device in ("mlu", "musa") and dtype in _UINT_DTYPES:
         pytest.skip(
-            "The `torch.musa` test cloning path does not support `uint16`, `uint32`, or `uint64`."
+            f"The `{device}` test cloning path does not support `uint16`, `uint32`, or `uint64`."
         )
 
     if dtype in _INT_DTYPES or dtype in _UINT_DTYPES: