Revert "Merge pull request #1056 from InfiniTensor/issue/1031"

This reverts commit 7f295448, reversing
changes made to e60985dc.

src/infiniop/ops/avg_pool1d/cuda/kernel.cuh

-#ifndef __INFINIOP_AVG_POOL1D_CUDA_KERNEL_CUH__
-#define __INFINIOP_AVG_POOL1D_CUDA_KERNEL_CUH__
-
-template <typename T>
-__device__ void avgPool1dKernel(
-    T *y,
-    const T *x,
-    size_t batch,
-    size_t channels,
-    size_t in_width,
-    size_t out_width,
-    size_t kernel_size,
-    size_t stride,
-    size_t padding,
-
-    ptrdiff_t y_stride_batch,
-    ptrdiff_t y_stride_channel,
-    ptrdiff_t y_stride_width,
-    ptrdiff_t x_stride_batch,
-    ptrdiff_t x_stride_channel,
-    ptrdiff_t x_stride_width) {
-
-    size_t total_elements = batch * channels * out_width;
-
-    for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
-         idx < total_elements;
-         idx += gridDim.x * blockDim.x) {
-
-        size_t ow = idx % out_width;
-        size_t temp = idx / out_width;
-        size_t c = temp % channels;
-        size_t b = temp / channels;
-
-        size_t y_offset = b * y_stride_batch + c * y_stride_channel + ow * y_stride_width;
-
-        long long start_w = static_cast<long long>(ow * stride) - padding;
-
-        T sum = 0;
-
-        for (size_t k = 0; k < kernel_size; ++k) {
-            long long iw = start_w + k;
-
-            if (iw >= 0 && iw < static_cast<long long>(in_width)) {
-                size_t x_offset = b * x_stride_batch + c * x_stride_channel + iw * x_stride_width;
-                sum += x[x_offset];
-            }
-        }
-
-#if defined(ENABLE_ILUVATAR_API)
-        // Iluvatar __half doesn't accept size_t directly.
-        y[y_offset] = sum / static_cast<T>(static_cast<double>(kernel_size));
-#else
-        y[y_offset] = sum / static_cast<T>(kernel_size);
-#endif
-    }
-}
-
-#endif

src/infiniop/ops/avg_pool1d/metax/avg_pool1d_metax.h

-#ifndef __INFINIOP_AVG_POOL1D_METAX_H__
-#define __INFINIOP_AVG_POOL1D_METAX_H__
-
-#include "../avg_pool1d.h"
-
-DESCRIPTOR(metax)
-
-#endif // __INFINIOP_AVG_POOL1D_METAX_H__

src/infiniop/ops/avg_pool1d/metax/avg_pool1d_metax.maca

-#include "../../../devices/metax/metax_common.h"
-#include "avg_pool1d_metax.h"
-#include "../../../devices/metax/metax_kernel_common.h"
-
-#include <type_traits>
-
-namespace op::avg_pool1d::metax {
-
-struct Descriptor::Opaque {
-    std::shared_ptr<device::metax::Handle::Internal> internal;
-};
-
-Descriptor::~Descriptor() {
-    delete _opaque;
-}
-
-infiniStatus_t Descriptor::create(
-    infiniopHandle_t handle_,
-    Descriptor **desc_ptr,
-    infiniopTensorDescriptor_t y_desc,
-    infiniopTensorDescriptor_t x_desc,
-    size_t kernel_size,
-    size_t stride,
-    size_t padding) {
-
-    auto handle = reinterpret_cast<device::metax::Handle *>(handle_);
-
-    auto info = AvgPool1dInfo::createAvgPool1dInfo(y_desc, x_desc, kernel_size, stride, padding);
-    CHECK_RESULT(info);
-
-    *desc_ptr = new Descriptor(
-        info.take(),
-        0,
-        new Opaque{handle->internal()},
-        handle->device,
-        handle->device_id);
-
-    return INFINI_STATUS_SUCCESS;
-}
-
-template <typename Tdata, typename Tcompute>
-__device__ __forceinline__ Tdata castToOutput(Tcompute val) {
-    if constexpr (std::is_same_v<Tdata, half>) {
-        return __float2half(static_cast<float>(val));
-    } else if constexpr (std::is_same_v<Tdata, cuda_bfloat16>) {
-        return __float2bfloat16(static_cast<float>(val));
-    } else {
-        return static_cast<Tdata>(val);
-    }
-}
-
-template <typename Tdata, typename Tcompute>
-INFINIOP_METAX_KERNEL avgPool1dGlobalKernel(
-    Tdata *y,
-    const Tdata *x,
-    size_t batch,
-    size_t channels,
-    size_t in_width,
-    size_t out_width,
-    size_t kernel_size,
-    size_t stride,
-    size_t padding,
-    ptrdiff_t y_stride_batch,
-    ptrdiff_t y_stride_channel,
-    ptrdiff_t y_stride_width,
-    ptrdiff_t x_stride_batch,
-    ptrdiff_t x_stride_channel,
-    ptrdiff_t x_stride_width) {
-
-    size_t total_elements = batch * channels * out_width;
-    Tcompute inv_kernel = Tcompute(1) / static_cast<Tcompute>(kernel_size);
-
-    for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
-         idx < total_elements;
-         idx += gridDim.x * blockDim.x) {
-
-        size_t ow = idx % out_width;
-        size_t temp = idx / out_width;
-        size_t c = temp % channels;
-        size_t b = temp / channels;
-
-        size_t y_offset = b * y_stride_batch + c * y_stride_channel + ow * y_stride_width;
-        size_t x_base = b * x_stride_batch + c * x_stride_channel;
-
-        long long start_w = static_cast<long long>(ow * stride) - static_cast<long long>(padding);
-        long long end_w = start_w + static_cast<long long>(kernel_size);
-        long long iw_start = start_w < 0 ? 0 : start_w;
-        long long iw_end = end_w > static_cast<long long>(in_width) ? static_cast<long long>(in_width) : end_w;
-
-        Tcompute sum = Tcompute(0);
-        if (iw_start < iw_end) {
-            size_t x_offset = x_base + static_cast<size_t>(iw_start) * x_stride_width;
-            for (long long iw = iw_start; iw < iw_end; ++iw) {
-                sum += static_cast<Tcompute>(x[x_offset]);
-                x_offset += x_stride_width;
-            }
-        }
-
-        y[y_offset] = castToOutput<Tdata, Tcompute>(sum * inv_kernel);
-    }
-}
-
-template <typename Tdata, typename Tcompute>
-infiniStatus_t calculateAvgPool1d(
-    const AvgPool1dInfo &info,
-    int max_threads_per_block,
-    Tdata *y,
-    const Tdata *x,
-    hcStream_t stream) {
-
-    size_t total_elements = info.batch * info.channels * info.out_width;
-
-    int block_size = 256;
-    if (max_threads_per_block > 0 && max_threads_per_block < block_size) {
-        block_size = max_threads_per_block;
-    }
-
-    size_t grid_size = (total_elements + block_size - 1) / block_size;
-    if (grid_size > 65535) {
-        grid_size = 65535;
-    }
-
-    avgPool1dGlobalKernel<Tdata, Tcompute><<<grid_size, block_size, 0, stream>>>(
-        y, x,
-        info.batch, info.channels, info.in_width, info.out_width,
-        info.kernel_size, info.stride, info.padding,
-        info.y_stride_batch, info.y_stride_channel, info.y_stride_width,
-        info.x_stride_batch, info.x_stride_channel, info.x_stride_width);
-
-    return INFINI_STATUS_SUCCESS;
-}
-
-#define CALCULATE(TDATA, TCOMPUTE) \
-    calculateAvgPool1d<TDATA, TCOMPUTE>( \
-        _info, \
-        _opaque->internal->maxThreadsPerBlock(), \
-        (TDATA *)y, \
-        (const TDATA *)x, \
-        (hcStream_t)stream)
-
-infiniStatus_t Descriptor::calculate(
-    void *workspace,
-    size_t workspace_size,
-    void *y,
-    const void *x,
-    void *stream) const {
-
-    (void)workspace;
-
-    if (workspace_size < _workspace_size) {
-        return INFINI_STATUS_INSUFFICIENT_WORKSPACE;
-    }
-
-    switch (_info.dtype) {
-    case INFINI_DTYPE_F16:
-        return CALCULATE(half, float);
-    case INFINI_DTYPE_BF16:
-        return CALCULATE(cuda_bfloat16, float);
-    case INFINI_DTYPE_F32:
-        return CALCULATE(float, float);
-    case INFINI_DTYPE_F64:
-        return CALCULATE(double, double);
-    default:
-        return INFINI_STATUS_BAD_TENSOR_DTYPE;
-    }
-}
-
-#undef CALCULATE
-
-} // namespace op::avg_pool1d::metax

src/infiniop/ops/avg_pool1d/moore/avg_pool1d_kernel.h

-#ifndef __INFINIOP_AVG_POOL1D_MOORE_KERNEL_H__
-#define __INFINIOP_AVG_POOL1D_MOORE_KERNEL_H__
-
-#include <type_traits>
-
-namespace op::avg_pool1d::moore {
-
-template <typename Tdata, typename Tcompute>
-__device__ __forceinline__ Tdata castToOutput(Tcompute val) {
-    if constexpr (std::is_same_v<Tdata, half>) {
-        return __float2half(static_cast<float>(val));
-    } else if constexpr (std::is_same_v<Tdata, cuda_bfloat16>) {
-        return __float2bfloat16_rn(static_cast<float>(val));
-    } else {
-        return static_cast<Tdata>(val);
-    }
-}
-
-template <typename Tdata, typename Tcompute>
-__device__ void avgPool1dKernel(
-    Tdata *y,
-    const Tdata *x,
-    size_t batch,
-    size_t channels,
-    size_t in_width,
-    size_t out_width,
-    size_t kernel_size,
-    size_t stride,
-    size_t padding,
-    ptrdiff_t y_stride_batch,
-    ptrdiff_t y_stride_channel,
-    ptrdiff_t y_stride_width,
-    ptrdiff_t x_stride_batch,
-    ptrdiff_t x_stride_channel,
-    ptrdiff_t x_stride_width) {
-
-    size_t total_elements = batch * channels * out_width;
-    Tcompute inv_kernel = Tcompute(1) / static_cast<Tcompute>(kernel_size);
-
-    for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
-         idx < total_elements;
-         idx += gridDim.x * blockDim.x) {
-
-        size_t ow = idx % out_width;
-        size_t temp = idx / out_width;
-        size_t c = temp % channels;
-        size_t b = temp / channels;
-
-        size_t y_offset = b * y_stride_batch + c * y_stride_channel + ow * y_stride_width;
-        size_t x_base = b * x_stride_batch + c * x_stride_channel;
-
-        long long start_w = static_cast<long long>(ow * stride) - static_cast<long long>(padding);
-        long long end_w = start_w + static_cast<long long>(kernel_size);
-        long long iw_start = start_w < 0 ? 0 : start_w;
-        long long iw_end = end_w > static_cast<long long>(in_width) ? static_cast<long long>(in_width) : end_w;
-
-        Tcompute sum = Tcompute(0);
-        if (iw_start < iw_end) {
-            size_t x_offset = x_base + static_cast<size_t>(iw_start) * x_stride_width;
-            for (long long iw = iw_start; iw < iw_end; ++iw) {
-                sum += static_cast<Tcompute>(x[x_offset]);
-                x_offset += x_stride_width;
-            }
-        }
-
-        y[y_offset] = castToOutput<Tdata, Tcompute>(sum * inv_kernel);
-    }
-}
-
-} // namespace op::avg_pool1d::moore
-
-#endif // __INFINIOP_AVG_POOL1D_MOORE_KERNEL_H__

src/infiniop/ops/avg_pool1d/moore/avg_pool1d_moore.h

-#ifndef __INFINIOP_AVG_POOL1D_MOORE_H__
-#define __INFINIOP_AVG_POOL1D_MOORE_H__
-
-#include "../avg_pool1d.h"
-
-DESCRIPTOR(moore)
-
-#endif // __INFINIOP_AVG_POOL1D_MOORE_H__

src/infiniop/ops/avg_pool1d/moore/avg_pool1d_moore.mu

-#include "../../../devices/moore/moore_common.h"
-#include "avg_pool1d_moore.h"
-
-#include "../../../devices/moore/moore_kernel_common.h"
-
-#include "avg_pool1d_kernel.h"
-
-namespace op::avg_pool1d::moore {
-
-struct Descriptor::Opaque {
-    std::shared_ptr<device::moore::Handle::Internal> internal;
-};
-
-Descriptor::~Descriptor() {
-    delete _opaque;
-}
-
-infiniStatus_t Descriptor::create(
-    infiniopHandle_t handle_,
-    Descriptor **desc_ptr,
-    infiniopTensorDescriptor_t y_desc,
-    infiniopTensorDescriptor_t x_desc,
-    size_t kernel_size,
-    size_t stride,
-    size_t padding) {
-
-    auto handle = reinterpret_cast<device::moore::Handle *>(handle_);
-
-    auto info = AvgPool1dInfo::createAvgPool1dInfo(y_desc, x_desc, kernel_size, stride, padding);
-    CHECK_RESULT(info);
-
-    *desc_ptr = new Descriptor(
-        info.take(),
-        0,
-        new Opaque{handle->internal()},
-        handle->device,
-        handle->device_id);
-
-    return INFINI_STATUS_SUCCESS;
-}
-
-template <typename Tdata, typename Tcompute>
-INFINIOP_MOORE_KERNEL avgPool1dGlobalKernel(
-    Tdata *y,
-    const Tdata *x,
-    size_t batch,
-    size_t channels,
-    size_t in_width,
-    size_t out_width,
-    size_t kernel_size,
-    size_t stride,
-    size_t padding,
-    ptrdiff_t y_stride_batch,
-    ptrdiff_t y_stride_channel,
-    ptrdiff_t y_stride_width,
-    ptrdiff_t x_stride_batch,
-    ptrdiff_t x_stride_channel,
-    ptrdiff_t x_stride_width) {
-
-    avgPool1dKernel<Tdata, Tcompute>(
-        y, x,
-        batch, channels, in_width, out_width,
-        kernel_size, stride, padding,
-        y_stride_batch, y_stride_channel, y_stride_width,
-        x_stride_batch, x_stride_channel, x_stride_width);
-}
-
-template <typename Tdata, typename Tcompute>
-infiniStatus_t calculateAvgPool1d(
-    const AvgPool1dInfo &info,
-    int max_threads_per_block,
-    Tdata *y,
-    const Tdata *x,
-    musaStream_t stream) {
-
-    size_t total_elements = info.batch * info.channels * info.out_width;
-
-    int block_size = 256;
-    if (max_threads_per_block > 0 && max_threads_per_block < block_size) {
-        block_size = max_threads_per_block;
-    }
-
-    size_t grid_size = (total_elements + block_size - 1) / block_size;
-    if (grid_size > 65535) {
-        grid_size = 65535;
-    }
-
-    avgPool1dGlobalKernel<Tdata, Tcompute><<<grid_size, block_size, 0, stream>>>(
-        y, x,
-        info.batch, info.channels, info.in_width, info.out_width,
-        info.kernel_size, info.stride, info.padding,
-        info.y_stride_batch, info.y_stride_channel, info.y_stride_width,
-        info.x_stride_batch, info.x_stride_channel, info.x_stride_width);
-
-    return INFINI_STATUS_SUCCESS;
-}
-
-#define CALCULATE(TDATA, TCOMPUTE) \
-    calculateAvgPool1d<TDATA, TCOMPUTE>(\
-        _info,\
-        _opaque->internal->maxThreadsPerBlock(),\
-        (TDATA *)y,\
-        (const TDATA *)x,\
-        (musaStream_t)stream)
-
-infiniStatus_t Descriptor::calculate(
-    void *workspace,
-    size_t workspace_size,
-    void *y,
-    const void *x,
-    void *stream) const {
-
-    (void)workspace;
-
-    if (workspace_size < _workspace_size) {
-        return INFINI_STATUS_INSUFFICIENT_WORKSPACE;
-    }
-
-    switch (_info.dtype) {
-    case INFINI_DTYPE_F16:
-        return CALCULATE(half, float);
-    case INFINI_DTYPE_BF16:
-        return CALCULATE(cuda_bfloat16, float);
-    case INFINI_DTYPE_F32:
-        return CALCULATE(float, float);
-    case INFINI_DTYPE_F64:
-        return CALCULATE(double, double);
-    default:
-        return INFINI_STATUS_BAD_TENSOR_DTYPE;
-    }
-}
-
-#undef CALCULATE
-
-} // namespace op::avg_pool1d::moore

src/infiniop/ops/avg_pool1d/nvidia/avg_pool1d_nvidia.cu

-#include "../../../devices/nvidia/nvidia_common.cuh"
-#include "../../../devices/nvidia/nvidia_kernel_common.cuh"
-#include "../cuda/kernel.cuh"
-#include "avg_pool1d_nvidia.cuh"
-
-template <typename T>
-__global__ void avgPool1dGlobalKernel(
-    T *y,
-    const T *x,
-    size_t batch,
-    size_t channels,
-    size_t in_width,
-    size_t out_width,
-    size_t kernel_size,
-    size_t stride,
-    size_t padding,
-    ptrdiff_t y_stride_batch,
-    ptrdiff_t y_stride_channel,
-    ptrdiff_t y_stride_width,
-    ptrdiff_t x_stride_batch,
-    ptrdiff_t x_stride_channel,
-    ptrdiff_t x_stride_width) {
-
-    avgPool1dKernel<T>(
-        y, x,
-        batch, channels, in_width, out_width,
-        kernel_size, stride, padding,
-        y_stride_batch, y_stride_channel, y_stride_width,
-        x_stride_batch, x_stride_channel, x_stride_width);
-}
-
-namespace op::avg_pool1d::nvidia {
-
-struct Descriptor::Opaque {
-    std::shared_ptr<device::nvidia::Handle::Internal> internal;
-};
-
-Descriptor::~Descriptor() {
-    delete _opaque;
-}
-
-infiniStatus_t Descriptor::create(
-    infiniopHandle_t handle_,
-    Descriptor **desc_ptr,
-    infiniopTensorDescriptor_t y_desc,
-    infiniopTensorDescriptor_t x_desc,
-    size_t kernel_size,
-    size_t stride,
-    size_t padding) {
-
-    auto handle = reinterpret_cast<device::nvidia::Handle *>(handle_);
-
-    auto info = AvgPool1dInfo::createAvgPool1dInfo(y_desc, x_desc, kernel_size, stride, padding);
-    CHECK_RESULT(info);
-
-    *desc_ptr = new Descriptor(
-        info.take(),
-        0,
-        new Opaque{reinterpret_cast<device::nvidia::Handle *>(handle)->internal()},
-        handle->device,
-        handle->device_id);
-
-    return INFINI_STATUS_SUCCESS;
-}
-
-template <typename T>
-infiniStatus_t calculateAvgPool1d(
-    const AvgPool1dInfo &info,
-    int max_threads_per_block,
-    T *y,
-    const T *x,
-    cudaStream_t stream) {
-
-    size_t total_elements = info.batch * info.channels * info.out_width;
-
-    int block_size = 256;
-    if (max_threads_per_block > 0 && max_threads_per_block < 256) {
-        block_size = max_threads_per_block;
-    }
-
-    size_t grid_size = (total_elements + block_size - 1) / block_size;
-    if (grid_size > 65535) {
-        grid_size = 65535;
-    }
-
-    avgPool1dGlobalKernel<T><<<grid_size, block_size, 0, stream>>>(
-        y, x,
-        info.batch, info.channels, info.in_width, info.out_width,
-        info.kernel_size, info.stride, info.padding,
-        info.y_stride_batch, info.y_stride_channel, info.y_stride_width,
-        info.x_stride_batch, info.x_stride_channel, info.x_stride_width);
-
-    return INFINI_STATUS_SUCCESS;
-}
-
-#define CALCULATE(TDATA)                                        \
-    calculateAvgPool1d(_info,                                   \
-                       _opaque->internal->maxThreadsPerBlock(), \
-                       (TDATA *)y,                              \
-                       (const TDATA *)x,                        \
-                       (cudaStream_t)stream)
-
-infiniStatus_t Descriptor::calculate(
-    void *workspace,
-    size_t workspace_size,
-    void *y,
-    const void *x,
-    void *stream) const {
-
-    switch (_info.dtype) {
-    case INFINI_DTYPE_F16:
-        return CALCULATE(half);
-    case INFINI_DTYPE_BF16:
-        return CALCULATE(cuda_bfloat16);
-    case INFINI_DTYPE_F32:
-        return CALCULATE(float);
-    case INFINI_DTYPE_F64:
-        return CALCULATE(double);
-    default:
-        return INFINI_STATUS_BAD_TENSOR_DTYPE;
-    }
-}
-
-#undef CALCULATE
-
-} // namespace op::avg_pool1d::nvidia

src/infiniop/ops/avg_pool1d/nvidia/avg_pool1d_nvidia.cuh

-#ifndef __INFINIOP_AVG_POOL1D_CUDA_H__
-#define __INFINIOP_AVG_POOL1D_CUDA_H__
-
-#include "../avg_pool1d.h"
-
-DESCRIPTOR(nvidia)
-
-#endif

src/infiniop/ops/avg_pool1d/operator.cc

-#include "../../operator.h"
-#include "../../handle.h"
-#include "infiniop/ops/avg_pool1d.h"
-
-#ifdef ENABLE_CPU_API
-#include "cpu/avg_pool1d_cpu.h"
-#endif
-#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API)
-#include "nvidia/avg_pool1d_nvidia.cuh"
-#endif
-#ifdef ENABLE_ASCEND_API
-#include "ascend/avg_pool1d_ascend.h"
-#endif
-#ifdef ENABLE_CAMBRICON_API
-#include "bang/avg_pool1d_bang.h"
-#endif
-#ifdef ENABLE_METAX_API
-#include "metax/avg_pool1d_metax.h"
-#endif
-#ifdef ENABLE_KUNLUN_API
-#include "kunlun/avg_pool1d_kunlun.h"
-#endif
-#ifdef ENABLE_MOORE_API
-#include "moore/avg_pool1d_moore.h"
-#endif
-
-__INFINI_C infiniStatus_t infiniopCreateAvgPool1dDescriptor(
-    infiniopHandle_t handle,
-    infiniopAvgPool1dDescriptor_t *desc_ptr,
-    infiniopTensorDescriptor_t y,
-    infiniopTensorDescriptor_t x,
-    size_t kernel_size,
-    size_t stride,
-    size_t padding) {
-
-#define CREATE(CASE, NAMESPACE)                                                   \
-    case CASE:                                                                    \
-        return op::avg_pool1d::NAMESPACE::Descriptor::create(                     \
-            handle,                                                               \
-            reinterpret_cast<op::avg_pool1d::NAMESPACE::Descriptor **>(desc_ptr), \
-            y,                                                                    \
-            x,                                                                    \
-            kernel_size,                                                          \
-            stride,                                                               \
-            padding)
-
-    switch (handle->device) {
-#ifdef ENABLE_CPU_API
-        CREATE(INFINI_DEVICE_CPU, cpu);
-#endif
-#ifdef ENABLE_NVIDIA_API
-        CREATE(INFINI_DEVICE_NVIDIA, nvidia);
-#endif
-#ifdef ENABLE_ILUVATAR_API
-        CREATE(INFINI_DEVICE_ILUVATAR, nvidia);
-#endif
-#ifdef ENABLE_QY_API
-        CREATE(INFINI_DEVICE_QY, nvidia);
-#endif
-#ifdef ENABLE_HYGON_API
-        CREATE(INFINI_DEVICE_HYGON, nvidia);
-#endif
-#ifdef ENABLE_MOORE_API
-        CREATE(INFINI_DEVICE_MOORE, moore);
-#endif
-#ifdef ENABLE_METAX_API
-        CREATE(INFINI_DEVICE_METAX, metax);
-#endif
-#ifdef ENABLE_ASCEND_API
-        CREATE(INFINI_DEVICE_ASCEND, ascend);
-#endif
-#ifdef ENABLE_KUNLUN_API
-        CREATE(INFINI_DEVICE_KUNLUN, kunlun);
-#endif
-#ifdef ENABLE_CAMBRICON_API
-        CREATE(INFINI_DEVICE_CAMBRICON, bang);
-#endif
-    default:
-        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
-    }
-
-#undef CREATE
-}
-
-__INFINI_C infiniStatus_t infiniopGetAvgPool1dWorkspaceSize(infiniopAvgPool1dDescriptor_t desc,
-                                                            size_t *size) {
-#define GET(CASE, NAMESPACE)                                                                            \
-    case CASE:                                                                                          \
-        *size = reinterpret_cast<const op::avg_pool1d::NAMESPACE::Descriptor *>(desc)->workspaceSize(); \
-        return INFINI_STATUS_SUCCESS
-
-    switch (desc->device_type) {
-#ifdef ENABLE_CPU_API
-        GET(INFINI_DEVICE_CPU, cpu);
-#endif
-#ifdef ENABLE_NVIDIA_API
-        GET(INFINI_DEVICE_NVIDIA, nvidia);
-#endif
-#ifdef ENABLE_ILUVATAR_API
-        GET(INFINI_DEVICE_ILUVATAR, nvidia);
-#endif
-#ifdef ENABLE_QY_API
-        GET(INFINI_DEVICE_QY, nvidia);
-#endif
-#ifdef ENABLE_HYGON_API
-        GET(INFINI_DEVICE_HYGON, nvidia);
-#endif
-#ifdef ENABLE_MOORE_API
-        GET(INFINI_DEVICE_MOORE, moore);
-#endif
-#ifdef ENABLE_METAX_API
-        GET(INFINI_DEVICE_METAX, metax);
-#endif
-#ifdef ENABLE_KUNLUN_API
-        GET(INFINI_DEVICE_KUNLUN, kunlun);
-#endif
-#ifdef ENABLE_CAMBRICON_API
-        GET(INFINI_DEVICE_CAMBRICON, bang);
-#endif
-#ifdef ENABLE_ASCEND_API
-        GET(INFINI_DEVICE_ASCEND, ascend);
-#endif
-    default:
-        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
-    }
-
-#undef GET
-}
-
-__INFINI_C infiniStatus_t infiniopAvgPool1d(
-    infiniopAvgPool1dDescriptor_t desc,
-    void *workspace,
-    size_t workspace_size,
-    void *y,
-    const void *x,
-    void *stream) {
-
-#define CALCULATE(CASE, NAMESPACE)                                                   \
-    case CASE:                                                                       \
-        return reinterpret_cast<const op::avg_pool1d::NAMESPACE::Descriptor *>(desc) \
-            ->calculate(workspace, workspace_size, y, x, stream)
-
-    switch (desc->device_type) {
-#ifdef ENABLE_CPU_API
-        CALCULATE(INFINI_DEVICE_CPU, cpu);
-#endif
-#ifdef ENABLE_NVIDIA_API
-        CALCULATE(INFINI_DEVICE_NVIDIA, nvidia);
-#endif
-#ifdef ENABLE_ILUVATAR_API
-        CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia);
-#endif
-#ifdef ENABLE_QY_API
-        CALCULATE(INFINI_DEVICE_QY, nvidia);
-#endif
-#ifdef ENABLE_HYGON_API
-        CALCULATE(INFINI_DEVICE_HYGON, nvidia);
-#endif
-#ifdef ENABLE_MOORE_API
-        CALCULATE(INFINI_DEVICE_MOORE, moore);
-#endif
-#ifdef ENABLE_METAX_API
-        CALCULATE(INFINI_DEVICE_METAX, metax);
-#endif
-#ifdef ENABLE_KUNLUN_API
-        CALCULATE(INFINI_DEVICE_KUNLUN, kunlun);
-#endif
-#ifdef ENABLE_CAMBRICON_API
-        CALCULATE(INFINI_DEVICE_CAMBRICON, bang);
-#endif
-#ifdef ENABLE_ASCEND_API
-        CALCULATE(INFINI_DEVICE_ASCEND, ascend);
-#endif
-    default:
-        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
-    }
-
-#undef CALCULATE
-}
-
-__INFINI_C infiniStatus_t
-infiniopDestroyAvgPool1dDescriptor(infiniopAvgPool1dDescriptor_t desc) {
-
-#define DELETE(CASE, NAMESPACE)                                                       \
-    case CASE:                                                                        \
-        delete reinterpret_cast<const op::avg_pool1d::NAMESPACE::Descriptor *>(desc); \
-        return INFINI_STATUS_SUCCESS;
-
-    switch (desc->device_type) {
-#ifdef ENABLE_CPU_API
-        DELETE(INFINI_DEVICE_CPU, cpu);
-#endif
-#ifdef ENABLE_NVIDIA_API
-        DELETE(INFINI_DEVICE_NVIDIA, nvidia);
-#endif
-#ifdef ENABLE_ILUVATAR_API
-        DELETE(INFINI_DEVICE_ILUVATAR, nvidia);
-#endif
-#ifdef ENABLE_QY_API
-        DELETE(INFINI_DEVICE_QY, nvidia);
-#endif
-#ifdef ENABLE_HYGON_API
-        DELETE(INFINI_DEVICE_HYGON, nvidia);
-#endif
-#ifdef ENABLE_MOORE_API
-        DELETE(INFINI_DEVICE_MOORE, moore);
-#endif
-#ifdef ENABLE_METAX_API
-        DELETE(INFINI_DEVICE_METAX, metax);
-#endif
-#ifdef ENABLE_KUNLUN_API
-        DELETE(INFINI_DEVICE_KUNLUN, kunlun);
-#endif
-#ifdef ENABLE_CAMBRICON_API
-        DELETE(INFINI_DEVICE_CAMBRICON, bang);
-#endif
-#ifdef ENABLE_ASCEND_API
-        DELETE(INFINI_DEVICE_ASCEND, ascend);
-#endif
-    default:
-        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
-    }
-
-#undef DELETE
-}

src/infiniop/ops/cross_entropy/cpu/cross_entropy_cpu.cc

-#include "cross_entropy_cpu.h"
-#include "../../../devices/cpu/common_cpu.h"
-#include "../../../reduce/cpu/reduce.h"
-#include <algorithm>
-#include <cmath>
-
-namespace op::cross_entropy::cpu {
-
-Descriptor::~Descriptor() = default;
-
-infiniStatus_t Descriptor::create(
-    infiniopHandle_t handle,
-    Descriptor **desc_ptr,
-    infiniopTensorDescriptor_t y_desc,
-    infiniopTensorDescriptor_t x_desc,
-    infiniopTensorDescriptor_t target_desc) {
-
-    auto x_dtype = x_desc->dtype();
-    auto t_dtype = target_desc->dtype();
-
-    CHECK_DTYPE(x_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16);
-    CHECK_DTYPE(t_dtype, INFINI_DTYPE_I32, INFINI_DTYPE_I64);
-
-    CrossEntropyInfo info{};
-    info.dtype = x_dtype;
-    info.target_dtype = t_dtype;
-
-    info.outer_size = target_desc->numel();
-
-    info.vocab_size = x_desc->shape().back();
-
-    info.x_stride = static_cast<ptrdiff_t>(info.vocab_size);
-
-    *desc_ptr = new Descriptor(nullptr, info, 0, handle->device, handle->device_id);
-    return INFINI_STATUS_SUCCESS;
-}
-
-template <typename T, typename Tidx>
-infiniStatus_t cross_entropy_kernel(const CrossEntropyInfo *info,
-                                    T *y, const T *x, const void *target) {
-    const Tidx *label = reinterpret_cast<const Tidx *>(target);
-
-#pragma omp parallel for
-    for (ptrdiff_t i = 0; i < ptrdiff_t(info->outer_size); ++i) {
-        const T *row = x + i * info->x_stride;
-        Tidx idx = label[i];
-
-        if (idx < 0 || static_cast<size_t>(idx) >= info->vocab_size) {
-            y[i] = utils::cast<T>(0.f);
-            continue;
-        }
-
-        float max_val = op::common_cpu::reduce_op::max(row, info->vocab_size, 1);
-
-        float sum_exp = 0.f;
-        for (size_t j = 0; j < info->vocab_size; ++j) {
-            sum_exp += std::exp(utils::cast<float>(row[j]) - max_val);
-        }
-
-        float log_term = std::log(sum_exp) + max_val;
-        float target_logit = utils::cast<float>(row[idx]);
-        y[i] = utils::cast<T>(log_term - target_logit);
-    }
-    return INFINI_STATUS_SUCCESS;
-}
-
-template <typename T>
-infiniStatus_t dispatch_target_type(const CrossEntropyInfo *info,
-                                    T *y, const T *x, const void *target) {
-
-    if (info->target_dtype == INFINI_DTYPE_I32) {
-        return cross_entropy_kernel<T, int32_t>(info, y, x, target);
-    } else if (info->target_dtype == INFINI_DTYPE_I64) {
-        return cross_entropy_kernel<T, int64_t>(info, y, x, target);
-    }
-    return INFINI_STATUS_BAD_TENSOR_DTYPE;
-}
-
-infiniStatus_t Descriptor::calculate(
-    void *workspace,
-    size_t workspace_size,
-    void *y,
-    const void *x,
-    const void *target,
-    void *stream) const {
-
-    switch (_info.dtype) {
-    case INFINI_DTYPE_F16:
-        return dispatch_target_type(&_info, (fp16_t *)y, (const fp16_t *)x, target);
-    case INFINI_DTYPE_BF16:
-        return dispatch_target_type(&_info, (bf16_t *)y, (const bf16_t *)x, target);
-    case INFINI_DTYPE_F32:
-        return dispatch_target_type(&_info, (float *)y, (const float *)x, target);
-    default:
-        return INFINI_STATUS_BAD_TENSOR_DTYPE;
-    }
-}
-
-} // namespace op::cross_entropy::cpu

src/infiniop/ops/cross_entropy/cpu/cross_entropy_cpu.h

-#ifndef __CROSS_ENTROPY_CPU_H__
-#define __CROSS_ENTROPY_CPU_H__
-
-#include "../cross_entropy.h"
-
-DESCRIPTOR(cpu)
-
-#endif

src/infiniop/ops/cross_entropy/cross_entropy.h

-#ifndef CROSS_ENTROPY_H
-#define CROSS_ENTROPY_H
-
-#include "../../operator.h"
-#include "info.h"
-
-#define DESCRIPTOR(NAMESPACE)                                                 \
-    namespace op::cross_entropy::NAMESPACE {                                  \
-    class Descriptor final : public InfiniopDescriptor {                      \
-        struct Opaque;                                                        \
-        Opaque *_opaque;                                                      \
-        CrossEntropyInfo _info;                                               \
-        size_t _workspace_size;                                               \
-                                                                              \
-        Descriptor(Opaque *opaque,                                            \
-                   CrossEntropyInfo info,                                     \
-                   size_t workspace_size,                                     \
-                   infiniDevice_t device_type,                                \
-                   int device_id)                                             \
-            : InfiniopDescriptor{device_type, device_id},                     \
-              _opaque(opaque),                                                \
-              _info(info),                                                    \
-              _workspace_size(workspace_size) {}                              \
-                                                                              \
-    public:                                                                   \
-        ~Descriptor();                                                        \
-        size_t workspaceSize() const { return _workspace_size; }              \
-        static infiniStatus_t create(infiniopHandle_t handle,                 \
-                                     Descriptor **desc_ptr,                   \
-                                     infiniopTensorDescriptor_t y_desc,       \
-                                     infiniopTensorDescriptor_t x_desc,       \
-                                     infiniopTensorDescriptor_t target_desc); \
-        infiniStatus_t calculate(void *workspace,                             \
-                                 size_t workspace_size,                       \
-                                 void *y,                                     \
-                                 const void *x,                               \
-                                 const void *target,                          \
-                                 void *stream) const;                         \
-    };                                                                        \
-    }
-
-#endif

src/infiniop/ops/cross_entropy/cuda/kernel.cuh

-#ifndef __CROSS_ENTROPY_KERNEL_CUH__
-#define __CROSS_ENTROPY_KERNEL_CUH__
-
-#include "../../../devices/nvidia/nvidia_common.cuh"
-#include "../../../reduce/cuda/reduce.cuh"
-
-template <unsigned int BLOCK_SIZE, typename Tdata, typename Tidx, typename Tcompute = float>
-__device__ void crossEntropyKernel(
-    Tdata *y_,
-    const Tdata *x_,
-    const void *target_,
-    size_t outer_size,
-    size_t vocab_size,
-    ptrdiff_t x_stride) {
-
-    size_t row_idx = blockIdx.x;
-    if (row_idx >= outer_size) {
-        return;
-    }
-
-    const Tdata *x = x_ + row_idx * x_stride;
-    const Tidx *target = reinterpret_cast<const Tidx *>(target_);
-
-    Tidx label = target[row_idx];
-
-    Tdata max_val_raw = op::common_cuda::reduce_op::max<BLOCK_SIZE, Tdata>(x, vocab_size);
-    __shared__ Tcompute max_val_shared;
-    if (threadIdx.x == 0) {
-        max_val_shared = static_cast<Tcompute>(max_val_raw);
-    }
-    __syncthreads();
-    Tcompute max_val = max_val_shared;
-
-    Tcompute thread_sum = 0.0f;
-    for (size_t col = threadIdx.x; col < vocab_size; col += BLOCK_SIZE) {
-        Tcompute val = static_cast<Tcompute>(x[col]);
-        thread_sum += expf(val - max_val);
-    }
-
-    for (int offset = warpSize / 2; offset > 0; offset /= 2) {
-        thread_sum += __shfl_down_sync(0xffffffff, thread_sum, offset);
-    }
-
-    static __shared__ Tcompute shared_sum[32];
-    int lane = threadIdx.x % warpSize;
-    int warp = threadIdx.x / warpSize;
-
-    if (lane == 0) {
-        shared_sum[warp] = thread_sum;
-    }
-    __syncthreads();
-
-    Tcompute block_sum = 0.0f;
-    if (warp == 0) {
-
-        if (lane < (BLOCK_SIZE + warpSize - 1) / warpSize) {
-            block_sum = shared_sum[lane];
-        }
-        for (int offset = warpSize / 2; offset > 0; offset /= 2) {
-            block_sum += __shfl_down_sync(0xffffffff, block_sum, offset);
-        }
-    }
-
-    if (threadIdx.x == 0) {
-        Tcompute log_term = logf(block_sum) + max_val;
-
-        Tcompute target_logit = 0.0f;
-
-        if (label >= 0 && static_cast<size_t>(label) < vocab_size) {
-            target_logit = static_cast<Tcompute>(x[label]);
-        } else {
-
-            log_term = 0.0f;
-        }
-
-        y_[row_idx] = static_cast<Tdata>(log_term - target_logit);
-    }
-}
-
-#endif

src/infiniop/ops/cross_entropy/info.h

-#ifndef CROSS_ENTROPY_INFO_H
-#define CROSS_ENTROPY_INFO_H
-#include "../../../utils.h"
-#include "../../tensor.h"
-#include <vector>
-
-#include <cstddef>
-
-struct CrossEntropyInfo {
-    int dtype;
-    int target_dtype;
-    size_t outer_size;
-    size_t vocab_size;
-    ptrdiff_t x_stride;
-};
-
-#endif

src/infiniop/ops/cross_entropy/metax/cross_entropy_metax.h

-#ifndef __CROSS_ENTROPY_METAX_H__
-#define __CROSS_ENTROPY_METAX_H__
-
-#include "../cross_entropy.h"
-
-DESCRIPTOR(metax)
-
-#endif // __CROSS_ENTROPY_METAX_H__

src/infiniop/ops/cross_entropy/metax/cross_entropy_metax.maca

-#include "../../../devices/metax/metax_common.h"
-#include "cross_entropy_metax.h"
-#include "../../../devices/metax/metax_kernel_common.h"
-
-#include <cub/block/block_reduce.cuh>
-
-#include "../../../reduce/cuda/reduce.cuh"
-
-#include <cmath>
-
-namespace {
-
-template <unsigned int BLOCK_SIZE, typename Tdata, typename Tidx, typename Tcompute>
-__device__ void crossEntropyKernel(
-    Tdata *y_,
-    const Tdata *x_,
-    const void *target_,
-    size_t outer_size,
-    size_t vocab_size,
-    ptrdiff_t x_stride) {
-
-    size_t row_idx = blockIdx.x;
-    if (row_idx >= outer_size) {
-        return;
-    }
-
-    const Tdata *x = x_ + row_idx * x_stride;
-    const Tidx *target = reinterpret_cast<const Tidx *>(target_);
-
-    Tidx label = target[row_idx];
-
-    Tdata max_val_raw = op::common_cuda::reduce_op::max<BLOCK_SIZE, Tdata>(x, vocab_size);
-    __shared__ Tcompute max_val_shared;
-    if (threadIdx.x == 0) {
-        max_val_shared = static_cast<Tcompute>(max_val_raw);
-    }
-    __syncthreads();
-
-    Tcompute max_val = max_val_shared;
-
-    Tcompute thread_sum = Tcompute(0);
-    for (size_t col = threadIdx.x; col < vocab_size; col += BLOCK_SIZE) {
-        Tcompute val = static_cast<Tcompute>(x[col]);
-        thread_sum += expf(val - max_val);
-    }
-
-    using BlockReduce = cub::BlockReduce<Tcompute, BLOCK_SIZE>;
-    __shared__ typename BlockReduce::TempStorage temp_storage;
-    Tcompute block_sum = BlockReduce(temp_storage).Sum(thread_sum);
-
-    if (threadIdx.x == 0) {
-        if (label < 0 || static_cast<size_t>(label) >= vocab_size) {
-            y_[row_idx] = static_cast<Tdata>(0.0f);
-            return;
-        }
-        Tcompute log_term = logf(block_sum) + max_val;
-        Tcompute target_logit = static_cast<Tcompute>(x[label]);
-        y_[row_idx] = static_cast<Tdata>(log_term - target_logit);
-    }
-}
-
-template <unsigned int BLOCK_SIZE, typename Tdata, typename Tidx, typename Tcompute>
-INFINIOP_METAX_KERNEL crossEntropy(
-    Tdata *y, const Tdata *x, const void *target,
-    size_t outer_size, size_t vocab_size, ptrdiff_t x_stride) {
-    crossEntropyKernel<BLOCK_SIZE, Tdata, Tidx, Tcompute>(
-        y, x, target, outer_size, vocab_size, x_stride);
-}
-
-} // namespace
-
-namespace op::cross_entropy::metax {
-
-struct Descriptor::Opaque {
-    std::shared_ptr<device::metax::Handle::Internal> internal;
-};
-
-Descriptor::~Descriptor() {
-    delete _opaque;
-}
-
-infiniStatus_t Descriptor::create(
-    infiniopHandle_t handle,
-    Descriptor **desc_ptr,
-    infiniopTensorDescriptor_t y_desc,
-    infiniopTensorDescriptor_t x_desc,
-    infiniopTensorDescriptor_t target_desc) {
-
-    (void)y_desc;
-
-    auto x_dtype = x_desc->dtype();
-    auto t_dtype = target_desc->dtype();
-
-    CHECK_DTYPE(x_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32);
-    CHECK_DTYPE(t_dtype, INFINI_DTYPE_I32, INFINI_DTYPE_I64);
-
-    CrossEntropyInfo info{};
-    info.dtype = x_dtype;
-    info.target_dtype = t_dtype;
-    info.vocab_size = x_desc->shape().back();
-    info.outer_size = target_desc->numel();
-    info.x_stride = static_cast<ptrdiff_t>(info.vocab_size);
-
-    *desc_ptr = new Descriptor(
-        new Opaque{reinterpret_cast<device::metax::Handle *>(handle)->internal()},
-        info, 0, handle->device, handle->device_id);
-    return INFINI_STATUS_SUCCESS;
-}
-
-template <unsigned int BLOCK_SIZE>
-infiniStatus_t launchKernel(void *y, const void *x, const void *target,
-                            const CrossEntropyInfo &info, hcStream_t stream) {
-    dim3 grid(static_cast<uint32_t>(info.outer_size), 1, 1);
-
-    if (info.target_dtype == INFINI_DTYPE_I64) {
-        if (info.dtype == INFINI_DTYPE_F16) {
-            crossEntropy<BLOCK_SIZE, half, int64_t, float>
-                <<<grid, BLOCK_SIZE, 0, stream>>>(
-                    (half *)y, (const half *)x, target,
-                    info.outer_size, info.vocab_size, info.x_stride);
-        } else if (info.dtype == INFINI_DTYPE_BF16) {
-            crossEntropy<BLOCK_SIZE, cuda_bfloat16, int64_t, float>
-                <<<grid, BLOCK_SIZE, 0, stream>>>(
-                    (cuda_bfloat16 *)y, (const cuda_bfloat16 *)x, target,
-                    info.outer_size, info.vocab_size, info.x_stride);
-        } else if (info.dtype == INFINI_DTYPE_F32) {
-            crossEntropy<BLOCK_SIZE, float, int64_t, float>
-                <<<grid, BLOCK_SIZE, 0, stream>>>(
-                    (float *)y, (const float *)x, target,
-                    info.outer_size, info.vocab_size, info.x_stride);
-        } else {
-            return INFINI_STATUS_BAD_TENSOR_DTYPE;
-        }
-    } else if (info.target_dtype == INFINI_DTYPE_I32) {
-        if (info.dtype == INFINI_DTYPE_F16) {
-            crossEntropy<BLOCK_SIZE, half, int32_t, float>
-                <<<grid, BLOCK_SIZE, 0, stream>>>(
-                    (half *)y, (const half *)x, target,
-                    info.outer_size, info.vocab_size, info.x_stride);
-        } else if (info.dtype == INFINI_DTYPE_BF16) {
-            crossEntropy<BLOCK_SIZE, cuda_bfloat16, int32_t, float>
-                <<<grid, BLOCK_SIZE, 0, stream>>>(
-                    (cuda_bfloat16 *)y, (const cuda_bfloat16 *)x, target,
-                    info.outer_size, info.vocab_size, info.x_stride);
-        } else if (info.dtype == INFINI_DTYPE_F32) {
-            crossEntropy<BLOCK_SIZE, float, int32_t, float>
-                <<<grid, BLOCK_SIZE, 0, stream>>>(
-                    (float *)y, (const float *)x, target,
-                    info.outer_size, info.vocab_size, info.x_stride);
-        } else {
-            return INFINI_STATUS_BAD_TENSOR_DTYPE;
-        }
-    } else {
-        return INFINI_STATUS_BAD_TENSOR_DTYPE;
-    }
-
-    return INFINI_STATUS_SUCCESS;
-}
-
-infiniStatus_t Descriptor::calculate(
-    void *workspace,
-    size_t workspace_size,
-    void *y,
-    const void *x,
-    const void *target,
-    void *stream_) const {
-
-    (void)workspace;
-
-    if (workspace_size < _workspace_size) {
-        return INFINI_STATUS_INSUFFICIENT_WORKSPACE;
-    }
-
-    auto stream = reinterpret_cast<hcStream_t>(stream_);
-    int max_threads = _opaque->internal->maxThreadsPerBlock();
-
-    if (max_threads >= METAX_BLOCK_SIZE_1024) {
-        CHECK_STATUS(launchKernel<METAX_BLOCK_SIZE_1024>(y, x, target, _info, stream));
-    } else if (max_threads >= METAX_BLOCK_SIZE_512) {
-        CHECK_STATUS(launchKernel<METAX_BLOCK_SIZE_512>(y, x, target, _info, stream));
-    } else {
-        CHECK_STATUS(launchKernel<256>(y, x, target, _info, stream));
-    }
-
-    return INFINI_STATUS_SUCCESS;
-}
-
-} // namespace op::cross_entropy::metax

src/infiniop/ops/cross_entropy/moore/cross_entropy_kernel.h

-#ifndef __CROSS_ENTROPY_KERNEL_CUH__
-#define __CROSS_ENTROPY_KERNEL_CUH__
-
-template <unsigned int BLOCK_SIZE, typename Tdata, typename Tidx, typename Tcompute>
-__device__ void crossEntropyKernel(
-    Tdata *y_,
-    const Tdata *x_,
-    const void *target_,
-    size_t outer_size,
-    size_t vocab_size,
-    ptrdiff_t x_stride) {
-
-    size_t row_idx = blockIdx.x;
-    if (row_idx >= outer_size) {
-        return;
-    }
-
-    const Tdata *x = x_ + row_idx * x_stride;
-    const Tidx *target = reinterpret_cast<const Tidx *>(target_);
-
-    Tidx label = target[row_idx];
-
-    Tdata max_val_raw = op::common_cuda::reduce_op::max<BLOCK_SIZE, Tdata>(x, vocab_size);
-    __shared__ Tcompute max_val_shared;
-    if (threadIdx.x == 0) {
-        max_val_shared = static_cast<Tcompute>(max_val_raw);
-    }
-    __syncthreads();
-
-    Tcompute max_val = max_val_shared;
-
-    Tcompute thread_sum = Tcompute(0);
-    for (size_t col = threadIdx.x; col < vocab_size; col += BLOCK_SIZE) {
-        Tcompute val = static_cast<Tcompute>(x[col]);
-        thread_sum += expf(val - max_val);
-    }
-
-    using BlockReduce = cub::BlockReduce<Tcompute, BLOCK_SIZE>;
-    __shared__ typename BlockReduce::TempStorage temp_storage;
-    Tcompute block_sum = BlockReduce(temp_storage).Sum(thread_sum);
-
-    if (threadIdx.x == 0) {
-        if (label < 0 || static_cast<size_t>(label) >= vocab_size) {
-            y_[row_idx] = static_cast<Tdata>(0.0f);
-            return;
-        }
-        Tcompute log_term = logf(block_sum) + max_val;
-        Tcompute target_logit = static_cast<Tcompute>(x[label]);
-        y_[row_idx] = static_cast<Tdata>(log_term - target_logit);
-    }
-}
-
-#endif

src/infiniop/ops/cross_entropy/moore/cross_entropy_moore.h

-#ifndef __CROSS_ENTROPY_MOORE_H__
-#define __CROSS_ENTROPY_MOORE_H__
-
-#include "../cross_entropy.h"
-
-DESCRIPTOR(moore)
-
-#endif

src/infiniop/ops/cross_entropy/moore/cross_entropy_moore.mu

-#include "../../../devices/moore/moore_common.h"
-#include "cross_entropy_moore.h"
-
-#include <cub/block/block_reduce.cuh>
-#include "../../../devices/moore/moore_kernel_common.h"
-
-#include "../../../reduce/cuda/reduce.cuh"
-
-#include "cross_entropy_kernel.h"
-
-template <unsigned int BLOCK_SIZE, typename Tdata, typename Tidx, typename Tcompute>
-INFINIOP_MOORE_KERNEL crossEntropy(
-    Tdata *y, const Tdata *x, const void *target,
-    size_t outer_size, size_t vocab_size, ptrdiff_t x_stride) {
-    crossEntropyKernel<BLOCK_SIZE, Tdata, Tidx, Tcompute>(
-        y, x, target, outer_size, vocab_size, x_stride);
-}
-
-namespace op::cross_entropy::moore {
-
-struct Descriptor::Opaque {
-    std::shared_ptr<device::moore::Handle::Internal> internal;
-};
-
-Descriptor::~Descriptor() {
-    delete _opaque;
-}
-
-infiniStatus_t Descriptor::create(
-    infiniopHandle_t handle,
-    Descriptor **desc_ptr,
-    infiniopTensorDescriptor_t y_desc,
-    infiniopTensorDescriptor_t x_desc,
-    infiniopTensorDescriptor_t target_desc) {
-
-    (void)y_desc;
-
-    auto x_dtype = x_desc->dtype();
-    auto t_dtype = target_desc->dtype();
-
-    CHECK_DTYPE(x_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32);
-    CHECK_DTYPE(t_dtype, INFINI_DTYPE_I32, INFINI_DTYPE_I64);
-
-    CrossEntropyInfo info{};
-    info.dtype = x_dtype;
-    info.target_dtype = t_dtype;
-    info.vocab_size = x_desc->shape().back();
-    info.outer_size = target_desc->numel();
-    info.x_stride = static_cast<ptrdiff_t>(info.vocab_size);
-
-    *desc_ptr = new Descriptor(
-        new Opaque{reinterpret_cast<device::moore::Handle *>(handle)->internal()},
-        info, 0, handle->device, handle->device_id);
-    return INFINI_STATUS_SUCCESS;
-}
-
-template <unsigned int BLOCK_SIZE>
-infiniStatus_t launchKernel(void *y, const void *x, const void *target,
-                            const CrossEntropyInfo &info, musaStream_t stream) {
-    dim3 grid(static_cast<uint32_t>(info.outer_size), 1, 1);
-
-    if (info.target_dtype == INFINI_DTYPE_I64) {
-        if (info.dtype == INFINI_DTYPE_F16) {
-            crossEntropy<BLOCK_SIZE, half, int64_t, float>
-                <<<grid, BLOCK_SIZE, 0, stream>>>(
-                    (half *)y, (const half *)x, target,
-                    info.outer_size, info.vocab_size, info.x_stride);
-        } else if (info.dtype == INFINI_DTYPE_BF16) {
-            crossEntropy<BLOCK_SIZE, __mt_bfloat16, int64_t, float>
-                <<<grid, BLOCK_SIZE, 0, stream>>>(
-                    (__mt_bfloat16 *)y, (const __mt_bfloat16 *)x, target,
-                    info.outer_size, info.vocab_size, info.x_stride);
-        } else if (info.dtype == INFINI_DTYPE_F32) {
-            crossEntropy<BLOCK_SIZE, float, int64_t, float>
-                <<<grid, BLOCK_SIZE, 0, stream>>>(
-                    (float *)y, (const float *)x, target,
-                    info.outer_size, info.vocab_size, info.x_stride);
-        } else {
-            return INFINI_STATUS_BAD_TENSOR_DTYPE;
-        }
-    } else if (info.target_dtype == INFINI_DTYPE_I32) {
-        if (info.dtype == INFINI_DTYPE_F16) {
-            crossEntropy<BLOCK_SIZE, half, int32_t, float>
-                <<<grid, BLOCK_SIZE, 0, stream>>>(
-                    (half *)y, (const half *)x, target,
-                    info.outer_size, info.vocab_size, info.x_stride);
-        } else if (info.dtype == INFINI_DTYPE_BF16) {
-            crossEntropy<BLOCK_SIZE, __mt_bfloat16, int32_t, float>
-                <<<grid, BLOCK_SIZE, 0, stream>>>(
-                    (__mt_bfloat16 *)y, (const __mt_bfloat16 *)x, target,
-                    info.outer_size, info.vocab_size, info.x_stride);
-        } else if (info.dtype == INFINI_DTYPE_F32) {
-            crossEntropy<BLOCK_SIZE, float, int32_t, float>
-                <<<grid, BLOCK_SIZE, 0, stream>>>(
-                    (float *)y, (const float *)x, target,
-                    info.outer_size, info.vocab_size, info.x_stride);
-        } else {
-            return INFINI_STATUS_BAD_TENSOR_DTYPE;
-        }
-    } else {
-        return INFINI_STATUS_BAD_TENSOR_DTYPE;
-    }
-
-    return INFINI_STATUS_SUCCESS;
-}
-
-infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size,
-                                     void *y,
-                                     const void *x,
-                                     const void *target,
-                                     void *stream_) const {
-    musaStream_t stream = (musaStream_t)stream_;
-    (void)workspace;
-
-    if (workspace_size < _workspace_size) {
-        return INFINI_STATUS_INSUFFICIENT_WORKSPACE;
-    }
-
-    if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_1024) {
-        CHECK_STATUS(launchKernel<MOORE_BLOCK_SIZE_1024>(y, x, target, _info, stream));
-    } else if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_512) {
-        CHECK_STATUS(launchKernel<MOORE_BLOCK_SIZE_512>(y, x, target, _info, stream));
-    } else {
-        return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED;
-    }
-    return INFINI_STATUS_SUCCESS;
-}
-
-} // namespace op::cross_entropy::moore

src/infiniop/ops/cross_entropy/nvidia/cross_entropy_nvidia.cu

-#include "../../../devices/nvidia/nvidia_common.cuh"
-#include "../../../devices/nvidia/nvidia_kernel_common.cuh"
-#include "../cuda/kernel.cuh"
-#include "cross_entropy_nvidia.cuh"
-
-template <unsigned int BLOCK_SIZE, typename Tdata, typename Tidx, typename Tcompute = float>
-INFINIOP_CUDA_KERNEL crossEntropy(
-    Tdata *y, const Tdata *x, const void *target,
-    size_t outer_size, size_t vocab_size, ptrdiff_t x_stride) {
-
-    crossEntropyKernel<BLOCK_SIZE, Tdata, Tidx, Tcompute>(
-        y, x, target, outer_size, vocab_size, x_stride);
-}
-
-namespace op::cross_entropy::nvidia {
-
-struct Descriptor::Opaque {
-    std::shared_ptr<device::nvidia::Handle::Internal> internal;
-};
-
-Descriptor::~Descriptor() {
-    delete _opaque;
-}
-
-infiniStatus_t Descriptor::create(
-    infiniopHandle_t handle,
-    Descriptor **desc_ptr,
-    infiniopTensorDescriptor_t y_desc,
-    infiniopTensorDescriptor_t x_desc,
-    infiniopTensorDescriptor_t target_desc) {
-
-    auto x_dtype = x_desc->dtype();
-    auto t_dtype = target_desc->dtype();
-
-    CrossEntropyInfo info;
-    info.dtype = x_dtype;
-    info.target_dtype = t_dtype;
-
-    info.vocab_size = x_desc->shape().back();
-    info.outer_size = target_desc->numel();
-    info.x_stride = static_cast<ptrdiff_t>(info.vocab_size);
-
-    auto internal = reinterpret_cast<device::nvidia::Handle *>(handle)->internal();
-
-    *desc_ptr = new Descriptor(
-        new Opaque{internal},
-        info, 0, handle->device, handle->device_id);
-    return INFINI_STATUS_SUCCESS;
-}
-
-template <unsigned int BLOCK_SIZE>
-infiniStatus_t launchKernel(void *y, const void *x, const void *target,
-                            const CrossEntropyInfo &info, cudaStream_t stream) {
-
-    dim3 grid(static_cast<uint32_t>(info.outer_size), 1, 1);
-
-    if (info.target_dtype == INFINI_DTYPE_I64) {
-        if (info.dtype == INFINI_DTYPE_F16) {
-            crossEntropy<BLOCK_SIZE, half, int64_t>
-                <<<grid, BLOCK_SIZE, 0, stream>>>((half *)y, (const half *)x, target, info.outer_size, info.vocab_size, info.x_stride);
-        } else if (info.dtype == INFINI_DTYPE_BF16) {
-            crossEntropy<BLOCK_SIZE, __nv_bfloat16, int64_t>
-                <<<grid, BLOCK_SIZE, 0, stream>>>((__nv_bfloat16 *)y, (const __nv_bfloat16 *)x, target, info.outer_size, info.vocab_size, info.x_stride);
-        } else if (info.dtype == INFINI_DTYPE_F32) {
-            crossEntropy<BLOCK_SIZE, float, int64_t>
-                <<<grid, BLOCK_SIZE, 0, stream>>>((float *)y, (const float *)x, target, info.outer_size, info.vocab_size, info.x_stride);
-        }
-    } else if (info.target_dtype == INFINI_DTYPE_I32) {
-
-        if (info.dtype == INFINI_DTYPE_F16) {
-            crossEntropy<BLOCK_SIZE, half, int32_t>
-                <<<grid, BLOCK_SIZE, 0, stream>>>((half *)y, (const half *)x, target, info.outer_size, info.vocab_size, info.x_stride);
-        } else if (info.dtype == INFINI_DTYPE_BF16) {
-            crossEntropy<BLOCK_SIZE, __nv_bfloat16, int32_t>
-                <<<grid, BLOCK_SIZE, 0, stream>>>((__nv_bfloat16 *)y, (const __nv_bfloat16 *)x, target, info.outer_size, info.vocab_size, info.x_stride);
-        } else if (info.dtype == INFINI_DTYPE_F32) {
-            crossEntropy<BLOCK_SIZE, float, int32_t>
-                <<<grid, BLOCK_SIZE, 0, stream>>>((float *)y, (const float *)x, target, info.outer_size, info.vocab_size, info.x_stride);
-        }
-    } else {
-        return INFINI_STATUS_BAD_TENSOR_DTYPE;
-    }
-
-    return INFINI_STATUS_SUCCESS;
-}
-
-infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size,
-                                     void *y,
-                                     const void *x,
-                                     const void *target,
-                                     void *stream_) const {
-    cudaStream_t stream = (cudaStream_t)stream_;
-
-    int max_threads = _opaque->internal->maxThreadsPerBlock();
-
-    if (max_threads >= 1024) {
-        CHECK_STATUS(launchKernel<1024>(y, x, target, _info, stream));
-    } else if (max_threads >= 512) {
-        CHECK_STATUS(launchKernel<512>(y, x, target, _info, stream));
-    } else {
-        CHECK_STATUS(launchKernel<256>(y, x, target, _info, stream));
-    }
-
-    return INFINI_STATUS_SUCCESS;
-}
-
-} // namespace op::cross_entropy::nvidia