From 7c07b4d53568e01498d4e982fd69b1a696e43683 Mon Sep 17 00:00:00 2001 From: Li Baoming <41820386+baominghelly@users.noreply.github.com> Date: Thu, 3 Sep 2026 11:09:15 +0800 Subject: [PATCH] feat(ascend): add fill provider --- src/native/ascend/ops/fill/kernel.h | 113 ++++++++++++++++++++++++++++ tests/test_fill.py | 35 ++++++++- 2 files changed, 147 insertions(+), 1 deletion(-) create mode 100644 src/native/ascend/ops/fill/kernel.h diff --git a/src/native/ascend/ops/fill/kernel.h b/src/native/ascend/ops/fill/kernel.h new file mode 100644 index 000000000..c067ab4a7 --- /dev/null +++ b/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 + +#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 : 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(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(stream_); + auto t_value = value_cache_.get(const_cast(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_ diff --git a/tests/test_fill.py b/tests/test_fill.py index da624d725..0f527c52f 100644 --- a/tests/test_fill.py +++ b/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)