From 57db3d5f3b5d5600893f2bbd968a270c7c9d9bed Mon Sep 17 00:00:00 2001 From: rubik Date: Wed, 22 Jul 2026 15:40:47 +0800 Subject: [PATCH 1/2] feat(op): Add Metax conv and softmax operators Implement convolution and softmax operators for Metax backend. - Conv: Uses mcDNN library for standard paths plus custom MACA kernels for 2D direct conv and 3D patch embed optimizations. - Softmax: Implements custom kernel using warp/block reduction for high throughput. - Types: Supports FP16, BF16, and FP32 for both operators. --- src/infiniop/ops/conv/metax/conv_metax.cc | 323 ++++++++++++++++++ src/infiniop/ops/conv/metax/conv_metax.h | 28 ++ .../ops/conv/metax/conv_metax_kernel.maca | 284 +++++++++++++++ src/infiniop/ops/conv/operator.cc | 16 + .../ops/softmax/metax/softmax_metax.h | 8 + .../ops/softmax/metax/softmax_metax.maca | 188 ++++++++++ src/infiniop/ops/softmax/operator.cc | 16 + 7 files changed, 863 insertions(+) create mode 100644 src/infiniop/ops/conv/metax/conv_metax.cc create mode 100644 src/infiniop/ops/conv/metax/conv_metax.h create mode 100644 src/infiniop/ops/conv/metax/conv_metax_kernel.maca create mode 100644 src/infiniop/ops/softmax/metax/softmax_metax.h create mode 100644 src/infiniop/ops/softmax/metax/softmax_metax.maca diff --git a/src/infiniop/ops/conv/metax/conv_metax.cc b/src/infiniop/ops/conv/metax/conv_metax.cc new file mode 100644 index 000000000..7778501cd --- /dev/null +++ b/src/infiniop/ops/conv/metax/conv_metax.cc @@ -0,0 +1,323 @@ +#include "conv_metax.h" +#include "../../../devices/metax/metax_common.h" +#include "mcdnn/mcdnn.h" // Metax DNN Library Header +#include +#include + +namespace op::conv::metax { + +// --------------------------------------------------------------------------- +// Helper Functions +// --------------------------------------------------------------------------- + +/** + * @brief Maps InfiniCore data types to mcDNN data types. + * @param dtype The InfiniCore data type. + * @return The corresponding mcDNN data type. + */ +inline mcdnnDataType_t toMcDNNDataType(infiniDtype_t dtype) { + switch (dtype) { + case INFINI_DTYPE_F32: + return MCDNN_DATA_FLOAT; + case INFINI_DTYPE_F16: + return MCDNN_DATA_HALF; + case INFINI_DTYPE_BF16: + return MCDNN_DATA_BFLOAT16; + default: + return MCDNN_DATA_FLOAT; + } +} + +/** + * @brief Checks mcDNN API return status and converts to InfiniCore status. + */ +inline infiniStatus_t checkMcDNNStatus(mcdnnStatus_t status) { + if (status == MCDNN_STATUS_SUCCESS) { + return INFINI_STATUS_SUCCESS; + } + // Log the error if a logging mechanism is available + return INFINI_STATUS_INTERNAL_ERROR; +} + +struct Descriptor::Opaque { + mcdnnHandle_t mcdnn_handle = nullptr; + mcdnnConvolutionDescriptor_t conv_desc = nullptr; + mcdnnFilterDescriptor_t filter_desc = nullptr; + mcdnnTensorDescriptor_t input_desc = nullptr; + mcdnnTensorDescriptor_t output_desc = nullptr; + mcdnnTensorDescriptor_t bias_desc = nullptr; + mcdnnActivationDescriptor_t activation_desc = nullptr; + mcdnnConvolutionFwdAlgo_t algo = MCDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM; + bool direct_conv2d_kernel = false; + bool patch_embed_kernel = false; + + ~Opaque() { + if (mcdnn_handle) { + mcdnnDestroy(mcdnn_handle); + } + if (conv_desc) { + mcdnnDestroyConvolutionDescriptor(conv_desc); + } + if (filter_desc) { + mcdnnDestroyFilterDescriptor(filter_desc); + } + if (input_desc) { + mcdnnDestroyTensorDescriptor(input_desc); + } + if (output_desc) { + mcdnnDestroyTensorDescriptor(output_desc); + } + if (bias_desc) { + mcdnnDestroyTensorDescriptor(bias_desc); + } + if (activation_desc) { + mcdnnDestroyActivationDescriptor(activation_desc); + } + } +}; + +Descriptor::~Descriptor() = default; + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t w_desc, + infiniopTensorDescriptor_t b_desc, + const void *pads, + const void *strides, + const void *dilations, + size_t n) { + + auto handle = reinterpret_cast(handle_); + auto dtype = y_desc->dtype(); + + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); + + auto result = ConvInfo::create(handle_, y_desc, x_desc, w_desc, b_desc, + pads, strides, dilations, n); + CHECK_RESULT(result); + ConvInfo info = result.take(); + + size_t workspace_size = 0; + auto opaque = new Opaque(); + mcdnnStatus_t status; + + const bool can_use_direct_conv2d_kernel = info.ndim() == 2; + + if (can_use_direct_conv2d_kernel) { + opaque->direct_conv2d_kernel = true; + *desc_ptr = new Descriptor( + dtype, std::move(info), workspace_size, + opaque, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; + } + + const bool can_use_patch_embed_kernel = info.ndim() == 3 && info.output_dim(0) == 1 && info.output_dim(1) == 1 && info.output_dim(2) == 1 && info.input_dim(0) == info.kernel_dim(0) && info.input_dim(1) == info.kernel_dim(1) && info.input_dim(2) == info.kernel_dim(2) && info.pad_info(0) == 0 && info.pad_info(1) == 0 && info.pad_info(2) == 0 && info.stride_info(0) == static_cast(info.kernel_dim(0)) && info.stride_info(1) == static_cast(info.kernel_dim(1)) && info.stride_info(2) == static_cast(info.kernel_dim(2)) && info.dilation_info(0) == 1 && info.dilation_info(1) == 1 && info.dilation_info(2) == 1; + + if (can_use_patch_embed_kernel) { + opaque->patch_embed_kernel = true; + *desc_ptr = new Descriptor( + dtype, std::move(info), workspace_size, + opaque, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; + } + + do { + // 1. Create mcDNN Handle + status = mcdnnCreate(&opaque->mcdnn_handle); + if (status != MCDNN_STATUS_SUCCESS) { + break; + } + + // 2. Create Convolution Descriptor + std::vector pad_vec(info.ndim()); + std::vector stride_vec(info.ndim()); + std::vector dilation_vec(info.ndim()); + for (size_t i = 0; i < info.ndim(); ++i) { + pad_vec[i] = static_cast(info.pad_info(i)); + stride_vec[i] = static_cast(info.stride_info(i)); + dilation_vec[i] = static_cast(info.dilation_info(i)); + } + + status = mcdnnCreateConvolutionDescriptor(&opaque->conv_desc); + if (status != MCDNN_STATUS_SUCCESS) { + break; + } + + status = mcdnnSetConvolutionNdDescriptor( + opaque->conv_desc, info.ndim(), pad_vec.data(), stride_vec.data(), + dilation_vec.data(), MCDNN_CONVOLUTION, toMcDNNDataType(dtype)); + if (status != MCDNN_STATUS_SUCCESS) { + break; + } + + // 3. Create Filter Descriptor + std::vector filter_dim(info.ndim() + 2); + filter_dim[0] = info.out_channels(); + filter_dim[1] = info.in_channels(); // groups = 1 + for (size_t i = 0; i < info.ndim(); ++i) { + filter_dim[i + 2] = info.kernel_dim(i); + } + + status = mcdnnCreateFilterDescriptor(&opaque->filter_desc); + if (status != MCDNN_STATUS_SUCCESS) { + break; + } + + status = mcdnnSetFilterNdDescriptor( + opaque->filter_desc, toMcDNNDataType(dtype), MCDNN_TENSOR_NCHW, + info.ndim() + 2, filter_dim.data()); + if (status != MCDNN_STATUS_SUCCESS) { + break; + } + + // 4. Create Tensor Descriptors + std::vector input_dim(info.ndim() + 2); + input_dim[0] = info.batch(); + input_dim[1] = info.in_channels(); + for (size_t i = 0; i < info.ndim(); ++i) { + input_dim[i + 2] = info.input_dim(i); + } + + status = mcdnnCreateTensorDescriptor(&opaque->input_desc); + if (status != MCDNN_STATUS_SUCCESS) { + break; + } + + status = mcdnnSetTensorNdDescriptor(opaque->input_desc, toMcDNNDataType(dtype), + info.ndim() + 2, input_dim.data(), nullptr); + if (status != MCDNN_STATUS_SUCCESS) { + break; + } + + std::vector output_dim(info.ndim() + 2); + output_dim[0] = info.batch(); + output_dim[1] = info.out_channels(); + for (size_t i = 0; i < info.ndim(); ++i) { + output_dim[i + 2] = info.output_dim(i); + } + + status = mcdnnCreateTensorDescriptor(&opaque->output_desc); + if (status != MCDNN_STATUS_SUCCESS) { + break; + } + + status = mcdnnSetTensorNdDescriptor(opaque->output_desc, toMcDNNDataType(dtype), + info.ndim() + 2, output_dim.data(), nullptr); + if (status != MCDNN_STATUS_SUCCESS) { + break; + } + + if (b_desc != nullptr) { + std::vector bias_dim(info.ndim() + 2, 1); + bias_dim[1] = info.out_channels(); // 1xCx1x1 + + status = mcdnnCreateTensorDescriptor(&opaque->bias_desc); + if (status != MCDNN_STATUS_SUCCESS) { + break; + } + + status = mcdnnSetTensorNdDescriptor( + opaque->bias_desc, + toMcDNNDataType(dtype), + info.ndim() + 2, + bias_dim.data(), + nullptr); + if (status != MCDNN_STATUS_SUCCESS) { + break; + } + } + + // 5. Get Workspace Size + status = mcdnnGetConvolutionForwardWorkspaceSize( + opaque->mcdnn_handle, + opaque->input_desc, opaque->filter_desc, opaque->conv_desc, + opaque->output_desc, opaque->algo, &workspace_size); + + if (status != MCDNN_STATUS_SUCCESS) { + workspace_size = 0; + } + + // Success: Create Descriptor + *desc_ptr = new Descriptor( + dtype, std::move(info), workspace_size, + opaque, handle->device, handle->device_id); + + return INFINI_STATUS_SUCCESS; + + } while (0); + + // Error Handling: Cleanup + if (opaque) { + // auto clean mcdnn_handle and descriptors + delete opaque; + } + return INFINI_STATUS_INTERNAL_ERROR; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *y, + const void *x, + const void *w, + const void *bias, + void *stream) const { + + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + if (_opaque->direct_conv2d_kernel) { + return launchDirectConv2d(_dtype, _info, y, x, w, bias, stream); + } + + if (_opaque->patch_embed_kernel) { + return launchPatchEmbedConv3d(_dtype, _info, y, x, w, bias, stream); + } + + mcdnnStatus_t status; + float alpha = 1.0f; + float beta = 0.0f; + + // --------------------------------------------------------------- + // Step 1: y = conv(x, w) + // --------------------------------------------------------------- + status = mcdnnConvolutionForward( + _opaque->mcdnn_handle, + &alpha, + _opaque->input_desc, x, + _opaque->filter_desc, w, + _opaque->conv_desc, + _opaque->algo, + workspace, workspace_size, + &beta, + _opaque->output_desc, y); + + CHECK_STATUS(checkMcDNNStatus(status)); + + // --------------------------------------------------------------- + // Step 2: y = y + bias + // --------------------------------------------------------------- + if (bias != nullptr) { + // mcdnnAddTensor: y = alpha * bias + beta * y + // we need do: y = 1.0 * bias + 1.0 * y_old + float alpha_add = 1.0f; + float beta_add = 1.0f; + + status = mcdnnAddTensor( + _opaque->mcdnn_handle, + &alpha_add, + _opaque->bias_desc, bias, + &beta_add, + _opaque->output_desc, y); + + CHECK_STATUS(checkMcDNNStatus(status)); + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::conv::metax diff --git a/src/infiniop/ops/conv/metax/conv_metax.h b/src/infiniop/ops/conv/metax/conv_metax.h new file mode 100644 index 000000000..75c4a546b --- /dev/null +++ b/src/infiniop/ops/conv/metax/conv_metax.h @@ -0,0 +1,28 @@ +#ifndef __CONV_METAX_H__ +#define __CONV_METAX_H__ + +#include "../conv.h" + +DESCRIPTOR(metax) + +namespace op::conv::metax { +infiniStatus_t launchDirectConv2d( + infiniDtype_t dtype, + const ConvInfo &info, + void *y, + const void *x, + const void *w, + const void *bias, + void *stream); + +infiniStatus_t launchPatchEmbedConv3d( + infiniDtype_t dtype, + const ConvInfo &info, + void *y, + const void *x, + const void *w, + const void *bias, + void *stream); +} // namespace op::conv::metax + +#endif // __CONV_METAX_H__ diff --git a/src/infiniop/ops/conv/metax/conv_metax_kernel.maca b/src/infiniop/ops/conv/metax/conv_metax_kernel.maca new file mode 100644 index 000000000..1e8b09e7b --- /dev/null +++ b/src/infiniop/ops/conv/metax/conv_metax_kernel.maca @@ -0,0 +1,284 @@ +#include "../../../devices/metax/metax_common.h" +#include "../../../devices/metax/metax_kernel_common.h" +#include "conv_metax.h" + +namespace op::conv::metax { +namespace { + +template +__device__ __forceinline__ float load_as_float(const T *ptr, ptrdiff_t offset) { + return static_cast(ptr[offset]); +} + +template <> +__device__ __forceinline__ float load_as_float(const half *ptr, ptrdiff_t offset) { + return __half2float(ptr[offset]); +} + +template <> +__device__ __forceinline__ float load_as_float(const cuda_bfloat16 *ptr, ptrdiff_t offset) { + return __bfloat162float(ptr[offset]); +} + +template +__device__ __forceinline__ T cast_from_float(float x) { + return static_cast(x); +} + +template <> +__device__ __forceinline__ half cast_from_float(float x) { + return __float2half(x); +} + +template <> +__device__ __forceinline__ cuda_bfloat16 cast_from_float(float x) { + return __float2bfloat16(x); +} + +template +__global__ void direct_conv2d_kernel( + T *y, + const T *x, + const T *w, + const T *bias, + size_t batch, + size_t in_channels, + size_t out_channels, + size_t in_h, + size_t in_w, + size_t kernel_h, + size_t kernel_w, + size_t out_h, + size_t out_w, + size_t pad_h, + size_t pad_w, + ptrdiff_t stride_h, + ptrdiff_t stride_w, + size_t dilation_h, + size_t dilation_w) { + const size_t out_idx = static_cast(blockIdx.x); + const size_t total_outputs = batch * out_channels * out_h * out_w; + if (out_idx >= total_outputs) { + return; + } + + size_t tmp = out_idx; + const size_t ow = tmp % out_w; + tmp /= out_w; + const size_t oh = tmp % out_h; + tmp /= out_h; + const size_t oc = tmp % out_channels; + const size_t n = tmp / out_channels; + + const size_t reduce_elems = in_channels * kernel_h * kernel_w; + float acc = 0.0f; + for (size_t i = threadIdx.x; i < reduce_elems; i += BLOCK_SIZE) { + const size_t kw = i % kernel_w; + const size_t kh = (i / kernel_w) % kernel_h; + const size_t ic = i / (kernel_h * kernel_w); + + const ptrdiff_t ih = static_cast(oh) * stride_h + + static_cast(kh * dilation_h) + - static_cast(pad_h); + const ptrdiff_t iw = static_cast(ow) * stride_w + + static_cast(kw * dilation_w) + - static_cast(pad_w); + if (ih < 0 || iw < 0 || ih >= static_cast(in_h) || iw >= static_cast(in_w)) { + continue; + } + + const size_t x_idx = ((n * in_channels + ic) * in_h + static_cast(ih)) * in_w + static_cast(iw); + const size_t w_idx = ((oc * in_channels + ic) * kernel_h + kh) * kernel_w + kw; + acc += load_as_float(x, static_cast(x_idx)) * load_as_float(w, static_cast(w_idx)); + } + + __shared__ float reduce[BLOCK_SIZE]; + reduce[threadIdx.x] = acc; + __syncthreads(); + + for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + reduce[threadIdx.x] += reduce[threadIdx.x + stride]; + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + float value = reduce[0]; + if (bias != nullptr) { + value += load_as_float(bias, static_cast(oc)); + } + y[out_idx] = cast_from_float(value); + } +} + +template +__global__ void patch_embed_conv3d_kernel( + T *y, + const T *x, + const T *w, + const T *bias, + size_t batch, + size_t in_channels, + size_t out_channels, + size_t kt, + size_t kh, + size_t kw) { + const size_t out_idx = static_cast(blockIdx.x); + const size_t total_outputs = batch * out_channels; + if (out_idx >= total_outputs) { + return; + } + + const size_t n = out_idx / out_channels; + const size_t oc = out_idx - n * out_channels; + const size_t patch_elems = in_channels * kt * kh * kw; + + float acc = 0.0f; + for (size_t i = threadIdx.x; i < patch_elems; i += BLOCK_SIZE) { + acc += load_as_float(x, static_cast(n * patch_elems + i)) * load_as_float(w, static_cast(oc * patch_elems + i)); + } + + __shared__ float reduce[BLOCK_SIZE]; + reduce[threadIdx.x] = acc; + __syncthreads(); + + for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + reduce[threadIdx.x] += reduce[threadIdx.x + stride]; + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + float value = reduce[0]; + if (bias != nullptr) { + value += load_as_float(bias, static_cast(oc)); + } + y[out_idx] = cast_from_float(value); + } +} + +template +infiniStatus_t launch_typed_direct_conv2d( + void *y, + const void *x, + const void *w, + const void *bias, + const ConvInfo &info, + hcStream_t stream) { + constexpr int BLOCK_SIZE = 256; + const size_t total_outputs = info.batch() * info.out_channels() * info.output_dim(0) * info.output_dim(1); + if (total_outputs == 0) { + return INFINI_STATUS_SUCCESS; + } + + direct_conv2d_kernel<<(total_outputs), BLOCK_SIZE, 0, stream>>>( + static_cast(y), + static_cast(x), + static_cast(w), + static_cast(bias), + info.batch(), + info.in_channels(), + info.out_channels(), + info.input_dim(0), + info.input_dim(1), + info.kernel_dim(0), + info.kernel_dim(1), + info.output_dim(0), + info.output_dim(1), + info.pad_info(0), + info.pad_info(1), + info.stride_info(0), + info.stride_info(1), + info.dilation_info(0), + info.dilation_info(1)); + CHECK_METAX(hcGetLastError()); + return INFINI_STATUS_SUCCESS; +} + +template +infiniStatus_t launch_typed_patch_embed_conv3d( + void *y, + const void *x, + const void *w, + const void *bias, + size_t batch, + size_t in_channels, + size_t out_channels, + size_t kt, + size_t kh, + size_t kw, + hcStream_t stream) { + constexpr int BLOCK_SIZE = 256; + const size_t total_outputs = batch * out_channels; + if (total_outputs == 0) { + return INFINI_STATUS_SUCCESS; + } + + patch_embed_conv3d_kernel<<(total_outputs), BLOCK_SIZE, 0, stream>>>( + static_cast(y), + static_cast(x), + static_cast(w), + static_cast(bias), + batch, + in_channels, + out_channels, + kt, + kh, + kw); + CHECK_METAX(hcGetLastError()); + return INFINI_STATUS_SUCCESS; +} + +} // namespace + +infiniStatus_t launchDirectConv2d( + infiniDtype_t dtype, + const ConvInfo &info, + void *y, + const void *x, + const void *w, + const void *bias, + void *stream) { + auto hc_stream = reinterpret_cast(stream); + switch (dtype) { + case INFINI_DTYPE_F16: + return launch_typed_direct_conv2d(y, x, w, bias, info, hc_stream); + case INFINI_DTYPE_BF16: + return launch_typed_direct_conv2d(y, x, w, bias, info, hc_stream); + case INFINI_DTYPE_F32: + return launch_typed_direct_conv2d(y, x, w, bias, info, hc_stream); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +infiniStatus_t launchPatchEmbedConv3d( + infiniDtype_t dtype, + const ConvInfo &info, + void *y, + const void *x, + const void *w, + const void *bias, + void *stream) { + auto hc_stream = reinterpret_cast(stream); + switch (dtype) { + case INFINI_DTYPE_F16: + return launch_typed_patch_embed_conv3d( + y, x, w, bias, info.batch(), info.in_channels(), info.out_channels(), + info.kernel_dim(0), info.kernel_dim(1), info.kernel_dim(2), hc_stream); + case INFINI_DTYPE_BF16: + return launch_typed_patch_embed_conv3d( + y, x, w, bias, info.batch(), info.in_channels(), info.out_channels(), + info.kernel_dim(0), info.kernel_dim(1), info.kernel_dim(2), hc_stream); + case INFINI_DTYPE_F32: + return launch_typed_patch_embed_conv3d( + y, x, w, bias, info.batch(), info.in_channels(), info.out_channels(), + info.kernel_dim(0), info.kernel_dim(1), info.kernel_dim(2), hc_stream); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +} // namespace op::conv::metax diff --git a/src/infiniop/ops/conv/operator.cc b/src/infiniop/ops/conv/operator.cc index 5d27c9232..5536ab268 100644 --- a/src/infiniop/ops/conv/operator.cc +++ b/src/infiniop/ops/conv/operator.cc @@ -13,6 +13,10 @@ #include "nvidia/conv_nvidia.cuh" #endif +#ifdef ENABLE_METAX_API +#include "metax/conv_metax.h" +#endif + __INFINI_C __export infiniStatus_t infiniopCreateConvDescriptor(infiniopHandle_t handle, infiniopConvDescriptor_t *desc_ptr, infiniopTensorDescriptor_t y_desc, @@ -58,6 +62,9 @@ __INFINI_C __export infiniStatus_t infiniopCreateConvDescriptor(infiniopHandle_t #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); #endif +#ifdef ENABLE_METAX_API + CREATE(INFINI_DEVICE_METAX, metax); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -98,6 +105,9 @@ infiniopGetConvWorkspaceSize( #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); #endif +#ifdef ENABLE_METAX_API + GET(INFINI_DEVICE_METAX, metax); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -146,6 +156,9 @@ __INFINI_C infiniStatus_t infiniopConv( #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); #endif +#ifdef ENABLE_METAX_API + CALCULATE(INFINI_DEVICE_METAX, metax); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -182,6 +195,9 @@ infiniopDestroyConvDescriptor(infiniopConvDescriptor_t desc) { #ifdef ENABLE_CAMBRICON_API DELETE(INFINI_DEVICE_CAMBRICON, bang); #endif +#ifdef ENABLE_METAX_API + DELETE(INFINI_DEVICE_METAX, metax); +#endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/softmax/metax/softmax_metax.h b/src/infiniop/ops/softmax/metax/softmax_metax.h new file mode 100644 index 000000000..d05575647 --- /dev/null +++ b/src/infiniop/ops/softmax/metax/softmax_metax.h @@ -0,0 +1,8 @@ +#ifndef __SOFTMAX_METAX_H__ +#define __SOFTMAX_METAX_H__ + +#include "../softmax.h" + +DESCRIPTOR(metax) + +#endif // __SOFTMAX_METAX_H__ diff --git a/src/infiniop/ops/softmax/metax/softmax_metax.maca b/src/infiniop/ops/softmax/metax/softmax_metax.maca new file mode 100644 index 000000000..f778440b6 --- /dev/null +++ b/src/infiniop/ops/softmax/metax/softmax_metax.maca @@ -0,0 +1,188 @@ +#include "../../../devices/metax/metax_common.h" +#include "../../../devices/metax/metax_handle.h" +#include "../../../devices/metax/metax_kernel_common.h" +#include "softmax_metax.h" +#include +#include + +namespace op::softmax::metax { +namespace { + +__device__ __forceinline__ float to_float(float val) { return val; } +__device__ __forceinline__ float to_float(double val) { return static_cast(val); } +__device__ __forceinline__ float to_float(__half val) { return __half2float(val); } +__device__ __forceinline__ float to_float(cuda_bfloat16 val) { return __bfloat162float(val); } + +template +__device__ __forceinline__ T from_float(float val) { return static_cast(val); } + +template <> +__device__ __forceinline__ __half from_float<__half>(float val) { return __float2half(val); } + +template <> +__device__ __forceinline__ cuda_bfloat16 from_float(float val) { return __float2bfloat16(val); } + +template +__device__ __forceinline__ T warp_reduce_max(T val) { + for (int offset = 16; offset > 0; offset /= 2) { + T shuffled = __shfl_down_sync(0xffffffff, val, offset); + val = val > shuffled ? val : shuffled; + } + return val; +} + +template +__device__ __forceinline__ T warp_reduce_sum(T val) { + for (int offset = 16; offset > 0; offset /= 2) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return val; +} + +__device__ __forceinline__ float block_reduce_max(float val) { + __shared__ float shared[32]; + int lane = threadIdx.x & 31; + int wid = threadIdx.x >> 5; + val = warp_reduce_max(val); + if (lane == 0) { + shared[wid] = val; + } + __syncthreads(); + val = (threadIdx.x < ((blockDim.x + 31) >> 5)) ? shared[lane] : -INFINITY; + if (wid == 0) { + val = warp_reduce_max(val); + } + return val; +} + +__device__ __forceinline__ float block_reduce_sum(float val) { + __shared__ float shared[32]; + int lane = threadIdx.x & 31; + int wid = threadIdx.x >> 5; + val = warp_reduce_sum(val); + if (lane == 0) { + shared[wid] = val; + } + __syncthreads(); + val = (threadIdx.x < ((blockDim.x + 31) >> 5)) ? shared[lane] : 0.0f; + if (wid == 0) { + val = warp_reduce_sum(val); + } + return val; +} + +template +__global__ void softmax_kernel( + T *__restrict__ output, + const T *__restrict__ input, + size_t othersize, + size_t dimsize, + ptrdiff_t stride) { + const size_t other_idx = static_cast(blockIdx.x); + if (other_idx >= othersize) { + return; + } + + const size_t inner_idx = other_idx % static_cast(stride); + const size_t outer_idx = other_idx / static_cast(stride); + const size_t base_offset = outer_idx * dimsize * static_cast(stride) + inner_idx; + + float local_max = -INFINITY; + for (size_t i = threadIdx.x; i < dimsize; i += blockDim.x) { + const float val = to_float(input[base_offset + i * stride]); + local_max = local_max > val ? local_max : val; + } + + __shared__ float s_max; + __shared__ float s_sum; + const float max_val = block_reduce_max(local_max); + if (threadIdx.x == 0) { + s_max = max_val; + } + __syncthreads(); + + float local_sum = 0.0f; + for (size_t i = threadIdx.x; i < dimsize; i += blockDim.x) { + local_sum += expf(to_float(input[base_offset + i * stride]) - s_max); + } + + const float sum_val = block_reduce_sum(local_sum); + if (threadIdx.x == 0) { + s_sum = sum_val; + } + __syncthreads(); + + const float inv_sum = 1.0f / s_sum; + for (size_t i = threadIdx.x; i < dimsize; i += blockDim.x) { + const size_t idx = base_offset + i * stride; + output[idx] = from_float(expf(to_float(input[idx]) - s_max) * inv_sum); + } +} + +template +void launch_kernel(void *output, const void *input, const SoftmaxInfo &info, void *stream) { + auto hc_stream = reinterpret_cast(stream); + unsigned int threads_per_block = 256; + if (info.dimsize < 256) { + threads_per_block = 128; + } + if (info.dimsize < 128) { + threads_per_block = 64; + } + if (info.dimsize < 64) { + threads_per_block = 32; + } + softmax_kernel<<>>( + reinterpret_cast(output), + reinterpret_cast(input), + info.othersize, + info.dimsize, + info.stride); +} + +} // namespace + +struct Descriptor::Opaque {}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc, + int axis) { + auto handle = reinterpret_cast(handle_); + auto info = SoftmaxInfo::create(y_desc, x_desc, axis); + CHECK_RESULT(info); + *desc_ptr = new Descriptor(new Opaque(), info.take(), 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream) const { + (void)workspace; + (void)workspace_size; + switch (_info.dtype) { + case INFINI_DTYPE_F16: + launch_kernel<__half>(y, x, _info, stream); + break; + case INFINI_DTYPE_BF16: + launch_kernel(y, x, _info, stream); + break; + case INFINI_DTYPE_F32: + launch_kernel(y, x, _info, stream); + break; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::softmax::metax diff --git a/src/infiniop/ops/softmax/operator.cc b/src/infiniop/ops/softmax/operator.cc index 1664e257f..040c6160c 100644 --- a/src/infiniop/ops/softmax/operator.cc +++ b/src/infiniop/ops/softmax/operator.cc @@ -10,6 +10,10 @@ #include "nvidia/softmax_nvidia.cuh" #endif +#ifdef ENABLE_METAX_API +#include "metax/softmax_metax.h" +#endif + __INFINI_C infiniStatus_t infiniopCreateSoftmaxDescriptor( infiniopHandle_t handle, infiniopSoftmaxDescriptor_t *desc_ptr, @@ -43,6 +47,9 @@ __INFINI_C infiniStatus_t infiniopCreateSoftmaxDescriptor( #endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); +#endif +#ifdef ENABLE_METAX_API + CREATE(INFINI_DEVICE_METAX, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -74,6 +81,9 @@ __INFINI_C infiniStatus_t infiniopGetSoftmaxWorkspaceSize(infiniopSoftmaxDescrip #endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); +#endif +#ifdef ENABLE_METAX_API + GET(INFINI_DEVICE_METAX, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -110,6 +120,9 @@ __INFINI_C infiniStatus_t infiniopSoftmax( #endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); +#endif +#ifdef ENABLE_METAX_API + CALCULATE(INFINI_DEVICE_METAX, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -141,6 +154,9 @@ __INFINI_C infiniStatus_t infiniopDestroySoftmaxDescriptor(infiniopSoftmaxDescri #endif #ifdef ENABLE_CAMBRICON_API DESTROY(INFINI_DEVICE_CAMBRICON, bang); +#endif +#ifdef ENABLE_METAX_API + DESTROY(INFINI_DEVICE_METAX, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; From 548363fe580bde0fb9e1172476c2e6e02f364339 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Thu, 23 Jul 2026 17:11:01 +0800 Subject: [PATCH 2/2] feat: support conv on maca and hpcc --- src/infiniop/devices/metax/metax_ht2mc.h | 17 +++ src/infiniop/ops/conv/metax/conv_metax.cc | 131 +++++++++++----------- 2 files changed, 82 insertions(+), 66 deletions(-) diff --git a/src/infiniop/devices/metax/metax_ht2mc.h b/src/infiniop/devices/metax/metax_ht2mc.h index 9e0da85ef..51a4824e6 100644 --- a/src/infiniop/devices/metax/metax_ht2mc.h +++ b/src/infiniop/devices/metax/metax_ht2mc.h @@ -79,6 +79,7 @@ #define hcdnnCreateTensorDescriptor mcdnnCreateTensorDescriptor #define hcdnnCreateTensorDescriptor mcdnnCreateTensorDescriptor #define hcdnnCreatePoolingDescriptor mcdnnCreatePoolingDescriptor +#define hcdnnDestroy mcdnnDestroy #define hcdnnCreate mcdnnCreate #define HCDNN_TENSOR_NCHW MCDNN_TENSOR_NCHW #define HCDNN_STATUS_SUCCESS MCDNN_STATUS_SUCCESS @@ -153,4 +154,20 @@ #define HCBLAS_POINTER_MODE_HOST MCBLAS_POINTER_MODE_HOST #define __hpcc_fp8_e4m3 __maca_fp8_e4m3 #define __hpcc_bfloat16 __maca_bfloat16 +#define hcdnnActivationDescriptor_t mcdnnActivationDescriptor_t +#define hcdnnConvolutionDescriptor_t mcdnnConvolutionDescriptor_t +#define hcdnnConvolutionFwdAlgo_t mcdnnConvolutionFwdAlgo_t +#define hcdnnFilterDescriptor_t mcdnnFilterDescriptor_t +#define hcdnnCreateConvolutionDescriptor mcdnnCreateConvolutionDescriptor +#define hcdnnDestroyConvolutionDescriptor mcdnnDestroyConvolutionDescriptor +#define hcdnnSetConvolutionNdDescriptor mcdnnSetConvolutionNdDescriptor +#define hcdnnCreateFilterDescriptor mcdnnCreateFilterDescriptor +#define hcdnnDestroyFilterDescriptor mcdnnDestroyFilterDescriptor +#define hcdnnSetFilterNdDescriptor mcdnnSetFilterNdDescriptor +#define hcdnnGetConvolutionForwardWorkspaceSize mcdnnGetConvolutionForwardWorkspaceSize +#define hcdnnConvolutionForward mcdnnConvolutionForward +#define hcdnnAddTensor mcdnnAddTensor +#define hcdnnDestroyActivationDescriptor mcdnnDestroyActivationDescriptor +#define HCDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM MCDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM +#define HCDNN_CONVOLUTION MCDNN_CONVOLUTION #endif diff --git a/src/infiniop/ops/conv/metax/conv_metax.cc b/src/infiniop/ops/conv/metax/conv_metax.cc index 7778501cd..3262a2e32 100644 --- a/src/infiniop/ops/conv/metax/conv_metax.cc +++ b/src/infiniop/ops/conv/metax/conv_metax.cc @@ -1,6 +1,5 @@ #include "conv_metax.h" #include "../../../devices/metax/metax_common.h" -#include "mcdnn/mcdnn.h" // Metax DNN Library Header #include #include @@ -11,28 +10,28 @@ namespace op::conv::metax { // --------------------------------------------------------------------------- /** - * @brief Maps InfiniCore data types to mcDNN data types. + * @brief Maps InfiniCore data types to hcDNN data types. * @param dtype The InfiniCore data type. - * @return The corresponding mcDNN data type. + * @return The corresponding hcDNN data type. */ -inline mcdnnDataType_t toMcDNNDataType(infiniDtype_t dtype) { +inline hcdnnDataType_t toHcDNNDataType(infiniDtype_t dtype) { switch (dtype) { case INFINI_DTYPE_F32: - return MCDNN_DATA_FLOAT; + return HCDNN_DATA_FLOAT; case INFINI_DTYPE_F16: - return MCDNN_DATA_HALF; + return HCDNN_DATA_HALF; case INFINI_DTYPE_BF16: - return MCDNN_DATA_BFLOAT16; + return HCDNN_DATA_BFLOAT16; default: - return MCDNN_DATA_FLOAT; + return HCDNN_DATA_FLOAT; } } /** - * @brief Checks mcDNN API return status and converts to InfiniCore status. + * @brief Checks hcDNN API return status and converts to InfiniCore status. */ -inline infiniStatus_t checkMcDNNStatus(mcdnnStatus_t status) { - if (status == MCDNN_STATUS_SUCCESS) { +inline infiniStatus_t checkHcDNNStatus(hcdnnStatus_t status) { + if (status == HCDNN_STATUS_SUCCESS) { return INFINI_STATUS_SUCCESS; } // Log the error if a logging mechanism is available @@ -40,38 +39,38 @@ inline infiniStatus_t checkMcDNNStatus(mcdnnStatus_t status) { } struct Descriptor::Opaque { - mcdnnHandle_t mcdnn_handle = nullptr; - mcdnnConvolutionDescriptor_t conv_desc = nullptr; - mcdnnFilterDescriptor_t filter_desc = nullptr; - mcdnnTensorDescriptor_t input_desc = nullptr; - mcdnnTensorDescriptor_t output_desc = nullptr; - mcdnnTensorDescriptor_t bias_desc = nullptr; - mcdnnActivationDescriptor_t activation_desc = nullptr; - mcdnnConvolutionFwdAlgo_t algo = MCDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM; + hcdnnHandle_t hcdnn_handle = nullptr; + hcdnnConvolutionDescriptor_t conv_desc = nullptr; + hcdnnFilterDescriptor_t filter_desc = nullptr; + hcdnnTensorDescriptor_t input_desc = nullptr; + hcdnnTensorDescriptor_t output_desc = nullptr; + hcdnnTensorDescriptor_t bias_desc = nullptr; + hcdnnActivationDescriptor_t activation_desc = nullptr; + hcdnnConvolutionFwdAlgo_t algo = HCDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM; bool direct_conv2d_kernel = false; bool patch_embed_kernel = false; ~Opaque() { - if (mcdnn_handle) { - mcdnnDestroy(mcdnn_handle); + if (hcdnn_handle) { + hcdnnDestroy(hcdnn_handle); } if (conv_desc) { - mcdnnDestroyConvolutionDescriptor(conv_desc); + hcdnnDestroyConvolutionDescriptor(conv_desc); } if (filter_desc) { - mcdnnDestroyFilterDescriptor(filter_desc); + hcdnnDestroyFilterDescriptor(filter_desc); } if (input_desc) { - mcdnnDestroyTensorDescriptor(input_desc); + hcdnnDestroyTensorDescriptor(input_desc); } if (output_desc) { - mcdnnDestroyTensorDescriptor(output_desc); + hcdnnDestroyTensorDescriptor(output_desc); } if (bias_desc) { - mcdnnDestroyTensorDescriptor(bias_desc); + hcdnnDestroyTensorDescriptor(bias_desc); } if (activation_desc) { - mcdnnDestroyActivationDescriptor(activation_desc); + hcdnnDestroyActivationDescriptor(activation_desc); } } }; @@ -102,7 +101,7 @@ infiniStatus_t Descriptor::create( size_t workspace_size = 0; auto opaque = new Opaque(); - mcdnnStatus_t status; + hcdnnStatus_t status; const bool can_use_direct_conv2d_kernel = info.ndim() == 2; @@ -125,9 +124,9 @@ infiniStatus_t Descriptor::create( } do { - // 1. Create mcDNN Handle - status = mcdnnCreate(&opaque->mcdnn_handle); - if (status != MCDNN_STATUS_SUCCESS) { + // 1. Create hcDNN Handle + status = hcdnnCreate(&opaque->hcdnn_handle); + if (status != HCDNN_STATUS_SUCCESS) { break; } @@ -141,15 +140,15 @@ infiniStatus_t Descriptor::create( dilation_vec[i] = static_cast(info.dilation_info(i)); } - status = mcdnnCreateConvolutionDescriptor(&opaque->conv_desc); - if (status != MCDNN_STATUS_SUCCESS) { + status = hcdnnCreateConvolutionDescriptor(&opaque->conv_desc); + if (status != HCDNN_STATUS_SUCCESS) { break; } - status = mcdnnSetConvolutionNdDescriptor( + status = hcdnnSetConvolutionNdDescriptor( opaque->conv_desc, info.ndim(), pad_vec.data(), stride_vec.data(), - dilation_vec.data(), MCDNN_CONVOLUTION, toMcDNNDataType(dtype)); - if (status != MCDNN_STATUS_SUCCESS) { + dilation_vec.data(), HCDNN_CONVOLUTION, toHcDNNDataType(dtype)); + if (status != HCDNN_STATUS_SUCCESS) { break; } @@ -161,15 +160,15 @@ infiniStatus_t Descriptor::create( filter_dim[i + 2] = info.kernel_dim(i); } - status = mcdnnCreateFilterDescriptor(&opaque->filter_desc); - if (status != MCDNN_STATUS_SUCCESS) { + status = hcdnnCreateFilterDescriptor(&opaque->filter_desc); + if (status != HCDNN_STATUS_SUCCESS) { break; } - status = mcdnnSetFilterNdDescriptor( - opaque->filter_desc, toMcDNNDataType(dtype), MCDNN_TENSOR_NCHW, + status = hcdnnSetFilterNdDescriptor( + opaque->filter_desc, toHcDNNDataType(dtype), HCDNN_TENSOR_NCHW, info.ndim() + 2, filter_dim.data()); - if (status != MCDNN_STATUS_SUCCESS) { + if (status != HCDNN_STATUS_SUCCESS) { break; } @@ -181,14 +180,14 @@ infiniStatus_t Descriptor::create( input_dim[i + 2] = info.input_dim(i); } - status = mcdnnCreateTensorDescriptor(&opaque->input_desc); - if (status != MCDNN_STATUS_SUCCESS) { + status = hcdnnCreateTensorDescriptor(&opaque->input_desc); + if (status != HCDNN_STATUS_SUCCESS) { break; } - status = mcdnnSetTensorNdDescriptor(opaque->input_desc, toMcDNNDataType(dtype), + status = hcdnnSetTensorNdDescriptor(opaque->input_desc, toHcDNNDataType(dtype), info.ndim() + 2, input_dim.data(), nullptr); - if (status != MCDNN_STATUS_SUCCESS) { + if (status != HCDNN_STATUS_SUCCESS) { break; } @@ -199,14 +198,14 @@ infiniStatus_t Descriptor::create( output_dim[i + 2] = info.output_dim(i); } - status = mcdnnCreateTensorDescriptor(&opaque->output_desc); - if (status != MCDNN_STATUS_SUCCESS) { + status = hcdnnCreateTensorDescriptor(&opaque->output_desc); + if (status != HCDNN_STATUS_SUCCESS) { break; } - status = mcdnnSetTensorNdDescriptor(opaque->output_desc, toMcDNNDataType(dtype), + status = hcdnnSetTensorNdDescriptor(opaque->output_desc, toHcDNNDataType(dtype), info.ndim() + 2, output_dim.data(), nullptr); - if (status != MCDNN_STATUS_SUCCESS) { + if (status != HCDNN_STATUS_SUCCESS) { break; } @@ -214,29 +213,29 @@ infiniStatus_t Descriptor::create( std::vector bias_dim(info.ndim() + 2, 1); bias_dim[1] = info.out_channels(); // 1xCx1x1 - status = mcdnnCreateTensorDescriptor(&opaque->bias_desc); - if (status != MCDNN_STATUS_SUCCESS) { + status = hcdnnCreateTensorDescriptor(&opaque->bias_desc); + if (status != HCDNN_STATUS_SUCCESS) { break; } - status = mcdnnSetTensorNdDescriptor( + status = hcdnnSetTensorNdDescriptor( opaque->bias_desc, - toMcDNNDataType(dtype), + toHcDNNDataType(dtype), info.ndim() + 2, bias_dim.data(), nullptr); - if (status != MCDNN_STATUS_SUCCESS) { + if (status != HCDNN_STATUS_SUCCESS) { break; } } // 5. Get Workspace Size - status = mcdnnGetConvolutionForwardWorkspaceSize( - opaque->mcdnn_handle, + status = hcdnnGetConvolutionForwardWorkspaceSize( + opaque->hcdnn_handle, opaque->input_desc, opaque->filter_desc, opaque->conv_desc, opaque->output_desc, opaque->algo, &workspace_size); - if (status != MCDNN_STATUS_SUCCESS) { + if (status != HCDNN_STATUS_SUCCESS) { workspace_size = 0; } @@ -251,7 +250,7 @@ infiniStatus_t Descriptor::create( // Error Handling: Cleanup if (opaque) { - // auto clean mcdnn_handle and descriptors + // auto clean hcdnn_handle and descriptors delete opaque; } return INFINI_STATUS_INTERNAL_ERROR; @@ -278,15 +277,15 @@ infiniStatus_t Descriptor::calculate( return launchPatchEmbedConv3d(_dtype, _info, y, x, w, bias, stream); } - mcdnnStatus_t status; + hcdnnStatus_t status; float alpha = 1.0f; float beta = 0.0f; // --------------------------------------------------------------- // Step 1: y = conv(x, w) // --------------------------------------------------------------- - status = mcdnnConvolutionForward( - _opaque->mcdnn_handle, + status = hcdnnConvolutionForward( + _opaque->hcdnn_handle, &alpha, _opaque->input_desc, x, _opaque->filter_desc, w, @@ -296,25 +295,25 @@ infiniStatus_t Descriptor::calculate( &beta, _opaque->output_desc, y); - CHECK_STATUS(checkMcDNNStatus(status)); + CHECK_STATUS(checkHcDNNStatus(status)); // --------------------------------------------------------------- // Step 2: y = y + bias // --------------------------------------------------------------- if (bias != nullptr) { - // mcdnnAddTensor: y = alpha * bias + beta * y + // hcdnnAddTensor: y = alpha * bias + beta * y // we need do: y = 1.0 * bias + 1.0 * y_old float alpha_add = 1.0f; float beta_add = 1.0f; - status = mcdnnAddTensor( - _opaque->mcdnn_handle, + status = hcdnnAddTensor( + _opaque->hcdnn_handle, &alpha_add, _opaque->bias_desc, bias, &beta_add, _opaque->output_desc, y); - CHECK_STATUS(checkMcDNNStatus(status)); + CHECK_STATUS(checkHcDNNStatus(status)); } return INFINI_STATUS_SUCCESS;