From 9639aed76355f1b9bc2d8647584849890b1b2ef1 Mon Sep 17 00:00:00 2001 From: Li Baoming <1508269885@qq.com> Date: Mon, 31 Aug 2026 06:55:36 +0000 Subject: [PATCH] feat(cambricon): add mul provider --- src/native/cambricon/kernel_utils.h | 40 +++++++ src/native/cambricon/ops/mul/kernel.h | 79 +++++++++++++ src/native/cambricon/ops/mul/kernel.mlu | 144 ++++++++++++++++++++++++ tests/test_mul.py | 4 +- 4 files changed, 265 insertions(+), 2 deletions(-) create mode 100644 src/native/cambricon/kernel_utils.h create mode 100644 src/native/cambricon/ops/mul/kernel.h create mode 100644 src/native/cambricon/ops/mul/kernel.mlu diff --git a/src/native/cambricon/kernel_utils.h b/src/native/cambricon/kernel_utils.h new file mode 100644 index 000000000..3f252ad62 --- /dev/null +++ b/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 + +#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(coordinate) * strides[dim]; + } + return offset; +} + +} // namespace infini::ops::cambricon::kernel_utils + +#endif // __BANG__ + +#endif // INFINI_OPS_CAMBRICON_KERNEL_UTILS_H_ diff --git a/src/native/cambricon/ops/mul/kernel.h b/src/native/cambricon/ops/mul/kernel.h new file mode 100644 index 000000000..6cd9090eb --- /dev/null +++ b/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 + +#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 +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 : 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(stream_ ? stream_ : 0); + void* workspace = workspace_ ? workspace_ : default_workspace_.get(); + + using SupportedTypes = + List; + DispatchFunc( + out_type_, + [&](auto tag) { + using T = typename decltype(tag)::type; + MulUnion(workspace, queue, core_per_cluster_, cluster_count_, + static_cast(input.data()), + static_cast(other.data()), + static_cast(out.data()), out_shape_.data(), + input_strides_.data(), other_strides_.data(), + out_strides_.data(), output_size_, + static_cast(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_ diff --git a/src/native/cambricon/ops/mul/kernel.mlu b/src/native/cambricon/ops/mul/kernel.mlu new file mode 100644 index 000000000..189f13f53 --- /dev/null +++ b/src/native/cambricon/ops/mul/kernel.mlu @@ -0,0 +1,144 @@ +#include +#include +#include +#include + +#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 +__mlu_device__ void Multiply(const T* input, const T* other, T* out, + std::size_t count) { + if constexpr (std::is_same_v) { + __bang_mul(reinterpret_cast(out), + reinterpret_cast(input), + reinterpret_cast(other), count); + } else if constexpr (std::is_same_v || + std::is_same_v) { + __bang_mul(out, input, other, count); + } else { + for (std::size_t i = 0; i < count; ++i) { + out[i] = static_cast(input[i] * other[i]); + } + } +} + +template +__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(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 +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(workspace); + auto* device_shape = reinterpret_cast(bytes); + auto* device_input_strides = + reinterpret_cast(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(out_shape), + ndim * sizeof(std::size_t), queue, cnrtMemcpyHostToDev)); + CNRT_CHECK(cnrtMemcpyAsync( + device_input_strides, const_cast(input_strides), + ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev)); + CNRT_CHECK(cnrtMemcpyAsync( + device_other_strides, const_cast(other_strides), + ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev)); + CNRT_CHECK( + cnrtMemcpyAsync(device_out_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}; + MulKernel<<>>( + 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(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 diff --git a/tests/test_mul.py b/tests/test_mul.py index cdfb67222..c9cb674a8 100644 --- a/tests/test_mul.py +++ b/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: