From 2cc81ee5879e9d83ef8140a03160273d68be5831 Mon Sep 17 00:00:00 2001 From: Li Baoming <1508269885@qq.com> Date: Tue, 1 Sep 2026 03:29:32 +0000 Subject: [PATCH] feat(cambricon): add gelu provider --- src/native/cambricon/ops/gelu/kernel.h | 147 +++++++++++++++++++ src/native/cambricon/ops/gelu/kernel.mlu | 172 +++++++++++++++++++++++ tests/test_gelu.py | 2 + 3 files changed, 321 insertions(+) create mode 100644 src/native/cambricon/ops/gelu/kernel.h create mode 100644 src/native/cambricon/ops/gelu/kernel.mlu diff --git a/src/native/cambricon/ops/gelu/kernel.h b/src/native/cambricon/ops/gelu/kernel.h new file mode 100644 index 000000000..30e354733 --- /dev/null +++ b/src/native/cambricon/ops/gelu/kernel.h @@ -0,0 +1,147 @@ +#ifndef INFINI_OPS_CAMBRICON_GELU_CNNL_H_ +#define INFINI_OPS_CAMBRICON_GELU_CNNL_H_ + +#include +#include + +#include "base/gelu.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" +#include "native/cambricon/data_type_.h" + +namespace infini::ops { + +template +void GeluUnion(void* workspace, int core_per_cluster, int cluster_count, + cnrtQueue_t queue, const void* input, void* out, + const size_t* input_shape, const ptrdiff_t* input_strides, + const size_t* out_shape, const ptrdiff_t* out_strides, + size_t output_size, int ndim, bool approximate, + bool input_contiguous, bool out_contiguous); + +template <> +class Operator : public Gelu { + public: + Operator(const Tensor input, const std::string approximate, Tensor out) + : Gelu{input, approximate, out} { + assert(input_shape_ == out_shape_ && + "`CambriconGelu` requires matching input and output shapes."); + assert(input_type_ == out_type_ && + "`CambriconGelu` requires matching input and output dtypes."); + assert(input.device() == out.device() && + "`CambriconGelu` requires input and output on the same device."); + assert((input_type_ == DataType::kFloat16 || + input_type_ == DataType::kBFloat16 || + input_type_ == DataType::kFloat32) && + "`CambriconGelu` supports float16, bfloat16, and float32 only."); + assert(!out.HasBroadcastDim() && + "`CambriconGelu` output must not have broadcast dimensions."); + assert(std::all_of(input_strides_.begin(), input_strides_.end(), + [](auto stride) { return stride >= 0; }) && + "`CambriconGelu` does not support negative input strides."); + assert(std::all_of(out_strides_.begin(), out_strides_.end(), + [](auto stride) { return stride >= 0; }) && + "`CambriconGelu` does not support negative output strides."); + + if (output_size_ == 0) { + return; + } + + if (approximate_ == "tanh" || input_type_ == DataType::kBFloat16) { + cnrt_utils::GetLaunchConfig(input.device(), &core_per_cluster_, + &cluster_count_); + const auto workspace_size = workspace_size_in_bytes(); + if (workspace_size != 0) { + CNRT_CHECK(cnrtMalloc(&default_workspace_, workspace_size)); + } + return; + } + + cnnl_handle_ = cnnl_utils::CreateHandle(); + input_desc_ = cnnl_utils::MakeTensorDescriptor(input_type_, input_shape_, + input_strides_); + out_desc_ = + cnnl_utils::MakeTensorDescriptor(out_type_, out_shape_, out_strides_); + + INFINI_OPS_CNNL_CHECK(cnnlCreateActivationDescriptor(&activation_desc_)); + const cnnlActivationMode_t mode = CNNL_ACTIVATION_GELU; + const cnnlComputationPreference_t preference = + CNNL_COMPUTATION_HIGH_PRECISION; + const cnnlNanPropagation_t nan_propagation = CNNL_PROPAGATE_NAN; + const bool use_approximation = false; + INFINI_OPS_CNNL_CHECK(cnnlSetActivationDescAttr( + activation_desc_, CNNL_ACTIVATION_MODE, &mode, sizeof(mode))); + INFINI_OPS_CNNL_CHECK( + cnnlSetActivationDescAttr(activation_desc_, CNNL_ACTIVATION_PREFERENCE, + &preference, sizeof(preference))); + INFINI_OPS_CNNL_CHECK( + cnnlSetActivationDescAttr(activation_desc_, CNNL_ACTIVATION_NAN_PROP, + &nan_propagation, sizeof(nan_propagation))); + INFINI_OPS_CNNL_CHECK(cnnlSetActivationDescAttr( + activation_desc_, CNNL_ACTIVATION_APPROXIMATE, &use_approximation, + sizeof(use_approximation))); + } + + ~Operator() { + if (default_workspace_) { + (void)cnrtFree(default_workspace_); + } + if (activation_desc_) { + (void)cnnlDestroyActivationDescriptor(activation_desc_); + } + } + + void operator()(const Tensor input, const std::string approximate, + Tensor out) const override { + assert(approximate == approximate_ && + "`CambriconGelu` attributes changed after descriptor creation."); + if (output_size_ == 0) { + return; + } + + if (approximate_ == "tanh" || input_type_ == DataType::kBFloat16) { + auto queue = static_cast(stream_ ? stream_ : 0); + auto* workspace = workspace_ ? workspace_ : default_workspace_; + DispatchFunc< + Device::Type::kCambricon, + List>( + {out_type_}, + [&](auto tag) { + using T = typename decltype(tag)::type; + GeluUnion(workspace, core_per_cluster_, cluster_count_, queue, + input.data(), out.data(), input_shape_.data(), + input_strides_.data(), out_shape_.data(), + out_strides_.data(), output_size_, + static_cast(ndim_), approximate_ == "tanh", + is_input_contiguous_, is_out_contiguous_); + }, + "CambriconGelu::operator() - output dispatch"); + return; + } + + INFINI_OPS_CNNL_CHECK(cnnlSetQueue( + cnnl_handle_.get(), static_cast(stream_ ? stream_ : 0))); + INFINI_OPS_CNNL_CHECK(cnnlActivationForward( + cnnl_handle_.get(), activation_desc_, nullptr, input_desc_.get(), + input.data(), nullptr, out_desc_.get(), out.data())); + } + + std::size_t workspace_size_in_bytes() const override { + return (approximate_ == "tanh" || input_type_ == DataType::kBFloat16) + ? ndim_ * (2 * sizeof(size_t) + 2 * sizeof(ptrdiff_t)) + : 0; + } + + private: + cnnl_utils::Handle cnnl_handle_{}; + cnnl_utils::TensorDescriptor input_desc_{}; + cnnl_utils::TensorDescriptor out_desc_{}; + cnnlActivationDescriptor_t activation_desc_{nullptr}; + void* default_workspace_{nullptr}; + int core_per_cluster_{0}; + int cluster_count_{0}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/gelu/kernel.mlu b/src/native/cambricon/ops/gelu/kernel.mlu new file mode 100644 index 000000000..b79775873 --- /dev/null +++ b/src/native/cambricon/ops/gelu/kernel.mlu @@ -0,0 +1,172 @@ +#include +#include +#include + +#include "kernel.h" + +namespace infini::ops { +namespace { + +__nram__ char gelu_nram_buffer[NRAM_MAX_SIZE] __attribute__((aligned(128))); + +__mlu_device__ ptrdiff_t LogicalToOffset(size_t logical_index, int ndim, + const size_t* shape, + const ptrdiff_t* strides) { + ptrdiff_t offset = 0; + for (int dim = ndim - 1; dim >= 0; --dim) { + const size_t coordinate = logical_index % shape[dim]; + logical_index /= shape[dim]; + offset += static_cast(coordinate) * strides[dim]; + } + return offset; +} + +template +__mlu_device__ void ComputeGelu(const T* input, T* output, float* values, + size_t count, bool approximate) { + if constexpr (std::is_same_v) { + __bang_half2float(values, reinterpret_cast(const_cast(input)), + count); + } else if constexpr (std::is_same_v) { + __bang_bfloat162float(values, const_cast(input), count); + } else { + __memcpy(values, const_cast(input), count * sizeof(float), NRAM2NRAM); + } + + constexpr float kSqrtTwoOverPi = 0.7978845608028654F; + constexpr float kCubicCoefficient = 0.044715F; + for (size_t i = 0; i < count; ++i) { + const float x = values[i]; + if (approximate) { + const float inner = kSqrtTwoOverPi * (x + kCubicCoefficient * x * x * x); + values[i] = 0.5F * x * (1.0F + tanhf(inner)); + } else { + constexpr float kInverseSqrtTwo = 0.7071067811865475F; + values[i] = 0.5F * x * (1.0F + erff(x * kInverseSqrtTwo)); + } + } + + if constexpr (std::is_same_v) { + __bang_float2half(reinterpret_cast(output), values, count); + } else if constexpr (std::is_same_v) { + __bang_float2bfloat16(output, values, count); + } else { + __memcpy(output, values, count * sizeof(float), NRAM2NRAM); + } +} + +template +__mlu_global__ void GeluKernel(const T* input, T* output, + const size_t* input_shape, + const ptrdiff_t* input_strides, + const size_t* out_shape, + const ptrdiff_t* out_strides, size_t output_size, + int ndim, bool approximate, + bool input_contiguous, bool out_contiguous) { + const size_t elements_per_task = (output_size + taskDim - 1) / taskDim; + const size_t begin = taskId * elements_per_task; + const size_t end = begin + elements_per_task < output_size + ? begin + elements_per_task + : output_size; + if (begin >= end) { + return; + } + + size_t block_size = NRAM_MAX_SIZE / (2 * sizeof(T) + sizeof(float)); + block_size = block_size / 64 * 64; + if (block_size == 0) { + block_size = 1; + } + auto* input_buffer = reinterpret_cast(gelu_nram_buffer); + auto* output_buffer = input_buffer + block_size; + auto* values = reinterpret_cast(output_buffer + block_size); + + size_t processed = begin; + while (processed < end) { + const size_t current = + block_size < end - processed ? block_size : end - processed; + if (input_contiguous) { + __memcpy(input_buffer, input + processed, current * sizeof(T), + GDRAM2NRAM); + } else { + for (size_t i = 0; i < current; ++i) { + const size_t logical = processed + i; + input_buffer[i] = + input[LogicalToOffset(logical, ndim, input_shape, input_strides)]; + } + } + + ComputeGelu(input_buffer, output_buffer, values, current, approximate); + if (out_contiguous) { + __memcpy(output + processed, output_buffer, current * sizeof(T), + NRAM2GDRAM); + } else { + for (size_t i = 0; i < current; ++i) { + const size_t logical = processed + i; + output[LogicalToOffset(logical, ndim, out_shape, out_strides)] = + output_buffer[i]; + } + } + processed += current; + } +} + +} // namespace + +template +void GeluUnion(void* workspace, int core_per_cluster, int cluster_count, + cnrtQueue_t queue, const void* input, void* out, + const size_t* input_shape, const ptrdiff_t* input_strides, + const size_t* out_shape, const ptrdiff_t* out_strides, + size_t output_size, int ndim, bool approximate, + bool input_contiguous, bool out_contiguous) { + size_t* device_input_shape = nullptr; + size_t* device_out_shape = nullptr; + ptrdiff_t* device_input_strides = nullptr; + ptrdiff_t* device_out_strides = nullptr; + if (ndim != 0) { + auto* workspace_bytes = static_cast(workspace); + device_input_shape = reinterpret_cast(workspace_bytes); + device_out_shape = device_input_shape + ndim; + device_input_strides = + reinterpret_cast(device_out_shape + ndim); + device_out_strides = device_input_strides + ndim; + + CNRT_CHECK( + cnrtMemcpyAsync(device_input_shape, const_cast(input_shape), + ndim * sizeof(size_t), queue, cnrtMemcpyHostToDev)); + CNRT_CHECK(cnrtMemcpyAsync(device_out_shape, const_cast(out_shape), + ndim * sizeof(size_t), queue, + cnrtMemcpyHostToDev)); + CNRT_CHECK(cnrtMemcpyAsync( + device_input_strides, const_cast(input_strides), + ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev)); + CNRT_CHECK( + cnrtMemcpyAsync(device_out_strides, const_cast(out_strides), + ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev)); + } + + cnrtDim3_t kernel_dim; + kernel_dim.x = core_per_cluster; + kernel_dim.y = cluster_count; + kernel_dim.z = 1; + (void)cnrtGetLastError(); + GeluKernel<<>>( + static_cast(input), static_cast(out), device_input_shape, + device_input_strides, device_out_shape, device_out_strides, output_size, + ndim, approximate, input_contiguous, out_contiguous); + CNRT_CHECK(cnrtGetLastError()); +} + +#define INSTANTIATE_GELU(T) \ + template void GeluUnion(void*, int, int, cnrtQueue_t, const void*, void*, \ + const size_t*, const ptrdiff_t*, const size_t*, \ + const ptrdiff_t*, size_t, int, bool, bool, bool) + +INSTANTIATE_GELU(__half); +INSTANTIATE_GELU(__bang_bfloat16); +INSTANTIATE_GELU(float); + +#undef INSTANTIATE_GELU + +} // namespace infini::ops diff --git a/tests/test_gelu.py b/tests/test_gelu.py index c943bf34f..710e4eb11 100644 --- a/tests/test_gelu.py +++ b/tests/test_gelu.py @@ -42,6 +42,8 @@ def test_gelu( ): if device == "musa" and dtype == torch.float64: pytest.skip("MUSA does not support float64 GELU") + if device == "mlu" and dtype == torch.float64: + pytest.skip("Cambricon CNNL does not support float64 GELU") input = randn_strided(shape, input_strides, dtype=dtype, device=device) out = (