From 6d389ab68ae080bd947624a95fb23e108847645a Mon Sep 17 00:00:00 2001 From: InfiniTensor integration candidate Date: Thu, 5 Nov 2026 23:15:23 +0800 Subject: [PATCH] feat(cambricon): extend native model operator coverage --- include/infini/convolution_layout.h | 13 + src/CMakeLists.txt | 19 +- src/base/gather_axis.h | 20 ++ src/base/grid_sample.h | 19 ++ src/base/instance_norm.h | 18 ++ src/base/masked_scatter.h | 19 ++ src/base/resize2d.h | 23 ++ src/base/roi_align.h | 20 ++ src/base/scatter_nd.h | 18 ++ src/base/tile.h | 16 ++ src/convolution_layout.cc | 14 + src/native/cambricon/activation_cnnl.h | 50 ++++ src/native/cambricon/cnnl_utils.h | 12 +- src/native/cambricon/common.h | 6 + src/native/cambricon/compare_cnnl.h | 64 +++++ src/native/cambricon/data_type_.h | 24 +- src/native/cambricon/host_bang_compat.h | 4 + src/native/cambricon/op_tensor.h | 54 ++++ src/native/cambricon/ops/abs/kernel.h | 18 ++ src/native/cambricon/ops/add/kernel.h | 27 +- src/native/cambricon/ops/add/kernel.mlu | 88 +++++-- src/native/cambricon/ops/amax/kernel.h | 31 +++ src/native/cambricon/ops/argmax/kernel.h | 69 ++++- src/native/cambricon/ops/cast/kernel.h | 84 ++++++ src/native/cambricon/ops/cat/kernel.h | 97 +++++++ src/native/cambricon/ops/clip/kernel.h | 74 ++++++ src/native/cambricon/ops/convolution/kernel.h | 246 ++++++++++++++++++ src/native/cambricon/ops/cos/kernel.h | 18 ++ src/native/cambricon/ops/div/kernel.h | 123 +++++++++ src/native/cambricon/ops/embedding/kernel.mlu | 4 +- src/native/cambricon/ops/eq/kernel.h | 23 ++ src/native/cambricon/ops/erf/kernel.h | 18 ++ src/native/cambricon/ops/fill/kernel.mlu | 2 +- .../ops/fused_add_rms_norm/kernel.mlu | 6 +- src/native/cambricon/ops/ge/kernel.h | 23 ++ src/native/cambricon/ops/gelu/kernel.mlu | 4 +- src/native/cambricon/ops/grid_sample/kernel.h | 80 ++++++ .../cambricon/ops/index_select/kernel.h | 83 ++++++ .../cambricon/ops/instance_norm/kernel.h | 116 +++++++++ src/native/cambricon/ops/layer_norm/kernel.h | 91 +++++++ src/native/cambricon/ops/less/kernel.h | 23 ++ src/native/cambricon/ops/log/kernel.h | 24 ++ src/native/cambricon/ops/logical_not/kernel.h | 18 ++ .../cambricon/ops/masked_scatter/kernel.h | 93 +++++++ src/native/cambricon/ops/matmul/kernel.h | 163 ++++++++++++ src/native/cambricon/ops/max_pool2d/kernel.h | 84 ++++++ src/native/cambricon/ops/maximum/kernel.h | 51 ++++ src/native/cambricon/ops/mean/kernel.h | 45 ++++ src/native/cambricon/ops/mul/kernel.h | 23 +- src/native/cambricon/ops/mul/kernel.mlu | 4 +- src/native/cambricon/ops/neg/kernel.h | 18 ++ src/native/cambricon/ops/norm/kernel.h | 37 +++ src/native/cambricon/ops/pow/kernel.h | 71 +++++ src/native/cambricon/ops/relu/kernel.mlu | 2 +- src/native/cambricon/ops/resize2d/kernel.h | 152 +++++++++++ src/native/cambricon/ops/rms_norm/kernel.mlu | 47 +++- src/native/cambricon/ops/rms_norm/rms_norm.h | 14 +- src/native/cambricon/ops/roi_align/kernel.h | 72 +++++ .../cambricon/ops/rotary_embedding/kernel.mlu | 4 +- src/native/cambricon/ops/rsqrt/kernel.h | 18 ++ src/native/cambricon/ops/scatter/kernel.h | 115 ++++++++ src/native/cambricon/ops/scatter_nd/kernel.h | 95 +++++++ .../cambricon/ops/scatter_nd/kernel.mlu | 73 ++++++ src/native/cambricon/ops/sign/kernel.h | 18 ++ .../cambricon/ops/silu_and_mul/kernel.h | 2 +- .../cambricon/ops/silu_and_mul/kernel.mlu | 4 +- src/native/cambricon/ops/sin/kernel.h | 18 ++ src/native/cambricon/ops/sqrt/kernel.h | 18 ++ src/native/cambricon/ops/sum/kernel.h | 38 +++ src/native/cambricon/ops/tanh/kernel.h | 18 ++ src/native/cambricon/ops/tile/kernel.h | 40 +++ .../kernel.mlu | 2 +- src/native/cambricon/ops/where/kernel.h | 68 +++++ src/native/cambricon/reduce_cnnl.h | 97 +++++++ src/native/cambricon/scatter_index.h | 19 ++ src/native/cambricon/unary_cnnl.h | 39 +++ 76 files changed, 3295 insertions(+), 70 deletions(-) create mode 100644 include/infini/convolution_layout.h create mode 100644 src/base/gather_axis.h create mode 100644 src/base/grid_sample.h create mode 100644 src/base/instance_norm.h create mode 100644 src/base/masked_scatter.h create mode 100644 src/base/resize2d.h create mode 100644 src/base/roi_align.h create mode 100644 src/base/scatter_nd.h create mode 100644 src/base/tile.h create mode 100644 src/convolution_layout.cc create mode 100644 src/native/cambricon/activation_cnnl.h create mode 100644 src/native/cambricon/compare_cnnl.h create mode 100644 src/native/cambricon/op_tensor.h create mode 100644 src/native/cambricon/ops/abs/kernel.h create mode 100644 src/native/cambricon/ops/amax/kernel.h create mode 100644 src/native/cambricon/ops/cast/kernel.h create mode 100644 src/native/cambricon/ops/cat/kernel.h create mode 100644 src/native/cambricon/ops/clip/kernel.h create mode 100644 src/native/cambricon/ops/convolution/kernel.h create mode 100644 src/native/cambricon/ops/cos/kernel.h create mode 100644 src/native/cambricon/ops/div/kernel.h create mode 100644 src/native/cambricon/ops/eq/kernel.h create mode 100644 src/native/cambricon/ops/erf/kernel.h create mode 100644 src/native/cambricon/ops/ge/kernel.h create mode 100644 src/native/cambricon/ops/grid_sample/kernel.h create mode 100644 src/native/cambricon/ops/index_select/kernel.h create mode 100644 src/native/cambricon/ops/instance_norm/kernel.h create mode 100644 src/native/cambricon/ops/layer_norm/kernel.h create mode 100644 src/native/cambricon/ops/less/kernel.h create mode 100644 src/native/cambricon/ops/log/kernel.h create mode 100644 src/native/cambricon/ops/logical_not/kernel.h create mode 100644 src/native/cambricon/ops/masked_scatter/kernel.h create mode 100644 src/native/cambricon/ops/matmul/kernel.h create mode 100644 src/native/cambricon/ops/max_pool2d/kernel.h create mode 100644 src/native/cambricon/ops/maximum/kernel.h create mode 100644 src/native/cambricon/ops/mean/kernel.h create mode 100644 src/native/cambricon/ops/neg/kernel.h create mode 100644 src/native/cambricon/ops/norm/kernel.h create mode 100644 src/native/cambricon/ops/pow/kernel.h create mode 100644 src/native/cambricon/ops/resize2d/kernel.h create mode 100644 src/native/cambricon/ops/roi_align/kernel.h create mode 100644 src/native/cambricon/ops/rsqrt/kernel.h create mode 100644 src/native/cambricon/ops/scatter/kernel.h create mode 100644 src/native/cambricon/ops/scatter_nd/kernel.h create mode 100644 src/native/cambricon/ops/scatter_nd/kernel.mlu create mode 100644 src/native/cambricon/ops/sign/kernel.h create mode 100644 src/native/cambricon/ops/sin/kernel.h create mode 100644 src/native/cambricon/ops/sqrt/kernel.h create mode 100644 src/native/cambricon/ops/sum/kernel.h create mode 100644 src/native/cambricon/ops/tanh/kernel.h create mode 100644 src/native/cambricon/ops/tile/kernel.h create mode 100644 src/native/cambricon/ops/where/kernel.h create mode 100644 src/native/cambricon/reduce_cnnl.h create mode 100644 src/native/cambricon/scatter_index.h create mode 100644 src/native/cambricon/unary_cnnl.h diff --git a/include/infini/convolution_layout.h b/include/infini/convolution_layout.h new file mode 100644 index 000000000..ed5a1f3fc --- /dev/null +++ b/include/infini/convolution_layout.h @@ -0,0 +1,13 @@ +#ifndef INFINI_OPS_CONVOLUTION_LAYOUT_H_ +#define INFINI_OPS_CONVOLUTION_LAYOUT_H_ + +#include + +namespace infini::ops { + +// Logical NCHW tensors may use packed channels-last strides. +bool SupportsConvolutionChannelsLast(const infini::rt::Device& device); + +} // namespace infini::ops + +#endif // INFINI_OPS_CONVOLUTION_LAYOUT_H_ diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index e685b4f12..673293d57 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -434,9 +434,23 @@ if(WITH_CAMBRICON) if(CNCC_COMPILER) message(STATUS "Found cncc: ${CNCC_COMPILER}") separate_arguments(INFINI_MLU_HOST_ABI_FLAGS NATIVE_COMMAND "${CMAKE_CXX_FLAGS}") + set(INFINIOPS_CAMBRICON_MLU_ARCH "mtp_592" CACHE STRING "BANG target architecture") + option(INFINIOPS_CAMBRICON_BANG_BF16 "Enable BF16 in native BANG kernels" ON) + target_compile_definitions(infiniops PRIVATE + INFINIOPS_CAMBRICON_BANG_BF16=$) + set(_cncc_host_includes "") + foreach(_dir IN LISTS CMAKE_CXX_IMPLICIT_INCLUDE_DIRECTORIES) + if(_dir MATCHES "/c[+][+]/") + list(APPEND _cncc_host_includes -isystem "${_dir}") + endif() + endforeach() set(MLU_COMPILE_OPTS - -c --bang-mlu-arch=mtp_592 -O3 -fPIC -Wall -Werror -std=c++17 -pthread + --bang-mlu-arch=${INFINIOPS_CAMBRICON_MLU_ARCH} + -DINFINIOPS_CAMBRICON_NATIVE=1 + -DINFINIOPS_CAMBRICON_BANG_BF16=$ + -O3 -fPIC -Wall -Werror -Wno-error=return-type -std=c++17 -pthread -I${CMAKE_CURRENT_SOURCE_DIR} -I${NEUWARE_HOME}/include + ${_cncc_host_includes} ${INFINI_RT_INCLUDE_FLAGS} ${INFINI_MLU_HOST_ABI_FLAGS} -idirafter /usr/local/neuware/lib/clang/11.1.0/include @@ -1490,6 +1504,9 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS) set_source_files_properties(${ILUVATAR_OPERATOR_DISPATCH_OBJECTS} PROPERTIES EXTERNAL_OBJECT TRUE GENERATED TRUE) target_sources(infiniops PRIVATE ${ILUVATAR_OPERATOR_DISPATCH_OBJECTS}) + elseif(WITH_CAMBRICON AND NOT GENERATE_PYTHON_BINDINGS) + # C++ consumers need dispatch entry points without the Python extension. + target_sources(infiniops PRIVATE ${OPERATOR_DISPATCH_SOURCES}) endif() endif() endif() diff --git a/src/base/gather_axis.h b/src/base/gather_axis.h new file mode 100644 index 000000000..2bee6252e --- /dev/null +++ b/src/base/gather_axis.h @@ -0,0 +1,20 @@ +#ifndef INFINI_OPS_BASE_GATHER_AXIS_H_ +#define INFINI_OPS_BASE_GATHER_AXIS_H_ + +#include "operator.h" + +namespace infini::ops { + +// Select a one-dimensional list of indices, wrapping negatives on the axis. +class GatherAxis : public Operator { + public: + GatherAxis(const Tensor input, const Tensor index, const int64_t dim, + Tensor out) {} + + virtual void operator()(const Tensor input, const Tensor index, + const int64_t dim, Tensor out) const = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/grid_sample.h b/src/base/grid_sample.h new file mode 100644 index 000000000..21ef365c1 --- /dev/null +++ b/src/base/grid_sample.h @@ -0,0 +1,19 @@ +#ifndef INFINI_OPS_BASE_GRID_SAMPLE_H_ +#define INFINI_OPS_BASE_GRID_SAMPLE_H_ + +#include "operator.h" + +namespace infini::ops { + +// Spatial sampling in NHWC layout; the last grid dimension stores (x, y). +class GridSample : public Operator { + public: + GridSample(const Tensor input, const Tensor grid, bool nearest, + bool reflection, bool align_corners, Tensor out) {} + virtual void operator()(const Tensor input, const Tensor grid, bool nearest, + bool reflection, bool align_corners, Tensor out) const = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/instance_norm.h b/src/base/instance_norm.h new file mode 100644 index 000000000..83b738f98 --- /dev/null +++ b/src/base/instance_norm.h @@ -0,0 +1,18 @@ +#ifndef INFINI_OPS_BASE_INSTANCE_NORM_H_ +#define INFINI_OPS_BASE_INSTANCE_NORM_H_ + +#include "operator.h" + +namespace infini::ops { + +class InstanceNorm : public Operator { + public: + InstanceNorm(const Tensor input, const Tensor scale, const Tensor bias, + float epsilon, Tensor out) {} + virtual void operator()(const Tensor input, const Tensor scale, const Tensor bias, + float epsilon, Tensor out) const = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/masked_scatter.h b/src/base/masked_scatter.h new file mode 100644 index 000000000..6314d1b67 --- /dev/null +++ b/src/base/masked_scatter.h @@ -0,0 +1,19 @@ +#ifndef INFINI_OPS_BASE_MASKED_SCATTER_H_ +#define INFINI_OPS_BASE_MASKED_SCATTER_H_ + +#include "operator.h" + +namespace infini::ops { + +// Replace true positions in row-major order with consecutive source values. +class MaskedScatter : public Operator { + public: + MaskedScatter(const Tensor input, const Tensor mask, const Tensor source, + Tensor out) {} + virtual void operator()(const Tensor input, const Tensor mask, + const Tensor source, Tensor out) const = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/resize2d.h b/src/base/resize2d.h new file mode 100644 index 000000000..559493b55 --- /dev/null +++ b/src/base/resize2d.h @@ -0,0 +1,23 @@ +#ifndef INFINI_OPS_BASE_RESIZE2D_H_ +#define INFINI_OPS_BASE_RESIZE2D_H_ + +#include "operator.h" + +namespace infini::ops { + +class Resize2d : public Operator { + public: + enum Mode : int64_t { kNearest, kLinear }; + enum Coordinates : int64_t { kHalfPixel, kPytorchHalfPixel, kAlignCorners, kAsymmetric }; + enum Rounding : int64_t { kPreferFloor, kPreferCeil, kFloor, kCeil }; + + Resize2d(const Tensor input, int64_t mode, int64_t coordinates, + int64_t rounding, float height_scale, float width_scale, Tensor out) {} + virtual void operator()(const Tensor input, int64_t mode, int64_t coordinates, + int64_t rounding, float height_scale, float width_scale, + Tensor out) const = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/roi_align.h b/src/base/roi_align.h new file mode 100644 index 000000000..9c4a5f18b --- /dev/null +++ b/src/base/roi_align.h @@ -0,0 +1,20 @@ +#ifndef INFINI_OPS_BASE_ROI_ALIGN_H_ +#define INFINI_OPS_BASE_ROI_ALIGN_H_ + +#include "operator.h" + +namespace infini::ops { + +// Average pooling over NHWC feature maps. Boxes are [batch, x1, y1, x2, y2]. +class RoiAlign : public Operator { + public: + RoiAlign(const Tensor input, const Tensor boxes, int64_t sampling_ratio, + float spatial_scale, bool aligned, Tensor out) {} + virtual void operator()(const Tensor input, const Tensor boxes, + int64_t sampling_ratio, float spatial_scale, + bool aligned, Tensor out) const = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/scatter_nd.h b/src/base/scatter_nd.h new file mode 100644 index 000000000..50eae239a --- /dev/null +++ b/src/base/scatter_nd.h @@ -0,0 +1,18 @@ +#ifndef INFINI_OPS_BASE_SCATTER_ND_H_ +#define INFINI_OPS_BASE_SCATTER_ND_H_ + +#include "operator.h" + +namespace infini::ops { + +class ScatterNd : public Operator { + public: + ScatterNd(const Tensor input, const Tensor indices, const Tensor updates, + bool add, Tensor out) {} + virtual void operator()(const Tensor input, const Tensor indices, + const Tensor updates, bool add, Tensor out) const = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/tile.h b/src/base/tile.h new file mode 100644 index 000000000..f84f1c2b8 --- /dev/null +++ b/src/base/tile.h @@ -0,0 +1,16 @@ +#ifndef INFINI_OPS_BASE_TILE_H_ +#define INFINI_OPS_BASE_TILE_H_ + +#include "operator.h" + +namespace infini::ops { + +class Tile : public Operator { + public: + Tile(const Tensor input, Tensor out) {} + virtual void operator()(const Tensor input, Tensor out) const = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/convolution_layout.cc b/src/convolution_layout.cc new file mode 100644 index 000000000..3c8ce2e0a --- /dev/null +++ b/src/convolution_layout.cc @@ -0,0 +1,14 @@ +#include + +namespace infini::ops { + +bool SupportsConvolutionChannelsLast(const infini::rt::Device& device) { +#ifdef WITH_CAMBRICON + return device.type() == infini::rt::Device::Type::kCambricon; +#else + (void)device; + return false; +#endif +} + +} // namespace infini::ops diff --git a/src/native/cambricon/activation_cnnl.h b/src/native/cambricon/activation_cnnl.h new file mode 100644 index 000000000..e05bf7246 --- /dev/null +++ b/src/native/cambricon/activation_cnnl.h @@ -0,0 +1,50 @@ +#ifndef INFINI_OPS_CAMBRICON_ACTIVATION_CNNL_H_ +#define INFINI_OPS_CAMBRICON_ACTIVATION_CNNL_H_ + +#include + +#include "native/cambricon/cnnl_utils.h" + +namespace infini::ops::cnnl_utils { + +template +class Activation : public Base { + public: + Activation(const Tensor input, Tensor out) : Base{input, out} { + if (input.shape() != out.shape() || input.dtype() != out.dtype() || + input.device() != out.device() || !input.IsContiguous() || + !out.IsContiguous() || + (out.dtype() != DataType::kFloat32 && out.dtype() != DataType::kFloat16)) + throw std::invalid_argument("Cambricon activation requires matching FP32 or FP16 tensors"); + if (out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + descriptor_ = cnnl_utils::MakeTensorDescriptor( + input.dtype(), Tensor::Shape(input.shape()), Tensor::Strides(input.strides())); + cnnlActivationDescriptor_t value = nullptr; + INFINI_OPS_CNNL_CHECK(cnnlCreateActivationDescriptor(&value)); + operation_.reset(value); + INFINI_OPS_CNNL_CHECK(cnnlSetActivationDescriptor_v5( + value, Mode, CNNL_ACTIVATION_HIGH_PRECISION, + CNNL_PROPAGATE_NAN, 0.0f, 0.0f, 0.0f, 0.0f, true)); + } + + void operator()(const Tensor input, Tensor out) const override { + if (!descriptor_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(this->stream_))); + const float one = 1.0f, zero = 0.0f; + INFINI_OPS_CNNL_CHECK(cnnlActivationForward( + handle_.get(), operation_.get(), &one, descriptor_.get(), input.data(), + &zero, descriptor_.get(), out.data())); + } + + private: + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor descriptor_; + std::unique_ptr, + decltype(&cnnlDestroyActivationDescriptor)> + operation_{nullptr, cnnlDestroyActivationDescriptor}; +}; + +} // namespace infini::ops::cnnl_utils + +#endif diff --git a/src/native/cambricon/cnnl_utils.h b/src/native/cambricon/cnnl_utils.h index e5c78754c..619b2b1fd 100644 --- a/src/native/cambricon/cnnl_utils.h +++ b/src/native/cambricon/cnnl_utils.h @@ -4,6 +4,8 @@ #include #include #include +#include +#include #include #include "native/cambricon/common.h" @@ -12,8 +14,9 @@ #define INFINI_OPS_CNNL_CHECK(call) \ do { \ const auto cnnl_status = (call); \ - assert(cnnl_status == CNNL_STATUS_SUCCESS && "`" #call "` failed."); \ - (void)cnnl_status; \ + if (cnnl_status != CNNL_STATUS_SUCCESS) \ + throw std::runtime_error(std::string(#call ": ") + \ + cnnlGetErrorString(cnnl_status)); \ } while (false) namespace infini::ops::cnnl_utils { @@ -92,6 +95,11 @@ inline TensorDescriptor MakeTensorDescriptor(DataType dtype, return desc; } +inline TensorDescriptor MakeTensorDescriptor(const Tensor& tensor) { + return MakeTensorDescriptor(tensor.dtype(), Tensor::Shape(tensor.shape()), + Tensor::Strides(tensor.strides())); +} + } // namespace infini::ops::cnnl_utils #endif diff --git a/src/native/cambricon/common.h b/src/native/cambricon/common.h index fc8ede07a..9a859e8de 100644 --- a/src/native/cambricon/common.h +++ b/src/native/cambricon/common.h @@ -38,10 +38,16 @@ namespace infini::ops::cnnl_utils { inline cnnlDataType_t GetDataType(DataType dtype) { switch (dtype) { + case DataType::kBool: + return CNNL_DTYPE_BOOL; case DataType::kInt8: return CNNL_DTYPE_INT8; case DataType::kUInt8: return CNNL_DTYPE_UINT8; + case DataType::kInt16: + return CNNL_DTYPE_INT16; + case DataType::kUInt32: + return CNNL_DTYPE_UINT32; case DataType::kInt32: return CNNL_DTYPE_INT32; case DataType::kInt64: diff --git a/src/native/cambricon/compare_cnnl.h b/src/native/cambricon/compare_cnnl.h new file mode 100644 index 000000000..8e2babd9b --- /dev/null +++ b/src/native/cambricon/compare_cnnl.h @@ -0,0 +1,64 @@ +#ifndef INFINI_OPS_CAMBRICON_COMPARE_CNNL_H_ +#define INFINI_OPS_CAMBRICON_COMPARE_CNNL_H_ + +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops::cnnl_utils { + +template +class Compare : public Base { + public: + Compare(const Tensor input, const Tensor other, Tensor out) + : Base{input, other, out} { + if (input.dtype() != other.dtype() || out.dtype() != DataType::kBool || + input.device() != out.device() || other.device() != out.device() || + !input.IsContiguous() || !other.IsContiguous() || !out.IsContiguous()) + throw std::invalid_argument("Cambricon comparison requires matching inputs and BOOL output"); + if (out.numel() == 0) return; + handle_ = CreateHandle(); + lhs_ = MakeTensorDescriptor(input); + rhs_ = MakeTensorDescriptor(other); + out_ = MakeTensorDescriptor(out); + INFINI_OPS_CNNL_CHECK(cnnlGetLogicOpWorkspaceSize( + handle_.get(), lhs_.get(), rhs_.get(), out_.get(), &workspace_size_)); + scratch_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_); + if (workspace_size_ && !scratch_) + throw std::runtime_error("Cambricon comparison workspace allocation failed"); + } + Compare(const Tensor input, const double other, Tensor out) + : Base{input, other, out} { Unsupported(); } + Compare(Tensor input, const Tensor other) + : Base{input, other, input} { Unsupported(); } + Compare(Tensor input, const double other) + : Base{input, other, input} { Unsupported(); } + + std::size_t workspace_size_in_bytes() const override { return workspace_size_; } + + void operator()(const Tensor input, const Tensor other, Tensor out) const override { + if (!out_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(this->stream_))); + auto* scratch = this->workspace_ && this->workspace_size_in_bytes_ >= workspace_size_ + ? this->workspace_ : scratch_.get(); + INFINI_OPS_CNNL_CHECK(cnnlLogicOp( + handle_.get(), Kind, lhs_.get(), input.data(), rhs_.get(), other.data(), + scratch, workspace_size_, out_.get(), out.data())); + } + void operator()(const Tensor, const double, Tensor) const override { Unsupported(); } + void operator()(Tensor, const Tensor) const override { Unsupported(); } + void operator()(Tensor, const double) const override { Unsupported(); } + + private: + [[noreturn]] static void Unsupported() { + throw std::invalid_argument("Cambricon comparison requires two tensors and a separate BOOL output"); + } + + Handle handle_; + TensorDescriptor lhs_, rhs_, out_; + cnrt_utils::DeviceBuffer scratch_; + std::size_t workspace_size_ = 0; +}; + +} // namespace infini::ops::cnnl_utils + +#endif diff --git a/src/native/cambricon/data_type_.h b/src/native/cambricon/data_type_.h index af029413a..762472fc9 100644 --- a/src/native/cambricon/data_type_.h +++ b/src/native/cambricon/data_type_.h @@ -2,20 +2,40 @@ #define INFINI_OPS_CAMBRICON_DATA_TYPE__H_ #include +#include "common/traits.h" +#include "data_type.h" + +namespace infini::ops { +template +using CambriconBangFloatTypes = + List; +} + // Native arithmetic is confined to genuine BANG compilation. The public RT // type map uses the same complete storage types in host and native objects. -#if defined(__BANGCC__) +#if defined(__BANGCC__) || defined(INFINIOPS_CAMBRICON_NATIVE) #include +#if __has_include() #include +#endif namespace infini::ops::cambricon { +#if __has_include() +using NativeHalf = __half; +#else +using NativeHalf = half; +#endif template struct NativeScalar { using type = StorageT; }; template <> struct NativeScalar { - using type = __half; + using type = NativeHalf; }; template <> struct NativeScalar { diff --git a/src/native/cambricon/host_bang_compat.h b/src/native/cambricon/host_bang_compat.h index f9c5a3cbc..c44eea19a 100644 --- a/src/native/cambricon/host_bang_compat.h +++ b/src/native/cambricon/host_bang_compat.h @@ -8,6 +8,7 @@ // inline semantics. The vendor include guards then keep later cnrt/cnnl // includes from emitting the non-inline definitions. #if !defined(__BANGCC__) && !defined(__CUDACC__) +#if __has_include() #pragma push_macro("__BANGCC__") #pragma push_macro("__mlu_host__") #pragma push_macro("__mlu_func__") @@ -18,7 +19,9 @@ #pragma pop_macro("__mlu_func__") #pragma pop_macro("__mlu_host__") #pragma pop_macro("__BANGCC__") +#endif +#if __has_include() #pragma push_macro("__BANGCC__") #pragma push_macro("__mlu_host__") #pragma push_macro("__mlu_func__") @@ -30,5 +33,6 @@ #pragma pop_macro("__mlu_host__") #pragma pop_macro("__BANGCC__") #endif +#endif #endif diff --git a/src/native/cambricon/op_tensor.h b/src/native/cambricon/op_tensor.h new file mode 100644 index 000000000..a5c4730a7 --- /dev/null +++ b/src/native/cambricon/op_tensor.h @@ -0,0 +1,54 @@ +#ifndef INFINI_OPS_CAMBRICON_OP_TENSOR_H_ +#define INFINI_OPS_CAMBRICON_OP_TENSOR_H_ + +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops::cnnl_utils { + +class OpTensor { + public: + OpTensor(cnnlOpTensorDesc_t kind, const Tensor& lhs, const Tensor& rhs, + const Tensor& out) + : handle_(CreateHandle()), + lhs_(MakeTensorDescriptor(lhs)), rhs_(MakeTensorDescriptor(rhs)), out_(MakeTensorDescriptor(out)) { + cnnlOpTensorDescriptor_t desc = nullptr; + INFINI_OPS_CNNL_CHECK(cnnlCreateOpTensorDescriptor(&desc)); + operation_.reset(desc); + INFINI_OPS_CNNL_CHECK(cnnlSetOpTensorDescriptor( + desc, kind, GetDataType(out.dtype()), CNNL_PROPAGATE_NAN)); + INFINI_OPS_CNNL_CHECK(cnnlGetOpTensorWorkspaceSize( + handle_.get(), lhs_.get(), rhs_.get(), out_.get(), &workspace_size_)); + scratch_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_); + if (workspace_size_ && !scratch_) + throw std::runtime_error("Cambricon OpTensor workspace allocation failed"); + } + + size_t WorkspaceSize() const { return workspace_size_; } + + void Run(const Tensor& lhs, const Tensor& rhs, Tensor out, float alpha, + cnrtQueue_t queue, void* workspace, size_t capacity) const { + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), queue)); + auto* scratch = workspace && capacity >= workspace_size_ + ? workspace : scratch_.get(); + const float one = 1.0f, zero = 0.0f; + INFINI_OPS_CNNL_CHECK(cnnlOpTensor( + handle_.get(), operation_.get(), &one, lhs_.get(), lhs.data(), &alpha, + rhs_.get(), rhs.data(), scratch, workspace_size_, &zero, out_.get(), + out.data())); + } + + private: + + Handle handle_; + TensorDescriptor lhs_, rhs_, out_; + std::unique_ptr, + decltype(&cnnlDestroyOpTensorDescriptor)> + operation_{nullptr, cnnlDestroyOpTensorDescriptor}; + cnrt_utils::DeviceBuffer scratch_; + size_t workspace_size_ = 0; +}; + +} // namespace infini::ops::cnnl_utils + +#endif diff --git a/src/native/cambricon/ops/abs/kernel.h b/src/native/cambricon/ops/abs/kernel.h new file mode 100644 index 000000000..17759362e --- /dev/null +++ b/src/native/cambricon/ops/abs/kernel.h @@ -0,0 +1,18 @@ +#ifndef INFINI_OPS_CAMBRICON_ABS_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_ABS_KERNEL_H_ + +#include "base/abs.h" +#include "native/cambricon/unary_cnnl.h" + +namespace infini::ops { + +template <> +class Operator + : public cnnl_utils::Unary { + public: + Operator(const Tensor input, Tensor out) : Unary{input, out} {} +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/add/kernel.h b/src/native/cambricon/ops/add/kernel.h index 8a701d39b..1afc3d9d1 100644 --- a/src/native/cambricon/ops/add/kernel.h +++ b/src/native/cambricon/ops/add/kernel.h @@ -4,6 +4,7 @@ #include "base/add.h" #include "native/cambricon/common.h" #include "native/cambricon/data_type_.h" +#include "native/cambricon/op_tensor.h" namespace infini::ops { @@ -22,6 +23,14 @@ class Operator : public Add { Operator(const Tensor input, const Tensor other, const double alpha, Tensor out) : Add{input, other, alpha, out} { + if (output_size_ == 0) return; + if ((out_type_ == DataType::kFloat16 || out_type_ == DataType::kFloat32) && + input_type_ == out_type_ && other_type_ == out_type_ && + input.device() == out.device() && other.device() == out.device() && + input.IsContiguous() && other.IsContiguous() && out.IsContiguous()) { + cnnl_ = std::make_unique(CNNL_OP_TENSOR_ADD, input, other, out); + return; + } cnrt_utils::GetLaunchConfig(input.device(), &core_per_cluster, &cluster_count); cnrtMalloc(&default_workspace_, workspace_size_in_bytes()); @@ -32,7 +41,13 @@ class Operator : public Add { void operator()(const Tensor input, const Tensor other, const double alpha, Tensor out) const override { + if (output_size_ == 0) return; auto queue = static_cast(stream_ ? stream_ : 0); + if (cnnl_) { + cnnl_->Run(input, other, out, static_cast(alpha), queue, + workspace_, workspace_size_in_bytes_); + return; + } auto workspace{workspace_ ? workspace_ : default_workspace_}; bool fast_path = is_input_contiguous_ && is_other_contiguous_ && @@ -40,8 +55,7 @@ class Operator : public Add { other_shape_ == out_shape_; DispatchFunc>( + CambriconBangFloatTypes>( {out_type_}, [&](auto tag) { using T = typename decltype(tag)::type; @@ -55,15 +69,22 @@ class Operator : public Add { "CambriconAdd::operator() - output dispatch"); } - ~Operator() { cnrtFree(default_workspace_); } + ~Operator() { + if (default_workspace_) cnrtFree(default_workspace_); + } std::size_t workspace_size_in_bytes() const override { + if (cnnl_) return cnnl_->WorkspaceSize(); + if (output_size_ == 0) return 0; return ndim_ * (3 * sizeof(size_t) + 3 * sizeof(ptrdiff_t)); } void* default_workspace_{nullptr}; int core_per_cluster = 0; int cluster_count = 0; + + private: + std::unique_ptr cnnl_; }; } // namespace infini::ops diff --git a/src/native/cambricon/ops/add/kernel.mlu b/src/native/cambricon/ops/add/kernel.mlu index 2cbda223d..90eb4f2d2 100644 --- a/src/native/cambricon/ops/add/kernel.mlu +++ b/src/native/cambricon/ops/add/kernel.mlu @@ -7,30 +7,65 @@ namespace infini::ops { template using AddAlpha = std::conditional_t, int64_t, float>; +__mlu_device__ uint32_t MultiplyHigh32(uint32_t a, uint32_t b) { + const uint32_t a0 = a & 0xffffu, a1 = a >> 16; + const uint32_t b0 = b & 0xffffu, b1 = b >> 16; + const uint32_t t = a1 * b0 + ((a0 * b0) >> 16); + return a1 * b1 + (t >> 16) + ((a0 * b1 + (t & 0xffffu)) >> 16); +} + template __mlu_device__ void BangAdd(const T* src1, T* src2, T* dst, size_t n, - AddAlpha alpha) { - if (alpha != static_cast>(1)) { - if constexpr (std::is_same_v) { - __bang_mul_scalar(reinterpret_cast(src2), - reinterpret_cast(src2), alpha, n); - } else if constexpr (std::is_same_v || - std::is_same_v) { - __bang_mul_scalar(src2, src2, alpha, n); - } else { - const T scalar = static_cast(alpha); - for (size_t i = 0; i < n; ++i) { - src2[i] *= scalar; + AddAlpha alpha, uint32_t alpha_hi) { + if constexpr (std::is_same_v) { + // Preserve all 64 bits on older BANG scalar datapaths. + const auto* a = reinterpret_cast(src1); + const auto* b = reinterpret_cast(src2); + auto* y = reinterpret_cast(dst); + const auto alpha_lo = static_cast(alpha); + for (size_t i = 0; i < n; ++i) { + const auto j = 2 * i; + const uint32_t product = b[j] * alpha_lo; + const uint32_t low = a[j] + product; + y[j] = low; + y[j + 1] = a[j + 1] + (low < a[j]) + b[j + 1] * alpha_lo + + b[j] * alpha_hi + MultiplyHigh32(b[j], alpha_lo); + } + } else { + if (alpha != static_cast>(1)) { + if constexpr (std::is_same_v) { + __bang_mul_scalar(reinterpret_cast(src2), + reinterpret_cast(src2), alpha, n); + } else if constexpr (std::is_same_v || + std::is_same_v) { + __bang_mul_scalar(src2, src2, alpha, n); + } else { + const T scalar = static_cast(alpha); + for (size_t i = 0; i < n; ++i) { + src2[i] *= scalar; + } } } + + if constexpr (std::is_same_v) { + __bang_add(reinterpret_cast(dst), + reinterpret_cast(src1), + reinterpret_cast(src2), n); + } else { + __bang_add(dst, src1, src2, n); + } } +} - if constexpr (std::is_same_v) { - __bang_add(reinterpret_cast(dst), - reinterpret_cast(src1), - reinterpret_cast(src2), n); +template +__mlu_device__ void CopyElement(T* dst, const T* src) { + if constexpr (std::is_same_v) { + auto* out = reinterpret_cast(dst); + const auto* in = reinterpret_cast(src); + out[0] = in[0]; + out[1] = in[1]; } else { - __bang_add(dst, src1, src2, n); + *dst = *src; } } @@ -40,7 +75,7 @@ __mlu_global__ void AddKernel( const ptrdiff_t* out_strides, const size_t* input_shape, const ptrdiff_t* input_strides, const size_t* other_shape, const ptrdiff_t* other_strides, size_t output_size, int ndim, - bool fast_path, bool out_contiguous, AddAlpha alpha) { + bool fast_path, bool out_contiguous, AddAlpha alpha, uint32_t alpha_hi) { size_t elements_per_task = (output_size + taskDim - 1) / taskDim; size_t start = taskId * elements_per_task; size_t end = start + elements_per_task; @@ -69,7 +104,7 @@ __mlu_global__ void AddKernel( GDRAM2NRAM); __memcpy(other_buf, other + start + processed, curr * sizeof(T), GDRAM2NRAM); - BangAdd(input_buf, other_buf, output_buf, curr, alpha); + BangAdd(input_buf, other_buf, output_buf, curr, alpha, alpha_hi); __memcpy(output + start + processed, output_buf, curr * sizeof(T), NRAM2GDRAM); @@ -96,7 +131,7 @@ __mlu_global__ void AddKernel( size_t c = coord < input_shape[d] ? coord : 0; offset += static_cast(c) * input_strides[d]; } - input_buf[i] = input[offset]; + CopyElement(input_buf + i, input + offset); } // Compute `other` offset. @@ -109,11 +144,11 @@ __mlu_global__ void AddKernel( size_t c = coord < other_shape[d] ? coord : 0; offset += static_cast(c) * other_strides[d]; } - other_buf[i] = other[offset]; + CopyElement(other_buf + i, other + offset); } } - BangAdd(input_buf, other_buf, output_buf, curr, alpha); + BangAdd(input_buf, other_buf, output_buf, curr, alpha, alpha_hi); if (out_contiguous) { __memcpy(output + start + processed, output_buf, curr * sizeof(T), @@ -128,7 +163,7 @@ __mlu_global__ void AddKernel( offset += static_cast(coord) * out_strides[d]; tmp /= out_shape[d]; } - output[offset] = output_buf[i]; + CopyElement(output + offset, output_buf + i); } } @@ -186,10 +221,13 @@ void AddUnion(void* workspace, int core_per_cluster, int cluster_count, cnrtMemcpyAsync(mlu_other_strides, const_cast(other_strides), ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev)); + uint32_t alpha_hi = 0; + if constexpr (std::is_same_v) + alpha_hi = static_cast(static_cast(static_cast(alpha)) >> 32); AddKernel<<>>( input_, other_, out_, mlu_out_shape, mlu_out_strides, mlu_input_shape, mlu_input_strides, mlu_other_shape, mlu_other_strides, output_size, ndim, - fast_path, out_contiguous, static_cast>(alpha)); + fast_path, out_contiguous, static_cast>(alpha), alpha_hi); } template void AddUnion(void*, int, int, cnrtQueue_t, const void*, @@ -199,12 +237,14 @@ template void AddUnion(void*, int, int, cnrtQueue_t, const void*, const ptrdiff_t*, size_t, int, bool, bool, double); +#if INFINIOPS_CAMBRICON_BANG_BF16 template void AddUnion(void*, int, int, cnrtQueue_t, const void*, const void*, void*, const size_t*, const ptrdiff_t*, const size_t*, const ptrdiff_t*, const size_t*, const ptrdiff_t*, size_t, int, bool, bool, double); +#endif template void AddUnion(void*, int, int, cnrtQueue_t, const void*, const void*, void*, const size_t*, diff --git a/src/native/cambricon/ops/amax/kernel.h b/src/native/cambricon/ops/amax/kernel.h new file mode 100644 index 000000000..fd6230bc6 --- /dev/null +++ b/src/native/cambricon/ops/amax/kernel.h @@ -0,0 +1,31 @@ +#ifndef INFINI_OPS_CAMBRICON_AMAX_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_AMAX_KERNEL_H_ + +#include "base/amax.h" +#include "native/cambricon/reduce_cnnl.h" + +namespace infini::ops { + +template <> +class Operator : public Amax { + public: + Operator(const Tensor input, const std::vector dim, const bool keepdim, + Tensor out) + : Amax{input, dim, keepdim, out}, + plan_{CNNL_REDUCE_MAX, input, dim, keepdim, out} {} + + std::size_t workspace_size_in_bytes() const override { return plan_.WorkspaceSize(); } + + void operator()(const Tensor input, const std::vector, const bool, + Tensor out) const override { + plan_.Run(input, out, static_cast(stream_), + workspace_, workspace_size_in_bytes_); + } + + private: + cnnl_utils::ReducePlan plan_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/argmax/kernel.h b/src/native/cambricon/ops/argmax/kernel.h index ae31554eb..5ffafd64c 100644 --- a/src/native/cambricon/ops/argmax/kernel.h +++ b/src/native/cambricon/ops/argmax/kernel.h @@ -5,7 +5,9 @@ #include #include "base/argmax.h" +#include "native/cambricon/cnnl_utils.h" #include "native/cambricon/common.h" +#include "native/cambricon/cnrt_utils.h" namespace infini::ops { void CambriconArgmaxLaunch(cnrtQueue_t, const void*, void*, int32_t, std::size_t, uint32_t); @@ -14,10 +16,26 @@ class Operator : public Argmax { public: static bool supports(const Tensor x, const std::optional dim, bool keepdim, const Tensor y) { + if (dim) { + const auto rank = static_cast(x.ndim()); + if (x.device().type() != Device::Type::kCambricon || + x.device() != y.device() || !x.IsContiguous() || !y.IsContiguous() || + (x.dtype() != DataType::kFloat32 && x.dtype() != DataType::kFloat16 && + x.dtype() != DataType::kInt32) || + y.dtype() != DataType::kInt64 || *dim < -rank || *dim >= rank) + return false; + const auto axis = static_cast(*dim < 0 ? *dim + rank : *dim); + if (x.shape()[axis] == 0) return false; + Tensor::Shape expected; + for (size_t i = 0; i < x.ndim(); ++i) { + if (i != axis || keepdim) expected.push_back(i == axis ? 1 : x.shape()[i]); + } + return expected == y.shape(); + } if (x.device().type() != Device::Type::kCambricon || x.device().index() < 0 || x.dtype() != DataType::kFloat32 || !x.numel() || x.numel() > INT_MAX || !x.IsContiguous() || !x.data() || - dim || keepdim || x.device() != y.device() || y.ndim() != 0 || + keepdim || x.device() != y.device() || y.ndim() != 0 || y.dtype() != DataType::kInt64 || !y.IsContiguous() || !y.data()) return false; const auto a = reinterpret_cast(x.data()), @@ -28,7 +46,34 @@ class Operator : public Argmax { } Operator(const Tensor x, const std::optional dim, bool keepdim, Tensor y) - : Argmax(x, dim, keepdim, y) {} + : Argmax(x, dim, keepdim, y) { + if (!dim) return; + if (!supports(x, dim, keepdim, y)) + throw std::invalid_argument("Unsupported Cambricon axis Argmax"); + if (!y.numel()) return; + const auto rank = static_cast(x.ndim()); + axis_ = static_cast(*dim < 0 ? *dim + rank : *dim); + auto top_shape = input_shape_; + top_shape[axis_] = 1; + Tensor::Strides top_strides(top_shape.size()); + Tensor::Stride stride = 1; + for (size_t i = top_shape.size(); i-- > 0;) { + top_strides[i] = stride; + stride *= top_shape[i]; + } + handle_ = cnnl_utils::CreateHandle(); + input_ = cnnl_utils::MakeTensorDescriptor(input_type_, input_shape_, input_strides_); + values_ = cnnl_utils::MakeTensorDescriptor(input_type_, top_shape, top_strides); + indices_ = cnnl_utils::MakeTensorDescriptor(out_type_, top_shape, top_strides); + INFINI_OPS_CNNL_CHECK(cnnlGetTopKTensorWorkspaceSize( + handle_.get(), input_.get(), 1, axis_, true, values_.get(), + indices_.get(), &workspace_size_)); + workspace_own_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_); + values_storage_ = cnrt_utils::AllocateDeviceBuffer(y.numel() * x.element_size()); + if ((workspace_size_ && !workspace_own_) || !values_storage_) + throw std::runtime_error("Cambricon Argmax workspace allocation failed"); + } + std::size_t workspace_size_in_bytes() const override { return workspace_size_; } void operator()(const Tensor x, const std::optional dim, bool keepdim, Tensor y) const override { if (!supports(x, dim, keepdim, y) || x.shape() != input_shape_ || @@ -37,12 +82,30 @@ class Operator : public Argmax { y.dtype() != out_type_ || x.device().index() != device_index_ || dim != dim_ || keepdim != keepdim_) throw std::invalid_argument( - "Cambricon Argmax prepared FP32 metadata changed"); + "Cambricon Argmax prepared metadata changed"); + if (dim) { + if (!indices_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + auto* scratch = workspace_ && workspace_size_in_bytes_ >= workspace_size_ + ? workspace_ : workspace_own_.get(); + INFINI_OPS_CNNL_CHECK(cnnlTopKTensor_v3( + handle_.get(), input_.get(), x.data(), 1, axis_, true, true, true, + scratch, workspace_size_, values_.get(), values_storage_.get(), + indices_.get(), y.data())); + return; + } CambriconArgmaxLaunch(static_cast(stream_), x.data(), y.data(), static_cast(x.numel()), sizeof(float), 0x7f800000U); CNRT_CHECK(cnrtGetLastError()); } + + private: + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor input_, values_, indices_; + cnrt_utils::DeviceBuffer workspace_own_, values_storage_; + std::size_t workspace_size_ = 0; + int axis_ = 0; }; } // namespace infini::ops #endif diff --git a/src/native/cambricon/ops/cast/kernel.h b/src/native/cambricon/ops/cast/kernel.h new file mode 100644 index 000000000..11c507d0e --- /dev/null +++ b/src/native/cambricon/ops/cast/kernel.h @@ -0,0 +1,84 @@ +#ifndef INFINI_OPS_CAMBRICON_CAST_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_CAST_KERNEL_H_ + +#include + +#include "base/cast.h" +#include "native/cambricon/cnnl_utils.h" + +namespace infini::ops { + +template <> +class Operator : public Cast { + public: + Operator(const Tensor input, Tensor out) : Cast{input, out} { + if (input_shape_ != out_shape_ || input.device() != out.device() || + !is_input_contiguous_ || !is_out_contiguous_) + throw std::invalid_argument("Cambricon Cast requires matching contiguous tensor shapes"); + cast_type_ = CastType(input_dtype_, out_dtype_); + if (output_size_ == 0) return; + handle_ = cnnl_utils::CreateHandle(); + input_ = cnnl_utils::MakeTensorDescriptor(input_dtype_, input_shape_, input_strides_); + output_ = cnnl_utils::MakeTensorDescriptor(out_dtype_, out_shape_, out_strides_); + } + + void operator()(const Tensor input, Tensor out) const override { + if (!output_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + INFINI_OPS_CNNL_CHECK(cnnlCastDataType( + handle_.get(), input_.get(), input.data(), cast_type_, output_.get(), out.data())); + } + + private: + static cnnlCastDataType_t CastType(DataType input, DataType out) { + struct Conversion { DataType src, dst; cnnlCastDataType_t cnnl; }; + static constexpr Conversion conversions[] = { + {DataType::kFloat32, DataType::kFloat16, CNNL_CAST_FLOAT_TO_HALF}, + {DataType::kFloat32, DataType::kInt64, CNNL_CAST_FLOAT_TO_INT64}, + {DataType::kFloat32, DataType::kInt32, CNNL_CAST_FLOAT_TO_INT32}, + {DataType::kFloat32, DataType::kInt16, CNNL_CAST_FLOAT_TO_INT16}, + {DataType::kFloat32, DataType::kInt8, CNNL_CAST_FLOAT_TO_INT8}, + {DataType::kFloat32, DataType::kBFloat16, CNNL_CAST_FLOAT_TO_BFLOAT16}, + {DataType::kInt32, DataType::kFloat32, CNNL_CAST_INT32_TO_FLOAT}, + {DataType::kInt32, DataType::kInt8, CNNL_CAST_INT32_TO_INT8}, + {DataType::kInt32, DataType::kInt16, CNNL_CAST_INT32_TO_INT16}, + {DataType::kInt32, DataType::kInt64, CNNL_CAST_INT32_TO_INT64}, + {DataType::kInt16, DataType::kFloat32, CNNL_CAST_INT16_TO_FLOAT}, + {DataType::kInt16, DataType::kInt32, CNNL_CAST_INT16_TO_INT32}, + {DataType::kInt8, DataType::kFloat32, CNNL_CAST_INT8_TO_FLOAT}, + {DataType::kInt8, DataType::kInt16, CNNL_CAST_INT8_TO_INT16}, + {DataType::kInt8, DataType::kInt32, CNNL_CAST_INT8_TO_INT32}, + {DataType::kUInt8, DataType::kFloat32, CNNL_CAST_UINT8_TO_FLOAT}, + {DataType::kUInt8, DataType::kInt32, CNNL_CAST_UINT8_TO_INT32}, + {DataType::kUInt8, DataType::kInt64, CNNL_CAST_UINT8_TO_INT64}, + {DataType::kInt64, DataType::kInt32, CNNL_CAST_INT64_TO_INT32}, + {DataType::kInt64, DataType::kUInt32, CNNL_CAST_INT64_TO_UINT32}, + {DataType::kInt64, DataType::kFloat32, CNNL_CAST_INT64_TO_FLOAT}, + {DataType::kUInt32, DataType::kInt64, CNNL_CAST_UINT32_TO_INT64}, + {DataType::kFloat16, DataType::kFloat32, CNNL_CAST_HALF_TO_FLOAT}, + {DataType::kBFloat16, DataType::kFloat32, CNNL_CAST_BFLOAT16_TO_FLOAT}, + {DataType::kFloat32, DataType::kBool, CNNL_CAST_FLOAT_TO_BOOL}, + {DataType::kFloat16, DataType::kBool, CNNL_CAST_HALF_TO_BOOL}, + {DataType::kBFloat16, DataType::kBool, CNNL_CAST_BFLOAT16_TO_BOOL}, + {DataType::kInt32, DataType::kBool, CNNL_CAST_INT32_TO_BOOL}, + {DataType::kInt64, DataType::kBool, CNNL_CAST_INT64_TO_BOOL}, + {DataType::kBool, DataType::kFloat32, CNNL_CAST_BOOL_TO_FLOAT}, + {DataType::kBool, DataType::kFloat16, CNNL_CAST_BOOL_TO_HALF}, + {DataType::kBool, DataType::kBFloat16, CNNL_CAST_BOOL_TO_BFLOAT16}, + {DataType::kBool, DataType::kInt32, CNNL_CAST_BOOL_TO_INT32}, + {DataType::kBool, DataType::kInt64, CNNL_CAST_BOOL_TO_INT64}, + }; + for (const auto& entry : conversions) { + if (entry.src == input && entry.dst == out) + return entry.cnnl; + } + throw std::invalid_argument("Unsupported Cambricon Cast dtype pair"); + } + cnnlCastDataType_t cast_type_{}; + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor input_, output_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/cat/kernel.h b/src/native/cambricon/ops/cat/kernel.h new file mode 100644 index 000000000..a342a1012 --- /dev/null +++ b/src/native/cambricon/ops/cat/kernel.h @@ -0,0 +1,97 @@ +#ifndef INFINI_OPS_CAMBRICON_CAT_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_CAT_KERNEL_H_ + +#include +#include + +#include "base/cat.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops { + +template <> +class Operator : public Cat { + public: + Operator(const std::vector tensors, int64_t dim, Tensor out) + : Cat{tensors, CheckedAxis(tensors, dim, out), out} { + if (out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + output_ = cnnl_utils::MakeTensorDescriptor(out_type_, out_shape_, out_strides_); + for (size_t i = 0; i < tensors.size(); ++i) { + const auto& input = tensors[i]; + if (input.numel() == 0) continue; + indices_.push_back(i); + inputs_.push_back(cnnl_utils::MakeTensorDescriptor( + input.dtype(), Tensor::Shape(input.shape()), Tensor::Strides(input.strides()))); + descriptors_.push_back(inputs_.back().get()); + } + if (indices_.size() == 1) return; + INFINI_OPS_CNNL_CHECK(cnnlGetConcatWorkspaceSize( + handle_.get(), indices_.size(), &workspace_size_)); + workspace_own_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_); + if (workspace_size_ && !workspace_own_) + throw std::runtime_error("Cambricon Cat workspace allocation failed"); + } + + std::size_t workspace_size_in_bytes() const override { return workspace_size_; } + + void operator()(const std::vector tensors, int64_t, Tensor out) const override { + if (!output_) return; + if (indices_.size() == 1) { + const auto* data = tensors[indices_.front()].data(); + if (data != out.data() && + cnrtMemcpyAsync(out.data(), const_cast(data), out.numel() * out.element_size(), + static_cast(stream_), cnrtMemcpyDevToDev) != cnrtSuccess) + throw std::runtime_error("Cambricon Cat single-input copy failed"); + return; + } + std::vector buffers; + buffers.reserve(indices_.size()); + for (auto index : indices_) buffers.push_back(tensors[index].data()); + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + auto* scratch = workspace_ && workspace_size_in_bytes_ >= workspace_size_ + ? workspace_ : workspace_own_.get(); + INFINI_OPS_CNNL_CHECK(cnnlConcat( + handle_.get(), buffers.size(), dim_, descriptors_.data(), buffers.data(), + scratch, workspace_size_, output_.get(), out.data())); + } + + private: + static int64_t CheckedAxis(const std::vector& tensors, int64_t axis, const Tensor& out) { + const auto rank = static_cast(out.ndim()); + if (axis < 0) axis += rank; + if (tensors.empty() || tensors.size() > std::numeric_limits::max() || + axis < 0 || axis >= rank || !out.IsContiguous()) + throw std::invalid_argument("Invalid Cambricon Cat input list or axis"); + size_t extent = 0; + for (const auto& tensor : tensors) { + if (tensor.ndim() != out.ndim()) + throw std::invalid_argument("Cambricon Cat input rank differs from output"); + if (tensor.dtype() != out.dtype() || tensor.device() != out.device()) + throw std::invalid_argument("Cambricon Cat input dtype or device differs from output"); + if (tensor.numel() != 0 && !tensor.IsContiguous()) + throw std::invalid_argument("Cambricon Cat requires contiguous nonempty inputs"); + for (int64_t i = 0; i < rank; ++i) { + if (i != axis && tensor.size(i) != out.size(i)) + throw std::invalid_argument("Inconsistent Cambricon Cat dimensions"); + } + extent += tensor.size(axis); + } + if (extent != out.size(axis)) + throw std::invalid_argument("Invalid Cambricon Cat output shape"); + return axis; + } + + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor output_; + std::vector inputs_; + std::vector descriptors_; + std::vector indices_; + cnrt_utils::DeviceBuffer workspace_own_; + std::size_t workspace_size_ = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/clip/kernel.h b/src/native/cambricon/ops/clip/kernel.h new file mode 100644 index 000000000..6bb6dc042 --- /dev/null +++ b/src/native/cambricon/ops/clip/kernel.h @@ -0,0 +1,74 @@ +#ifndef INFINI_OPS_CAMBRICON_CLIP_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_CLIP_KERNEL_H_ + +#include + +#include "base/clip.h" +#include "native/cambricon/cnnl_utils.h" + +namespace infini::ops { + +template <> +class Operator : public Clip { + public: + Operator(const Tensor input, const std::optional min, + const std::optional max, Tensor out) + : Clip{input, min, max, out} { + Init(input, out); + } + + Operator(const Tensor input, const std::optional min, + const std::optional max, Tensor out) + : Clip{input, min, max, out} { + if (!min || !max || min->numel() != 1 || max->numel() != 1 || + min->dtype() != input.dtype() || max->dtype() != input.dtype() || + min->device() != input.device() || max->device() != input.device()) + throw std::invalid_argument("Cambricon Clip device bounds must be two same-type scalars"); + Init(input, out); + } + + void operator()(const Tensor input, const std::optional min, + const std::optional max, Tensor out) const override { + if (!descriptor_) return; + const auto infinity = std::numeric_limits::infinity(); + const float low = min.value_or(-infinity), high = max.value_or(infinity); + if (input.dtype() == DataType::kFloat16) { + const auto low_half = ::infini::rt::Float16::FromFloat(low).bits; + const auto high_half = ::infini::rt::Float16::FromFloat(high).bits; + Run(input, CNNL_POINTER_MODE_HOST, &low_half, &high_half, out); + } else { + Run(input, CNNL_POINTER_MODE_HOST, &low, &high, out); + } + } + + void operator()(const Tensor input, const std::optional min, + const std::optional max, Tensor out) const override { + if (descriptor_) Run(input, CNNL_POINTER_MODE_DEVICE, min->data(), max->data(), out); + } + + private: + void Init(const Tensor& input, const Tensor& out) { + if (input.shape() != out.shape() || input.dtype() != out.dtype() || + input.device() != out.device() || !input.IsContiguous() || !out.IsContiguous() || + (out.dtype() != DataType::kFloat32 && out.dtype() != DataType::kFloat16)) + throw std::invalid_argument("Cambricon Clip requires matching FP32 or FP16 tensors"); + if (out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + descriptor_ = cnnl_utils::MakeTensorDescriptor( + input.dtype(), Tensor::Shape(input.shape()), Tensor::Strides(input.strides())); + } + + void Run(const Tensor& input, cnnlPointerMode_t pointer_mode, + const void* min, const void* max, Tensor out) const { + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + INFINI_OPS_CNNL_CHECK(cnnlClip_v2(handle_.get(), pointer_mode, descriptor_.get(), + input.data(), min, max, descriptor_.get(), out.data())); + } + + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor descriptor_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/convolution/kernel.h b/src/native/cambricon/ops/convolution/kernel.h new file mode 100644 index 000000000..8726a96ec --- /dev/null +++ b/src/native/cambricon/ops/convolution/kernel.h @@ -0,0 +1,246 @@ +#ifndef INFINI_OPS_CAMBRICON_CONVOLUTION_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_CONVOLUTION_KERNEL_H_ + +#include +#include +#include + +#include "base/convolution.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops { + +template <> +class Operator : public Convolution { + public: + Operator(const Tensor input, const Tensor weight, std::optional bias, + std::vector stride, std::vector padding, + std::vector dilation, bool transposed, + std::vector output_padding, int64_t groups, Tensor out) + : Convolution{Validate(input, weight, bias, stride, padding, dilation, + transposed, output_padding, groups, out), + weight, bias, stride, padding, dilation, transposed, + output_padding, groups, out}, + transposed_(transposed) { + if (out.numel() == 0) return; + const auto rank = input.ndim(); + std::vector forward{0}, backward{0, static_cast(rank - 1)}; + for (size_t i = 2; i < rank; ++i) { + forward.push_back(static_cast(i)); + backward.push_back(static_cast(i - 1)); + } + forward.push_back(1); + const auto channels_first = rank == 4 ? CNNL_LAYOUT_NCHW : CNNL_LAYOUT_NCDHW; + const auto channels_last = rank == 4 ? CNNL_LAYOUT_NHWC : CNNL_LAYOUT_NDHWC; + const auto dtype = cnnl_utils::GetDataType(input.dtype()); + handle_ = cnnl_utils::CreateHandle(); + if (!IsChannelsLast(input)) input_first_ = MakeDescriptor(input, dtype, channels_first); + if (!IsChannelsLast(weight)) weight_first_ = MakeDescriptor(weight, dtype, channels_first); + if (!IsChannelsLast(out)) output_first_ = MakeDescriptor(out, dtype, channels_first); + input_last_ = MakeDescriptor(input, dtype, channels_last, forward); + weight_last_ = MakeDescriptor(weight, dtype, channels_last, forward); + output_last_ = MakeDescriptor(out, dtype, channels_last, forward); + if (bias) { + bias_.reset(CreateTensorDescriptor()); + std::vector shape(rank, 1); + shape.back() = bias->numel(); + INFINI_OPS_CNNL_CHECK(cnnlSetTensorDescriptor_v2( + bias_.get(), channels_last, dtype, rank, shape.data())); + } + + MakeTranspose(to_last_, forward); + MakeTranspose(to_first_, backward); + cnnlConvolutionDescriptor_t conv = nullptr; + INFINI_OPS_CNNL_CHECK(cnnlCreateConvolutionDescriptor(&conv)); + conv_.reset(conv); + std::vector pad, step, dilate; + for (size_t i = 0; i < rank - 2; ++i) { + pad.push_back(static_cast(padding[i])); + pad.push_back(static_cast(padding[i])); + step.push_back(static_cast(stride[i])); + dilate.push_back(static_cast(dilation[i])); + } + INFINI_OPS_CNNL_CHECK(cnnlSetConvolutionDescriptor( + conv_.get(), rank, pad.data(), step.data(), dilate.data(), groups, + CNNL_DTYPE_FLOAT)); + if (transposed_) { + INFINI_OPS_CNNL_CHECK(cnnlGetConvolutionBackwardDataAlgorithm( + handle_.get(), input_last_.get(), weight_last_.get(), conv_.get(), + output_last_.get(), CNNL_CONVOLUTION_BWD_DATA_FASTEST, &transpose_algorithm_)); + INFINI_OPS_CNNL_CHECK(cnnlGetConvolutionBackwardDataWorkspaceSize( + handle_.get(), weight_last_.get(), input_last_.get(), conv_.get(), + output_last_.get(), transpose_algorithm_, &conv_workspace_)); + } else { + INFINI_OPS_CNNL_CHECK(cnnlGetConvolutionForwardAlgorithm( + handle_.get(), conv_.get(), input_last_.get(), weight_last_.get(), + output_last_.get(), CNNL_CONVOLUTION_FWD_FASTEST, &algorithm_)); + INFINI_OPS_CNNL_CHECK(cnnlGetConvolutionForwardWorkspaceSize( + handle_.get(), input_last_.get(), weight_last_.get(), output_last_.get(), + bias_.get(), conv_.get(), algorithm_, &conv_workspace_)); + } + if (input_first_) input_workspace_ = TransposeWorkspace(input_first_.get(), to_last_.get()); + if (weight_first_) weight_workspace_ = TransposeWorkspace(weight_first_.get(), to_last_.get()); + if (output_first_) output_workspace_ = TransposeWorkspace(output_last_.get(), to_first_.get()); + scratch_size_ = std::max({conv_workspace_, input_workspace_, weight_workspace_, output_workspace_}); + const auto bytes = ::infini::rt::kDataTypeToSize.at(input.dtype()); + if (input_first_) data_last_ = Allocate(input.numel() * bytes); + if (weight_first_) filter_last_ = Allocate(weight.numel() * bytes); + if (output_first_) result_last_ = Allocate(out.numel() * bytes); + scratch_ = Allocate(scratch_size_); + } + + std::size_t workspace_size_in_bytes() const override { return scratch_size_; } + + void operator()(const Tensor input, const Tensor weight, std::optional bias, + const std::vector, const std::vector, + const std::vector, bool, const std::vector, + int64_t, Tensor out) const override { + if (!output_last_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + auto* scratch = workspace_ && workspace_size_in_bytes_ >= scratch_size_ + ? workspace_ : scratch_.get(); + if (input_first_) INFINI_OPS_CNNL_CHECK(cnnlTranspose_v2( + handle_.get(), to_last_.get(), input_first_.get(), input.data(), + input_last_.get(), data_last_.get(), scratch, input_workspace_)); + if (weight_first_) INFINI_OPS_CNNL_CHECK(cnnlTranspose_v2( + handle_.get(), to_last_.get(), weight_first_.get(), weight.data(), + weight_last_.get(), filter_last_.get(), scratch, weight_workspace_)); + const auto* input_data = input_first_ ? data_last_.get() : input.data(); + const auto* weight_data = weight_first_ ? filter_last_.get() : weight.data(); + auto* output_data = output_first_ ? result_last_.get() : out.data(); + if (transposed_) { + INFINI_OPS_CNNL_CHECK(cnnlConvolutionBackwardData( + handle_.get(), nullptr, weight_last_.get(), weight_data, + input_last_.get(), input_data, conv_.get(), transpose_algorithm_, + scratch, conv_workspace_, nullptr, output_last_.get(), output_data)); + } else { + INFINI_OPS_CNNL_CHECK(cnnlConvolutionForward( + handle_.get(), conv_.get(), algorithm_, nullptr, input_last_.get(), + input_data, weight_last_.get(), weight_data, bias_.get(), + bias ? bias->data() : nullptr, scratch, conv_workspace_, nullptr, + output_last_.get(), output_data)); + } + if (output_first_) INFINI_OPS_CNNL_CHECK(cnnlTranspose_v2( + handle_.get(), to_first_.get(), output_last_.get(), result_last_.get(), + output_first_.get(), out.data(), scratch, output_workspace_)); + } + + private: + using Transpose = std::unique_ptr, + decltype(&cnnlDestroyTransposeDescriptor)>; + + static bool IsChannelsLast(const Tensor& tensor) { + const auto rank = tensor.ndim(); + if (rank != 4 && rank != 5) return false; + int64_t stride = 1; + for (size_t i = 0; i < rank; ++i) { + const size_t axis = i == 0 ? 1 : i == rank - 1 ? 0 : rank - i; + if (tensor.shape()[axis] > 1 && tensor.stride(axis) != stride) return false; + stride *= tensor.shape()[axis]; + } + return true; + } + + static bool IsSupportedLayout(const Tensor& tensor) { + return tensor.IsContiguous() || IsChannelsLast(tensor); + } + + static Tensor Validate(const Tensor& input, const Tensor& weight, + const std::optional& bias, + const std::vector& stride, + const std::vector& padding, + const std::vector& dilation, bool transposed, + const std::vector& output_padding, + int64_t groups, const Tensor& out) { + const auto rank = input.ndim(); + if ((rank != 4 && rank != 5) || weight.ndim() != rank || out.ndim() != rank || + stride.size() != rank - 2 || padding.size() != rank - 2 || + dilation.size() != rank - 2 || output_padding.size() != rank - 2 || + (transposed && (rank != 4 || bias)) || + groups < 1 || groups > std::numeric_limits::max() || + (input.dtype() != DataType::kFloat32 && input.dtype() != DataType::kFloat16) || + input.dtype() != weight.dtype() || input.dtype() != out.dtype() || + input.device() != weight.device() || input.device() != out.device() || + !IsSupportedLayout(input) || !IsSupportedLayout(weight) || !IsSupportedLayout(out) || + (bias && (bias->ndim() != 1 || bias->numel() != weight.shape()[0] || + bias->dtype() != out.dtype() || bias->device() != out.device() || + !bias->IsContiguous()))) + throw std::invalid_argument("Unsupported Cambricon Convolution tensor metadata"); + for (size_t i = 0; i < rank - 2; ++i) { + const auto limit = std::numeric_limits::max(); + if (stride[i] <= 0 || stride[i] > limit || dilation[i] <= 0 || + dilation[i] > limit || padding[i] < 0 || padding[i] > limit || + output_padding[i] < 0 || + (transposed ? output_padding[i] >= stride[i] : output_padding[i] != 0)) + throw std::invalid_argument("Unsupported Cambricon Convolution spatial parameters"); + } + const auto input_channels = input.shape()[1], output_channels = out.shape()[1]; + if (input_channels % groups || output_channels % groups || + (transposed ? (weight.shape()[0] != input_channels || + weight.shape()[1] != output_channels / groups) + : (weight.shape()[0] != output_channels || + weight.shape()[1] != input_channels / groups)) || + out.shape()[0] != input.shape()[0]) + throw std::invalid_argument("Invalid Cambricon Convolution channels"); + return input; + } + + static cnnlTensorDescriptor_t CreateTensorDescriptor() { + cnnlTensorDescriptor_t desc = nullptr; + INFINI_OPS_CNNL_CHECK(cnnlCreateTensorDescriptor(&desc)); + return desc; + } + + static cnnl_utils::TensorDescriptor MakeDescriptor( + const Tensor& tensor, cnnlDataType_t dtype, cnnlTensorLayout_t layout, + const std::vector& permutation = {}) { + cnnl_utils::TensorDescriptor result{CreateTensorDescriptor()}; + std::vector dims(tensor.ndim()); + for (size_t i = 0; i < dims.size(); ++i) + dims[i] = tensor.shape()[permutation.empty() ? i : permutation[i]]; + INFINI_OPS_CNNL_CHECK(cnnlSetTensorDescriptor_v2( + result.get(), layout, dtype, dims.size(), dims.data())); + return result; + } + + static void MakeTranspose(Transpose& descriptor, const std::vector& permutation) { + cnnlTransposeDescriptor_t value = nullptr; + INFINI_OPS_CNNL_CHECK(cnnlCreateTransposeDescriptor(&value)); + descriptor.reset(value); + INFINI_OPS_CNNL_CHECK(cnnlSetTransposeDescriptor( + value, permutation.size(), permutation.data())); + } + + std::size_t TransposeWorkspace(cnnlTensorDescriptor_t input, + cnnlTransposeDescriptor_t transpose) const { + std::size_t bytes = 0; + INFINI_OPS_CNNL_CHECK(cnnlGetTransposeWorkspaceSize(handle_.get(), input, transpose, &bytes)); + return bytes; + } + + static cnrt_utils::DeviceBuffer Allocate(std::size_t bytes) { + auto result = cnrt_utils::AllocateDeviceBuffer(bytes); + if (bytes && !result) throw std::runtime_error("Cambricon Convolution allocation failed"); + return result; + } + + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor input_first_, weight_first_, output_first_; + cnnl_utils::TensorDescriptor input_last_, weight_last_, output_last_, bias_; + Transpose to_last_{nullptr, cnnlDestroyTransposeDescriptor}; + Transpose to_first_{nullptr, cnnlDestroyTransposeDescriptor}; + std::unique_ptr, + decltype(&cnnlDestroyConvolutionDescriptor)> + conv_{nullptr, cnnlDestroyConvolutionDescriptor}; + cnnlConvolutionForwardAlgo_t algorithm_{}; + cnnlConvolutionBwdDataAlgo_t transpose_algorithm_{}; + bool transposed_; + cnrt_utils::DeviceBuffer data_last_, filter_last_, result_last_, scratch_; + std::size_t input_workspace_ = 0, weight_workspace_ = 0; + std::size_t output_workspace_ = 0, conv_workspace_ = 0, scratch_size_ = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/cos/kernel.h b/src/native/cambricon/ops/cos/kernel.h new file mode 100644 index 000000000..c63230ea3 --- /dev/null +++ b/src/native/cambricon/ops/cos/kernel.h @@ -0,0 +1,18 @@ +#ifndef INFINI_OPS_CAMBRICON_COS_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_COS_KERNEL_H_ + +#include "base/cos.h" +#include "native/cambricon/unary_cnnl.h" + +namespace infini::ops { + +template <> +class Operator + : public cnnl_utils::Unary { + public: + Operator(const Tensor input, Tensor out) : Unary{input, out} {} +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/div/kernel.h b/src/native/cambricon/ops/div/kernel.h new file mode 100644 index 000000000..3f32755d8 --- /dev/null +++ b/src/native/cambricon/ops/div/kernel.h @@ -0,0 +1,123 @@ +#ifndef INFINI_OPS_CAMBRICON_DIV_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_DIV_KERNEL_H_ + +#include "base/div.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops { + +template <> +class Operator : public Div { + public: + Operator(const Tensor input, const Tensor other, Tensor out) + : Operator{input, other, std::nullopt, out} {} + + Operator(const Tensor input, const Tensor other, + const std::optional rounding_mode, Tensor out) + : Div{input, other, rounding_mode, out} { + RequireTrueDivision(rounding_mode); + if (input.dtype() != out.dtype() || other.dtype() != out.dtype() || + input.device() != out.device() || other.device() != out.device() || + !input.IsContiguous() || !other.IsContiguous() || !out.IsContiguous() || + (out.dtype() != DataType::kFloat32 && out.dtype() != DataType::kFloat16)) + throw std::invalid_argument("Cambricon Div requires matching FP32 or FP16 tensors"); + if (out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + lhs_ = Describe(input, out.ndim()); + rhs_ = Describe(other, out.ndim()); + out_ = Describe(out, out.ndim()); + INFINI_OPS_CNNL_CHECK(cnnlGetDivWorkspaceSize( + handle_.get(), lhs_.get(), rhs_.get(), out_.get(), &workspace_size_)); + scratch_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_); + if (workspace_size_ && !scratch_) + throw std::runtime_error("Cambricon Div workspace allocation failed"); + } + + Operator(Tensor input, const Tensor other) + : Operator{input, other, input} {} + + Operator(Tensor input, const Tensor other, + const std::optional rounding_mode) + : Operator{input, other, rounding_mode, input} {} + + Operator(const Tensor input, const double other, Tensor out) + : Div{input, other, out} { + UnsupportedScalar(); + } + + Operator(const Tensor input, const double other, + const std::optional rounding_mode, Tensor out) + : Div{input, other, rounding_mode, out} { + UnsupportedScalar(); + } + + std::size_t workspace_size_in_bytes() const override { return workspace_size_; } + + void operator()(const Tensor input, const Tensor other, Tensor out) const override { + if (!out_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + auto* scratch = workspace_ && workspace_size_in_bytes_ >= workspace_size_ + ? workspace_ : scratch_.get(); + INFINI_OPS_CNNL_CHECK(cnnlDiv_v2( + handle_.get(), CNNL_COMPUTATION_HIGH_PRECISION, lhs_.get(), input.data(), + rhs_.get(), other.data(), scratch, workspace_size_, out_.get(), out.data())); + } + + void operator()(const Tensor input, const Tensor other, + const std::optional rounding_mode, Tensor out) const override { + RequireTrueDivision(rounding_mode); + (*this)(input, other, out); + } + + void operator()(Tensor input, const Tensor other) const override { + (*this)(input, other, input); + } + + void operator()(Tensor input, const Tensor other, + const std::optional rounding_mode) const override { + (*this)(input, other, rounding_mode, input); + } + + void operator()(const Tensor, const double, Tensor) const override { + UnsupportedScalar(); + } + + void operator()(const Tensor, const double, + const std::optional, Tensor) const override { + UnsupportedScalar(); + } + + private: + static void RequireTrueDivision(const std::optional& rounding_mode) { + if (rounding_mode) + throw std::invalid_argument("Cambricon Div does not support rounding modes"); + } + + [[noreturn]] static void UnsupportedScalar() { + throw std::invalid_argument("Cambricon Div requires a tensor divisor"); + } + + static cnnl_utils::TensorDescriptor Describe(const Tensor& tensor, size_t rank) { + if (tensor.ndim() > rank) + throw std::invalid_argument("Cambricon Div operand rank exceeds output rank"); + auto desc = cnnl_utils::CreateTensorDescriptor(); + const auto dims = tensor.shape(); + std::vector shape(rank - tensor.ndim(), 1); + shape.insert(shape.end(), dims.begin(), dims.end()); + if (shape.empty()) shape.push_back(1); + INFINI_OPS_CNNL_CHECK(cnnlSetTensorDescriptor_v2( + desc.get(), CNNL_LAYOUT_ARRAY, cnnl_utils::GetDataType(tensor.dtype()), + static_cast(shape.size()), shape.data())); + return desc; + } + + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor lhs_, rhs_, out_; + cnrt_utils::DeviceBuffer scratch_; + std::size_t workspace_size_ = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/embedding/kernel.mlu b/src/native/cambricon/ops/embedding/kernel.mlu index a4f64adf1..bccfb22b7 100644 --- a/src/native/cambricon/ops/embedding/kernel.mlu +++ b/src/native/cambricon/ops/embedding/kernel.mlu @@ -70,7 +70,7 @@ __mlu_func__ ptrdiff_t LogicalOffset(size_t logical_index, size_t ndim, template __mlu_func__ float LoadFloat(const T* source) { - if constexpr (std::is_same::value) { + if constexpr (std::is_same::value) { return __half2float(*source); } else if constexpr (std::is_same::value) { return __bfloat162float(*source); @@ -253,7 +253,7 @@ void EmbeddingKernelLaunch( DispatchFunc< Device::Type::kCambricon, List, - List>( + CambriconBangFloatTypes<>>( {input_dtype, weight_dtype}, [&](auto input_tag, auto weight_tag) { using IndexT = typename decltype(input_tag)::type; diff --git a/src/native/cambricon/ops/eq/kernel.h b/src/native/cambricon/ops/eq/kernel.h new file mode 100644 index 000000000..078d87d1b --- /dev/null +++ b/src/native/cambricon/ops/eq/kernel.h @@ -0,0 +1,23 @@ +#ifndef INFINI_OPS_CAMBRICON_EQ_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_EQ_KERNEL_H_ + +#include "base/eq.h" +#include "native/cambricon/compare_cnnl.h" + +namespace infini::ops { + +template <> +class Operator + : public cnnl_utils::Compare { + public: + Operator(const Tensor input, const Tensor other, Tensor out) + : Compare{input, other, out} {} + Operator(const Tensor input, const double other, Tensor out) + : Compare{input, other, out} {} + Operator(Tensor input, const Tensor other) : Compare{input, other} {} + Operator(Tensor input, const double other) : Compare{input, other} {} +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/erf/kernel.h b/src/native/cambricon/ops/erf/kernel.h new file mode 100644 index 000000000..619271ce6 --- /dev/null +++ b/src/native/cambricon/ops/erf/kernel.h @@ -0,0 +1,18 @@ +#ifndef INFINI_OPS_CAMBRICON_ERF_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_ERF_KERNEL_H_ + +#include "base/erf.h" +#include "native/cambricon/unary_cnnl.h" + +namespace infini::ops { + +template <> +class Operator + : public cnnl_utils::Unary { + public: + Operator(const Tensor input, Tensor out) : Unary{input, out} {} +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/fill/kernel.mlu b/src/native/cambricon/ops/fill/kernel.mlu index bba39f1fc..4d95a3cd8 100644 --- a/src/native/cambricon/ops/fill/kernel.mlu +++ b/src/native/cambricon/ops/fill/kernel.mlu @@ -136,7 +136,7 @@ void FillUnion(void* workspace, cnrtQueue_t queue, int core_per_cluster, static_cast(cluster_count), 1}; T converted_host_value{}; if (device_value == nullptr) { - if constexpr (std::is_same_v || + if constexpr (std::is_same_v || std::is_same_v) { converted_host_value = static_cast(static_cast(host_value)); } else { diff --git a/src/native/cambricon/ops/fused_add_rms_norm/kernel.mlu b/src/native/cambricon/ops/fused_add_rms_norm/kernel.mlu index fc51b1928..38798b772 100644 --- a/src/native/cambricon/ops/fused_add_rms_norm/kernel.mlu +++ b/src/native/cambricon/ops/fused_add_rms_norm/kernel.mlu @@ -12,7 +12,7 @@ template __mlu_device__ void LoadAsFloat(float* dst, T* cache, const T* src, size_t count) { __memcpy(cache, src, count * sizeof(T), GDRAM2NRAM); - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { __bang_half2float(dst, reinterpret_cast(cache), count); } else if constexpr (std::is_same_v) { __bang_bfloat162float(dst, cache, count); @@ -23,7 +23,7 @@ __mlu_device__ void LoadAsFloat(float* dst, T* cache, const T* src, template __mlu_device__ void ConvertFromFloat(T* dst, const float* src, size_t count) { - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { __bang_float2half(reinterpret_cast(dst), const_cast(src), count); } else if constexpr (std::is_same_v) { @@ -35,7 +35,7 @@ __mlu_device__ void ConvertFromFloat(T* dst, const float* src, size_t count) { template __mlu_device__ void ConvertToFloat(float* dst, T* src, size_t count) { - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { __bang_half2float(dst, reinterpret_cast(src), count); } else if constexpr (std::is_same_v) { __bang_bfloat162float(dst, src, count); diff --git a/src/native/cambricon/ops/ge/kernel.h b/src/native/cambricon/ops/ge/kernel.h new file mode 100644 index 000000000..d228eda52 --- /dev/null +++ b/src/native/cambricon/ops/ge/kernel.h @@ -0,0 +1,23 @@ +#ifndef INFINI_OPS_CAMBRICON_GE_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_GE_KERNEL_H_ + +#include "base/ge.h" +#include "native/cambricon/compare_cnnl.h" + +namespace infini::ops { + +template <> +class Operator + : public cnnl_utils::Compare { + public: + Operator(const Tensor input, const Tensor other, Tensor out) + : Compare{input, other, out} {} + Operator(const Tensor input, const double other, Tensor out) + : Compare{input, other, out} {} + Operator(Tensor input, const Tensor other) : Compare{input, other} {} + Operator(Tensor input, const double other) : Compare{input, other} {} +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/gelu/kernel.mlu b/src/native/cambricon/ops/gelu/kernel.mlu index 46e563cc4..6e9d793a0 100644 --- a/src/native/cambricon/ops/gelu/kernel.mlu +++ b/src/native/cambricon/ops/gelu/kernel.mlu @@ -24,7 +24,7 @@ __mlu_device__ ptrdiff_t LogicalToOffset(size_t logical_index, int ndim, template __mlu_device__ void ComputeGelu(const T* input, T* output, float* values, size_t count, bool approximate) { - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { __bang_half2float(values, reinterpret_cast(const_cast(input)), count); } else if constexpr (std::is_same_v) { @@ -46,7 +46,7 @@ __mlu_device__ void ComputeGelu(const T* input, T* output, float* values, } } - if constexpr (std::is_same_v) { + 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); diff --git a/src/native/cambricon/ops/grid_sample/kernel.h b/src/native/cambricon/ops/grid_sample/kernel.h new file mode 100644 index 000000000..b4b23127c --- /dev/null +++ b/src/native/cambricon/ops/grid_sample/kernel.h @@ -0,0 +1,80 @@ +#ifndef INFINI_OPS_CAMBRICON_GRID_SAMPLE_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_GRID_SAMPLE_KERNEL_H_ + +#include "base/grid_sample.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops { + +template <> +class Operator : public GridSample { + public: + Operator(const Tensor input, const Tensor grid, bool nearest, + bool reflection, bool align_corners, Tensor out) + : GridSample{input, grid, nearest, reflection, align_corners, out} { + if (input.ndim() != 4 || grid.ndim() != 4 || out.ndim() != 4 || + input.dtype() != out.dtype() || grid.dtype() != out.dtype() || + (out.dtype() != DataType::kFloat32 && out.dtype() != DataType::kFloat16) || + input.device() != out.device() || grid.device() != out.device() || + !input.IsContiguous() || !grid.IsContiguous() || !out.IsContiguous() || + grid.size(3) != 2 || input.size(0) != grid.size(0) || + out.size(0) != grid.size(0) || out.size(1) != grid.size(1) || + out.size(2) != grid.size(2) || out.size(3) != input.size(3) || + (nearest && reflection)) + throw std::invalid_argument("Unsupported Cambricon GridSample parameters"); + if (out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + input_ = Describe(input); + grid_ = Describe(grid); + out_ = Describe(out); + cnnlGridSampleDescriptor_t desc = nullptr; + INFINI_OPS_CNNL_CHECK(cnnlCreateGridSampleDescriptor(&desc)); + sample_.reset(desc); + INFINI_OPS_CNNL_CHECK(cnnlSetGridSampleDescriptor( + desc, nearest ? CNNL_INTERP_NEAREST : CNNL_INTERP_BILINEAR, + reflection ? CNNL_GRIDSAMPLE_PADDING_REFLECTION : CNNL_GRIDSAMPLE_PADDING_ZEROS, + align_corners)); + INFINI_OPS_CNNL_CHECK(cnnlGetGridSampleForwardWorkspaceSize( + handle_.get(), input_.get(), grid_.get(), out_.get(), &workspace_size_)); + scratch_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_); + if (workspace_size_ && !scratch_) + throw std::runtime_error("Cambricon GridSample workspace allocation failed"); + } + + std::size_t workspace_size_in_bytes() const override { return workspace_size_; } + + void operator()(const Tensor input, const Tensor grid, bool, bool, bool, + Tensor out) const override { + if (!out_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + auto* scratch = workspace_ && workspace_size_in_bytes_ >= workspace_size_ + ? workspace_ : scratch_.get(); + INFINI_OPS_CNNL_CHECK(cnnlGridSampleForward( + handle_.get(), sample_.get(), input_.get(), input.data(), + grid_.get(), grid.data(), out_.get(), out.data(), scratch, workspace_size_)); + } + + private: + static cnnl_utils::TensorDescriptor Describe(const Tensor& tensor) { + auto desc = cnnl_utils::CreateTensorDescriptor(); + const auto shape = tensor.shape(); + const std::vector dims(shape.begin(), shape.end()); + INFINI_OPS_CNNL_CHECK(cnnlSetTensorDescriptor_v2( + desc.get(), CNNL_LAYOUT_NHWC, cnnl_utils::GetDataType(tensor.dtype()), + static_cast(dims.size()), dims.data())); + return desc; + } + + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor input_, grid_, out_; + std::unique_ptr, + decltype(&cnnlDestroyGridSampleDescriptor)> + sample_{nullptr, cnnlDestroyGridSampleDescriptor}; + cnrt_utils::DeviceBuffer scratch_; + std::size_t workspace_size_ = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/index_select/kernel.h b/src/native/cambricon/ops/index_select/kernel.h new file mode 100644 index 000000000..e32a005bd --- /dev/null +++ b/src/native/cambricon/ops/index_select/kernel.h @@ -0,0 +1,83 @@ +#ifndef INFINI_OPS_CAMBRICON_INDEX_SELECT_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_INDEX_SELECT_KERNEL_H_ + +#include + +#include "base/gather_axis.h" +#include "base/index_select.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" +#include "native/cambricon/scatter_index.h" + +namespace infini::ops { + +template +class CambriconAxisSelection : public Base { + public: + CambriconAxisSelection(const Tensor input, const Tensor index, int64_t dim, + Tensor out) + : Base{input, index, dim, out}, axis_(dim) { + const auto rank = static_cast(input.ndim()); + if (axis_ < 0) axis_ += rank; + if (axis_ < 0 || axis_ >= rank || index.ndim() != 1 || input.ndim() != out.ndim() || + (index.dtype() != DataType::kInt32 && index.dtype() != DataType::kInt64) || + input.dtype() != out.dtype() || input.device() != out.device() || + index.device() != out.device() || !input.IsContiguous() || + !index.IsContiguous() || !out.IsContiguous()) + throw std::invalid_argument("Cambricon IndexSelect: invalid tensor metadata"); + for (size_t i = 0; i < out.ndim(); ++i) { + const auto expected = i == static_cast(axis_) ? index.numel() : input.shape()[i]; + if (out.shape()[i] != expected) + throw std::invalid_argument("Cambricon IndexSelect: invalid output shape"); + } + if (out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + input_ = cnnl_utils::MakeTensorDescriptor(input); + index_ = cnnl_utils::MakeTensorDescriptor(index); + output_ = cnnl_utils::MakeTensorDescriptor(out); + if constexpr (WrapNegative) { + normalized_ = cnrt_utils::AllocateDeviceBuffer(index.numel() * index.element_size()); + if (!normalized_) + throw std::runtime_error("Cambricon GatherAxis index allocation failed"); + } + } + + void operator()(const Tensor input, const Tensor index, int64_t, + Tensor out) const override { + if (!output_) return; + auto queue = static_cast(this->stream_); + const void* indices = index.data(); + if constexpr (WrapNegative) { + NormalizeAxisIndices(indices, normalized_.get(), input.size(axis_), + index.numel(), index.dtype(), queue); + indices = normalized_.get(); + } + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), queue)); + INFINI_OPS_CNNL_CHECK(cnnlIndexSelect( + handle_.get(), static_cast(axis_), input_.get(), input.data(), + index_.get(), indices, output_.get(), out.data())); + } + + private: + + int64_t axis_; + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor input_, index_, output_; + cnrt_utils::DeviceBuffer normalized_; +}; + +template <> +class Operator + : public CambriconAxisSelection { + using CambriconAxisSelection::CambriconAxisSelection; +}; + +template <> +class Operator + : public CambriconAxisSelection { + using CambriconAxisSelection::CambriconAxisSelection; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/instance_norm/kernel.h b/src/native/cambricon/ops/instance_norm/kernel.h new file mode 100644 index 000000000..aa1d72f06 --- /dev/null +++ b/src/native/cambricon/ops/instance_norm/kernel.h @@ -0,0 +1,116 @@ +#ifndef INFINI_OPS_CAMBRICON_INSTANCE_NORM_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_INSTANCE_NORM_KERNEL_H_ + +#include +#include + +#include "base/instance_norm.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops { + +template <> +class Operator : public InstanceNorm { + public: + Operator(const Tensor input, const Tensor scale, const Tensor bias, + float epsilon, Tensor out) + : InstanceNorm{input, scale, bias, epsilon, out} { + const auto rank = input.ndim(); + if ((rank != 3 && rank != 4) || input.shape() != out.shape() || + input.dtype() != out.dtype() || + (out.dtype() != DataType::kFloat32 && out.dtype() != DataType::kFloat16) || + scale.ndim() != 1 || bias.ndim() != 1 || scale.size(0) != input.size(1) || + bias.size(0) != input.size(1) || scale.dtype() != input.dtype() || + bias.dtype() != input.dtype() || input.device() != out.device() || + scale.device() != out.device() || bias.device() != out.device() || + !input.IsContiguous() || !out.IsContiguous() || + !scale.IsContiguous() || !bias.IsContiguous() || !std::isfinite(epsilon) || epsilon < 0) + throw std::invalid_argument("Unsupported Cambricon InstanceNorm parameters"); + if (out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + const auto metadata = input.shape(); + const std::vector shape(metadata.begin(), metadata.end()); + std::vector forward = rank == 3 ? std::vector{0, 2, 1} : std::vector{0, 2, 3, 1}; + std::vector backward = rank == 3 ? forward : std::vector{0, 3, 1, 2}; + std::vector permuted; + for (auto dim : forward) permuted.push_back(shape[dim]); + const auto dtype = cnnl_utils::GetDataType(input.dtype()); + first_ = Describe(CNNL_LAYOUT_ARRAY, dtype, shape); + last_ = Describe(rank == 3 ? CNNL_LAYOUT_NLC : CNNL_LAYOUT_NHWC, dtype, permuted); + channel_ = Describe(CNNL_LAYOUT_ARRAY, dtype, {shape[1]}); + statistics_ = Describe(CNNL_LAYOUT_NC, dtype, {shape[0], shape[1]}); + SetupTranspose(to_last_, forward); + SetupTranspose(to_first_, backward); + INFINI_OPS_CNNL_CHECK(cnnlGetTransposeWorkspaceSize( + handle_.get(), first_.get(), to_last_.get(), &forward_bytes_)); + INFINI_OPS_CNNL_CHECK(cnnlGetTransposeWorkspaceSize( + handle_.get(), last_.get(), to_first_.get(), &backward_bytes_)); + INFINI_OPS_CNNL_CHECK(cnnlGetInstanceNormForwardWorkspaceSize( + handle_.get(), last_.get(), &normalization_bytes_)); + scratch_bytes_ = std::max({forward_bytes_, backward_bytes_, normalization_bytes_}); + input_last_ = Allocate(input.numel() * input.element_size()); + output_last_ = Allocate(out.numel() * out.element_size()); + const auto stats_bytes = static_cast(shape[0]) * shape[1] * sizeof(float); + saved_mean_ = Allocate(stats_bytes); + saved_inverse_ = Allocate(stats_bytes); + scratch_ = Allocate(scratch_bytes_); + } + + std::size_t workspace_size_in_bytes() const override { return scratch_bytes_; } + + void operator()(const Tensor input, const Tensor scale, const Tensor bias, + float epsilon, Tensor out) const override { + if (!first_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + auto* scratch = workspace_ && workspace_size_in_bytes_ >= scratch_bytes_ + ? workspace_ : scratch_.get(); + INFINI_OPS_CNNL_CHECK(cnnlTranspose_v2( + handle_.get(), to_last_.get(), first_.get(), input.data(), last_.get(), + input_last_.get(), scratch, forward_bytes_)); + INFINI_OPS_CNNL_CHECK(cnnlInstanceNormForward( + handle_.get(), epsilon, 0.0f, channel_.get(), scale.data(), last_.get(), + input_last_.get(), bias.data(), statistics_.get(), nullptr, nullptr, + scratch, normalization_bytes_, last_.get(), output_last_.get(), + saved_mean_.get(), saved_inverse_.get())); + INFINI_OPS_CNNL_CHECK(cnnlTranspose_v2( + handle_.get(), to_first_.get(), last_.get(), output_last_.get(), + first_.get(), out.data(), scratch, backward_bytes_)); + } + + private: + using Transpose = std::unique_ptr, + decltype(&cnnlDestroyTransposeDescriptor)>; + static cnnl_utils::TensorDescriptor Describe(cnnlTensorLayout_t layout, + cnnlDataType_t dtype, const std::vector& shape) { + auto desc = cnnl_utils::CreateTensorDescriptor(); + INFINI_OPS_CNNL_CHECK(cnnlSetTensorDescriptor_v2(desc.get(), layout, dtype, + shape.size(), shape.data())); + return desc; + } + + static void SetupTranspose(Transpose& desc, const std::vector& permutation) { + cnnlTransposeDescriptor_t value = nullptr; + INFINI_OPS_CNNL_CHECK(cnnlCreateTransposeDescriptor(&value)); + desc.reset(value); + INFINI_OPS_CNNL_CHECK(cnnlSetTransposeDescriptor(value, permutation.size(), permutation.data())); + } + + static cnrt_utils::DeviceBuffer Allocate(size_t bytes) { + auto result = cnrt_utils::AllocateDeviceBuffer(bytes); + if (bytes && !result) + throw std::runtime_error("Cambricon InstanceNorm allocation failed"); + return result; + } + + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor first_, last_, channel_, statistics_; + Transpose to_last_{nullptr, cnnlDestroyTransposeDescriptor}; + Transpose to_first_{nullptr, cnnlDestroyTransposeDescriptor}; + cnrt_utils::DeviceBuffer input_last_, output_last_, saved_mean_, saved_inverse_, scratch_; + std::size_t forward_bytes_ = 0, backward_bytes_ = 0, normalization_bytes_ = 0, scratch_bytes_ = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/layer_norm/kernel.h b/src/native/cambricon/ops/layer_norm/kernel.h new file mode 100644 index 000000000..02ab49ab0 --- /dev/null +++ b/src/native/cambricon/ops/layer_norm/kernel.h @@ -0,0 +1,91 @@ +#ifndef INFINI_OPS_CAMBRICON_LAYER_NORM_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_LAYER_NORM_KERNEL_H_ + +#include +#include + +#include "base/layer_norm.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops { + +template <> +class Operator : public LayerNorm { + public: + Operator(const Tensor input, const Tensor scale, std::optional bias, + float eps, Tensor out) + : LayerNorm{input, scale, bias, eps, out}, eps_(eps) { + const auto rank = static_cast(input.ndim()); + const auto axis = rank - 1; + if (axis < 0 || axis >= rank || !std::isfinite(eps) || eps < 0 || + input.shape() != out.shape() || + (input.dtype() != DataType::kFloat32 && input.dtype() != DataType::kFloat16) || + !Compatible(out, input) || !Compatible(scale, input) || + !input.IsContiguous() || (bias && (!Compatible(*bias, input) || bias->shape() != scale.shape()))) + throw std::invalid_argument("Invalid Cambricon LayerNorm tensor metadata"); + axis_ = static_cast(axis); + Tensor::Shape prefix, suffix; + for (int64_t i = 0; i < rank; ++i) { + if (i < axis) prefix.push_back(input.shape()[i]); + else suffix.push_back(input.shape()[i]); + } + if (Tensor::Shape(scale.shape()) != suffix) + throw std::invalid_argument("Cambricon LayerNorm scale must match the normalized shape"); + if (out.numel() == 0) return; + size_t stats_size = 1; + for (auto dim : prefix) stats_size *= dim; + Tensor::Strides stats_strides(prefix.size()); + Tensor::Stride stride = 1; + for (size_t i = prefix.size(); i-- > 0;) { + stats_strides[i] = stride; + stride *= prefix[i]; + } + handle_ = cnnl_utils::CreateHandle(); + input_ = cnnl_utils::MakeTensorDescriptor(input); + scale_ = cnnl_utils::MakeTensorDescriptor(scale); + output_ = cnnl_utils::MakeTensorDescriptor(out); + stats_ = cnnl_utils::MakeTensorDescriptor(input.dtype(), prefix, stats_strides); + INFINI_OPS_CNNL_CHECK(cnnlGetLayerNormOpWorkspaceSize( + handle_.get(), axis_, input_.get(), &workspace_size_)); + workspace_own_ = Allocate(workspace_size_); + const auto stats_bytes = stats_size * ::infini::rt::kDataTypeToSize.at(input.dtype()); + mean_ = Allocate(stats_bytes); + rstd_ = Allocate(stats_bytes); + } + + std::size_t workspace_size_in_bytes() const override { return workspace_size_; } + + void operator()(const Tensor input, const Tensor scale, std::optional bias, + float, Tensor out) const override { + if (!output_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + auto* scratch = workspace_ && workspace_size_in_bytes_ >= workspace_size_ + ? workspace_ : workspace_own_.get(); + INFINI_OPS_CNNL_CHECK(cnnlLayerNormForward( + handle_.get(), input_.get(), input.data(), axis_, scale_.get(), scale.data(), + bias ? bias->data() : nullptr, eps_, scratch, workspace_size_, + output_.get(), out.data(), stats_.get(), mean_.get(), rstd_.get())); + } + + private: + static bool Compatible(const Tensor& tensor, const Tensor& input) { + return tensor.dtype() == input.dtype() && tensor.device() == input.device() && + tensor.IsContiguous(); + } + static cnrt_utils::DeviceBuffer Allocate(std::size_t bytes) { + auto result = cnrt_utils::AllocateDeviceBuffer(bytes); + if (bytes && !result) throw std::runtime_error("Cambricon LayerNorm allocation failed"); + return result; + } + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor input_, scale_, output_, stats_; + cnrt_utils::DeviceBuffer workspace_own_, mean_, rstd_; + std::size_t workspace_size_ = 0; + int axis_ = 0; + float eps_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/less/kernel.h b/src/native/cambricon/ops/less/kernel.h new file mode 100644 index 000000000..e3defef0a --- /dev/null +++ b/src/native/cambricon/ops/less/kernel.h @@ -0,0 +1,23 @@ +#ifndef INFINI_OPS_CAMBRICON_LESS_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_LESS_KERNEL_H_ + +#include "base/less.h" +#include "native/cambricon/compare_cnnl.h" + +namespace infini::ops { + +template <> +class Operator + : public cnnl_utils::Compare { + public: + Operator(const Tensor input, const Tensor other, Tensor out) + : Compare{input, other, out} {} + Operator(const Tensor input, const double other, Tensor out) + : Compare{input, other, out} {} + Operator(Tensor input, const Tensor other) : Compare{input, other} {} + Operator(Tensor input, const double other) : Compare{input, other} {} +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/log/kernel.h b/src/native/cambricon/ops/log/kernel.h new file mode 100644 index 000000000..569343281 --- /dev/null +++ b/src/native/cambricon/ops/log/kernel.h @@ -0,0 +1,24 @@ +#ifndef INFINI_OPS_CAMBRICON_LOG_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_LOG_KERNEL_H_ + +#include "base/log.h" +#include "native/cambricon/unary_cnnl.h" + +namespace infini::ops { + +inline cnnlStatus_t NaturalLog(cnnlHandle_t handle, cnnlTensorDescriptor_t input_desc, + const void* input, cnnlTensorDescriptor_t out_desc, void* out) { + return cnnlLog_v2(handle, CNNL_COMPUTATION_HIGH_PRECISION, CNNL_LOG_E, + input_desc, input, out_desc, out); +} + +template <> +class Operator + : public cnnl_utils::Unary { + public: + Operator(const Tensor input, Tensor out) : Unary{input, out} {} +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/logical_not/kernel.h b/src/native/cambricon/ops/logical_not/kernel.h new file mode 100644 index 000000000..c29f978f4 --- /dev/null +++ b/src/native/cambricon/ops/logical_not/kernel.h @@ -0,0 +1,18 @@ +#ifndef INFINI_OPS_CAMBRICON_LOGICAL_NOT_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_LOGICAL_NOT_KERNEL_H_ + +#include "base/logical_not.h" +#include "native/cambricon/unary_cnnl.h" + +namespace infini::ops { + +template <> +class Operator + : public cnnl_utils::Unary { + public: + Operator(const Tensor input, Tensor out) : Unary{input, out} {} +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/masked_scatter/kernel.h b/src/native/cambricon/ops/masked_scatter/kernel.h new file mode 100644 index 000000000..7e68a6cc4 --- /dev/null +++ b/src/native/cambricon/ops/masked_scatter/kernel.h @@ -0,0 +1,93 @@ +#ifndef INFINI_OPS_CAMBRICON_MASKED_SCATTER_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_MASKED_SCATTER_KERNEL_H_ + +#include "base/masked_scatter.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops { + +template <> +class Operator : public MaskedScatter { + public: + Operator(const Tensor input, const Tensor mask, const Tensor source, Tensor out) + : MaskedScatter{input, mask, source, out} { + if (input.shape() != out.shape() || input.dtype() != out.dtype() || + source.dtype() != out.dtype() || mask.dtype() != DataType::kBool || + input.device() != out.device() || mask.device() != out.device() || + source.device() != out.device() || + !input.IsContiguous() || !mask.IsContiguous() || + !source.IsContiguous() || !out.IsContiguous() || mask.ndim() > input.ndim()) + throw std::invalid_argument("Unsupported Cambricon MaskedScatter tensors"); + for (size_t i = 0; i < mask.ndim(); ++i) { + const auto size = input.size(input.ndim() - mask.ndim() + i); + if (mask.size(i) != 1 && mask.size(i) != size) + throw std::invalid_argument("Cambricon MaskedScatter mask is not broadcastable"); + } + if (out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + const auto input_shape = input.shape(); + const auto mask_shape = mask.shape(); + input_ = Describe(input.dtype(), input_shape); + out_ = Describe(out.dtype(), out.shape()); + mask_ = Describe(mask.dtype(), mask_shape); + source_ = Describe(source.dtype(), Tensor::Shape{source.numel()}); + if (mask_shape != input_shape) { + expanded_mask_ = Describe(mask.dtype(), input_shape); + mask_bytes_ = (input.numel() + 255) & ~size_t{255}; + } + INFINI_OPS_CNNL_CHECK(cnnlGetMaskedWorkspaceSize( + handle_.get(), CNNL_MASKED_SCATTER, input_.get(), EffectiveMask(), + source_.get(), out_.get(), &masked_bytes_)); + scratch_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_in_bytes()); + if (workspace_size_in_bytes() && !scratch_) + throw std::runtime_error("Cambricon MaskedScatter workspace allocation failed"); + } + + std::size_t workspace_size_in_bytes() const override { return mask_bytes_ + masked_bytes_; } + + void operator()(const Tensor input, const Tensor mask, const Tensor source, + Tensor out) const override { + if (!out_) return; + auto* scratch = static_cast( + workspace_ && workspace_size_in_bytes_ >= workspace_size_in_bytes() + ? workspace_ : scratch_.get()); + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + const void* mask_data = mask.data(); + if (expanded_mask_) { + INFINI_OPS_CNNL_CHECK(cnnlExpand(handle_.get(), mask_.get(), mask_data, + expanded_mask_.get(), scratch)); + mask_data = scratch; + } + INFINI_OPS_CNNL_CHECK(cnnlMasked_v2( + handle_.get(), CNNL_MASKED_SCATTER, input_.get(), input.data(), + EffectiveMask(), mask_data, source_.get(), source.data(), + masked_bytes_ ? scratch + mask_bytes_ : nullptr, masked_bytes_, + out_.get(), out.data())); + } + + private: + template + static cnnl_utils::TensorDescriptor Describe(DataType dtype, const Shape& shape) { + auto desc = cnnl_utils::CreateTensorDescriptor(); + std::vector dims(shape.begin(), shape.end()); + if (dims.empty()) dims.push_back(1); + INFINI_OPS_CNNL_CHECK(cnnlSetTensorDescriptor_v2( + desc.get(), CNNL_LAYOUT_ARRAY, cnnl_utils::GetDataType(dtype), + static_cast(dims.size()), dims.data())); + return desc; + } + + cnnlTensorDescriptor_t EffectiveMask() const { + return expanded_mask_ ? expanded_mask_.get() : mask_.get(); + } + + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor input_, mask_, source_, out_, expanded_mask_; + cnrt_utils::DeviceBuffer scratch_; + std::size_t mask_bytes_ = 0, masked_bytes_ = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/matmul/kernel.h b/src/native/cambricon/ops/matmul/kernel.h new file mode 100644 index 000000000..368336836 --- /dev/null +++ b/src/native/cambricon/ops/matmul/kernel.h @@ -0,0 +1,163 @@ +#ifndef INFINI_OPS_CAMBRICON_MATMUL_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_MATMUL_KERNEL_H_ + +#include +#include +#include + +#include "base/matmul.h" +#include "native/cambricon/cnnl_utils.h" + +namespace infini::ops { + +template <> +class Operator : public Matmul { + public: + Operator(const Tensor input, const Tensor other, Tensor out) + : Matmul{input, other, out} { + if (input.device() != other.device() || input.device() != out.device() || + input_type_ != other_type_ || input_type_ != out_type_ || + (out_type_ != DataType::kFloat16 && out_type_ != DataType::kFloat32) || + input.ndim() < 2 || other.ndim() < 2 || out.ndim() < 2 || + !out.IsContiguous()) { + throw std::invalid_argument( + "Cambricon Matmul requires matching FP16 or FP32 matrix tensors"); + } + if (out.numel() == 0) return; + if (input.size(-1) == 0 || input.size(-1) != other.size(-2) || + input.size(-2) != out.size(-2) || other.size(-1) != out.size(-1)) { + throw std::invalid_argument("Unsupported Cambricon Matmul dimensions"); + } + CheckBatch(input, out); + CheckBatch(other, out); + + handle_ = Create(cnnlCreate, "create handle"); + descriptor_ = Create( + cnnlMatMulDescCreate, "create matmul descriptor"); + algo_ = Create( + cnnlMatMulAlgoCreate, "create algorithm"); + input_ = MatrixDescriptor(input, out.ndim(), CNNL_MATMUL_DESC_TRANSA); + other_ = MatrixDescriptor(other, out.ndim(), CNNL_MATMUL_DESC_TRANSB); + output_ = TensorDescriptor(Tensor::Shape(out.shape()), out_type_, out.ndim()); + const auto compute_type = CNNL_DTYPE_FLOAT; + const int32_t allow_tf32 = 0; + Check(cnnlSetMatMulDescAttr(descriptor_.get(), CNNL_MATMUL_DESC_COMPUTE_TYPE, + &compute_type, sizeof(compute_type)), "set accumulation type"); + Check(cnnlSetMatMulDescAttr(descriptor_.get(), CNNL_MATMUL_ALLOW_TF32, + &allow_tf32, sizeof(allow_tf32)), "disable TF32"); + + auto heuristic = Create( + cnnlCreateMatMulHeuristicResult, "create heuristic"); + auto result = heuristic.get(); + // Some CNNL releases fill the result without writing return_algo_count. + int count = -1; + Check(cnnlGetBatchMatMulExAlgoHeuristic(handle_.get(), descriptor_.get(), + input_.get(), other_.get(), output_.get(), + nullptr, 1, &result, &count), "select algorithm"); + if (count != -1 && count != 1) + throw std::runtime_error("Cambricon Matmul: no CNNL algorithm"); + Check(cnnlGetBatchMatMulExHeuristicResult(heuristic.get(), algo_.get(), &workspace_size_), + "get workspace size"); + if (workspace_size_ != 0) { + void* ptr = nullptr; + if (cnrtMalloc(&ptr, workspace_size_) != cnrtSuccess) + throw std::runtime_error("Cambricon Matmul: workspace allocation failed"); + scratch_.reset(ptr); + } + } + + std::size_t workspace_size_in_bytes() const override { return workspace_size_; } + + void operator()(const Tensor input, const Tensor other, Tensor out) const override { + if (!output_) return; + Check(cnnlSetQueue(handle_.get(), static_cast(stream_)), "set queue"); + const float alpha = 1.0f, beta = 0.0f; + auto* workspace = workspace_ && workspace_size_in_bytes_ >= workspace_size_ + ? workspace_ : scratch_.get(); + Check(cnnlBatchMatMulEx(handle_.get(), descriptor_.get(), algo_.get(), + &alpha, input_.get(), input.data(), other_.get(), + other.data(), &beta, output_.get(), out.data(), + workspace, workspace_size_), "execute"); + } + + private: + template + using Resource = std::unique_ptr, decltype(Destroy)>; + + template + static Resource Create(Creator creator, const char* operation) { + T value = nullptr; + Check(creator(&value), operation); + return Resource{value, Destroy}; + } + + static void Check(cnnlStatus_t status, const char* operation) { + if (status != CNNL_STATUS_SUCCESS) + throw std::runtime_error(std::string("Cambricon Matmul: ") + operation + + ": " + cnnlGetErrorString(status)); + } + + static void CheckBatch(const Tensor& input, const Tensor& out) { + if (input.ndim() > out.ndim()) + throw std::invalid_argument("Cambricon Matmul: invalid broadcast rank"); + const auto offset = out.ndim() - input.ndim(); + for (size_t i = 0; i + 2 < input.ndim(); ++i) { + if (input.shape()[i] != 1 && input.shape()[i] != out.shape()[offset + i]) + throw std::invalid_argument("Cambricon Matmul: incompatible batch dimensions"); + } + } + + static Resource + TensorDescriptor(const Tensor::Shape& shape, DataType dtype, size_t rank) { + auto result = Create( + cnnlCreateTensorDescriptor, "create tensor descriptor"); + std::vector dims(rank - shape.size(), 1); + dims.insert(dims.end(), shape.begin(), shape.end()); + Check(cnnlSetTensorDescriptor_v2(result.get(), CNNL_LAYOUT_ARRAY, + cnnl_utils::GetDataType(dtype), dims.size(), dims.data()), + "set tensor descriptor"); + return result; + } + + Resource + MatrixDescriptor(const Tensor& tensor, size_t rank, + cnnlMatMulDescAttribute_t attribute) { + Tensor::Shape shape(tensor.shape()); + const auto strides = tensor.strides(); + const size_t last = shape.size() - 1; + const auto rows = static_cast(shape[last - 1]); + const auto cols = static_cast(shape[last]); + const bool normal = (cols == 1 || strides[last] == 1) && + (rows == 1 || strides[last - 1] == cols); + const bool transposed = (rows == 1 || strides[last - 1] == 1) && + (cols == 1 || strides[last] == rows); + if (!normal && !transposed) + throw std::invalid_argument("Cambricon Matmul: unsupported matrix strides"); + auto stride = rows * cols; + for (size_t i = last - 1; i-- > 0;) { + if (shape[i] > 1 && strides[i] != stride) + throw std::invalid_argument("Cambricon Matmul: unsupported batch strides"); + stride *= shape[i]; + } + const int32_t transpose = normal ? 0 : 1; + if (transpose) std::swap(shape[last - 1], shape[last]); + Check(cnnlSetMatMulDescAttr(descriptor_.get(), attribute, &transpose, sizeof(transpose)), + "set transpose"); + // Preserve leading singleton dimensions even when the batch is one. + return TensorDescriptor(shape, tensor.dtype(), rank); + } + + Resource handle_{nullptr, cnnlDestroy}; + Resource descriptor_{nullptr, cnnlMatMulDescDestroy}; + Resource algo_{nullptr, cnnlMatMulAlgoDestroy}; + Resource + input_{nullptr, cnnlDestroyTensorDescriptor}, + other_{nullptr, cnnlDestroyTensorDescriptor}, + output_{nullptr, cnnlDestroyTensorDescriptor}; + std::size_t workspace_size_ = 0; + std::unique_ptr scratch_{nullptr, cnrtFree}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/max_pool2d/kernel.h b/src/native/cambricon/ops/max_pool2d/kernel.h new file mode 100644 index 000000000..963cd4f07 --- /dev/null +++ b/src/native/cambricon/ops/max_pool2d/kernel.h @@ -0,0 +1,84 @@ +#ifndef INFINI_OPS_CAMBRICON_MAX_POOL2D_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_MAX_POOL2D_KERNEL_H_ + +#include "base/max_pool2d.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops { + +template <> +class Operator : public MaxPool2d { + public: + Operator(const Tensor input, const std::vector kernel_size, + const std::vector stride, const std::vector padding, + const std::vector dilation, bool ceil_mode, Tensor out) + : MaxPool2d{input, kernel_size, stride, padding, dilation, ceil_mode, out} { + if (kernel_size.size() != 2 || stride.size() != 2 || + padding.size() != 2 || dilation.size() != 2 || + dilation[0] != 1 || dilation[1] != 1) + throw std::invalid_argument("Cambricon MaxPool requires 2D parameters and unit dilation"); + const auto kh = kernel_size[0], kw = kernel_size[1]; + const auto sh = stride[0], sw = stride[1]; + const auto ph = padding[0], pw = padding[1]; + if ((input.ndim() != 3 && input.ndim() != 4) || input.ndim() != out.ndim() || + input.dtype() != out.dtype() || + (input.dtype() != DataType::kFloat16 && input.dtype() != DataType::kFloat32) || + input.device() != out.device() || !input.IsContiguous() || + !out.IsContiguous() || kh <= 0 || kw <= 0 || sh <= 0 || sw <= 0 || + ph < 0 || pw < 0 || (input.ndim() == 3 && (kh != 1 || ph != 0 || sh != 1))) + throw std::invalid_argument("Unsupported Cambricon MaxPool tensor metadata"); + if (!out.numel()) return; + handle_ = cnnl_utils::CreateHandle(); + input_ = Describe(input); + output_ = Describe(out); + cnnlPoolingDescriptor_t desc = nullptr; + INFINI_OPS_CNNL_CHECK(cnnlCreatePoolingDescriptor(&desc)); + pool_.reset(desc); + INFINI_OPS_CNNL_CHECK(cnnlSetPooling2dDescriptor_v2( + desc, CNNL_POOLING_MAX, CNNL_PROPAGATE_NAN, kh, kw, ph, ph, pw, pw, + sh, sw, 1, 1, ceil_mode)); + INFINI_OPS_CNNL_CHECK(cnnlGetPoolingWorkspaceSize_v2( + handle_.get(), pool_.get(), input_.get(), output_.get(), &workspace_size_)); + scratch_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_); + } + + std::size_t workspace_size_in_bytes() const override { return workspace_size_; } + + void operator()(const Tensor input, const std::vector, + const std::vector, const std::vector, + const std::vector, bool, Tensor out) const override { + if (!pool_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + auto* scratch = workspace_ && workspace_size_in_bytes_ >= workspace_size_ + ? workspace_ : scratch_.get(); + const float one = 1.0f, zero = 0.0f; + INFINI_OPS_CNNL_CHECK(cnnlPoolingForward( + handle_.get(), pool_.get(), &one, input_.get(), input.data(), &zero, + output_.get(), out.data(), scratch, workspace_size_)); + } + + private: + static cnnl_utils::TensorDescriptor Describe(const Tensor& tensor) { + auto desc = cnnl_utils::CreateTensorDescriptor(); + const int dims[] = {static_cast(tensor.size(0)), + static_cast(tensor.size(1)), + tensor.ndim() == 4 ? static_cast(tensor.size(2)) : 1, + static_cast(tensor.size(tensor.ndim() - 1))}; + INFINI_OPS_CNNL_CHECK(cnnlSetTensorDescriptor( + desc.get(), CNNL_LAYOUT_NCHW, cnnl_utils::GetDataType(tensor.dtype()), 4, dims)); + return desc; + } + + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor input_, output_; + std::unique_ptr, + decltype(&cnnlDestroyPoolingDescriptor)> + pool_{nullptr, cnnlDestroyPoolingDescriptor}; + cnrt_utils::DeviceBuffer scratch_; + std::size_t workspace_size_ = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/maximum/kernel.h b/src/native/cambricon/ops/maximum/kernel.h new file mode 100644 index 000000000..178dd2d57 --- /dev/null +++ b/src/native/cambricon/ops/maximum/kernel.h @@ -0,0 +1,51 @@ +#ifndef INFINI_OPS_CAMBRICON_MAXIMUM_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_MAXIMUM_KERNEL_H_ + +#include "base/maximum.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops { + +template <> +class Operator : public Maximum { + public: + Operator(const Tensor input, const Tensor other, Tensor out) + : Maximum{input, other, out} { + if (input.dtype() != other.dtype() || input.dtype() != out.dtype() || + input.device() != other.device() || input.device() != out.device() || + !input.IsContiguous() || !other.IsContiguous() || !out.IsContiguous() || + (out.dtype() != DataType::kFloat32 && out.dtype() != DataType::kFloat16)) + throw std::invalid_argument("Cambricon Maximum requires matching FP32 or FP16 tensors"); + if (out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + lhs_ = cnnl_utils::MakeTensorDescriptor(input); + rhs_ = cnnl_utils::MakeTensorDescriptor(other); + out_ = cnnl_utils::MakeTensorDescriptor(out); + INFINI_OPS_CNNL_CHECK(cnnlGetMaximumWorkspaceSize(handle_.get(), out_.get(), &workspace_size_)); + scratch_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_); + if (workspace_size_ && !scratch_) + throw std::runtime_error("Cambricon Maximum workspace allocation failed"); + } + + std::size_t workspace_size_in_bytes() const override { return workspace_size_; } + + void operator()(const Tensor input, const Tensor other, Tensor out) const override { + if (!out_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + auto* scratch = workspace_ && workspace_size_in_bytes_ >= workspace_size_ + ? workspace_ : scratch_.get(); + INFINI_OPS_CNNL_CHECK(cnnlMaximum(handle_.get(), lhs_.get(), input.data(), + rhs_.get(), other.data(), out_.get(), out.data(), scratch, workspace_size_)); + } + + private: + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor lhs_, rhs_, out_; + cnrt_utils::DeviceBuffer scratch_; + std::size_t workspace_size_ = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/mean/kernel.h b/src/native/cambricon/ops/mean/kernel.h new file mode 100644 index 000000000..a30ee71a3 --- /dev/null +++ b/src/native/cambricon/ops/mean/kernel.h @@ -0,0 +1,45 @@ +#ifndef INFINI_OPS_CAMBRICON_MEAN_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_MEAN_KERNEL_H_ + +#include "base/mean.h" +#include "native/cambricon/reduce_cnnl.h" + +namespace infini::ops { + +template <> +class Operator : public Mean { + public: + Operator(const Tensor input, const std::optional> dim, + const bool keepdim, const std::optional dtype, Tensor out) + : Mean{input, dim, keepdim, dtype, out}, + plan_{ValidateMode(input, dtype), input, + dim.value_or(std::vector{}), keepdim, out} {} + + Operator(const Tensor input, const std::optional dtype, Tensor out) + : Operator{input, std::nullopt, false, dtype, out} {} + + std::size_t workspace_size_in_bytes() const override { return plan_.WorkspaceSize(); } + + void operator()(const Tensor input, const std::optional>, + const bool, const std::optional, Tensor out) const override { + (*this)(input, std::nullopt, out); + } + + void operator()(const Tensor input, const std::optional, Tensor out) const override { + plan_.Run(input, out, static_cast(stream_), + workspace_, workspace_size_in_bytes_); + } + + private: + static cnnlReduceOp_t ValidateMode(const Tensor& input, std::optional dtype) { + if (dtype && *dtype != input.dtype()) + throw std::invalid_argument("Cambricon Mean does not support mixed output dtype"); + return CNNL_REDUCE_AVG; + } + + cnnl_utils::ReducePlan plan_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/mul/kernel.h b/src/native/cambricon/ops/mul/kernel.h index 6cd9090eb..4d3352fac 100644 --- a/src/native/cambricon/ops/mul/kernel.h +++ b/src/native/cambricon/ops/mul/kernel.h @@ -9,6 +9,7 @@ #include "native/cambricon/cnrt_utils.h" #include "native/cambricon/common.h" #include "native/cambricon/data_type_.h" +#include "native/cambricon/op_tensor.h" namespace infini::ops { @@ -24,9 +25,16 @@ 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())} { + : Mul{input, other, out} { + if (!output_size_) return; + if ((out_type_ == DataType::kFloat32 || out_type_ == DataType::kFloat16) && + input.dtype() == out_type_ && other.dtype() == out_type_ && + input.device() == out.device() && other.device() == out.device() && + input.IsContiguous() && other.IsContiguous() && out.IsContiguous()) { + cnnl_ = std::make_unique(CNNL_OP_TENSOR_MUL, input, other, out); + return; + } + default_workspace_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_in_bytes()); cnrt_utils::GetLaunchConfig(input.device(), &core_per_cluster_, &cluster_count_); } @@ -41,10 +49,14 @@ class Operator : public Mul { is_out_contiguous_ && input.shape() == out.shape() && other.shape() == out.shape(); auto queue = static_cast(stream_ ? stream_ : 0); + if (cnnl_) { + cnnl_->Run(input, other, out, 1.0f, queue, workspace_, workspace_size_in_bytes_); + return; + } void* workspace = workspace_ ? workspace_ : default_workspace_.get(); using SupportedTypes = - List; DispatchFunc( @@ -63,10 +75,13 @@ class Operator : public Mul { } std::size_t workspace_size_in_bytes() const override { + if (cnnl_) return cnnl_->WorkspaceSize(); + if (!output_size_) return 0; return ndim_ * (sizeof(std::size_t) + 3 * sizeof(ptrdiff_t)); } private: + std::unique_ptr cnnl_; cnrt_utils::DeviceBuffer default_workspace_{}; int core_per_cluster_{0}; diff --git a/src/native/cambricon/ops/mul/kernel.mlu b/src/native/cambricon/ops/mul/kernel.mlu index 10719c3dc..988e4c266 100644 --- a/src/native/cambricon/ops/mul/kernel.mlu +++ b/src/native/cambricon/ops/mul/kernel.mlu @@ -14,7 +14,7 @@ __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) { + if constexpr (std::is_same_v) { __bang_mul(reinterpret_cast(out), reinterpret_cast(input), reinterpret_cast(other), count); @@ -136,7 +136,9 @@ void MulUnion(void* workspace, cnrtQueue_t queue, int core_per_cluster, int, bool, bool) INSTANTIATE_MUL(Float16); +#if INFINIOPS_CAMBRICON_BANG_BF16 INSTANTIATE_MUL(BFloat16); +#endif INSTANTIATE_MUL(float); INSTANTIATE_MUL(int16_t); INSTANTIATE_MUL(int32_t); diff --git a/src/native/cambricon/ops/neg/kernel.h b/src/native/cambricon/ops/neg/kernel.h new file mode 100644 index 000000000..e7bc087df --- /dev/null +++ b/src/native/cambricon/ops/neg/kernel.h @@ -0,0 +1,18 @@ +#ifndef INFINI_OPS_CAMBRICON_NEG_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_NEG_KERNEL_H_ + +#include "base/neg.h" +#include "native/cambricon/unary_cnnl.h" + +namespace infini::ops { + +template <> +class Operator + : public cnnl_utils::Unary { + public: + Operator(const Tensor input, Tensor out) : Unary{input, out} {} +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/norm/kernel.h b/src/native/cambricon/ops/norm/kernel.h new file mode 100644 index 000000000..915a3e141 --- /dev/null +++ b/src/native/cambricon/ops/norm/kernel.h @@ -0,0 +1,37 @@ +#ifndef INFINI_OPS_CAMBRICON_NORM_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_NORM_KERNEL_H_ + +#include "base/norm.h" +#include "native/cambricon/reduce_cnnl.h" + +namespace infini::ops { + +template <> +class Operator : public Norm { + public: + Operator(const Tensor input, std::optional p, + std::vector dim, bool keepdim, Tensor out) + : Norm{input, p, dim, keepdim, out}, + plan_{ValidateMode(p), input, dim, keepdim, out} {} + + std::size_t workspace_size_in_bytes() const override { return plan_.WorkspaceSize(); } + + void operator()(const Tensor input, std::optional, + std::vector, bool, Tensor out) const override { + plan_.Run(input, out, static_cast(stream_), + workspace_, workspace_size_in_bytes_); + } + + private: + static cnnlReduceOp_t ValidateMode(std::optional p) { + if (p.value_or(2.0) != 2.0) + throw std::invalid_argument("Cambricon Norm supports L2 only"); + return CNNL_REDUCE_NORM2; + } + + cnnl_utils::ReducePlan plan_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/pow/kernel.h b/src/native/cambricon/ops/pow/kernel.h new file mode 100644 index 000000000..8e2337d02 --- /dev/null +++ b/src/native/cambricon/ops/pow/kernel.h @@ -0,0 +1,71 @@ +#ifndef INFINI_OPS_CAMBRICON_POW_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_POW_KERNEL_H_ + +#include "base/pow.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops { + +template <> +class Operator : public Pow { + public: + Operator(const Tensor input, const Tensor exponent, Tensor out) + : Pow{input, exponent, out} { + if (input.dtype() != out.dtype() || exponent.dtype() != out.dtype() || + input.device() != out.device() || exponent.device() != out.device() || + !input.IsContiguous() || !exponent.IsContiguous() || !out.IsContiguous() || + (out.dtype() != DataType::kFloat32 && out.dtype() != DataType::kFloat16)) + throw std::invalid_argument("Cambricon Pow requires matching FP32 or FP16 tensors"); + if (out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + lhs_ = cnnl_utils::MakeTensorDescriptor(input); + rhs_ = cnnl_utils::MakeTensorDescriptor(exponent); + out_ = cnnl_utils::MakeTensorDescriptor(out); + INFINI_OPS_CNNL_CHECK(cnnlGetPowWorkspaceSize( + handle_.get(), lhs_.get(), rhs_.get(), out_.get(), &workspace_size_)); + scratch_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_); + if (workspace_size_ && !scratch_) + throw std::runtime_error("Cambricon Pow workspace allocation failed"); + } + + Operator(const Tensor input, const double exponent, Tensor out) + : Pow{input, exponent, out} { Unsupported(); } + Operator(Tensor input, const Tensor exponent) + : Operator{input, exponent, input} {} + Operator(Tensor input, const double exponent) + : Pow{input, exponent, input} { Unsupported(); } + + std::size_t workspace_size_in_bytes() const override { return workspace_size_; } + + void operator()(const Tensor input, const Tensor exponent, Tensor out) const override { + if (!out_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + auto* scratch = workspace_ && workspace_size_in_bytes_ >= workspace_size_ + ? workspace_ : scratch_.get(); + INFINI_OPS_CNNL_CHECK(cnnlPow_v2( + handle_.get(), lhs_.get(), input.data(), rhs_.get(), exponent.data(), + scratch, workspace_size_, out_.get(), out.data())); + } + + void operator()(const Tensor, const double, Tensor) const override { Unsupported(); } + void operator()(Tensor, const double) const override { Unsupported(); } + void operator()(Tensor input, const Tensor exponent) const override { + (*this)(input, exponent, input); + } + + private: + [[noreturn]] static void Unsupported() { + throw std::invalid_argument("Cambricon Pow requires a tensor exponent"); + } + + + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor lhs_, rhs_, out_; + cnrt_utils::DeviceBuffer scratch_; + std::size_t workspace_size_ = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/relu/kernel.mlu b/src/native/cambricon/ops/relu/kernel.mlu index beaeb3516..f07c5eb15 100644 --- a/src/native/cambricon/ops/relu/kernel.mlu +++ b/src/native/cambricon/ops/relu/kernel.mlu @@ -12,7 +12,7 @@ __nram__ char relu_nram_buffer[NRAM_MAX_SIZE] __attribute__((aligned(128))); template __mlu_device__ void ComputeRelu(const T* input, T* output, float* values, size_t count) { - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { __bang_half2float(values, reinterpret_cast(const_cast(input)), count); for (size_t i = 0; i < count; ++i) { diff --git a/src/native/cambricon/ops/resize2d/kernel.h b/src/native/cambricon/ops/resize2d/kernel.h new file mode 100644 index 000000000..60a6c9950 --- /dev/null +++ b/src/native/cambricon/ops/resize2d/kernel.h @@ -0,0 +1,152 @@ +#ifndef INFINI_OPS_CAMBRICON_RESIZE2D_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_RESIZE2D_KERNEL_H_ + +#include +#include + +#include "base/resize2d.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops { + +template <> +class Operator : public Resize2d { + public: + Operator(const Tensor input, int64_t mode, int64_t coordinates, + int64_t rounding, float height_scale, float width_scale, Tensor out) + : Resize2d{input, mode, coordinates, rounding, height_scale, width_scale, out} { + if (input.ndim() != 4 || out.ndim() != 4 || input.dtype() != out.dtype() || + (out.dtype() != DataType::kFloat32 && out.dtype() != DataType::kFloat16) || + input.device() != out.device() || !input.IsContiguous() || !out.IsContiguous() || + input.size(0) != out.size(0) || input.size(1) != out.size(1) || + input.size(2) == 0 || input.size(3) == 0 || + mode < kNearest || mode > kLinear || coordinates < kHalfPixel || + coordinates > kAsymmetric || rounding < kPreferFloor || rounding > kCeil || + !std::isfinite(height_scale) || !std::isfinite(width_scale) || + height_scale < 0 || width_scale < 0 || ((height_scale == 0) != (width_scale == 0))) + throw std::invalid_argument("Unsupported Cambricon Resize2d parameters"); + if (out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + nearest2x_ = mode == kNearest && coordinates == kAsymmetric && rounding == kFloor && + (height_scale == 0 || (height_scale == 2 && width_scale == 2)) && + out.size(2) == 2 * input.size(2) && out.size(3) == 2 * input.size(3); + if (nearest2x_) { + const auto n = static_cast(input.size(0)), c = static_cast(input.size(1)); + const auto h = static_cast(input.size(2)), w = static_cast(input.size(3)); + input_ = DescribeArray(input.dtype(), {n, c, h, w, 1}); + input_nhwc_ = DescribeArray(input.dtype(), {n, c, h, w, 2}); + out_nhwc_ = DescribeArray(input.dtype(), {n, c, h, 1, w * 2}); + out_ = DescribeArray(input.dtype(), {n, c, h, 2, w * 2}); + input_buffer_ = Allocate(input.numel() * input.element_size() * 2); + return; + } + static constexpr std::array coordinates_map{ + CNNL_INTERP_ALGO_9, CNNL_INTERP_ALGO_0, CNNL_INTERP_ALGO_2, CNNL_INTERP_ALGO_8}; + static constexpr std::array rounding_map{ + CNNL_INTERP_ROUND_PERFER_FLOOR, CNNL_INTERP_ROUND_PERFER_CEIL, + CNNL_INTERP_FLOOR, CNNL_INTERP_CEIL}; + input_ = Describe(input, false); + input_nhwc_ = Describe(input, true); + out_ = Describe(out, false); + out_nhwc_ = Describe(out, true); + SetupTranspose(to_nhwc_, {0, 2, 3, 1}); + SetupTranspose(to_nchw_, {0, 3, 1, 2}); + cnnlInterpDescriptor_t interp = nullptr; + INFINI_OPS_CNNL_CHECK(cnnlCreateInterpDescriptor(&interp)); + interp_.reset(interp); + const double scales[2] = {height_scale, width_scale}; + INFINI_OPS_CNNL_CHECK(cnnlSetInterpDescriptor_v2( + interp, input_nhwc_.get(), mode == kLinear ? CNNL_INTERP_BILINEAR : CNNL_INTERP_NEAREST, + coordinates_map[coordinates], height_scale > 0 ? scales : nullptr)); + INFINI_OPS_CNNL_CHECK(cnnlSetInterpDescriptorEx( + interp, input_nhwc_.get(), rounding_map[rounding], nullptr, nullptr, -0.75f, false)); + INFINI_OPS_CNNL_CHECK(cnnlGetTransposeWorkspaceSize( + handle_.get(), input_.get(), to_nhwc_.get(), &input_workspace_)); + INFINI_OPS_CNNL_CHECK(cnnlGetTransposeWorkspaceSize( + handle_.get(), out_nhwc_.get(), to_nchw_.get(), &output_workspace_)); + scratch_size_ = std::max(input_workspace_, output_workspace_); + input_buffer_ = Allocate(input.numel() * input.element_size()); + output_buffer_ = Allocate(out.numel() * out.element_size()); + scratch_ = Allocate(scratch_size_); + } + + std::size_t workspace_size_in_bytes() const override { return scratch_size_; } + + void operator()(const Tensor input, int64_t, int64_t, int64_t, float, float, + Tensor out) const override { + if (!out_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + if (nearest2x_) { + INFINI_OPS_CNNL_CHECK(cnnlTile(handle_.get(), input_.get(), input.data(), + input_nhwc_.get(), input_buffer_.get())); + INFINI_OPS_CNNL_CHECK(cnnlTile(handle_.get(), out_nhwc_.get(), input_buffer_.get(), + out_.get(), out.data())); + return; + } + auto* scratch = workspace_ && workspace_size_in_bytes_ >= scratch_size_ + ? workspace_ : scratch_.get(); + INFINI_OPS_CNNL_CHECK(cnnlTranspose_v2( + handle_.get(), to_nhwc_.get(), input_.get(), input.data(), + input_nhwc_.get(), input_buffer_.get(), scratch, input_workspace_)); + INFINI_OPS_CNNL_CHECK(cnnlInterp_v3( + handle_.get(), interp_.get(), input_nhwc_.get(), input_buffer_.get(), + out_nhwc_.get(), output_buffer_.get())); + INFINI_OPS_CNNL_CHECK(cnnlTranspose_v2( + handle_.get(), to_nchw_.get(), out_nhwc_.get(), output_buffer_.get(), + out_.get(), out.data(), scratch, output_workspace_)); + } + + private: + using Transpose = std::unique_ptr, + decltype(&cnnlDestroyTransposeDescriptor)>; + static cnnl_utils::TensorDescriptor DescribeArray( + DataType dtype, std::initializer_list dims) { + auto desc = cnnl_utils::CreateTensorDescriptor(); + INFINI_OPS_CNNL_CHECK(cnnlSetTensorDescriptor_v2( + desc.get(), CNNL_LAYOUT_ARRAY, cnnl_utils::GetDataType(dtype), dims.size(), dims.begin())); + return desc; + } + static cnnl_utils::TensorDescriptor Describe(const Tensor& tensor, bool nhwc) { + auto desc = cnnl_utils::CreateTensorDescriptor(); + std::array dims{}; + if (nhwc) + dims = {static_cast(tensor.size(0)), static_cast(tensor.size(2)), + static_cast(tensor.size(3)), static_cast(tensor.size(1))}; + else + for (int i = 0; i < 4; ++i) dims[i] = tensor.size(i); + INFINI_OPS_CNNL_CHECK(cnnlSetTensorDescriptor_v2( + desc.get(), nhwc ? CNNL_LAYOUT_NHWC : CNNL_LAYOUT_NCHW, + cnnl_utils::GetDataType(tensor.dtype()), dims.size(), dims.data())); + return desc; + } + + static void SetupTranspose(Transpose& desc, const std::array& permutation) { + cnnlTransposeDescriptor_t value = nullptr; + INFINI_OPS_CNNL_CHECK(cnnlCreateTransposeDescriptor(&value)); + desc.reset(value); + INFINI_OPS_CNNL_CHECK(cnnlSetTransposeDescriptor(value, permutation.size(), permutation.data())); + } + + static cnrt_utils::DeviceBuffer Allocate(size_t bytes) { + auto result = cnrt_utils::AllocateDeviceBuffer(bytes); + if (bytes && !result) + throw std::runtime_error("Cambricon Resize2d allocation failed"); + return result; + } + + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor input_, input_nhwc_, out_, out_nhwc_; + Transpose to_nhwc_{nullptr, cnnlDestroyTransposeDescriptor}; + Transpose to_nchw_{nullptr, cnnlDestroyTransposeDescriptor}; + std::unique_ptr, + decltype(&cnnlDestroyInterpDescriptor)> + interp_{nullptr, cnnlDestroyInterpDescriptor}; + cnrt_utils::DeviceBuffer input_buffer_, output_buffer_, scratch_; + std::size_t input_workspace_ = 0, output_workspace_ = 0, scratch_size_ = 0; + bool nearest2x_ = false; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/rms_norm/kernel.mlu b/src/native/cambricon/ops/rms_norm/kernel.mlu index 4d225039f..bc434c205 100644 --- a/src/native/cambricon/ops/rms_norm/kernel.mlu +++ b/src/native/cambricon/ops/rms_norm/kernel.mlu @@ -4,11 +4,23 @@ __nram__ char nram_buffer[NRAM_MAX_SIZE]; namespace infini::ops { +struct RmsNormMetadata { + size_t shape[kRmsNormInlineRank]; + ptrdiff_t output_strides[kRmsNormInlineRank]; + ptrdiff_t input_strides[kRmsNormInlineRank]; +}; + template __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output, - size_t* shape, ptrdiff_t* output_strides, - ptrdiff_t* input_strides, float epsilon, + RmsNormMetadata inline_metadata, const size_t* shape, + const ptrdiff_t* output_strides, + const ptrdiff_t* input_strides, float epsilon, int num_dims, int norm_dim_size) { + if (num_dims <= kRmsNormInlineRank) { + shape = inline_metadata.shape; + output_strides = inline_metadata.output_strides; + input_strides = inline_metadata.input_strides; + } // Calculate problem dimensions. int batch_volume = 1; for (int dim_idx = 0; dim_idx < num_dims - 1; ++dim_idx) { @@ -68,7 +80,7 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output, if (vector_size <= 128) { __memcpy(input_cache, input + input_offset, vector_size * sizeof(T), GDRAM2NRAM); - if constexpr (std::is_same::value) { + if constexpr (std::is_same::value) { __bang_half2float(float_buffer, reinterpret_cast(input_cache), vector_size); } else if constexpr (std::is_same::value) { @@ -98,7 +110,7 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output, processed_elements * input_strides[num_dims - 1], current_batch * sizeof(T), GDRAM2NRAM); - if constexpr (std::is_same::value) { + if constexpr (std::is_same::value) { __bang_half2float(float_buffer, reinterpret_cast(input_cache), current_batch); } else if constexpr (std::is_same::value) { @@ -136,7 +148,7 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output, GDRAM2NRAM); __memcpy(weight_cache, weight, vector_size * sizeof(TW), GDRAM2NRAM); - if constexpr (std::is_same::value) { + if constexpr (std::is_same::value) { __bang_half2float(float_buffer, reinterpret_cast(input_cache), vector_size); } else if constexpr (std::is_same::value) { @@ -146,7 +158,7 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output, NRAM2NRAM); } - if constexpr (std::is_same::value) { + if constexpr (std::is_same::value) { __bang_half2float(weight_float_buffer, reinterpret_cast(weight_cache), vector_size); } else if constexpr (std::is_same::value) { @@ -160,7 +172,7 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output, __bang_mul(float_buffer, float_buffer, weight_float_buffer, vector_size); __bang_mul_scalar(float_buffer, float_buffer, inv_rms, vector_size); - if constexpr (std::is_same::value) { + if constexpr (std::is_same::value) { __bang_float2half(reinterpret_cast(input_cache), float_buffer, vector_size); } else if constexpr (std::is_same::value) { @@ -187,7 +199,7 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output, __memcpy(weight_cache, weight + processed_elements, current_batch * sizeof(TW), GDRAM2NRAM); - if constexpr (std::is_same::value) { + if constexpr (std::is_same::value) { __bang_half2float(float_buffer, reinterpret_cast(input_cache), current_batch); } else if constexpr (std::is_same::value) { @@ -197,7 +209,7 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output, NRAM2NRAM); } - if constexpr (std::is_same::value) { + if constexpr (std::is_same::value) { __bang_half2float(weight_float_buffer, reinterpret_cast(weight_cache), current_batch); @@ -213,7 +225,7 @@ __mlu_global__ void RmsNorm(const T* input, const TW* weight, T* output, current_batch); __bang_mul_scalar(float_buffer, float_buffer, inv_rms, current_batch); - if constexpr (std::is_same::value) { + if constexpr (std::is_same::value) { __bang_float2half(reinterpret_cast(input_cache), float_buffer, current_batch); } else if constexpr (std::is_same::value) { @@ -269,6 +281,18 @@ void RmsNormUnion(void* workspace, int core_per_cluster, int cluster_count, auto y_ = reinterpret_cast(y); auto x_ = reinterpret_cast(x); auto w_ = reinterpret_cast(w); + if (ndim <= kRmsNormInlineRank) { + // By-value metadata requires no copies and can be recorded for replay. + RmsNormMetadata metadata{}; + for (int i = 0; i < ndim; ++i) { + metadata.shape[i] = shape[i]; + metadata.output_strides[i] = y_strides[i]; + metadata.input_strides[i] = x_strides[i]; + } + RmsNorm<<>>( + x_, w_, y_, metadata, nullptr, nullptr, nullptr, eps, ndim, dim_s); + return; + } char* tmp_device = reinterpret_cast(workspace); char* tmp_stride = tmp_device + ndim * sizeof(size_t); size_t* mlu_shape = (size_t*)tmp_device; @@ -287,7 +311,8 @@ void RmsNormUnion(void* workspace, int core_per_cluster, int cluster_count, cnrtMemcpyHostToDev)); RmsNorm<<>>( - x_, w_, y_, mlu_shape, mlu_y_strides, mlu_x_strides, eps, ndim, dim_s); + x_, w_, y_, RmsNormMetadata{}, mlu_shape, mlu_y_strides, mlu_x_strides, + eps, ndim, dim_s); } template void RmsNormUnion(void*, int, int, cnrtQueue_t, diff --git a/src/native/cambricon/ops/rms_norm/rms_norm.h b/src/native/cambricon/ops/rms_norm/rms_norm.h index f8e70a622..e2508629d 100644 --- a/src/native/cambricon/ops/rms_norm/rms_norm.h +++ b/src/native/cambricon/ops/rms_norm/rms_norm.h @@ -7,10 +7,13 @@ #include "base/rms_norm.h" #include "native/cambricon/common.h" +#include "native/cambricon/cnrt_utils.h" #include "native/cambricon/data_type_.h" namespace infini::ops { +inline constexpr int kRmsNormInlineRank = 8; + // TODO: Remove forward declaration. template void RmsNormUnion(void* workspace, int core_per_cluster, int cluster_count, @@ -31,13 +34,14 @@ class Operator : public RmsNorm { : RmsNorm{input, weight, eps, out} { cnrt_utils::GetLaunchConfig(input.device(), &core_per_cluster, &cluster_count); - cnrtMalloc(&default_workspace_, workspace_size_in_bytes()); + default_workspace_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_in_bytes()); } void operator()(const Tensor input, const Tensor weight, float eps, Tensor out) const override { auto queue = static_cast(stream_ ? stream_ : 0); - auto workspace{workspace_ ? workspace_ : default_workspace_}; + auto workspace{workspace_ && workspace_size_in_bytes_ >= workspace_size_in_bytes() + ? workspace_ : default_workspace_.get()}; DispatchFunc< Device::Type::kCambricon, @@ -56,13 +60,13 @@ class Operator : public RmsNorm { "CambriconRmsNorm::operator() - output dispatch"); } - ~Operator() { cnrtFree(default_workspace_); } - std::size_t workspace_size_in_bytes() const override { + if (ndim_ <= kRmsNormInlineRank) return 0; return ndim_ * (sizeof(size_t) + 2 * sizeof(ptrdiff_t)); } - void* default_workspace_{nullptr}; + private: + cnrt_utils::DeviceBuffer default_workspace_; int core_per_cluster = 0; int cluster_count = 0; }; diff --git a/src/native/cambricon/ops/roi_align/kernel.h b/src/native/cambricon/ops/roi_align/kernel.h new file mode 100644 index 000000000..58012fdfb --- /dev/null +++ b/src/native/cambricon/ops/roi_align/kernel.h @@ -0,0 +1,72 @@ +#ifndef INFINI_OPS_CAMBRICON_ROI_ALIGN_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_ROI_ALIGN_KERNEL_H_ + +#include +#include + +#include "base/roi_align.h" +#include "native/cambricon/cnnl_utils.h" + +namespace infini::ops { + +template <> +class Operator : public RoiAlign { + public: + Operator(const Tensor input, const Tensor boxes, int64_t sampling_ratio, + float spatial_scale, bool aligned, Tensor out) + : RoiAlign{input, boxes, sampling_ratio, spatial_scale, aligned, out} { + if (input.ndim() != 4 || boxes.ndim() != 2 || out.ndim() != 4 || + boxes.size(1) != 5 || out.size(0) != boxes.size(0) || + out.size(3) != input.size(3) || out.size(1) == 0 || out.size(2) == 0 || + out.size(1) > std::numeric_limits::max() || + out.size(2) > std::numeric_limits::max() || + sampling_ratio < 0 || sampling_ratio > std::numeric_limits::max() || + !std::isfinite(spatial_scale) || spatial_scale <= 0 || + input.device() != out.device() || boxes.device() != out.device() || + input.dtype() != out.dtype() || boxes.dtype() != out.dtype() || + (out.dtype() != DataType::kFloat32 && out.dtype() != DataType::kFloat16) || + !input.IsContiguous() || !boxes.IsContiguous() || !out.IsContiguous()) + throw std::invalid_argument("Unsupported Cambricon RoiAlign parameters"); + if (out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + input_ = Describe(input, CNNL_LAYOUT_NHWC); + boxes_ = Describe(boxes, CNNL_LAYOUT_ARRAY); + out_ = Describe(out, CNNL_LAYOUT_NHWC); + cnnlRoiAlignDescriptor_t desc = nullptr; + INFINI_OPS_CNNL_CHECK(cnnlCreateRoiAlignDescriptor(&desc)); + roi_.reset(desc); + INFINI_OPS_CNNL_CHECK(cnnlSetRoiAlignDescriptor_v2( + desc, out.size(1), out.size(2), sampling_ratio, spatial_scale, 1, aligned)); + } + + void operator()(const Tensor input, const Tensor boxes, int64_t, float, bool, + Tensor out) const override { + if (!out_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + INFINI_OPS_CNNL_CHECK(cnnlRoiAlign_v2( + handle_.get(), roi_.get(), input_.get(), input.data(), boxes_.get(), + boxes.data(), out_.get(), out.data(), nullptr, nullptr, nullptr, nullptr)); + } + + private: + static cnnl_utils::TensorDescriptor Describe(const Tensor& tensor, + cnnlTensorLayout_t layout) { + auto desc = cnnl_utils::CreateTensorDescriptor(); + const auto shape = tensor.shape(); + const std::vector dims(shape.begin(), shape.end()); + INFINI_OPS_CNNL_CHECK(cnnlSetTensorDescriptor_v2( + desc.get(), layout, cnnl_utils::GetDataType(tensor.dtype()), + static_cast(dims.size()), dims.data())); + return desc; + } + + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor input_, boxes_, out_; + std::unique_ptr, + decltype(&cnnlDestroyRoiAlignDescriptor)> + roi_{nullptr, cnnlDestroyRoiAlignDescriptor}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/rotary_embedding/kernel.mlu b/src/native/cambricon/ops/rotary_embedding/kernel.mlu index 416ecd9aa..61c197bba 100644 --- a/src/native/cambricon/ops/rotary_embedding/kernel.mlu +++ b/src/native/cambricon/ops/rotary_embedding/kernel.mlu @@ -12,7 +12,7 @@ __nram__ char rotary_embedding_nram_buffer[NRAM_MAX_SIZE]; template __mlu_device__ void ToFloat(float* output, const T* input, size_t count) { - if constexpr (std::is_same::value) { + if constexpr (std::is_same::value) { __bang_half2float(output, reinterpret_cast(input), count); } else if constexpr (std::is_same::value) { __bang_bfloat162float(output, input, count); @@ -23,7 +23,7 @@ __mlu_device__ void ToFloat(float* output, const T* input, size_t count) { template __mlu_device__ void FromFloat(T* output, const float* input, size_t count) { - if constexpr (std::is_same::value) { + if constexpr (std::is_same::value) { __bang_float2half(reinterpret_cast(output), input, count); } else if constexpr (std::is_same::value) { __bang_float2bfloat16(output, input, count); diff --git a/src/native/cambricon/ops/rsqrt/kernel.h b/src/native/cambricon/ops/rsqrt/kernel.h new file mode 100644 index 000000000..7cde86865 --- /dev/null +++ b/src/native/cambricon/ops/rsqrt/kernel.h @@ -0,0 +1,18 @@ +#ifndef INFINI_OPS_CAMBRICON_RSQRT_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_RSQRT_KERNEL_H_ + +#include "base/rsqrt.h" +#include "native/cambricon/unary_cnnl.h" + +namespace infini::ops { + +template <> +class Operator + : public cnnl_utils::Unary { + public: + Operator(const Tensor input, Tensor out) : Unary{input, out} {} +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/scatter/kernel.h b/src/native/cambricon/ops/scatter/kernel.h new file mode 100644 index 000000000..9bb2323d8 --- /dev/null +++ b/src/native/cambricon/ops/scatter/kernel.h @@ -0,0 +1,115 @@ +#ifndef INFINI_OPS_CAMBRICON_SCATTER_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_SCATTER_KERNEL_H_ + +#include "base/scatter.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" +#include "native/cambricon/scatter_index.h" + +namespace infini::ops { + +template <> +class Operator : public Scatter { + public: + Operator(const Tensor input, const Tensor index, const Tensor src, + const int64_t dim, Tensor out) + : Operator{input, index, src, dim, "none", out} {} + + Operator(const Tensor input, const Tensor index, const Tensor src, + const int64_t dim, const std::string reduce, Tensor out) + : Scatter{input, index, src, dim, reduce, out}, mode_(Mode(reduce)) { + const auto rank = static_cast(input.ndim()); + axis_ = dim < 0 ? dim + rank : dim; + if (axis_ < 0 || axis_ >= rank || input.shape() != out.shape() || + index.shape() != src.shape() || input.ndim() != index.ndim() || + input.dtype() != out.dtype() || src.dtype() != out.dtype() || + (index.dtype() != DataType::kInt32 && index.dtype() != DataType::kInt64) || + input.device() != out.device() || index.device() != out.device() || + src.device() != out.device() || !input.IsContiguous() || + !out.IsContiguous() || (index.numel() && (!index.IsContiguous() || !src.IsContiguous()))) + throw std::invalid_argument("Unsupported Cambricon Scatter tensor metadata"); + for (int64_t i = 0; i < rank; ++i) { + if (i != axis_ && index.size(i) > input.size(i)) + throw std::invalid_argument("Cambricon Scatter index extent exceeds input"); + } + if (index.numel() == 0 || out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + input_ = cnnl_utils::MakeTensorDescriptor(input); + index_ = cnnl_utils::MakeTensorDescriptor(index); + source_ = cnnl_utils::MakeTensorDescriptor(src); + output_ = cnnl_utils::MakeTensorDescriptor(out); + normalized_ = cnrt_utils::AllocateDeviceBuffer(index.numel() * index.element_size()); + INFINI_OPS_CNNL_CHECK(cnnlGetScatterWorkspaceSize( + handle_.get(), index_.get(), &workspace_size_)); + scratch_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_); + if (!normalized_ || (workspace_size_ && !scratch_)) + throw std::runtime_error("Cambricon Scatter allocation failed"); + } + + Operator(const Tensor input, const Tensor index, const double value, + const int64_t dim, Tensor out) + : Scatter{input, index, value, dim, out} { UnsupportedScalar(); } + Operator(const Tensor input, const Tensor index, const double value, + const int64_t dim, const std::string reduce, Tensor out) + : Scatter{input, index, value, dim, reduce, out} { UnsupportedScalar(); } + + std::size_t workspace_size_in_bytes() const override { return workspace_size_; } + + void operator()(const Tensor input, const Tensor index, const Tensor src, + const int64_t dim, Tensor out) const override { + (*this)(input, index, src, dim, "none", out); + } + + void operator()(const Tensor input, const Tensor index, const Tensor src, + const int64_t, const std::string, Tensor out) const override { + if (out.numel() == 0) return; + auto queue = static_cast(stream_); + if (index.numel() == 0) { + if (input.data() != out.data() && + cnrtMemcpyAsync(out.data(), const_cast(input.data()), + out.numel() * out.element_size(), queue, + cnrtMemcpyDevToDev) != cnrtSuccess) + throw std::runtime_error("Cambricon Scatter empty index copy failed"); + return; + } + NormalizeAxisIndices(index.data(), normalized_.get(), input.size(axis_), + index.numel(), index.dtype(), queue); + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), queue)); + auto* scratch = workspace_ && workspace_size_in_bytes_ >= workspace_size_ + ? workspace_ : scratch_.get(); + INFINI_OPS_CNNL_CHECK(cnnlScatter_v2( + handle_.get(), axis_, input_.get(), input.data(), index_.get(), + normalized_.get(), source_.get(), src.data(), scratch, workspace_size_, + output_.get(), out.data(), mode_)); + } + + void operator()(const Tensor, const Tensor, const double, const int64_t, + Tensor) const override { UnsupportedScalar(); } + void operator()(const Tensor, const Tensor, const double, const int64_t, + const std::string, Tensor) const override { UnsupportedScalar(); } + + private: + [[noreturn]] static void UnsupportedScalar() { + throw std::invalid_argument("Cambricon Scatter requires tensor updates"); + } + + static cnnlScatterMode_t Mode(const std::string& reduce) { + if (reduce.empty() || reduce == "none") return CNNL_SCATTER; + if (reduce == "add") return CNNL_SCATTER_ADD; + if (reduce == "max") return CNNL_SCATTER_MAX; + if (reduce == "min") return CNNL_SCATTER_MIN; + throw std::invalid_argument("Unsupported Cambricon Scatter reduction"); + } + + + cnnlScatterMode_t mode_; + int64_t axis_ = 0; + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor input_, index_, source_, output_; + cnrt_utils::DeviceBuffer normalized_, scratch_; + std::size_t workspace_size_ = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/scatter_nd/kernel.h b/src/native/cambricon/ops/scatter_nd/kernel.h new file mode 100644 index 000000000..2fb831949 --- /dev/null +++ b/src/native/cambricon/ops/scatter_nd/kernel.h @@ -0,0 +1,95 @@ +#ifndef INFINI_OPS_CAMBRICON_SCATTER_ND_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_SCATTER_ND_KERNEL_H_ + +#include "base/scatter_nd.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" +#include "native/cambricon/scatter_index.h" + +namespace infini::ops { + +template <> +class Operator : public ScatterNd { + public: + Operator(const Tensor input, const Tensor indices, const Tensor updates, + bool add, Tensor out) + : ScatterNd{input, indices, updates, add, out}, + mode_(add ? CNNL_SCATTERND_ADD : CNNL_SCATTERND_UPDATE) { + if (indices.ndim() == 0 || input.shape() != out.shape() || + input.dtype() != out.dtype() || updates.dtype() != out.dtype() || + (indices.dtype() != DataType::kInt32 && indices.dtype() != DataType::kInt64) || + input.device() != out.device() || indices.device() != out.device() || + updates.device() != out.device() || + !input.IsContiguous() || !indices.IsContiguous() || + !updates.IsContiguous() || !out.IsContiguous()) + throw std::invalid_argument("Unsupported Cambricon ScatterNd tensors"); + width_ = indices.size(-1); + if (width_ == 0 || width_ > input.ndim()) + throw std::invalid_argument("Cambricon ScatterNd has invalid coordinate width"); + const auto index_shape = indices.shape(); + const auto input_shape = input.shape(); + const auto update_shape = updates.shape(); + std::vector expected(index_shape.begin(), index_shape.end() - 1); + expected.insert(expected.end(), input_shape.begin() + width_, input_shape.end()); + if (expected.size() != update_shape.size() || + !std::equal(expected.begin(), expected.end(), update_shape.begin())) + throw std::invalid_argument("Cambricon ScatterNd has invalid update shape"); + if (indices.numel() == 0 || out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + input_ = Describe(input, false); + indices_ = Describe(indices, indices.ndim() == 1); + updates_ = Describe(updates, indices.ndim() == 1); + output_ = Describe(out, false); + dims_ = cnrt_utils::AllocateDeviceBuffer(width_ * sizeof(int64_t)); + normalized_ = cnrt_utils::AllocateDeviceBuffer( + indices.numel() * (indices.dtype() == DataType::kInt64 ? sizeof(int64_t) : sizeof(int32_t))); + if (!dims_ || !normalized_) + throw std::runtime_error("Cambricon ScatterNd index allocation failed"); + std::vector dims(input_shape.begin(), input_shape.begin() + width_); + if (cnrtMemcpy(dims_.get(), dims.data(), width_ * sizeof(int64_t), + cnrtMemcpyHostToDev) != cnrtSuccess) + throw std::runtime_error("Cambricon ScatterNd coordinate shape upload failed"); + } + + void operator()(const Tensor input, const Tensor indices, const Tensor updates, + bool, Tensor out) const override { + if (out.numel() == 0) return; + auto queue = static_cast(stream_); + if (indices.numel() == 0) { + if (input.data() != out.data() && + cnrtMemcpyAsync(out.data(), const_cast(input.data()), + out.numel() * out.element_size(), queue, + cnrtMemcpyDevToDev) != cnrtSuccess) + throw std::runtime_error("Cambricon ScatterNd empty index copy failed"); + return; + } + NormalizeScatterCoordinates(indices.data(), normalized_.get(), + static_cast(dims_.get()), indices.numel(), width_, indices.dtype(), queue); + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), queue)); + INFINI_OPS_CNNL_CHECK(cnnlScatterNd_v2( + handle_.get(), mode_, indices_.get(), normalized_.get(), updates_.get(), + updates.data(), input_.get(), input.data(), output_.get(), out.data())); + } + + private: + static cnnl_utils::TensorDescriptor Describe(const Tensor& tensor, bool prepend) { + auto desc = cnnl_utils::CreateTensorDescriptor(); + const auto shape = tensor.shape(); + std::vector dims(prepend ? 1 : 0, 1); + dims.insert(dims.end(), shape.begin(), shape.end()); + INFINI_OPS_CNNL_CHECK(cnnlSetTensorDescriptor_v2( + desc.get(), CNNL_LAYOUT_ARRAY, cnnl_utils::GetDataType(tensor.dtype()), + static_cast(dims.size()), dims.data())); + return desc; + } + + cnnlScatterNdMode_t mode_; + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor input_, indices_, updates_, output_; + cnrt_utils::DeviceBuffer dims_, normalized_; + size_t width_ = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/scatter_nd/kernel.mlu b/src/native/cambricon/ops/scatter_nd/kernel.mlu new file mode 100644 index 000000000..74376049e --- /dev/null +++ b/src/native/cambricon/ops/scatter_nd/kernel.mlu @@ -0,0 +1,73 @@ +#include "kernel.h" + +namespace infini::ops { +namespace { + +template +__mlu_global__ void NormalizeKernel(const T* input, T* output, + const int64_t* dims, size_t count, + size_t width, uint32_t axis_dim) { + constexpr size_t tile = 1024; + constexpr size_t max_cached_dims = 16; + __nram__ T values[tile]; + __nram__ int64_t cached_dims[max_cached_dims]; + if (dims && width <= max_cached_dims) + __memcpy(cached_dims, dims, width * sizeof(int64_t), GDRAM2NRAM); + const size_t per_task = (count + taskDim - 1) / taskDim; + const size_t begin = taskId * per_task; + const size_t end = std::min(begin + per_task, count); + for (size_t offset = begin; offset < end; offset += tile) { + const size_t length = std::min(tile, end - offset); + __memcpy(values, input + offset, length * sizeof(T), GDRAM2NRAM); + for (size_t j = 0; j < length; ++j) { + uint32_t dim = axis_dim; + if (dims) { + const auto d = (offset + j) % width; + dim = static_cast(width <= max_cached_dims ? cached_dims[d] : dims[d]); + } + if constexpr (std::is_same_v) { + auto* words = reinterpret_cast(values) + 2 * j; + if (words[1] & 0x80000000u) { + const uint32_t low = words[0]; + words[0] = low + dim; + words[1] += words[0] < low; + } + } else if (values[j] < 0) { + values[j] += static_cast(dim); + } + } + __memcpy(output + offset, values, length * sizeof(T), NRAM2GDRAM); + } +} + +void NormalizeIndices(const void* indices, void* normalized, const int64_t* dims, + size_t count, size_t width, uint32_t axis_dim, + DataType dtype, cnrtQueue_t queue) { + if (!count) return; + const cnrtDim3_t grid{static_cast(std::min(count, size_t{32})), 1, 1}; + if (dtype == DataType::kInt64) { + NormalizeKernel<<>>( + static_cast(indices), static_cast(normalized), + dims, count, width, axis_dim); + } else { + NormalizeKernel<<>>( + static_cast(indices), static_cast(normalized), + dims, count, width, axis_dim); + } +} + +} // namespace + +void NormalizeScatterCoordinates(const void* indices, void* normalized, + const int64_t* dims, size_t count, size_t width, + DataType dtype, cnrtQueue_t queue) { + NormalizeIndices(indices, normalized, dims, count, width, 0, dtype, queue); +} + +void NormalizeAxisIndices(const void* indices, void* normalized, + uint32_t dim, size_t count, DataType dtype, + cnrtQueue_t queue) { + NormalizeIndices(indices, normalized, nullptr, count, 1, dim, dtype, queue); +} + +} // namespace infini::ops diff --git a/src/native/cambricon/ops/sign/kernel.h b/src/native/cambricon/ops/sign/kernel.h new file mode 100644 index 000000000..d5ee1c6ae --- /dev/null +++ b/src/native/cambricon/ops/sign/kernel.h @@ -0,0 +1,18 @@ +#ifndef INFINI_OPS_CAMBRICON_SIGN_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_SIGN_KERNEL_H_ + +#include "base/sign.h" +#include "native/cambricon/unary_cnnl.h" + +namespace infini::ops { + +template <> +class Operator + : public cnnl_utils::Unary { + public: + Operator(const Tensor input, Tensor out) : Unary{input, out} {} +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/silu_and_mul/kernel.h b/src/native/cambricon/ops/silu_and_mul/kernel.h index ceb6a481b..5a2593508 100644 --- a/src/native/cambricon/ops/silu_and_mul/kernel.h +++ b/src/native/cambricon/ops/silu_and_mul/kernel.h @@ -39,7 +39,7 @@ class Operator : public SiluAndMul { auto* workspace = workspace_ ? workspace_ : default_workspace_; DispatchFunc< Device::Type::kCambricon, - List>( + CambriconBangFloatTypes<>>( out_type_, [&](auto tag) { using T = typename decltype(tag)::type; diff --git a/src/native/cambricon/ops/silu_and_mul/kernel.mlu b/src/native/cambricon/ops/silu_and_mul/kernel.mlu index b6c4d0228..47d4c1278 100644 --- a/src/native/cambricon/ops/silu_and_mul/kernel.mlu +++ b/src/native/cambricon/ops/silu_and_mul/kernel.mlu @@ -30,7 +30,7 @@ __mlu_device__ void ComputeSiluAndMul(const T* gate, const T* up, T* output, const float gate_value = gate[i]; output[i] = gate_value * up[i] / (1.0F + expf(-gate_value)); } - } else if constexpr (std::is_same_v) { + } else if constexpr (std::is_same_v) { auto* output_half = reinterpret_cast(output); const auto* gate_half = reinterpret_cast(gate); const auto* up_half = reinterpret_cast(up); @@ -156,7 +156,9 @@ void SiluAndMulUnion(void* workspace, int core_per_cluster, int cluster_count, size_t, int, bool, bool) INSTANTIATE_SILU_AND_MUL(Float16); +#if INFINIOPS_CAMBRICON_BANG_BF16 INSTANTIATE_SILU_AND_MUL(BFloat16); +#endif INSTANTIATE_SILU_AND_MUL(float); #undef INSTANTIATE_SILU_AND_MUL diff --git a/src/native/cambricon/ops/sin/kernel.h b/src/native/cambricon/ops/sin/kernel.h new file mode 100644 index 000000000..865e59431 --- /dev/null +++ b/src/native/cambricon/ops/sin/kernel.h @@ -0,0 +1,18 @@ +#ifndef INFINI_OPS_CAMBRICON_SIN_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_SIN_KERNEL_H_ + +#include "base/sin.h" +#include "native/cambricon/unary_cnnl.h" + +namespace infini::ops { + +template <> +class Operator + : public cnnl_utils::Unary { + public: + Operator(const Tensor input, Tensor out) : Unary{input, out} {} +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/sqrt/kernel.h b/src/native/cambricon/ops/sqrt/kernel.h new file mode 100644 index 000000000..f44dd3250 --- /dev/null +++ b/src/native/cambricon/ops/sqrt/kernel.h @@ -0,0 +1,18 @@ +#ifndef INFINI_OPS_CAMBRICON_SQRT_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_SQRT_KERNEL_H_ + +#include "base/sqrt.h" +#include "native/cambricon/unary_cnnl.h" + +namespace infini::ops { + +template <> +class Operator + : public cnnl_utils::Unary { + public: + Operator(const Tensor input, Tensor out) : Unary{input, out} {} +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/sum/kernel.h b/src/native/cambricon/ops/sum/kernel.h new file mode 100644 index 000000000..d229261c7 --- /dev/null +++ b/src/native/cambricon/ops/sum/kernel.h @@ -0,0 +1,38 @@ +#ifndef INFINI_OPS_CAMBRICON_SUM_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_SUM_KERNEL_H_ + +#include "base/sum.h" +#include "native/cambricon/reduce_cnnl.h" + +namespace infini::ops { + +template <> +class Operator : public Sum { + public: + Operator(const Tensor input, const std::optional> dim, + const bool keepdim, const std::optional dtype, Tensor out) + : Sum{input, dim, keepdim, dtype, out}, + plan_{ValidateMode(input, dtype), input, + dim.value_or(std::vector{}), keepdim, out} {} + + std::size_t workspace_size_in_bytes() const override { return plan_.WorkspaceSize(); } + + void operator()(const Tensor input, const std::optional>, + const bool, const std::optional, Tensor out) const override { + plan_.Run(input, out, static_cast(stream_), + workspace_, workspace_size_in_bytes_); + } + + private: + static cnnlReduceOp_t ValidateMode(const Tensor& input, std::optional dtype) { + if (dtype && *dtype != input.dtype()) + throw std::invalid_argument("Cambricon Sum does not support mixed output dtype"); + return CNNL_REDUCE_ADD; + } + + cnnl_utils::ReducePlan plan_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/tanh/kernel.h b/src/native/cambricon/ops/tanh/kernel.h new file mode 100644 index 000000000..76efa64a6 --- /dev/null +++ b/src/native/cambricon/ops/tanh/kernel.h @@ -0,0 +1,18 @@ +#ifndef INFINI_OPS_CAMBRICON_TANH_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_TANH_KERNEL_H_ + +#include "base/tanh.h" +#include "native/cambricon/activation_cnnl.h" + +namespace infini::ops { + +template <> +class Operator + : public cnnl_utils::Activation { + public: + Operator(const Tensor input, Tensor out) : Activation{input, out} {} +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/tile/kernel.h b/src/native/cambricon/ops/tile/kernel.h new file mode 100644 index 000000000..946aba6f2 --- /dev/null +++ b/src/native/cambricon/ops/tile/kernel.h @@ -0,0 +1,40 @@ +#ifndef INFINI_OPS_CAMBRICON_TILE_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_TILE_KERNEL_H_ + +#include "base/tile.h" +#include "native/cambricon/cnnl_utils.h" + +namespace infini::ops { + +template <> +class Operator : public Tile { + public: + Operator(const Tensor input, Tensor out) : Tile{input, out} { + if (input.dtype() != out.dtype() || input.device() != out.device() || + !input.IsContiguous() || !out.IsContiguous() || input.ndim() != out.ndim()) + throw std::invalid_argument("Cambricon Tile requires matching contiguous tensors"); + if (out.numel() == 0) return; + for (size_t i = 0; i < out.ndim(); ++i) { + if (input.size(i) == 0 || out.size(i) % input.size(i) != 0) + throw std::invalid_argument("Cambricon Tile output must be a multiple of input"); + } + handle_ = cnnl_utils::CreateHandle(); + input_ = cnnl_utils::MakeTensorDescriptor(input); + out_ = cnnl_utils::MakeTensorDescriptor(out); + } + + void operator()(const Tensor input, Tensor out) const override { + if (!out_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + INFINI_OPS_CNNL_CHECK(cnnlTile(handle_.get(), input_.get(), input.data(), + out_.get(), out.data())); + } + + private: + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor input_, out_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/ops/top_k_top_p_sampling_from_logits/kernel.mlu b/src/native/cambricon/ops/top_k_top_p_sampling_from_logits/kernel.mlu index 3300208f0..9f2d1f00f 100644 --- a/src/native/cambricon/ops/top_k_top_p_sampling_from_logits/kernel.mlu +++ b/src/native/cambricon/ops/top_k_top_p_sampling_from_logits/kernel.mlu @@ -16,7 +16,7 @@ template __mlu_device__ void LoadAsFloat(float* destination, T* cache, const T* source, std::size_t count) { __memcpy(cache, const_cast(source), count * sizeof(T), GDRAM2NRAM); - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { __bang_half2float(destination, reinterpret_cast(cache), count); } else if constexpr (std::is_same_v) { __bang_bfloat162float(destination, cache, count); diff --git a/src/native/cambricon/ops/where/kernel.h b/src/native/cambricon/ops/where/kernel.h new file mode 100644 index 000000000..efe6efb91 --- /dev/null +++ b/src/native/cambricon/ops/where/kernel.h @@ -0,0 +1,68 @@ +#ifndef INFINI_OPS_CAMBRICON_WHERE_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_WHERE_KERNEL_H_ + +#include + +#include "base/where.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops { + +template <> +class Operator : public Where { + public: + Operator(const Tensor condition, const Tensor input, const Tensor other, + Tensor out) : Where{condition, input, other, out} { + if (condition_type_ != DataType::kBool || input_type_ != out_type_ || + other_type_ != out_type_ || condition.device() != out.device() || + input.device() != out.device() || other.device() != out.device() || + !condition.IsContiguous() || !input.IsContiguous() || + !other.IsContiguous() || !out.IsContiguous()) { + throw std::invalid_argument("Cambricon Where requires a Boolean mask and contiguous tensors"); + } + for (const auto& value : {condition, input, other}) { + if (value.ndim() > out.ndim()) + throw std::invalid_argument("Cambricon Where: invalid broadcast rank"); + const auto offset = out.ndim() - value.ndim(); + for (size_t i = 0; i < value.ndim(); ++i) { + if (value.shape()[i] != 1 && value.shape()[i] != out.shape()[offset + i]) + throw std::invalid_argument("Cambricon Where: invalid broadcast dimensions"); + } + } + if (out.numel() == 0) return; + handle_ = cnnl_utils::CreateHandle(); + condition_ = cnnl_utils::MakeTensorDescriptor(condition_type_, condition_shape_, condition_strides_); + input_ = cnnl_utils::MakeTensorDescriptor(input_type_, input_shape_, input_strides_); + other_ = cnnl_utils::MakeTensorDescriptor(other_type_, other_shape_, other_strides_); + output_ = cnnl_utils::MakeTensorDescriptor(out_type_, out_shape_, out_strides_); + INFINI_OPS_CNNL_CHECK(cnnlGetSelectV2WorkspaceSize( + handle_.get(), condition_.get(), input_.get(), other_.get(), &workspace_size_)); + workspace_storage_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_); + if (workspace_size_ && !workspace_storage_) + throw std::runtime_error("Cambricon Where: workspace allocation failed"); + } + + std::size_t workspace_size_in_bytes() const override { return workspace_size_; } + + void operator()(const Tensor condition, const Tensor input, const Tensor other, + Tensor out) const override { + if (!output_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(stream_))); + auto* scratch = workspace_ && workspace_size_in_bytes_ >= workspace_size_ + ? workspace_ : workspace_storage_.get(); + INFINI_OPS_CNNL_CHECK(cnnlSelectV2( + handle_.get(), condition_.get(), condition.data(), input_.get(), input.data(), + other_.get(), other.data(), scratch, workspace_size_, output_.get(), out.data())); + } + + private: + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor condition_, input_, other_, output_; + cnrt_utils::DeviceBuffer workspace_storage_; + std::size_t workspace_size_ = 0; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/reduce_cnnl.h b/src/native/cambricon/reduce_cnnl.h new file mode 100644 index 000000000..b8996489f --- /dev/null +++ b/src/native/cambricon/reduce_cnnl.h @@ -0,0 +1,97 @@ +#ifndef INFINI_OPS_CAMBRICON_REDUCE_CNNL_H_ +#define INFINI_OPS_CAMBRICON_REDUCE_CNNL_H_ + +#include +#include + +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops::cnnl_utils { + +class ReducePlan { + public: + ReducePlan(cnnlReduceOp_t mode, const Tensor& input, + const std::vector& dim, bool keepdim, const Tensor& out) { + const bool floating = input.dtype() == DataType::kFloat32 || input.dtype() == DataType::kFloat16; + const bool integer_sum = mode == CNNL_REDUCE_ADD && + (input.dtype() == DataType::kInt32 || input.dtype() == DataType::kInt64); + if (input.ndim() == 0 || + input.device() != out.device() || input.dtype() != out.dtype() || + !(floating || integer_sum) || !input.IsContiguous() || !out.IsContiguous()) + throw std::invalid_argument("Cambricon Reduce requires contiguous inputs and a supported dtype"); + const auto rank = static_cast(input.ndim()); + std::set axes; + if (dim.empty()) { + for (int i = 0; i < rank; ++i) axes.insert(i); + } else { + for (auto axis : dim) { + if (axis < 0) axis += rank; + if (axis < 0 || axis >= rank || !axes.insert(static_cast(axis)).second) + throw std::invalid_argument("Cambricon Reduce: invalid axes"); + } + } + auto reduced_shape = Tensor::Shape(input.shape()); + Tensor::Shape expected; + for (int i = 0; i < rank; ++i) { + const bool reduce = axes.count(i) != 0; + if (reduce) reduced_shape[i] = 1; + if (keepdim || !reduce) expected.push_back(reduced_shape[i]); + } + if (expected != out.shape()) + throw std::invalid_argument("Cambricon Reduce: invalid output shape"); + if (out.numel() == 0) return; + if (input.numel() == 0) + throw std::invalid_argument("Cambricon Reduce: empty reduction not supported"); + Tensor::Strides reduced_strides(reduced_shape.size()); + Tensor::Stride stride = 1; + for (size_t i = reduced_shape.size(); i-- > 0;) { + reduced_strides[i] = stride; + stride *= reduced_shape[i]; + } + handle_ = cnnl_utils::CreateHandle(); + input_ = cnnl_utils::MakeTensorDescriptor(input.dtype(), Tensor::Shape(input.shape()), + Tensor::Strides(input.strides())); + output_ = cnnl_utils::MakeTensorDescriptor(out.dtype(), reduced_shape, reduced_strides); + cnnlReduceDescriptor_t descriptor = nullptr; + INFINI_OPS_CNNL_CHECK(cnnlCreateReduceDescriptor(&descriptor)); + descriptor_.reset(descriptor); + std::vector axis_list(axes.begin(), axes.end()); + INFINI_OPS_CNNL_CHECK(cnnlSetReduceDescriptor_v2( + descriptor_.get(), axis_list.data(), axis_list.size(), mode, + floating ? CNNL_DTYPE_FLOAT : cnnl_utils::GetDataType(input.dtype()), + CNNL_PROPAGATE_NAN, CNNL_REDUCE_NO_INDICES, + CNNL_32BIT_INDICES, 0.0)); + INFINI_OPS_CNNL_CHECK(cnnlGetReduceOpWorkspaceSize( + handle_.get(), input_.get(), output_.get(), descriptor_.get(), &workspace_size_)); + workspace_own_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_); + if (workspace_size_ && !workspace_own_) + throw std::runtime_error("Cambricon Reduce workspace allocation failed"); + } + + std::size_t WorkspaceSize() const { return workspace_size_; } + + void Run(const Tensor& input, Tensor out, cnrtQueue_t queue, + void* workspace, size_t capacity) const { + if (!output_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), queue)); + auto* scratch = workspace && capacity >= workspace_size_ + ? workspace : workspace_own_.get(); + INFINI_OPS_CNNL_CHECK(cnnlReduce( + handle_.get(), descriptor_.get(), scratch, workspace_size_, nullptr, + input_.get(), input.data(), 0, nullptr, nullptr, output_.get(), out.data())); + } + + private: + cnnl_utils::Handle handle_; + cnnl_utils::TensorDescriptor input_, output_; + std::unique_ptr, + decltype(&cnnlDestroyReduceDescriptor)> + descriptor_{nullptr, cnnlDestroyReduceDescriptor}; + cnrt_utils::DeviceBuffer workspace_own_; + std::size_t workspace_size_ = 0; +}; + +} // namespace infini::ops::cnnl_utils + +#endif diff --git a/src/native/cambricon/scatter_index.h b/src/native/cambricon/scatter_index.h new file mode 100644 index 000000000..1099103ed --- /dev/null +++ b/src/native/cambricon/scatter_index.h @@ -0,0 +1,19 @@ +#ifndef INFINI_OPS_CAMBRICON_SCATTER_INDEX_H_ +#define INFINI_OPS_CAMBRICON_SCATTER_INDEX_H_ + +#include "native/cambricon/common.h" +#include "tensor.h" + +namespace infini::ops { + +void NormalizeScatterCoordinates(const void* indices, void* normalized, + const int64_t* dims, size_t count, size_t width, + DataType dtype, cnrtQueue_t queue); + +void NormalizeAxisIndices(const void* indices, void* normalized, + uint32_t dim, size_t count, DataType dtype, + cnrtQueue_t queue); + +} // namespace infini::ops + +#endif diff --git a/src/native/cambricon/unary_cnnl.h b/src/native/cambricon/unary_cnnl.h new file mode 100644 index 000000000..7450db698 --- /dev/null +++ b/src/native/cambricon/unary_cnnl.h @@ -0,0 +1,39 @@ +#ifndef INFINI_OPS_CAMBRICON_UNARY_CNNL_H_ +#define INFINI_OPS_CAMBRICON_UNARY_CNNL_H_ + +#include "native/cambricon/cnnl_utils.h" + +namespace infini::ops::cnnl_utils { + +template +class Unary : public Base { + public: + Unary(const Tensor input, Tensor out) : Base{input, out} { + if (input.dtype() != out.dtype() || input.device() != out.device() || + input.shape() != out.shape() || + (Boolean ? out.dtype() != DataType::kBool + : out.dtype() != DataType::kFloat32 && out.dtype() != DataType::kFloat16)) + throw std::invalid_argument("Cambricon unary operation received unsupported tensor metadata"); + if (out.numel() == 0) return; + handle_ = CreateHandle(); + input_ = MakeTensorDescriptor( + input.dtype(), Tensor::Shape(input.shape()), Tensor::Strides(input.strides())); + out_ = MakeTensorDescriptor( + out.dtype(), Tensor::Shape(out.shape()), Tensor::Strides(out.strides())); + } + + void operator()(const Tensor input, Tensor out) const override { + if (!out_) return; + INFINI_OPS_CNNL_CHECK(cnnlSetQueue(handle_.get(), static_cast(this->stream_))); + INFINI_OPS_CNNL_CHECK(Execute(handle_.get(), input_.get(), input.data(), + out_.get(), out.data())); + } + + private: + Handle handle_; + TensorDescriptor input_, out_; +}; + +} // namespace infini::ops::cnnl_utils + +#endif