feat(ascend): add fill provider

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

@@ -0,0 +1,113 @@
+#ifndef INFINI_OPS_ASCEND_FILL_KERNEL_H_
+#define INFINI_OPS_ASCEND_FILL_KERNEL_H_
+
+#include <cassert>
+
+#include "acl/acl.h"
+#include "aclnn/aclnn_base.h"
+#include "aclnn_fill_scalar.h"
+#include "aclnn_fill_tensor.h"
+#include "base/fill.h"
+#include "native/ascend/common.h"
+#include "native/ascend/workspace_pool_.h"
+#include "operator.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Fill, Device::Type::kAscend> : public Fill {
+ public:
+  Operator(const Tensor input, const double value, Tensor out)
+      : Fill(input, value, out), out_cache_(out) {
+    Initialize(input, out);
+
+    // `aclCreateScalar` stores the address, so the backing value is a member.
+    scalar_storage_ = value;
+    scalar_ = aclCreateScalar(&scalar_storage_, ACL_DOUBLE);
+    assert(scalar_ != nullptr && "`AscendFill` failed to create scalar.");
+  }
+
+  Operator(const Tensor input, const Tensor value, Tensor out)
+      : Fill(input, value, out), value_cache_(value), out_cache_(out) {
+    assert(value.numel() == 1 &&
+           "`AscendFill` requires a scalar Tensor value.");
+    assert(value.device() == out.device() &&
+           "`AscendFill` requires Tensor value and output on the same "
+           "device.");
+    Initialize(input, out);
+  }
+
+  ~Operator() override {
+    if (scalar_ && ascend::IsAclRuntimeAlive()) {
+      aclDestroyScalar(scalar_);
+    }
+  }
+
+  void operator()(const Tensor input, const double /*value*/,
+                  Tensor out) const override {
+    (void)input;
+    if (out.numel() == 0) return;
+
+    auto stream = static_cast<aclrtStream>(stream_);
+    auto t_out = out_cache_.get(out.data());
+
+    aclOpExecutor* executor = nullptr;
+    uint64_t workspace_size = 0;
+    auto ret = aclnnInplaceFillScalarGetWorkspaceSize(
+        t_out, scalar_, &workspace_size, &executor);
+    assert(ret == ACL_SUCCESS &&
+           "`aclnnInplaceFillScalarGetWorkspaceSize` failed.");
+
+    auto& arena = ascend::GetWorkspacePool().Ensure(stream, workspace_size);
+    ret = aclnnInplaceFillScalar(arena.buf, workspace_size, executor, stream);
+    assert(ret == ACL_SUCCESS && "`aclnnInplaceFillScalar` failed.");
+  }
+
+  void operator()(const Tensor input, const Tensor value,
+                  Tensor out) const override {
+    (void)input;
+    if (out.numel() == 0) return;
+
+    auto stream = static_cast<aclrtStream>(stream_);
+    auto t_value = value_cache_.get(const_cast<void*>(value.data()));
+    auto t_out = out_cache_.get(out.data());
+
+    aclOpExecutor* executor = nullptr;
+    uint64_t workspace_size = 0;
+    auto ret = aclnnInplaceFillTensorGetWorkspaceSize(
+        t_out, t_value, &workspace_size, &executor);
+    assert(ret == ACL_SUCCESS &&
+           "`aclnnInplaceFillTensorGetWorkspaceSize` failed.");
+
+    auto& arena = ascend::GetWorkspacePool().Ensure(stream, workspace_size);
+    ret = aclnnInplaceFillTensor(arena.buf, workspace_size, executor, stream);
+    assert(ret == ACL_SUCCESS && "`aclnnInplaceFillTensor` failed.");
+  }
+
+ private:
+  static void Initialize(const Tensor input, const Tensor out) {
+    assert(input.shape() == out.shape() &&
+           "`AscendFill` requires input and output to have the same shape.");
+    assert(input.dtype() == out.dtype() &&
+           "`AscendFill` requires input and output to have the same dtype.");
+    assert(input.device() == out.device() &&
+           "`AscendFill` requires input and output on the same device.");
+    assert(!out.HasBroadcastDim() &&
+           "`AscendFill` output must not have broadcast dimensions.");
+    assert(out.ndim() <= 8 &&
+           "`AscendFill` does not support outputs with more than 8 "
+           "dimensions.");
+  }
+
+  mutable ascend::AclTensorCache value_cache_;
+
+  mutable ascend::AclTensorCache out_cache_;
+
+  double scalar_storage_ = 0.0;
+
+  aclScalar* scalar_ = nullptr;
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_ASCEND_FILL_KERNEL_H_

tests/test_fill.py

@@ -1,7 +1,7 @@
 import infini.ops
 import pytest
-import torch
 
+import torch
 from tests.utils import (
     Payload,
     empty_strided,
@@ -59,6 +59,39 @@ def test_fill(
     return Payload(_fill, _torch_fill, (input, value, out), {}, rtol=0, atol=0)
 
 
+@pytest.mark.auto_act_and_assert
+@pytest.mark.parametrize("shape, value_shape", (((13, 4), ()), ((3, 5), (1,))))
+@pytest.mark.parametrize(
+    "dtype, value",
+    (
+        (torch.uint8, 3),
+        (torch.int8, -3),
+        (torch.int16, -7),
+        (torch.int32, 11),
+        (torch.int64, -13),
+        (torch.float32, 2.5),
+        (torch.float16, -3.5),
+        (torch.bfloat16, 4.5),
+    ),
+)
+def test_fill_device_scalar(shape, value_shape, dtype, value, device):
+    if device != "npu":
+        pytest.skip("Device-scalar fill coverage is Ascend-specific")
+
+    input = _make_input(shape, None, dtype=dtype, device=device)
+    scalar = torch.full(value_shape, value, dtype=dtype, device=device)
+    out = torch.empty_like(input)
+
+    return Payload(
+        _fill,
+        _torch_fill,
+        (input, scalar, out),
+        {},
+        rtol=0,
+        atol=0,
+    )
+
+
 def _make_input(shape, strides, *, dtype, device):
     if dtype.is_floating_point:
         return randn_strided(shape, strides, dtype=dtype, device=device)