diff --git a/src/native/cambricon/ops/fill/kernel.h b/src/native/cambricon/ops/fill/kernel.h new file mode 100644 index 000000000..09ea2c789 --- /dev/null +++ b/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 +#include +#include +#include + +#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, ReducedFloatTypes>, + AllIntTypes>; + +template +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 : 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(stream_ ? stream_ : 0); + void* workspace = workspace_ ? workspace_ : default_workspace_.get(); + + DispatchFunc( + out_type_, + [&](auto tag) { + using T = typename decltype(tag)::type; + const T* effective_device_value = static_cast(device_value); + if constexpr (std::is_same_v || + std::is_same_v) { + if (effective_device_value == nullptr) { + const T converted_host_value = static_cast(host_value); + CNRT_CHECK(cnrtMemcpy(workspace, + const_cast(&converted_host_value), + sizeof(T), cnrtMemcpyHostToDev)); + effective_device_value = static_cast(workspace); + } + } + if constexpr (std::is_same_v || + std::is_same_v) { + FillRaw64Union(workspace, queue, core_per_cluster_, cluster_count_, + effective_device_value, out.data(), + out_shape_.data(), out_strides_.data(), output_size_, + static_cast(ndim_), is_out_contiguous_); + } else { + FillUnion(workspace, queue, core_per_cluster_, cluster_count_, + effective_device_value, host_value, + static_cast(out.data()), out_shape_.data(), + out_strides_.data(), output_size_, + static_cast(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_ diff --git a/src/native/cambricon/ops/fill/kernel.mlu b/src/native/cambricon/ops/fill/kernel.mlu new file mode 100644 index 000000000..0f508b9ca --- /dev/null +++ b/src/native/cambricon/ops/fill/kernel.mlu @@ -0,0 +1,170 @@ +#include +#include +#include +#include + +#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 +__mlu_device__ void FillBuffer(T* buffer, std::size_t count, + const T* device_value) { + if constexpr (std::is_same_v) { + 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 +__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(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(out); + for (std::size_t logical = range.begin; logical < range.end; ++logical) { + const ptrdiff_t offset = out_contiguous + ? static_cast(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(workspace) + sizeof(std::uint64_t); + std::size_t* device_shape = nullptr; + ptrdiff_t* device_strides = nullptr; + if (ndim != 0) { + device_shape = reinterpret_cast(bytes); + device_strides = reinterpret_cast(device_shape + ndim); + CNRT_CHECK(cnrtMemcpyAsync( + device_shape, const_cast(out_shape), + ndim * sizeof(std::size_t), queue, cnrtMemcpyHostToDev)); + CNRT_CHECK( + cnrtMemcpyAsync(device_strides, const_cast(out_strides), + ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev)); + } + + const cnrtDim3_t kernel_dim = {static_cast(core_per_cluster), + static_cast(cluster_count), 1}; + FillRaw64Kernel<<>>( + device_value, out, device_shape, device_strides, output_size, ndim, + out_contiguous); + CNRT_CHECK(cnrtGetLastError()); +} + +template +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(workspace); + auto* device_scalar = reinterpret_cast(bytes); + bytes += sizeof(std::uint64_t); + std::size_t* device_shape = nullptr; + ptrdiff_t* device_strides = nullptr; + if (ndim != 0) { + device_shape = reinterpret_cast(bytes); + device_strides = reinterpret_cast(device_shape + ndim); + CNRT_CHECK(cnrtMemcpyAsync( + device_shape, const_cast(out_shape), + ndim * sizeof(std::size_t), queue, cnrtMemcpyHostToDev)); + CNRT_CHECK( + cnrtMemcpyAsync(device_strides, const_cast(out_strides), + ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev)); + } + + const cnrtDim3_t kernel_dim = {static_cast(core_per_cluster), + static_cast(cluster_count), 1}; + if (device_value == nullptr) { + T converted_host_value; + if constexpr (std::is_same_v || + std::is_same_v) { + converted_host_value = static_cast(static_cast(host_value)); + } else { + converted_host_value = static_cast(host_value); + } + CNRT_CHECK(cnrtMemcpy(device_scalar, &converted_host_value, sizeof(T), + cnrtMemcpyHostToDev)); + device_value = device_scalar; + } + FillKernel<<>>( + device_value, out, device_shape, device_strides, output_size, ndim, + out_contiguous); + CNRT_CHECK(cnrtGetLastError()); +} + +#define INSTANTIATE_FILL(T) \ + template void FillUnion(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 diff --git a/tests/test_fill.py b/tests/test_fill.py index af2cf9b1a..da624d725 100644 --- a/tests/test_fill.py +++ b/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 = (