From 2ef786ca8fdbdfeb53313f3a18f356a54769d128 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Wed, 8 Jul 2026 13:00:07 +0800 Subject: [PATCH] feat: add hygon operators for w8a8 9g --- include/infinicore/ops/mul_scalar.hpp | 14 ++ include/infiniop.h | 1 + include/infiniop/ops/mul_scalar.h | 25 +++ python/infinicore/__init__.py | 2 + python/infinicore/ops/mul.py | 6 + python/infinicore/ops/mul_scalar.py | 12 ++ python/infinicore/tensor.py | 3 + src/infinicore/ops/mul_scalar/mul_scalar.cc | 27 +++ .../ops/mul_scalar/mul_scalar_infiniop.cc | 52 ++++++ src/infinicore/pybind11/ops.hpp | 2 + src/infinicore/pybind11/ops/mul_scalar.hpp | 26 +++ .../ops/mul_scalar/cpu/mul_scalar_cpu.cc | 98 +++++++++++ .../ops/mul_scalar/cpu/mul_scalar_cpu.h | 9 + src/infiniop/ops/mul_scalar/info.h | 47 ++++++ src/infiniop/ops/mul_scalar/mul_scalar.h | 49 ++++++ .../mul_scalar/nvidia/mul_scalar_nvidia.cu | 144 ++++++++++++++++ .../mul_scalar/nvidia/mul_scalar_nvidia.cuh | 9 + src/infiniop/ops/mul_scalar/operator.cc | 155 ++++++++++++++++++ .../quant/per_channel_quant_int8/operator.cc | 14 +- .../ops/scaled_mm/nvidia/int8_gemm_nvidia.cu | 4 +- src/infiniop/ops/scaled_mm/operator.cc | 14 +- xmake/hygon.lua | 11 +- 22 files changed, 718 insertions(+), 6 deletions(-) create mode 100644 include/infinicore/ops/mul_scalar.hpp create mode 100644 include/infiniop/ops/mul_scalar.h create mode 100644 python/infinicore/ops/mul_scalar.py create mode 100644 src/infinicore/ops/mul_scalar/mul_scalar.cc create mode 100644 src/infinicore/ops/mul_scalar/mul_scalar_infiniop.cc create mode 100644 src/infinicore/pybind11/ops/mul_scalar.hpp create mode 100644 src/infiniop/ops/mul_scalar/cpu/mul_scalar_cpu.cc create mode 100644 src/infiniop/ops/mul_scalar/cpu/mul_scalar_cpu.h create mode 100644 src/infiniop/ops/mul_scalar/info.h create mode 100644 src/infiniop/ops/mul_scalar/mul_scalar.h create mode 100644 src/infiniop/ops/mul_scalar/nvidia/mul_scalar_nvidia.cu create mode 100644 src/infiniop/ops/mul_scalar/nvidia/mul_scalar_nvidia.cuh create mode 100644 src/infiniop/ops/mul_scalar/operator.cc diff --git a/include/infinicore/ops/mul_scalar.hpp b/include/infinicore/ops/mul_scalar.hpp new file mode 100644 index 000000000..5ee5e0d0b --- /dev/null +++ b/include/infinicore/ops/mul_scalar.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(MulScalar, Tensor, const Tensor &, double); + +Tensor mul_scalar(const Tensor &a, double alpha); +void mul_scalar_(Tensor c, const Tensor &a, double alpha); + +} // namespace infinicore::op diff --git a/include/infiniop.h b/include/infiniop.h index 7794dfc5c..95bf75a0d 100644 --- a/include/infiniop.h +++ b/include/infiniop.h @@ -97,6 +97,7 @@ #include "infiniop/ops/moe_topk_softmax.h" #include "infiniop/ops/mrope.h" #include "infiniop/ops/mul.h" +#include "infiniop/ops/mul_scalar.h" #include "infiniop/ops/multi_margin_loss.h" #include "infiniop/ops/nrm2.h" #include "infiniop/ops/nsa_compress_paged_cache.h" diff --git a/include/infiniop/ops/mul_scalar.h b/include/infiniop/ops/mul_scalar.h new file mode 100644 index 000000000..9f90284ee --- /dev/null +++ b/include/infiniop/ops/mul_scalar.h @@ -0,0 +1,25 @@ +#ifndef __INFINIOP_MUL_SCALAR_API_H__ +#define __INFINIOP_MUL_SCALAR_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopMulScalarDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateMulScalarDescriptor(infiniopHandle_t handle, + infiniopMulScalarDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input); + +__INFINI_C __export infiniStatus_t infiniopGetMulScalarWorkspaceSize(infiniopMulScalarDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMulScalar(infiniopMulScalarDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + double alpha, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMulScalarDescriptor(infiniopMulScalarDescriptor_t desc); + +#endif diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py index cdefb904b..808995640 100644 --- a/python/infinicore/__init__.py +++ b/python/infinicore/__init__.py @@ -114,6 +114,7 @@ ) from infinicore.ops.mrope import mrope from infinicore.ops.mul import mul +from infinicore.ops.mul_scalar import mul_scalar from infinicore.ops.narrow import narrow from infinicore.ops.nrm2 import nrm2 from infinicore.ops.paged_attention import paged_attention @@ -234,6 +235,7 @@ "matmul", "equal", "mul", + "mul_scalar", "diff", "digamma", "dist", diff --git a/python/infinicore/ops/mul.py b/python/infinicore/ops/mul.py index e9e1d8c98..8ce703917 100644 --- a/python/infinicore/ops/mul.py +++ b/python/infinicore/ops/mul.py @@ -1,8 +1,14 @@ +import numbers + from infinicore.lib import _infinicore +from infinicore.ops.mul_scalar import mul_scalar from infinicore.tensor import Tensor def mul(input, other, *, out=None): + if isinstance(other, numbers.Real): + return mul_scalar(input, other, out=out) + if out is None: return Tensor(_infinicore.mul(input._underlying, other._underlying)) diff --git a/python/infinicore/ops/mul_scalar.py b/python/infinicore/ops/mul_scalar.py new file mode 100644 index 000000000..e7508a570 --- /dev/null +++ b/python/infinicore/ops/mul_scalar.py @@ -0,0 +1,12 @@ +from infinicore.lib import _infinicore +from infinicore.tensor import Tensor + + +def mul_scalar(input, alpha, *, out=None): + alpha = float(alpha) + if out is None: + return Tensor(_infinicore.mul_scalar(input._underlying, alpha)) + + _infinicore.mul_scalar_(out._underlying, input._underlying, alpha) + + return out diff --git a/python/infinicore/tensor.py b/python/infinicore/tensor.py index 4348887d6..bbe801f93 100644 --- a/python/infinicore/tensor.py +++ b/python/infinicore/tensor.py @@ -128,6 +128,9 @@ def __matmul__(self, other): def __mul__(self, other): return infinicore.mul(self, other) + def __rmul__(self, other): + return infinicore.mul(self, other) + def narrow(self, dim, start, length): return infinicore.narrow(self, dim, start, length) diff --git a/src/infinicore/ops/mul_scalar/mul_scalar.cc b/src/infinicore/ops/mul_scalar/mul_scalar.cc new file mode 100644 index 000000000..b5be25e7e --- /dev/null +++ b/src/infinicore/ops/mul_scalar/mul_scalar.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/mul_scalar.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MulScalar); + +MulScalar::MulScalar(Tensor c, const Tensor &a, double alpha) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a); + INFINICORE_GRAPH_OP_DISPATCH(c->device().getType(), c, a, alpha); +} + +void MulScalar::execute(Tensor c, const Tensor &a, double alpha) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(MulScalar, c, a, alpha); +} + +Tensor mul_scalar(const Tensor &a, double alpha) { + auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); + mul_scalar_(c, a, alpha); + return c; +} + +void mul_scalar_(Tensor c, const Tensor &a, double alpha) { + MulScalar::execute(c, a, alpha); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/mul_scalar/mul_scalar_infiniop.cc b/src/infinicore/ops/mul_scalar/mul_scalar_infiniop.cc new file mode 100644 index 000000000..13f3c615c --- /dev/null +++ b/src/infinicore/ops/mul_scalar/mul_scalar_infiniop.cc @@ -0,0 +1,52 @@ +#include "infinicore/ops/mul_scalar.hpp" + +#include "../infiniop_impl.hpp" + +namespace infinicore::op::mul_scalar_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, MulScalar, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace, c, a; + double alpha; +}; + +void *plan(Tensor c, const Tensor &a, double alpha) { + size_t seed = hash_combine(c, a, alpha); + + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, descriptor, MulScalar, + seed, c->desc(), a->desc()); + + INFINIOP_WORKSPACE_TENSOR(workspace, MulScalar, descriptor); + + return new PlannedMeta{ + descriptor, + graph::GraphTensor(workspace), + graph::GraphTensor(c), + graph::GraphTensor(a), + alpha}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + + INFINICORE_CHECK_ERROR(infiniopMulScalar( + planned->descriptor->desc, + planned->workspace->data(), + planned->workspace->numel(), + planned->c->data(), + planned->a->data(), + planned->alpha, + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MulScalar, &plan, &run, &cleanup); + +} // namespace infinicore::op::mul_scalar_impl::infiniop diff --git a/src/infinicore/pybind11/ops.hpp b/src/infinicore/pybind11/ops.hpp index 38aceb2fd..087261382 100644 --- a/src/infinicore/pybind11/ops.hpp +++ b/src/infinicore/pybind11/ops.hpp @@ -85,6 +85,7 @@ #include "ops/moe_topk_softmax.hpp" #include "ops/mrope.hpp" #include "ops/mul.hpp" +#include "ops/mul_scalar.hpp" #include "ops/multi_margin_loss.hpp" #include "ops/nrm2.hpp" #include "ops/pad.hpp" @@ -186,6 +187,7 @@ inline void bind(py::module &m) { bind_mamba_selective_scan(m); bind_kron(m); bind_mul(m); + bind_mul_scalar(m); bind_nrm2(m); bind_mha_kvcache(m); bind_mha_varlen(m); diff --git a/src/infinicore/pybind11/ops/mul_scalar.hpp b/src/infinicore/pybind11/ops/mul_scalar.hpp new file mode 100644 index 000000000..2bce862bc --- /dev/null +++ b/src/infinicore/pybind11/ops/mul_scalar.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/mul_scalar.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_mul_scalar(py::module &m) { + m.def("mul_scalar", + &op::mul_scalar, + py::arg("a"), + py::arg("alpha"), + R"doc(Multiply a tensor by a host scalar.)doc"); + + m.def("mul_scalar_", + &op::mul_scalar_, + py::arg("c"), + py::arg("a"), + py::arg("alpha"), + R"doc(Out-of-place tensor-scalar multiplication into c.)doc"); +} + +} // namespace infinicore::ops diff --git a/src/infiniop/ops/mul_scalar/cpu/mul_scalar_cpu.cc b/src/infiniop/ops/mul_scalar/cpu/mul_scalar_cpu.cc new file mode 100644 index 000000000..7049e5b6f --- /dev/null +++ b/src/infiniop/ops/mul_scalar/cpu/mul_scalar_cpu.cc @@ -0,0 +1,98 @@ +#include "mul_scalar_cpu.h" +#include "../../../devices/cpu/common_cpu.h" + +namespace op::mul_scalar::cpu { + +Descriptor::~Descriptor() = default; + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc) { + + auto handle = reinterpret_cast(handle_); + auto result = MulScalarInfo::create(output_desc, input_desc); + CHECK_RESULT(result); + + *desc_ptr = new Descriptor( + result.take(), + nullptr, + 0, + handle->device, + handle->device_id); + + return INFINI_STATUS_SUCCESS; +} + +template +inline T mulScalarHost(T value, double alpha) { + if constexpr (std::is_same_v || std::is_same_v) { + return utils::cast(utils::cast(value) * static_cast(alpha)); + } else { + return value * static_cast(alpha); + } +} + +template +infiniStatus_t calculateMulScalar(const MulScalarInfo &info, T *output, const T *input, double alpha) { + const auto &elementwise_info = info.elementwise_info; + const ptrdiff_t output_size = elementwise_info.getOutputSize(); + + if (info.contiguous) { +#pragma omp parallel for if (output_size > 1024) + for (ptrdiff_t i = 0; i < output_size; ++i) { + output[i] = mulScalarHost(input[i], alpha); + } + return INFINI_STATUS_SUCCESS; + } + +#pragma omp parallel for if (output_size > 1024) + for (ptrdiff_t i = 0; i < output_size; ++i) { + size_t out_idx = elementwise_info.isOutputContiguous() + ? i + : op::common_cpu::indexToOffset( + i, + elementwise_info.getNdim(), + elementwise_info.getOutputShape(), + elementwise_info.getOutputStrides()); + size_t in_idx = elementwise_info.getInputContiguous()[0] + ? i + : op::common_cpu::indexToOffset( + i, + elementwise_info.getNdim(), + elementwise_info.getInputShape(0), + elementwise_info.getInputStrides(0)); + output[out_idx] = mulScalarHost(input[in_idx], alpha); + } + + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *input, + double alpha, + void *stream) const { + + (void)workspace; + (void)workspace_size; + (void)stream; + + switch (_info.data_type) { + case INFINI_DTYPE_F16: + return calculateMulScalar(_info, static_cast(output), static_cast(input), alpha); + case INFINI_DTYPE_BF16: + return calculateMulScalar(_info, static_cast(output), static_cast(input), alpha); + case INFINI_DTYPE_F32: + return calculateMulScalar(_info, static_cast(output), static_cast(input), alpha); + case INFINI_DTYPE_F64: + return calculateMulScalar(_info, static_cast(output), static_cast(input), alpha); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +} // namespace op::mul_scalar::cpu diff --git a/src/infiniop/ops/mul_scalar/cpu/mul_scalar_cpu.h b/src/infiniop/ops/mul_scalar/cpu/mul_scalar_cpu.h new file mode 100644 index 000000000..19217a35c --- /dev/null +++ b/src/infiniop/ops/mul_scalar/cpu/mul_scalar_cpu.h @@ -0,0 +1,9 @@ +#ifndef __MUL_SCALAR_CPU_H__ +#define __MUL_SCALAR_CPU_H__ + +#include "../../../elementwise/cpu/elementwise_cpu.h" +#include "../mul_scalar.h" + +MUL_SCALAR_DESCRIPTOR(cpu) + +#endif // __MUL_SCALAR_CPU_H__ diff --git a/src/infiniop/ops/mul_scalar/info.h b/src/infiniop/ops/mul_scalar/info.h new file mode 100644 index 000000000..68ad91d31 --- /dev/null +++ b/src/infiniop/ops/mul_scalar/info.h @@ -0,0 +1,47 @@ +#ifndef __MUL_SCALAR_INFO_H__ +#define __MUL_SCALAR_INFO_H__ + +#include "../../../utils.h" +#include "../../elementwise/elementwise.h" +#include "../../tensor.h" + +class MulScalarInfo { +private: + MulScalarInfo( + op::elementwise::ElementwiseInfo elementwise_info_, + infiniDtype_t data_type_, + bool contiguous_) + : elementwise_info(std::move(elementwise_info_)), + data_type(data_type_), + contiguous(contiguous_) {} + +public: + op::elementwise::ElementwiseInfo elementwise_info; + infiniDtype_t data_type; + bool contiguous; + + size_t numel() const { return elementwise_info.getOutputSize(); } + + static utils::Result create( + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc) { + + CHECK_OR_RETURN(output_desc != nullptr, INFINI_STATUS_NULL_POINTER); + CHECK_OR_RETURN(input_desc != nullptr, INFINI_STATUS_NULL_POINTER); + + auto data_type = output_desc->dtype(); + CHECK_OR_RETURN(input_desc->dtype() == data_type, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_DTYPE(data_type, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); + CHECK_SAME_SHAPE(output_desc->shape(), input_desc->shape()); + + auto elementwise_info_result = op::elementwise::ElementwiseInfo::create(output_desc, {input_desc}); + CHECK_RESULT(elementwise_info_result); + + return utils::Result(MulScalarInfo( + elementwise_info_result.take(), + data_type, + output_desc->isContiguous() && input_desc->isContiguous())); + } +}; + +#endif // __MUL_SCALAR_INFO_H__ diff --git a/src/infiniop/ops/mul_scalar/mul_scalar.h b/src/infiniop/ops/mul_scalar/mul_scalar.h new file mode 100644 index 000000000..c0ca68e18 --- /dev/null +++ b/src/infiniop/ops/mul_scalar/mul_scalar.h @@ -0,0 +1,49 @@ +#ifndef __MUL_SCALAR_H__ +#define __MUL_SCALAR_H__ + +#include "../../operator.h" +#include "info.h" + +#include + +#define MUL_SCALAR_DESCRIPTOR(NAMESPACE) \ + \ + namespace op::mul_scalar::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + MulScalarInfo _info; \ + std::unique_ptr _device_info; \ + size_t _workspace_size; \ + \ + Descriptor( \ + MulScalarInfo info, \ + op::elementwise::NAMESPACE::DeviceImpl *device_info, \ + size_t workspace_size_, \ + infiniDevice_t device_type, \ + int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _info(std::move(info)), \ + _device_info(std::move(device_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 output_desc, \ + infiniopTensorDescriptor_t input_desc); \ + \ + infiniStatus_t calculate( \ + void *workspace, \ + size_t workspace_size, \ + void *output, \ + const void *input, \ + double alpha, \ + void *stream) const; \ + }; \ + } + +#endif // __MUL_SCALAR_H__ diff --git a/src/infiniop/ops/mul_scalar/nvidia/mul_scalar_nvidia.cu b/src/infiniop/ops/mul_scalar/nvidia/mul_scalar_nvidia.cu new file mode 100644 index 000000000..06ca80553 --- /dev/null +++ b/src/infiniop/ops/mul_scalar/nvidia/mul_scalar_nvidia.cu @@ -0,0 +1,144 @@ +#include "mul_scalar_nvidia.cuh" + +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" + +#include + +namespace { + +template +__device__ __forceinline__ T mulScalarValue(T value, double alpha) { + return value * static_cast(alpha); +} + +template <> +__device__ __forceinline__ half mulScalarValue(half value, double alpha) { + return __hmul(value, __float2half(static_cast(alpha))); +} + +template <> +__device__ __forceinline__ cuda_bfloat16 mulScalarValue(cuda_bfloat16 value, double alpha) { + return __hmul(value, __float2bfloat16_rn(static_cast(alpha))); +} + +template +__global__ void mulScalarKernel(size_t n, T *output, const T *input, double alpha) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + for (; idx < n; idx += stride) { + output[idx] = mulScalarValue(input[idx], alpha); + } +} + +template +infiniStatus_t launchMulScalarKernel(size_t numel, void *output, const void *input, double alpha, void *stream) { + if (numel == 0) { + return INFINI_STATUS_SUCCESS; + } + + constexpr int block = 256; + int grid = static_cast((numel + block - 1) / block); + grid = std::min(grid, 65535); + auto cuda_stream = reinterpret_cast(stream); + mulScalarKernel<<>>( + numel, + reinterpret_cast(output), + reinterpret_cast(input), + static_cast(alpha)); + auto err = cudaGetLastError(); + return err == cudaSuccess ? INFINI_STATUS_SUCCESS : INFINI_STATUS_INTERNAL_ERROR; +} + +struct MulScalarOp { + static constexpr size_t num_inputs = 1; + + template + __device__ __forceinline__ T operator()(const T &value, double alpha) const { + return mulScalarValue(value, alpha); + } +}; + +template +infiniStatus_t calculateMulScalar( + const MulScalarInfo &info, + op::elementwise::nvidia::DeviceImpl *device_info, + size_t workspace_size, + void *workspace, + void *output, + const void *input, + double alpha, + void *stream) { + + if (info.contiguous) { + return launchMulScalarKernel(info.numel(), output, input, alpha, stream); + } + + if (workspace_size < info.elementwise_info.getMetaMemSize() + sizeof(void *)) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + return device_info->calculate<256, MulScalarOp, T>( + info.elementwise_info, + workspace, + output, + {input}, + stream, + static_cast(alpha)); +} + +} // namespace + +namespace op::mul_scalar::nvidia { + +Descriptor::~Descriptor() = default; + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc) { + + auto handle = reinterpret_cast(handle_); + auto result = MulScalarInfo::create(output_desc, input_desc); + CHECK_RESULT(result); + + auto info = result.take(); + auto workspace_size = info.elementwise_info.getMetaMemSize() + sizeof(void *); + auto device_impl_result = op::elementwise::nvidia::DeviceImpl::create(handle->internal()); + CHECK_RESULT(device_impl_result); + + *desc_ptr = new Descriptor( + std::move(info), + device_impl_result.take(), + workspace_size, + handle->device, + handle->device_id); + + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *input, + double alpha, + void *stream) const { + + switch (_info.data_type) { + case INFINI_DTYPE_F16: + return calculateMulScalar(_info, _device_info.get(), workspace_size, workspace, output, input, alpha, stream); + case INFINI_DTYPE_BF16: + return calculateMulScalar(_info, _device_info.get(), workspace_size, workspace, output, input, alpha, stream); + case INFINI_DTYPE_F32: + return calculateMulScalar(_info, _device_info.get(), workspace_size, workspace, output, input, alpha, stream); + case INFINI_DTYPE_F64: + return calculateMulScalar(_info, _device_info.get(), workspace_size, workspace, output, input, alpha, stream); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +} // namespace op::mul_scalar::nvidia diff --git a/src/infiniop/ops/mul_scalar/nvidia/mul_scalar_nvidia.cuh b/src/infiniop/ops/mul_scalar/nvidia/mul_scalar_nvidia.cuh new file mode 100644 index 000000000..12f3553b6 --- /dev/null +++ b/src/infiniop/ops/mul_scalar/nvidia/mul_scalar_nvidia.cuh @@ -0,0 +1,9 @@ +#ifndef __MUL_SCALAR_NVIDIA_H__ +#define __MUL_SCALAR_NVIDIA_H__ + +#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" +#include "../mul_scalar.h" + +MUL_SCALAR_DESCRIPTOR(nvidia) + +#endif // __MUL_SCALAR_NVIDIA_H__ diff --git a/src/infiniop/ops/mul_scalar/operator.cc b/src/infiniop/ops/mul_scalar/operator.cc new file mode 100644 index 000000000..a8f50496f --- /dev/null +++ b/src/infiniop/ops/mul_scalar/operator.cc @@ -0,0 +1,155 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/mul_scalar.h" + +#ifdef ENABLE_CPU_API +#include "cpu/mul_scalar_cpu.h" +#endif +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) +#include "nvidia/mul_scalar_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreateMulScalarDescriptor( + infiniopHandle_t handle, + infiniopMulScalarDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc) { + +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::mul_scalar::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + output_desc, \ + input_desc) + + 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_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopGetMulScalarWorkspaceSize(infiniopMulScalarDescriptor_t desc, size_t *size) { + +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(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_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} + +__INFINI_C infiniStatus_t infiniopMulScalar( + infiniopMulScalarDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + double alpha, + void *stream) { + +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(workspace, workspace_size, output, input, alpha, 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_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroyMulScalarDescriptor(infiniopMulScalarDescriptor_t desc) { + +#define DELETE(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(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_ALI_API + DELETE(INFINI_DEVICE_ALI, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DELETE(INFINI_DEVICE_HYGON, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef DELETE +} diff --git a/src/infiniop/ops/quant/per_channel_quant_int8/operator.cc b/src/infiniop/ops/quant/per_channel_quant_int8/operator.cc index 3c5e7596c..273cf991b 100644 --- a/src/infiniop/ops/quant/per_channel_quant_int8/operator.cc +++ b/src/infiniop/ops/quant/per_channel_quant_int8/operator.cc @@ -2,7 +2,7 @@ #include "../../../handle.h" #include "infiniop/ops/quant/per_channel_quant_int8.h" -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || (defined(ENABLE_HYGON_API) && !defined(INFINIOP_DISABLE_HYGON_PER_CHANNEL_QUANT_INT8)) #include "nvidia/per_channel_quant_int8_nvidia.cuh" #endif #if defined(ENABLE_MOORE_API) @@ -31,6 +31,9 @@ __INFINI_C infiniStatus_t infiniopCreatePerChannelQuantI8Descriptor(infiniopHand #ifdef ENABLE_QY_API CREATE(INFINI_DEVICE_QY, nvidia) #endif +#if defined(ENABLE_HYGON_API) && !defined(INFINIOP_DISABLE_HYGON_PER_CHANNEL_QUANT_INT8) + CREATE(INFINI_DEVICE_HYGON, nvidia) +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore) #endif @@ -52,6 +55,9 @@ __INFINI_C infiniStatus_t infiniopGetPerChannelQuantI8WorkspaceSize(infiniopPerC #ifdef ENABLE_QY_API GET(INFINI_DEVICE_QY, nvidia) #endif +#if defined(ENABLE_HYGON_API) && !defined(INFINIOP_DISABLE_HYGON_PER_CHANNEL_QUANT_INT8) + GET(INFINI_DEVICE_HYGON, nvidia) +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore) #endif @@ -81,6 +87,9 @@ __INFINI_C infiniStatus_t infiniopPerChannelQuantI8(infiniopPerChannelQuantI8Des #ifdef ENABLE_QY_API QUANT(INFINI_DEVICE_QY, nvidia) #endif +#if defined(ENABLE_HYGON_API) && !defined(INFINIOP_DISABLE_HYGON_PER_CHANNEL_QUANT_INT8) + QUANT(INFINI_DEVICE_HYGON, nvidia) +#endif #ifdef ENABLE_MOORE_API QUANT(INFINI_DEVICE_MOORE, moore) #endif @@ -103,6 +112,9 @@ __INFINI_C infiniStatus_t infiniopDestroyPerChannelQuantI8Descriptor(infiniopPer #ifdef ENABLE_QY_API DESTROY(INFINI_DEVICE_QY, nvidia) #endif +#if defined(ENABLE_HYGON_API) && !defined(INFINIOP_DISABLE_HYGON_PER_CHANNEL_QUANT_INT8) + DESTROY(INFINI_DEVICE_HYGON, nvidia) +#endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore) #endif diff --git a/src/infiniop/ops/scaled_mm/nvidia/int8_gemm_nvidia.cu b/src/infiniop/ops/scaled_mm/nvidia/int8_gemm_nvidia.cu index 36351cf42..7532ac93e 100644 --- a/src/infiniop/ops/scaled_mm/nvidia/int8_gemm_nvidia.cu +++ b/src/infiniop/ops/scaled_mm/nvidia/int8_gemm_nvidia.cu @@ -64,7 +64,7 @@ infiniStatus_t Descriptor::create( return INFINI_STATUS_SUCCESS; } -#ifdef ENABLE_QY_API +#if defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) template infiniStatus_t Descriptor::launchKernel(const I8GemmInfo &info, Tdata *y, const Tdata *bias, const int8_t *x_packed, const float *x_scale, const int8_t *w_packed, const float *w_scale, void *stream_, void *workspace) const { cudaStream_t stream = (cudaStream_t)stream_; @@ -179,7 +179,7 @@ infiniStatus_t Descriptor::calculate( } else { return INFINI_STATUS_NOT_IMPLEMENTED; } -#elif defined ENABLE_QY_API +#elif defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #define CALCULATE_LINEAR(BLOCK_SIZE, TDATA) \ launchKernel(_info, (TDATA *)out, (const TDATA *)bias, (const int8_t *)a, (const float *)a_scale, (const int8_t *)b, (const float *)b_scale, stream, workspace) #define CALCULATE_LINEAR_WITH_BLOCK_SIZE(BLOCK_SIZE) \ diff --git a/src/infiniop/ops/scaled_mm/operator.cc b/src/infiniop/ops/scaled_mm/operator.cc index b852ab6c0..4b24330b6 100644 --- a/src/infiniop/ops/scaled_mm/operator.cc +++ b/src/infiniop/ops/scaled_mm/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/int8_gemm.h" -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/int8_gemm_nvidia.cuh" #endif @@ -36,6 +36,9 @@ __INFINI_C infiniStatus_t infiniopCreateI8GemmDescriptor(infiniopHandle_t handle #if defined(ENABLE_QY_API) CREATE(INFINI_DEVICE_QY, nvidia) #endif +#if defined(ENABLE_HYGON_API) + CREATE(INFINI_DEVICE_HYGON, nvidia) +#endif #if defined(ENABLE_MOORE_API) CREATE(INFINI_DEVICE_MOORE, moore) #endif @@ -57,6 +60,9 @@ __INFINI_C infiniStatus_t infiniopGetI8GemmWorkspaceSize(infiniopI8GemmDescripto #if defined(ENABLE_QY_API) GET(INFINI_DEVICE_QY, nvidia) #endif +#if defined(ENABLE_HYGON_API) + GET(INFINI_DEVICE_HYGON, nvidia) +#endif #if defined(ENABLE_MOORE_API) GET(INFINI_DEVICE_MOORE, moore) #endif @@ -87,6 +93,9 @@ __INFINI_C infiniStatus_t infiniopI8Gemm(infiniopI8GemmDescriptor_t desc, #if defined(ENABLE_QY_API) CACULATE(INFINI_DEVICE_QY, nvidia) #endif +#if defined(ENABLE_HYGON_API) + CACULATE(INFINI_DEVICE_HYGON, nvidia) +#endif #if defined(ENABLE_MOORE_API) CACULATE(INFINI_DEVICE_MOORE, moore) #endif @@ -108,6 +117,9 @@ __INFINI_C infiniStatus_t infiniopDestroyI8GemmDescriptor(infiniopI8GemmDescript #if defined(ENABLE_QY_API) DESTROY(INFINI_DEVICE_QY, nvidia) #endif +#if defined(ENABLE_HYGON_API) + DESTROY(INFINI_DEVICE_HYGON, nvidia) +#endif #if defined(ENABLE_MOORE_API) DESTROY(INFINI_DEVICE_MOORE, moore) #endif diff --git a/xmake/hygon.lua b/xmake/hygon.lua index cf4e970f8..b99198dbf 100644 --- a/xmake/hygon.lua +++ b/xmake/hygon.lua @@ -101,7 +101,12 @@ local function resolve_hygon_arch() end local HYGON_ARCH = resolve_hygon_arch() +local HYGON_SKIP_PER_CHANNEL_QUANT_INT8 = HYGON_ARCH == "gfx906" print("编译海光DCU架构: " .. HYGON_ARCH) +if HYGON_SKIP_PER_CHANNEL_QUANT_INT8 then + print("Hygon gfx906: skip per_channel_quant_int8 operator") + add_defines("INFINIOP_DISABLE_HYGON_PER_CHANNEL_QUANT_INT8") +end target("infiniop-hygon") set_kind("static") @@ -140,6 +145,9 @@ target("infiniop-hygon") -- 复用NVIDIA的CUDA实现，通过HIP兼容层 add_files("../src/infiniop/devices/nvidia/*.cu", "../src/infiniop/ops/*/nvidia/*.cu") + if not HYGON_SKIP_PER_CHANNEL_QUANT_INT8 then + add_files("../src/infiniop/ops/quant/per_channel_quant_int8/nvidia/*.cu") + end -- Keep platform-specific or currently unregistered NVIDIA sources out of the Hygon target. remove_files("../src/infiniop/ops/avg_pool3d/nvidia/*.cu") @@ -148,8 +156,7 @@ target("infiniop-hygon") remove_files("../src/infiniop/ops/dist/nvidia/*.cu") remove_files("../src/infiniop/ops/gptq_qyblas_gemm/nvidia/*.cu") remove_files("../src/infiniop/ops/histc/nvidia/*.cu") - remove_files("../src/infiniop/ops/quant*/nvidia/*.cu") - remove_files("../src/infiniop/ops/scaled_mm/nvidia/*.cu") + remove_files("../src/infiniop/ops/quant/per_tensor_quant_int8/nvidia/*.cu") if has_config("ninetoothed") then add_files("../build/ninetoothed/*.c", "../build/ninetoothed/*.cpp", {cxxflags = {"-Wno-return-type"}})