feat(cambricon): add fill provider

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

@@ -0,0 +1,149 @@
+#ifndef INFINI_OPS_CAMBRICON_FILL_KERNEL_H_
+#define INFINI_OPS_CAMBRICON_FILL_KERNEL_H_
+
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <type_traits>
+
+#include "base/fill.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 {
+
+using CambriconFillTypes =
+    ConcatType<ConcatType<List<DataType::kFloat32>, ReducedFloatTypes>,
+               AllIntTypes>;
+
+template <typename T>
+void FillUnion(void* workspace, cnrtQueue_t queue, int core_per_cluster,
+               int cluster_count, const T* device_value, double host_value,
+               T* out, const std::size_t* out_shape,
+               const ptrdiff_t* out_strides, std::size_t output_size, int ndim,
+               bool out_contiguous);
+
+void FillRaw64Union(void* workspace, cnrtQueue_t queue, int core_per_cluster,
+                    int cluster_count, const void* device_value, void* out,
+                    const std::size_t* out_shape, const ptrdiff_t* out_strides,
+                    std::size_t output_size, int ndim, bool out_contiguous);
+
+template <>
+class Operator<Fill, Device::Type::kCambricon> : public Fill {
+ public:
+  Operator(const Tensor input, const double value, Tensor out)
+      : Fill{input, value, out} {
+    Initialize(input, out);
+  }
+
+  Operator(const Tensor input, const Tensor value, Tensor out)
+      : Fill{input, value, out} {
+    assert(value.numel() == 1 &&
+           "`CambriconFill` requires a scalar Tensor value.");
+    assert(value.dtype() == out.dtype() &&
+           "`CambriconFill` requires Tensor value and output to have the same "
+           "dtype.");
+    assert(value.device() == out.device() &&
+           "`CambriconFill` requires Tensor value and output on the same "
+           "device.");
+    Initialize(input, out);
+  }
+
+  void operator()(const Tensor input, const double value,
+                  Tensor out) const override {
+    (void)input;
+    Run(nullptr, value, out);
+  }
+
+  void operator()(const Tensor input, const Tensor value,
+                  Tensor out) const override {
+    (void)input;
+    Run(value.data(), 0.0, out);
+  }
+
+  std::size_t workspace_size_in_bytes() const override {
+    return sizeof(std::uint64_t) +
+           ndim_ * (sizeof(std::size_t) + sizeof(ptrdiff_t));
+  }
+
+ private:
+  void Initialize(const Tensor input, const Tensor out) {
+    assert(input.shape() == out.shape() &&
+           "`CambriconFill` requires input and output to have the same "
+           "shape.");
+    assert(input.dtype() == out.dtype() &&
+           "`CambriconFill` requires input and output to have the same "
+           "dtype.");
+    assert(input.device() == out.device() &&
+           "`CambriconFill` requires input and output on the same device.");
+    assert(!out.HasBroadcastDim() &&
+           "`CambriconFill` output must not have broadcast dimensions.");
+
+    output_size_ = out.numel();
+    ndim_ = out.ndim();
+    is_out_contiguous_ = out.IsContiguous();
+    cnrt_utils::GetLaunchConfig(out.device(), &core_per_cluster_,
+                                &cluster_count_);
+    default_workspace_ =
+        cnrt_utils::AllocateDeviceBuffer(workspace_size_in_bytes());
+  }
+
+  void Run(const void* device_value, double host_value, Tensor out) const {
+    if (output_size_ == 0) {
+      return;
+    }
+
+    auto queue = static_cast<cnrtQueue_t>(stream_ ? stream_ : 0);
+    void* workspace = workspace_ ? workspace_ : default_workspace_.get();
+
+    DispatchFunc<Device::Type::kCambricon, CambriconFillTypes>(
+        out_type_,
+        [&](auto tag) {
+          using T = typename decltype(tag)::type;
+          const T* effective_device_value = static_cast<const T*>(device_value);
+          if constexpr (std::is_same_v<T, std::int64_t> ||
+                        std::is_same_v<T, std::uint64_t>) {
+            if (effective_device_value == nullptr) {
+              const T converted_host_value = static_cast<T>(host_value);
+              CNRT_CHECK(cnrtMemcpy(workspace,
+                                    const_cast<T*>(&converted_host_value),
+                                    sizeof(T), cnrtMemcpyHostToDev));
+              effective_device_value = static_cast<const T*>(workspace);
+            }
+          }
+          if constexpr (std::is_same_v<T, std::int64_t> ||
+                        std::is_same_v<T, std::uint64_t>) {
+            FillRaw64Union(workspace, queue, core_per_cluster_, cluster_count_,
+                           effective_device_value, out.data(),
+                           out_shape_.data(), out_strides_.data(), output_size_,
+                           static_cast<int>(ndim_), is_out_contiguous_);
+          } else {
+            FillUnion<T>(workspace, queue, core_per_cluster_, cluster_count_,
+                         effective_device_value, host_value,
+                         static_cast<T*>(out.data()), out_shape_.data(),
+                         out_strides_.data(), output_size_,
+                         static_cast<int>(ndim_), is_out_contiguous_);
+          }
+        },
+        "CambriconFill::operator()");
+  }
+
+  std::size_t output_size_{0};
+
+  std::size_t ndim_{0};
+
+  bool is_out_contiguous_{false};
+
+  cnrt_utils::DeviceBuffer default_workspace_{};
+
+  int core_per_cluster_{0};
+
+  int cluster_count_{0};
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_CAMBRICON_FILL_KERNEL_H_

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

@@ -0,0 +1,170 @@
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <type_traits>
+
+#include "kernel.h"
+#include "native/cambricon/kernel_utils.h"
+
+namespace infini::ops {
+namespace {
+
+__nram__ char fill_nram_buffer[NRAM_MAX_SIZE] __attribute__((aligned(128)));
+
+template <typename T>
+__mlu_device__ void FillBuffer(T* buffer, std::size_t count,
+                               const T* device_value) {
+  if constexpr (std::is_same_v<T, float>) {
+    const T value = device_value[0];
+    __bang_write_value(buffer, count, value);
+  } else {
+    const T value = device_value[0];
+    for (std::size_t i = 0; i < count; ++i) {
+      buffer[i] = value;
+    }
+  }
+}
+
+template <typename T>
+__mlu_global__ void FillKernel(const T* device_value, T* out,
+                               const std::size_t* out_shape,
+                               const ptrdiff_t* out_strides,
+                               std::size_t output_size, int ndim,
+                               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 / sizeof(T);
+  if (block_size >= 64) {
+    block_size = block_size / 64 * 64;
+  }
+  auto* buffer = reinterpret_cast<T*>(fill_nram_buffer);
+
+  for (std::size_t processed = range.begin; processed < range.end;) {
+    const std::size_t current = std::min(block_size, range.end - processed);
+    FillBuffer(buffer, current, device_value);
+
+    if (out_contiguous) {
+      __memcpy(out + processed, buffer, current * sizeof(T), NRAM2GDRAM);
+    } else {
+      for (std::size_t i = 0; i < current; ++i) {
+        const ptrdiff_t offset = cambricon::kernel_utils::LogicalToOffset(
+            processed + i, ndim, out_shape, out_strides);
+        out[offset] = buffer[i];
+      }
+    }
+    processed += current;
+  }
+}
+
+__mlu_global__ void FillRaw64Kernel(const void* device_value, void* out,
+                                    const std::size_t* out_shape,
+                                    const ptrdiff_t* out_strides,
+                                    std::size_t output_size, int ndim,
+                                    bool out_contiguous) {
+  constexpr std::size_t kElementSize = 8;
+  const auto range = cambricon::kernel_utils::GetTaskRange(output_size);
+  if (range.begin >= range.end) {
+    return;
+  }
+
+  auto* out_words = static_cast<std::int32_t*>(out);
+  for (std::size_t logical = range.begin; logical < range.end; ++logical) {
+    const ptrdiff_t offset = out_contiguous
+                                 ? static_cast<ptrdiff_t>(logical)
+                                 : cambricon::kernel_utils::LogicalToOffset(
+                                       logical, ndim, out_shape, out_strides);
+    __memcpy(out_words + 2 * offset, device_value, kElementSize, GDRAM2GDRAM);
+  }
+}
+
+}  // namespace
+
+void FillRaw64Union(void* workspace, cnrtQueue_t queue, int core_per_cluster,
+                    int cluster_count, const void* device_value, void* out,
+                    const std::size_t* out_shape, const ptrdiff_t* out_strides,
+                    std::size_t output_size, int ndim, bool out_contiguous) {
+  auto* bytes = static_cast<char*>(workspace) + sizeof(std::uint64_t);
+  std::size_t* device_shape = nullptr;
+  ptrdiff_t* device_strides = nullptr;
+  if (ndim != 0) {
+    device_shape = reinterpret_cast<std::size_t*>(bytes);
+    device_strides = reinterpret_cast<ptrdiff_t*>(device_shape + ndim);
+    CNRT_CHECK(cnrtMemcpyAsync(
+        device_shape, const_cast<std::size_t*>(out_shape),
+        ndim * sizeof(std::size_t), queue, cnrtMemcpyHostToDev));
+    CNRT_CHECK(
+        cnrtMemcpyAsync(device_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};
+  FillRaw64Kernel<<<kernel_dim, cnrtFuncTypeUnion1, queue>>>(
+      device_value, out, device_shape, device_strides, output_size, ndim,
+      out_contiguous);
+  CNRT_CHECK(cnrtGetLastError());
+}
+
+template <typename T>
+void FillUnion(void* workspace, cnrtQueue_t queue, int core_per_cluster,
+               int cluster_count, const T* device_value, double host_value,
+               T* out, const std::size_t* out_shape,
+               const ptrdiff_t* out_strides, std::size_t output_size, int ndim,
+               bool out_contiguous) {
+  auto* bytes = static_cast<char*>(workspace);
+  auto* device_scalar = reinterpret_cast<T*>(bytes);
+  bytes += sizeof(std::uint64_t);
+  std::size_t* device_shape = nullptr;
+  ptrdiff_t* device_strides = nullptr;
+  if (ndim != 0) {
+    device_shape = reinterpret_cast<std::size_t*>(bytes);
+    device_strides = reinterpret_cast<ptrdiff_t*>(device_shape + ndim);
+    CNRT_CHECK(cnrtMemcpyAsync(
+        device_shape, const_cast<std::size_t*>(out_shape),
+        ndim * sizeof(std::size_t), queue, cnrtMemcpyHostToDev));
+    CNRT_CHECK(
+        cnrtMemcpyAsync(device_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};
+  if (device_value == nullptr) {
+    T converted_host_value;
+    if constexpr (std::is_same_v<T, __half> ||
+                  std::is_same_v<T, __bang_bfloat16>) {
+      converted_host_value = static_cast<T>(static_cast<float>(host_value));
+    } else {
+      converted_host_value = static_cast<T>(host_value);
+    }
+    CNRT_CHECK(cnrtMemcpy(device_scalar, &converted_host_value, sizeof(T),
+                          cnrtMemcpyHostToDev));
+    device_value = device_scalar;
+  }
+  FillKernel<T><<<kernel_dim, cnrtFuncTypeUnion1, queue>>>(
+      device_value, out, device_shape, device_strides, output_size, ndim,
+      out_contiguous);
+  CNRT_CHECK(cnrtGetLastError());
+}
+
+#define INSTANTIATE_FILL(T)                                                  \
+  template void FillUnion<T>(void*, cnrtQueue_t, int, int, const T*, double, \
+                             T*, const std::size_t*, const ptrdiff_t*,       \
+                             std::size_t, int, bool)
+
+INSTANTIATE_FILL(__half);
+INSTANTIATE_FILL(__bang_bfloat16);
+INSTANTIATE_FILL(float);
+INSTANTIATE_FILL(int8_t);
+INSTANTIATE_FILL(int16_t);
+INSTANTIATE_FILL(int32_t);
+INSTANTIATE_FILL(uint8_t);
+INSTANTIATE_FILL(uint16_t);
+INSTANTIATE_FILL(uint32_t);
+
+#undef INSTANTIATE_FILL
+
+}  // namespace infini::ops

tests/test_fill.py

@@ -46,8 +46,8 @@ def test_fill(
     value,
     device,
 ):
-    if device == "musa" and dtype == torch.float64:
-        pytest.skip("MUSA does not support float64 fill")
+    if device in ("mlu", "musa") and dtype == torch.float64:
+        pytest.skip(f"{device.upper()} does not support float64 fill")
 
     input = _make_input(shape, input_strides, dtype=dtype, device=device)
     out = (