feat: add hygon operators for w8a8 9g

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

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"

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

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",

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))
 

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

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)
 

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

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> 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<PlannedMeta *>(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<PlannedMeta **>(planned_meta_ptr);
+    *planned_meta_ptr = nullptr;
+}
+
+INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MulScalar, &plan, &run, &cleanup);
+
+} // namespace infinicore::op::mul_scalar_impl::infiniop

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);

src/infinicore/pybind11/ops/mul_scalar.hpp

@@ -0,0 +1,26 @@
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+#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

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<device::cpu::Handle *>(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 <typename T>
+inline T mulScalarHost(T value, double alpha) {
+    if constexpr (std::is_same_v<T, fp16_t> || std::is_same_v<T, bf16_t>) {
+        return utils::cast<T>(utils::cast<float>(value) * static_cast<float>(alpha));
+    } else {
+        return value * static_cast<T>(alpha);
+    }
+}
+
+template <typename T>
+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<fp16_t *>(output), static_cast<const fp16_t *>(input), alpha);
+    case INFINI_DTYPE_BF16:
+        return calculateMulScalar(_info, static_cast<bf16_t *>(output), static_cast<const bf16_t *>(input), alpha);
+    case INFINI_DTYPE_F32:
+        return calculateMulScalar(_info, static_cast<float *>(output), static_cast<const float *>(input), alpha);
+    case INFINI_DTYPE_F64:
+        return calculateMulScalar(_info, static_cast<double *>(output), static_cast<const double *>(input), alpha);
+    default:
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+}
+
+} // namespace op::mul_scalar::cpu

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__

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<MulScalarInfo> 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>(MulScalarInfo(
+            elementwise_info_result.take(),
+            data_type,
+            output_desc->isContiguous() && input_desc->isContiguous()));
+    }
+};
+
+#endif // __MUL_SCALAR_INFO_H__