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;