From f8a4c480f67ee5962519b233d8af6f4476c19ae9 Mon Sep 17 00:00:00 2001 From: ShaneWu Date: Tue, 11 Aug 2026 11:10:43 +0800 Subject: [PATCH 1/2] feat: add aclnnMatmulAllReduce fusion in InfiniCore for Ascend RowParallelLinear --- include/infinicore/ops/linear_allreduce.hpp | 24 +++ .../ops/linear_allreduce/linear_allreduce.cc | 50 +++++ .../linear_allreduce_ascend.cc | 182 ++++++++++++++++++ 3 files changed, 256 insertions(+) create mode 100644 include/infinicore/ops/linear_allreduce.hpp create mode 100644 src/infinicore/ops/linear_allreduce/linear_allreduce.cc create mode 100644 src/infinicore/ops/linear_allreduce/linear_allreduce_ascend.cc diff --git a/include/infinicore/ops/linear_allreduce.hpp b/include/infinicore/ops/linear_allreduce.hpp new file mode 100644 index 000000000..beb20d25e --- /dev/null +++ b/include/infinicore/ops/linear_allreduce.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include "common/op.hpp" +#include +#include + +namespace infinicore::op { + +Tensor linear_allreduce( + Tensor input, + Tensor weight, + std::optional bias, + infinicclReduceOp_t op, + infinicclComm_t communicator); + +void linear_allreduce_( + Tensor output, + Tensor input, + Tensor weight, + std::optional bias, + infinicclReduceOp_t op, + infinicclComm_t communicator); + +} // namespace infinicore::op diff --git a/src/infinicore/ops/linear_allreduce/linear_allreduce.cc b/src/infinicore/ops/linear_allreduce/linear_allreduce.cc new file mode 100644 index 000000000..464e43116 --- /dev/null +++ b/src/infinicore/ops/linear_allreduce/linear_allreduce.cc @@ -0,0 +1,50 @@ +#include "infinicore/ops/linear_allreduce.hpp" +#include "infinicore/device.hpp" +#include "infinicore/ops/distributed/allreduce.hpp" +#include "infinicore/ops/linear.hpp" + +#if defined(ENABLE_ASCEND_API) +namespace infinicore::op::linear_allreduce_impl::ascend { +void linear_allreduce_impl( + Tensor output, Tensor input, Tensor weight, + std::optional bias, infinicclComm_t communicator); +} // namespace infinicore::op::linear_allreduce_impl::ascend +#endif + +namespace infinicore::op { + +Tensor linear_allreduce( + Tensor input, Tensor weight, std::optional bias, + infinicclReduceOp_t op, infinicclComm_t communicator) { +#if defined(ENABLE_ASCEND_API) + if (input->device().getType() == Device::Type::ASCEND) { + Size ndim = input->ndim(); + Size out_features = weight->shape()[0]; + auto out_shape = input->shape(); + out_shape[ndim - 1] = out_features; + auto out = Tensor::empty(out_shape, input->dtype(), input->device()); + linear_allreduce_impl::ascend::linear_allreduce_impl( + out, input, weight, bias, communicator); + return out; + } +#endif + auto output = linear(input, weight, bias); + return distributed::allreduce(output, op, communicator); +} + +void linear_allreduce_( + Tensor output, Tensor input, Tensor weight, + std::optional bias, infinicclReduceOp_t op, + infinicclComm_t communicator) { +#if defined(ENABLE_ASCEND_API) + if (input->device().getType() == Device::Type::ASCEND) { + linear_allreduce_impl::ascend::linear_allreduce_impl( + output, input, weight, bias, communicator); + return; + } +#endif + linear_(output, input, weight, bias); + distributed::allreduce_(output, output, op, communicator); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/linear_allreduce/linear_allreduce_ascend.cc b/src/infinicore/ops/linear_allreduce/linear_allreduce_ascend.cc new file mode 100644 index 000000000..d46e72aa6 --- /dev/null +++ b/src/infinicore/ops/linear_allreduce/linear_allreduce_ascend.cc @@ -0,0 +1,182 @@ +#if defined(ENABLE_ASCEND_API) + +#include "../../../infiniccl/infiniccl_impl.h" +#include "infinicore/context/context.hpp" +#include "infinicore/device.hpp" +#include "infinicore/ops/linear_allreduce.hpp" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace infinicore::op::linear_allreduce_impl::ascend { + +// ---- workspace pool ----------------------------------------- +// Per-stream reusable workspace. Once grown to the maximum +// required size, the buffer is never freed (leaked on growth). +// This avoids the MC2-notify-after-free bug: MC2 callbacks may +// still reference old workspace after stream sync completes. +struct WorkspaceBuf { + void *ptr = nullptr; + size_t cap = 0; +}; + +class WorkspacePool { + std::unordered_map bufs_; + std::mutex mtx_; + +public: + void *ensure(aclrtStream stream, size_t need) { + std::lock_guard lock(mtx_); + auto &b = bufs_[stream]; + if (need <= b.cap) { + return b.ptr; + } + void *new_ptr = nullptr; + aclError rc = aclrtMalloc(&new_ptr, need, ACL_MEM_MALLOC_HUGE_FIRST); + if (rc != ACL_SUCCESS || !new_ptr) { + fprintf(stderr, "[linear_allreduce/ascend] FATAL: aclrtMalloc(%zu MB) " + "failed rc=%d\n", + need / (1024 * 1024), (int)rc); + fflush(stderr); + throw std::runtime_error("[linear_allreduce/ascend] workspace alloc failed"); + } + b.ptr = new_ptr; + b.cap = need; + return b.ptr; + } +}; + +static WorkspacePool g_pool; + +static inline HcclComm get_hccl_comm(infinicclComm_t comm) { + return static_cast(comm->comm); +} + +static aclDataType to_acl_dtype(DataType dtype) { + switch (dtype) { + case DataType::F16: + return ACL_FLOAT16; + case DataType::BF16: + return ACL_BF16; + default: + throw std::runtime_error( + "[linear_allreduce/ascend] unsupported dtype: " + std::to_string(static_cast(dtype)) + ". aclnnMatmulAllReduce only supports F16/BF16"); + } +} + +void linear_allreduce_impl( + Tensor output, Tensor input, Tensor weight, + std::optional bias, + infinicclComm_t communicator) { + infinicore::context::setDevice(input->device()); + + auto atype = input->dtype(); + if (atype != DataType::F16 && atype != DataType::BF16) { + throw std::runtime_error( + "[linear_allreduce/ascend] unsupported activation dtype: " + std::to_string(static_cast(atype)) + ". aclnnMatmulAllReduce only supports F16/BF16"); + } + + auto w_perm = weight->permute({1, 0}); + Tensor weight_w = w_perm->is_contiguous() ? Tensor(w_perm) : w_perm->contiguous(); + + auto in_shape = input->shape(); + auto wt_shape = weight_w->shape(); + auto out_shape = output->shape(); + + std::vector in_dims(in_shape.begin(), in_shape.end()); + std::vector wt_dims(wt_shape.begin(), wt_shape.end()); + std::vector out_dims(out_shape.begin(), out_shape.end()); + + aclTensor *x1_acl = aclCreateTensor( + in_dims.data(), in_dims.size(), + to_acl_dtype(atype), + nullptr, 0, ACL_FORMAT_ND, + in_dims.data(), in_dims.size(), + const_cast(reinterpret_cast(input->data()))); + + aclTensor *x2_acl = aclCreateTensor( + wt_dims.data(), wt_dims.size(), + to_acl_dtype(atype), + nullptr, 0, ACL_FORMAT_ND, + wt_dims.data(), wt_dims.size(), + const_cast(reinterpret_cast(weight_w->data()))); + + aclTensor *bias_acl = nullptr; + if (bias.has_value()) { + Tensor bias_w = bias.value()->is_contiguous() ? Tensor(bias.value()) + : bias.value()->contiguous(); + auto bias_shape = bias_w->shape(); + std::vector bias_dims(bias_shape.begin(), bias_shape.end()); + bias_acl = aclCreateTensor( + bias_dims.data(), bias_dims.size(), + to_acl_dtype(atype), + nullptr, 0, ACL_FORMAT_ND, + bias_dims.data(), bias_dims.size(), + const_cast(reinterpret_cast(bias_w->data()))); + } + + aclTensor *out_acl = aclCreateTensor( + out_dims.data(), out_dims.size(), + to_acl_dtype(atype), + nullptr, 0, ACL_FORMAT_ND, + out_dims.data(), out_dims.size(), + const_cast(reinterpret_cast(output->data()))); + + HcclComm hccl_comm = get_hccl_comm(communicator); + char group_name[COMM_NAME_MAX_LENGTH] = {}; + HcclGetCommName(hccl_comm, group_name); + + uint64_t workspace_size = 0; + aclOpExecutor *executor = nullptr; + aclnnStatus ret = aclnnMatmulAllReduceGetWorkspaceSize( + x1_acl, x2_acl, bias_acl, + group_name, "sum", + 0, 1, + out_acl, + &workspace_size, &executor); + + if (ret != 0) { + if (bias_acl) { + aclDestroyTensor(bias_acl); + } + aclDestroyTensor(x1_acl); + aclDestroyTensor(x2_acl); + aclDestroyTensor(out_acl); + const char *err = aclGetRecentErrMsg(); + throw std::runtime_error( + std::string("[linear_allreduce/ascend] GetWorkspaceSize failed: ") + std::to_string(ret) + ", msg: " + (err ? err : "(null)")); + } + + aclrtStream stream = static_cast( + infinicore::context::getStream()); + void *workspace = g_pool.ensure(stream, (size_t)workspace_size); + + ret = aclnnMatmulAllReduce( + workspace, workspace_size, executor, stream); + + if (bias_acl) { + aclDestroyTensor(bias_acl); + } + aclDestroyTensor(x1_acl); + aclDestroyTensor(x2_acl); + aclDestroyTensor(out_acl); + + if (ret != 0) { + const char *err = aclGetRecentErrMsg(); + throw std::runtime_error( + std::string("[linear_allreduce/ascend] execution failed: ") + std::to_string(ret) + ", msg: " + (err ? err : "(null)")); + } +} + +} // namespace infinicore::op::linear_allreduce_impl::ascend + +#endif // ENABLE_ASCEND_API From a830fcfc44223730aef7445d778c089f34b34a6c Mon Sep 17 00:00:00 2001 From: wooway777 Date: Thu, 13 Aug 2026 19:39:54 +0800 Subject: [PATCH 2/2] refactor: move matmul all-reduce to infiniop --- include/infiniccl.h | 7 + include/infinicore/ops.hpp | 5 +- include/infinicore/ops/linear_allreduce.hpp | 16 +- include/infiniop.h | 3 +- include/infiniop/ops/matmul_all_reduce.h | 34 ++++ src/infiniccl/ascend/infiniccl_ascend.cc | 14 ++ .../cambricon/infiniccl_cambricon.cc | 7 + src/infiniccl/cuda/infiniccl_cuda.cu | 7 + src/infiniccl/infiniccl.cc | 31 +++ src/infiniccl/infiniccl_impl.h | 28 +++ src/infiniccl/kunlun/infiniccl_kunlun.cc | 7 + src/infiniccl/metax/infiniccl_metax.cc | 7 + src/infiniccl/moore/infiniccl_moore.cc | 7 + .../ops/linear_allreduce/linear_allreduce.cc | 109 +++++++---- .../linear_allreduce_ascend.cc | 182 ------------------ .../linear_allreduce_infiniop.cc | 86 +++++++++ .../ascend/matmul_all_reduce_ascend.cc | 130 +++++++++++++ .../ascend/matmul_all_reduce_ascend.h | 8 + .../ops/matmul_all_reduce/matmul_all_reduce.h | 48 +++++ .../ops/matmul_all_reduce/operator.cc | 91 +++++++++ 20 files changed, 602 insertions(+), 225 deletions(-) create mode 100644 include/infiniop/ops/matmul_all_reduce.h delete mode 100644 src/infinicore/ops/linear_allreduce/linear_allreduce_ascend.cc create mode 100644 src/infinicore/ops/linear_allreduce/linear_allreduce_infiniop.cc create mode 100644 src/infiniop/ops/matmul_all_reduce/ascend/matmul_all_reduce_ascend.cc create mode 100644 src/infiniop/ops/matmul_all_reduce/ascend/matmul_all_reduce_ascend.h create mode 100644 src/infiniop/ops/matmul_all_reduce/matmul_all_reduce.h create mode 100644 src/infiniop/ops/matmul_all_reduce/operator.cc diff --git a/include/infiniccl.h b/include/infiniccl.h index cced30879..b0f98c214 100644 --- a/include/infiniccl.h +++ b/include/infiniccl.h @@ -21,6 +21,13 @@ typedef struct { char internal[INFINICCL_UNIQUE_ID_BYTES]; } infinicclUniqueId_t; +#define INFINICCL_COMM_NAME_MAX_LENGTH 128 + +__INFINI_C __export infiniStatus_t infinicclGetCommName( + infinicclComm_t comm, + char *comm_name, + size_t comm_name_size); + __INFINI_C __export infiniStatus_t infinicclCommInitAll( infiniDevice_t device_type, infinicclComm_t *comms, diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index 63f16a9c9..5e93e1457 100644 --- a/include/infinicore/ops.hpp +++ b/include/infinicore/ops.hpp @@ -44,10 +44,11 @@ #include "ops/gelutanh.hpp" #include "ops/hardswish.hpp" #include "ops/hardtanh.hpp" -#include "ops/kv_caching.hpp" #include "ops/kimi_delta_attention.hpp" +#include "ops/kv_caching.hpp" #include "ops/layer_norm.hpp" #include "ops/linear.hpp" +#include "ops/linear_allreduce.hpp" #include "ops/linear_mxfp4.hpp" #include "ops/mamba_selective_scan.hpp" #include "ops/matmul.hpp" @@ -73,8 +74,8 @@ #include "ops/relu.hpp" #include "ops/rms_norm.hpp" #include "ops/rope.hpp" -#include "ops/rotary_embedding.hpp" #include "ops/rot.hpp" +#include "ops/rotary_embedding.hpp" #include "ops/rotg.hpp" #include "ops/rotm.hpp" #include "ops/rotmg.hpp" diff --git a/include/infinicore/ops/linear_allreduce.hpp b/include/infinicore/ops/linear_allreduce.hpp index beb20d25e..82aea3f18 100644 --- a/include/infinicore/ops/linear_allreduce.hpp +++ b/include/infinicore/ops/linear_allreduce.hpp @@ -1,24 +1,30 @@ #pragma once +#include "../graph/graph.hpp" #include "common/op.hpp" #include #include namespace infinicore::op { +INFINICORE_GRAPH_OP_CLASS( + LinearAllReduce, + Tensor, + const Tensor &, + const Tensor &, + const std::optional &, + infinicclComm_t); + Tensor linear_allreduce( Tensor input, Tensor weight, std::optional bias, - infinicclReduceOp_t op, infinicclComm_t communicator); -void linear_allreduce_( - Tensor output, +Tensor linear_allreduce_packed( Tensor input, - Tensor weight, + Tensor packed_weight, std::optional bias, - infinicclReduceOp_t op, infinicclComm_t communicator); } // namespace infinicore::op diff --git a/include/infiniop.h b/include/infiniop.h index 38d6d1bfd..9f632e27f 100644 --- a/include/infiniop.h +++ b/include/infiniop.h @@ -78,10 +78,10 @@ #include "infiniop/ops/index_copy.h" #include "infiniop/ops/inner.h" #include "infiniop/ops/int8_gemm.h" +#include "infiniop/ops/kimi_delta_attention.h" #include "infiniop/ops/kron.h" #include "infiniop/ops/kthvalue.h" #include "infiniop/ops/kv_caching.h" -#include "infiniop/ops/kimi_delta_attention.h" #include "infiniop/ops/layer_norm.h" #include "infiniop/ops/ldexp.h" #include "infiniop/ops/lerp.h" @@ -96,6 +96,7 @@ #include "infiniop/ops/lp_norm.h" #include "infiniop/ops/mamba_selective_scan.h" #include "infiniop/ops/masked_select.h" +#include "infiniop/ops/matmul_all_reduce.h" #include "infiniop/ops/matrix_power.h" #include "infiniop/ops/moe_align.h" #include "infiniop/ops/moe_fused_dense.h" diff --git a/include/infiniop/ops/matmul_all_reduce.h b/include/infiniop/ops/matmul_all_reduce.h new file mode 100644 index 000000000..1234353d3 --- /dev/null +++ b/include/infiniop/ops/matmul_all_reduce.h @@ -0,0 +1,34 @@ +#ifndef __INFINIOP_MATMUL_ALL_REDUCE_API_H__ +#define __INFINIOP_MATMUL_ALL_REDUCE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopMatmulAllReduceDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateMatmulAllReduceDescriptor( + infiniopHandle_t handle, + infiniopMatmulAllReduceDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t weight_desc, + infiniopTensorDescriptor_t bias_desc, + const char *group_name); + +__INFINI_C __export infiniStatus_t infiniopGetMatmulAllReduceWorkspaceSize( + infiniopMatmulAllReduceDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMatmulAllReduce( + infiniopMatmulAllReduceDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + const void *weight, + const void *bias, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMatmulAllReduceDescriptor( + infiniopMatmulAllReduceDescriptor_t desc); + +#endif diff --git a/src/infiniccl/ascend/infiniccl_ascend.cc b/src/infiniccl/ascend/infiniccl_ascend.cc index 76f4c961b..302ac6f9c 100644 --- a/src/infiniccl/ascend/infiniccl_ascend.cc +++ b/src/infiniccl/ascend/infiniccl_ascend.cc @@ -54,6 +54,20 @@ inline HcclReduceOp getHcclRedOp(infinicclReduceOp_t op) { namespace infiniccl::ascend { +infiniStatus_t getCommName( + infinicclComm_t comm, + char *comm_name, + size_t comm_name_size) { + if (comm == nullptr || comm_name == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + if (comm_name_size < COMM_NAME_MAX_LENGTH) { + return INFINI_STATUS_BAD_PARAM; + } + CHECK_HCCL(HcclGetCommName(getHcclComm(comm), comm_name)); + return INFINI_STATUS_SUCCESS; +} + infiniStatus_t commInitAll( infinicclComm_t *comms, int ndevice, diff --git a/src/infiniccl/cambricon/infiniccl_cambricon.cc b/src/infiniccl/cambricon/infiniccl_cambricon.cc index 4bd1ca402..5f0ab536a 100644 --- a/src/infiniccl/cambricon/infiniccl_cambricon.cc +++ b/src/infiniccl/cambricon/infiniccl_cambricon.cc @@ -52,6 +52,13 @@ inline cnclReduceOp_t getCnclRedOp(infinicclReduceOp_t op) { namespace infiniccl::cambricon { +infiniStatus_t getCommName( + infinicclComm_t comm, + char *comm_name, + size_t comm_name_size) { + return getCommNameFromHandle(comm, comm_name, comm_name_size); +} + infiniStatus_t commInitAll( infinicclComm_t *comms, int ndevice, diff --git a/src/infiniccl/cuda/infiniccl_cuda.cu b/src/infiniccl/cuda/infiniccl_cuda.cu index e9a177dda..9b305afd0 100644 --- a/src/infiniccl/cuda/infiniccl_cuda.cu +++ b/src/infiniccl/cuda/infiniccl_cuda.cu @@ -63,6 +63,13 @@ inline ncclComm_t getNcclComm(infinicclComm_t comm) { namespace infiniccl::cuda { +infiniStatus_t getCommName( + infinicclComm_t comm, + char *comm_name, + size_t comm_name_size) { + return getCommNameFromHandle(comm, comm_name, comm_name_size); +} + infiniStatus_t commInitAll( infinicclComm_t *comms, int ndevice, diff --git a/src/infiniccl/infiniccl.cc b/src/infiniccl/infiniccl.cc index f348e355f..ec48b9b23 100644 --- a/src/infiniccl/infiniccl.cc +++ b/src/infiniccl/infiniccl.cc @@ -8,6 +8,37 @@ #include "./metax/infiniccl_metax.h" #include "./moore/infiniccl_moore.h" +__INFINI_C infiniStatus_t infinicclGetCommName( + infinicclComm_t comm, + char *comm_name, + size_t comm_name_size) { + if (comm == nullptr || comm_name == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + +#define GET_COMM_NAME(CASE_, NAMESPACE_) \ + case CASE_: \ + return infiniccl::NAMESPACE_::getCommName( \ + comm, comm_name, comm_name_size) + + switch (comm->device_type) { + GET_COMM_NAME(INFINI_DEVICE_NVIDIA, cuda); + GET_COMM_NAME(INFINI_DEVICE_ILUVATAR, cuda); + GET_COMM_NAME(INFINI_DEVICE_QY, cuda); + GET_COMM_NAME(INFINI_DEVICE_HYGON, cuda); + GET_COMM_NAME(INFINI_DEVICE_ASCEND, ascend); + GET_COMM_NAME(INFINI_DEVICE_CAMBRICON, cambricon); + GET_COMM_NAME(INFINI_DEVICE_METAX, metax); + GET_COMM_NAME(INFINI_DEVICE_MOORE, moore); + GET_COMM_NAME(INFINI_DEVICE_KUNLUN, kunlun); + GET_COMM_NAME(INFINI_DEVICE_ALI, cuda); + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef GET_COMM_NAME +} + __INFINI_C infiniStatus_t infinicclCommInitAll( infiniDevice_t device_type, infinicclComm_t *comms, diff --git a/src/infiniccl/infiniccl_impl.h b/src/infiniccl/infiniccl_impl.h index 790355bd9..6120451ff 100644 --- a/src/infiniccl/infiniccl_impl.h +++ b/src/infiniccl/infiniccl_impl.h @@ -3,6 +3,8 @@ #include "infiniccl.h" +#include + struct InfinicclComm { infiniDevice_t device_type; int device_id; // the actual device ID, not rank number @@ -11,8 +13,34 @@ struct InfinicclComm { int world_size = 1; }; +namespace infiniccl { +inline infiniStatus_t getCommNameFromHandle( + infinicclComm_t comm, + char *comm_name, + size_t comm_name_size) { + if (comm == nullptr || comm_name == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + const int size = std::snprintf( + comm_name, + comm_name_size, + "infiniccl-%d-%p", + static_cast(comm->device_type), + comm->comm); + if (size < 0 || static_cast(size) >= comm_name_size) { + return INFINI_STATUS_BAD_PARAM; + } + return INFINI_STATUS_SUCCESS; +} +} // namespace infiniccl + #define INFINICCL_DEVICE_API(NAMSPACE, IMPL) \ namespace infiniccl::NAMSPACE { \ + infiniStatus_t getCommName( \ + infinicclComm_t comm, \ + char *comm_name, \ + size_t comm_name_size) IMPL; \ + \ infiniStatus_t commInitAll( \ infinicclComm_t *comms, \ int ndevice, \ diff --git a/src/infiniccl/kunlun/infiniccl_kunlun.cc b/src/infiniccl/kunlun/infiniccl_kunlun.cc index 0fa19c4aa..65f74524b 100644 --- a/src/infiniccl/kunlun/infiniccl_kunlun.cc +++ b/src/infiniccl/kunlun/infiniccl_kunlun.cc @@ -56,6 +56,13 @@ inline BKCLOp getBkclRedOp(infinicclReduceOp_t op) { namespace infiniccl::kunlun { +infiniStatus_t getCommName( + infinicclComm_t comm, + char *comm_name, + size_t comm_name_size) { + return getCommNameFromHandle(comm, comm_name, comm_name_size); +} + infiniStatus_t commInitAll( infinicclComm_t *comms, int ndevice, diff --git a/src/infiniccl/metax/infiniccl_metax.cc b/src/infiniccl/metax/infiniccl_metax.cc index 4b9f6303f..6ababd2ee 100644 --- a/src/infiniccl/metax/infiniccl_metax.cc +++ b/src/infiniccl/metax/infiniccl_metax.cc @@ -60,6 +60,13 @@ inline hcclComm_t getHcclComm(infinicclComm_t comm) { namespace infiniccl::metax { +infiniStatus_t getCommName( + infinicclComm_t comm, + char *comm_name, + size_t comm_name_size) { + return getCommNameFromHandle(comm, comm_name, comm_name_size); +} + infiniStatus_t commInitAll( infinicclComm_t *comms, int ndevice, diff --git a/src/infiniccl/moore/infiniccl_moore.cc b/src/infiniccl/moore/infiniccl_moore.cc index 7c2faed65..ea4100de9 100644 --- a/src/infiniccl/moore/infiniccl_moore.cc +++ b/src/infiniccl/moore/infiniccl_moore.cc @@ -59,6 +59,13 @@ inline mcclComm_t getMcclComm(infinicclComm_t comm) { namespace infiniccl::moore { +infiniStatus_t getCommName( + infinicclComm_t comm, + char *comm_name, + size_t comm_name_size) { + return getCommNameFromHandle(comm, comm_name, comm_name_size); +} + infiniStatus_t commInitAll( infinicclComm_t *comms, int ndevice, diff --git a/src/infinicore/ops/linear_allreduce/linear_allreduce.cc b/src/infinicore/ops/linear_allreduce/linear_allreduce.cc index 464e43116..c89c91b6e 100644 --- a/src/infinicore/ops/linear_allreduce/linear_allreduce.cc +++ b/src/infinicore/ops/linear_allreduce/linear_allreduce.cc @@ -1,50 +1,89 @@ #include "infinicore/ops/linear_allreduce.hpp" -#include "infinicore/device.hpp" #include "infinicore/ops/distributed/allreduce.hpp" #include "infinicore/ops/linear.hpp" -#if defined(ENABLE_ASCEND_API) -namespace infinicore::op::linear_allreduce_impl::ascend { -void linear_allreduce_impl( - Tensor output, Tensor input, Tensor weight, - std::optional bias, infinicclComm_t communicator); -} // namespace infinicore::op::linear_allreduce_impl::ascend -#endif +#include "../../utils.hpp" namespace infinicore::op { -Tensor linear_allreduce( - Tensor input, Tensor weight, std::optional bias, - infinicclReduceOp_t op, infinicclComm_t communicator) { -#if defined(ENABLE_ASCEND_API) - if (input->device().getType() == Device::Type::ASCEND) { - Size ndim = input->ndim(); - Size out_features = weight->shape()[0]; - auto out_shape = input->shape(); - out_shape[ndim - 1] = out_features; - auto out = Tensor::empty(out_shape, input->dtype(), input->device()); - linear_allreduce_impl::ascend::linear_allreduce_impl( - out, input, weight, bias, communicator); - return out; +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(LinearAllReduce); + +LinearAllReduce::LinearAllReduce( + Tensor output, + const Tensor &input, + const Tensor &weight, + const std::optional &bias, + infinicclComm_t communicator) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input, weight); + if (bias) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, *bias); } -#endif - auto output = linear(input, weight, bias); - return distributed::allreduce(output, op, communicator); + INFINICORE_GRAPH_OP_DISPATCH( + output->device().getType(), output, input, weight, bias, communicator); } -void linear_allreduce_( - Tensor output, Tensor input, Tensor weight, - std::optional bias, infinicclReduceOp_t op, +void LinearAllReduce::execute( + Tensor output, + const Tensor &input, + const Tensor &weight, + const std::optional &bias, infinicclComm_t communicator) { -#if defined(ENABLE_ASCEND_API) - if (input->device().getType() == Device::Type::ASCEND) { - linear_allreduce_impl::ascend::linear_allreduce_impl( - output, input, weight, bias, communicator); - return; + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + LinearAllReduce, output, input, weight, bias, communicator); +} + +static Tensor linear_allreduce_impl( + Tensor input, + Tensor weight, + std::optional bias, + infinicclComm_t communicator, + bool weight_is_packed) { + Size in_features = weight->shape()[weight_is_packed ? 0 : 1]; + Size out_features = weight->shape()[weight_is_packed ? 1 : 0]; + auto output_shape = input->shape(); + output_shape.back() = out_features; + + const bool aclnn_supported = input->dtype() == DataType::F16 || input->dtype() == DataType::BF16; + if (input->device().getType() == Device::Type::ASCEND && aclnn_supported) { + auto output = Tensor::empty(output_shape, input->dtype(), input->device()); + Size rows = 1; + for (Size i = 0; i + 1 < input->ndim(); ++i) { + rows *= input->size(i); + } + auto input_matrix = input->view({rows, in_features}); + auto output_matrix = output->view({rows, out_features}); + auto weight_matrix = weight_is_packed + ? weight + : weight->permute({1, 0}); + LinearAllReduce::execute( + output_matrix, input_matrix, weight_matrix, bias, communicator); + return output; } -#endif - linear_(output, input, weight, bias); - distributed::allreduce_(output, output, op, communicator); + + auto output = weight_is_packed + ? linear_packed(input, weight, bias) + : linear(input, weight, bias); + distributed::allreduce_( + output, output, INFINICCL_SUM, communicator); + return output; +} + +Tensor linear_allreduce( + Tensor input, + Tensor weight, + std::optional bias, + infinicclComm_t communicator) { + return linear_allreduce_impl( + input, weight, bias, communicator, false); +} + +Tensor linear_allreduce_packed( + Tensor input, + Tensor packed_weight, + std::optional bias, + infinicclComm_t communicator) { + return linear_allreduce_impl( + input, packed_weight, bias, communicator, true); } } // namespace infinicore::op diff --git a/src/infinicore/ops/linear_allreduce/linear_allreduce_ascend.cc b/src/infinicore/ops/linear_allreduce/linear_allreduce_ascend.cc deleted file mode 100644 index d46e72aa6..000000000 --- a/src/infinicore/ops/linear_allreduce/linear_allreduce_ascend.cc +++ /dev/null @@ -1,182 +0,0 @@ -#if defined(ENABLE_ASCEND_API) - -#include "../../../infiniccl/infiniccl_impl.h" -#include "infinicore/context/context.hpp" -#include "infinicore/device.hpp" -#include "infinicore/ops/linear_allreduce.hpp" - -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -namespace infinicore::op::linear_allreduce_impl::ascend { - -// ---- workspace pool ----------------------------------------- -// Per-stream reusable workspace. Once grown to the maximum -// required size, the buffer is never freed (leaked on growth). -// This avoids the MC2-notify-after-free bug: MC2 callbacks may -// still reference old workspace after stream sync completes. -struct WorkspaceBuf { - void *ptr = nullptr; - size_t cap = 0; -}; - -class WorkspacePool { - std::unordered_map bufs_; - std::mutex mtx_; - -public: - void *ensure(aclrtStream stream, size_t need) { - std::lock_guard lock(mtx_); - auto &b = bufs_[stream]; - if (need <= b.cap) { - return b.ptr; - } - void *new_ptr = nullptr; - aclError rc = aclrtMalloc(&new_ptr, need, ACL_MEM_MALLOC_HUGE_FIRST); - if (rc != ACL_SUCCESS || !new_ptr) { - fprintf(stderr, "[linear_allreduce/ascend] FATAL: aclrtMalloc(%zu MB) " - "failed rc=%d\n", - need / (1024 * 1024), (int)rc); - fflush(stderr); - throw std::runtime_error("[linear_allreduce/ascend] workspace alloc failed"); - } - b.ptr = new_ptr; - b.cap = need; - return b.ptr; - } -}; - -static WorkspacePool g_pool; - -static inline HcclComm get_hccl_comm(infinicclComm_t comm) { - return static_cast(comm->comm); -} - -static aclDataType to_acl_dtype(DataType dtype) { - switch (dtype) { - case DataType::F16: - return ACL_FLOAT16; - case DataType::BF16: - return ACL_BF16; - default: - throw std::runtime_error( - "[linear_allreduce/ascend] unsupported dtype: " + std::to_string(static_cast(dtype)) + ". aclnnMatmulAllReduce only supports F16/BF16"); - } -} - -void linear_allreduce_impl( - Tensor output, Tensor input, Tensor weight, - std::optional bias, - infinicclComm_t communicator) { - infinicore::context::setDevice(input->device()); - - auto atype = input->dtype(); - if (atype != DataType::F16 && atype != DataType::BF16) { - throw std::runtime_error( - "[linear_allreduce/ascend] unsupported activation dtype: " + std::to_string(static_cast(atype)) + ". aclnnMatmulAllReduce only supports F16/BF16"); - } - - auto w_perm = weight->permute({1, 0}); - Tensor weight_w = w_perm->is_contiguous() ? Tensor(w_perm) : w_perm->contiguous(); - - auto in_shape = input->shape(); - auto wt_shape = weight_w->shape(); - auto out_shape = output->shape(); - - std::vector in_dims(in_shape.begin(), in_shape.end()); - std::vector wt_dims(wt_shape.begin(), wt_shape.end()); - std::vector out_dims(out_shape.begin(), out_shape.end()); - - aclTensor *x1_acl = aclCreateTensor( - in_dims.data(), in_dims.size(), - to_acl_dtype(atype), - nullptr, 0, ACL_FORMAT_ND, - in_dims.data(), in_dims.size(), - const_cast(reinterpret_cast(input->data()))); - - aclTensor *x2_acl = aclCreateTensor( - wt_dims.data(), wt_dims.size(), - to_acl_dtype(atype), - nullptr, 0, ACL_FORMAT_ND, - wt_dims.data(), wt_dims.size(), - const_cast(reinterpret_cast(weight_w->data()))); - - aclTensor *bias_acl = nullptr; - if (bias.has_value()) { - Tensor bias_w = bias.value()->is_contiguous() ? Tensor(bias.value()) - : bias.value()->contiguous(); - auto bias_shape = bias_w->shape(); - std::vector bias_dims(bias_shape.begin(), bias_shape.end()); - bias_acl = aclCreateTensor( - bias_dims.data(), bias_dims.size(), - to_acl_dtype(atype), - nullptr, 0, ACL_FORMAT_ND, - bias_dims.data(), bias_dims.size(), - const_cast(reinterpret_cast(bias_w->data()))); - } - - aclTensor *out_acl = aclCreateTensor( - out_dims.data(), out_dims.size(), - to_acl_dtype(atype), - nullptr, 0, ACL_FORMAT_ND, - out_dims.data(), out_dims.size(), - const_cast(reinterpret_cast(output->data()))); - - HcclComm hccl_comm = get_hccl_comm(communicator); - char group_name[COMM_NAME_MAX_LENGTH] = {}; - HcclGetCommName(hccl_comm, group_name); - - uint64_t workspace_size = 0; - aclOpExecutor *executor = nullptr; - aclnnStatus ret = aclnnMatmulAllReduceGetWorkspaceSize( - x1_acl, x2_acl, bias_acl, - group_name, "sum", - 0, 1, - out_acl, - &workspace_size, &executor); - - if (ret != 0) { - if (bias_acl) { - aclDestroyTensor(bias_acl); - } - aclDestroyTensor(x1_acl); - aclDestroyTensor(x2_acl); - aclDestroyTensor(out_acl); - const char *err = aclGetRecentErrMsg(); - throw std::runtime_error( - std::string("[linear_allreduce/ascend] GetWorkspaceSize failed: ") + std::to_string(ret) + ", msg: " + (err ? err : "(null)")); - } - - aclrtStream stream = static_cast( - infinicore::context::getStream()); - void *workspace = g_pool.ensure(stream, (size_t)workspace_size); - - ret = aclnnMatmulAllReduce( - workspace, workspace_size, executor, stream); - - if (bias_acl) { - aclDestroyTensor(bias_acl); - } - aclDestroyTensor(x1_acl); - aclDestroyTensor(x2_acl); - aclDestroyTensor(out_acl); - - if (ret != 0) { - const char *err = aclGetRecentErrMsg(); - throw std::runtime_error( - std::string("[linear_allreduce/ascend] execution failed: ") + std::to_string(ret) + ", msg: " + (err ? err : "(null)")); - } -} - -} // namespace infinicore::op::linear_allreduce_impl::ascend - -#endif // ENABLE_ASCEND_API diff --git a/src/infinicore/ops/linear_allreduce/linear_allreduce_infiniop.cc b/src/infinicore/ops/linear_allreduce/linear_allreduce_infiniop.cc new file mode 100644 index 000000000..35369b07a --- /dev/null +++ b/src/infinicore/ops/linear_allreduce/linear_allreduce_infiniop.cc @@ -0,0 +1,86 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/linear_allreduce.hpp" + +#include +#include + +namespace infinicore::op::linear_allreduce_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, MatmulAllReduce, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace; + graph::GraphTensor output; + graph::GraphTensor input; + graph::GraphTensor weight; + std::optional bias; +}; + +void *plan( + Tensor output, + const Tensor &input, + const Tensor &weight, + const std::optional &bias, + infinicclComm_t communicator) { + std::array group_name{}; + INFINICORE_CHECK_ERROR(infinicclGetCommName( + communicator, group_name.data(), group_name.size())); + + size_t seed = hash_combine( + output, input, weight, bias, std::string(group_name.data())); + auto bias_desc = bias ? (*bias)->desc() : nullptr; + + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, descriptor, MatmulAllReduce, + seed, output->desc(), input->desc(), weight->desc(), + bias_desc, group_name.data()); + + INFINIOP_WORKSPACE_TENSOR( + workspace, MatmulAllReduce, descriptor); + + std::optional graph_bias; + if (bias) { + graph_bias.emplace(*bias); + } + return new PlannedMeta{ + descriptor, + graph::GraphTensor(workspace), + graph::GraphTensor(output), + graph::GraphTensor(input), + graph::GraphTensor(weight), + std::move(graph_bias)}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + const void *bias = planned->bias + ? (*planned->bias)->data() + : nullptr; + INFINICORE_CHECK_ERROR(infiniopMatmulAllReduce( + planned->descriptor->desc, + planned->workspace->data(), + planned->workspace->numel(), + planned->output->data(), + planned->input->data(), + planned->weight->data(), + bias, + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + LinearAllReduce::plan_dispatcher().registerDevice( + Device::Type::ASCEND, &plan); + LinearAllReduce::run_dispatcher().registerDevice( + Device::Type::ASCEND, &run); + LinearAllReduce::cleanup_dispatcher().registerDevice( + Device::Type::ASCEND, &cleanup); + return true; +}(); + +} // namespace infinicore::op::linear_allreduce_impl::infiniop diff --git a/src/infiniop/ops/matmul_all_reduce/ascend/matmul_all_reduce_ascend.cc b/src/infiniop/ops/matmul_all_reduce/ascend/matmul_all_reduce_ascend.cc new file mode 100644 index 000000000..d14d2f450 --- /dev/null +++ b/src/infiniop/ops/matmul_all_reduce/ascend/matmul_all_reduce_ascend.cc @@ -0,0 +1,130 @@ +#include "matmul_all_reduce_ascend.h" +#include "../../../devices/ascend/common_ascend.h" + +#include + +namespace op::matmul_all_reduce::ascend { + +struct Descriptor::Opaque { + aclnnTensorDescriptor_t output; + aclnnTensorDescriptor_t input; + aclnnTensorDescriptor_t weight; + aclnnTensorDescriptor_t bias; + aclOpExecutor *executor; + + ~Opaque() { + delete output; + delete input; + delete weight; + delete bias; + aclDestroyAclOpExecutor(executor); + } +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t weight_desc, + infiniopTensorDescriptor_t bias_desc, + const char *group_name) { + if (desc_ptr == nullptr || group_name == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + + auto dtype = input_desc->dtype(); + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16); + CHECK_API_OR(output_desc->dtype() == dtype, true, + return INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_API_OR(weight_desc->dtype() == dtype, true, + return INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_API_OR(bias_desc == nullptr || bias_desc->dtype() == dtype, true, + return INFINI_STATUS_BAD_TENSOR_DTYPE); + + CHECK_API_OR(input_desc->ndim() == 2, true, + return INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_API_OR(weight_desc->ndim() == 2, true, + return INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_API_OR(output_desc->ndim() == 2, true, + return INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_API_OR(bias_desc == nullptr || bias_desc->ndim() == 1, true, + return INFINI_STATUS_BAD_TENSOR_SHAPE); + + const auto &input_shape = input_desc->shape(); + const auto &weight_shape = weight_desc->shape(); + const auto &output_shape = output_desc->shape(); + CHECK_API_OR(input_shape[1] == weight_shape[0], true, + return INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_API_OR(output_shape[0] == input_shape[0], true, + return INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_API_OR(output_shape[1] == weight_shape[1], true, + return INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_API_OR(bias_desc == nullptr || bias_desc->shape()[0] == output_shape[1], + true, return INFINI_STATUS_BAD_TENSOR_SHAPE); + + auto output = new aclnnTensorDescriptor(output_desc); + auto input = new aclnnTensorDescriptor(input_desc); + auto weight = new aclnnTensorDescriptor(weight_desc); + auto bias = bias_desc == nullptr + ? nullptr + : new aclnnTensorDescriptor(bias_desc); + + uint64_t workspace_size = 0; + aclOpExecutor *executor = nullptr; + CHECK_ACL(aclnnMatmulAllReduceGetWorkspaceSize( + input->tensor, + weight->tensor, + bias == nullptr ? nullptr : bias->tensor, + group_name, + "sum", + 0, + 1, + output->tensor, + &workspace_size, + &executor)); + CHECK_ACL(aclSetAclOpExecutorRepeatable(executor)); + + auto handle = reinterpret_cast(handle_); + *desc_ptr = new Descriptor( + workspace_size, + new Opaque{output, input, weight, bias, executor}, + handle->device, + handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *input, + const void *weight, + const void *bias, + void *stream) const { + if (workspace_size < workspaceSize()) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + CHECK_ACL(AclSetTensorAddr( + _opaque->executor, 0, _opaque->input->tensor, + const_cast(input))); + CHECK_ACL(AclSetTensorAddr( + _opaque->executor, 1, _opaque->weight->tensor, + const_cast(weight))); + if (_opaque->bias != nullptr) { + CHECK_ACL(AclSetTensorAddr( + _opaque->executor, 2, _opaque->bias->tensor, + const_cast(bias))); + } + CHECK_ACL(AclSetTensorAddr( + _opaque->executor, 3, _opaque->output->tensor, output)); + CHECK_ACL(aclnnMatmulAllReduce( + workspace, workspace_size, _opaque->executor, stream)); + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::matmul_all_reduce::ascend diff --git a/src/infiniop/ops/matmul_all_reduce/ascend/matmul_all_reduce_ascend.h b/src/infiniop/ops/matmul_all_reduce/ascend/matmul_all_reduce_ascend.h new file mode 100644 index 000000000..547c95f63 --- /dev/null +++ b/src/infiniop/ops/matmul_all_reduce/ascend/matmul_all_reduce_ascend.h @@ -0,0 +1,8 @@ +#ifndef __MATMUL_ALL_REDUCE_ASCEND_H__ +#define __MATMUL_ALL_REDUCE_ASCEND_H__ + +#include "../matmul_all_reduce.h" + +DESCRIPTOR(ascend) + +#endif diff --git a/src/infiniop/ops/matmul_all_reduce/matmul_all_reduce.h b/src/infiniop/ops/matmul_all_reduce/matmul_all_reduce.h new file mode 100644 index 000000000..8357d3a58 --- /dev/null +++ b/src/infiniop/ops/matmul_all_reduce/matmul_all_reduce.h @@ -0,0 +1,48 @@ +#ifndef __MATMUL_ALL_REDUCE_H__ +#define __MATMUL_ALL_REDUCE_H__ + +#include "../../operator.h" + +#define DESCRIPTOR(NAMESPACE) \ + \ + namespace op::matmul_all_reduce::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ + size_t _workspace_size; \ + \ + Descriptor( \ + size_t workspace_size, \ + Opaque *opaque, \ + infiniDevice_t device_type, \ + int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _opaque(opaque), \ + _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, \ + infiniopTensorDescriptor_t weight_desc, \ + infiniopTensorDescriptor_t bias_desc, \ + const char *group_name); \ + \ + infiniStatus_t calculate( \ + void *workspace, \ + size_t workspace_size, \ + void *output, \ + const void *input, \ + const void *weight, \ + const void *bias, \ + void *stream) const; \ + }; \ + } + +#endif diff --git a/src/infiniop/ops/matmul_all_reduce/operator.cc b/src/infiniop/ops/matmul_all_reduce/operator.cc new file mode 100644 index 000000000..05f424562 --- /dev/null +++ b/src/infiniop/ops/matmul_all_reduce/operator.cc @@ -0,0 +1,91 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/matmul_all_reduce.h" + +#ifdef ENABLE_ASCEND_API +#include "ascend/matmul_all_reduce_ascend.h" +#endif + +__INFINI_C infiniStatus_t infiniopCreateMatmulAllReduceDescriptor( + infiniopHandle_t handle, + infiniopMatmulAllReduceDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t weight_desc, + infiniopTensorDescriptor_t bias_desc, + const char *group_name) { + switch (handle->device) { +#ifdef ENABLE_ASCEND_API + case INFINI_DEVICE_ASCEND: + return op::matmul_all_reduce::ascend::Descriptor::create( + handle, + reinterpret_cast( + desc_ptr), + output_desc, input_desc, weight_desc, bias_desc, group_name); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +} + +__INFINI_C infiniStatus_t infiniopGetMatmulAllReduceWorkspaceSize( + infiniopMatmulAllReduceDescriptor_t desc, + size_t *size) { + if (desc == nullptr || size == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + switch (desc->device_type) { +#ifdef ENABLE_ASCEND_API + case INFINI_DEVICE_ASCEND: + *size = reinterpret_cast< + const op::matmul_all_reduce::ascend::Descriptor *>(desc) + ->workspaceSize(); + return INFINI_STATUS_SUCCESS; +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +} + +__INFINI_C infiniStatus_t infiniopMatmulAllReduce( + infiniopMatmulAllReduceDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + const void *weight, + const void *bias, + void *stream) { + if (desc == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + switch (desc->device_type) { +#ifdef ENABLE_ASCEND_API + case INFINI_DEVICE_ASCEND: + return reinterpret_cast< + const op::matmul_all_reduce::ascend::Descriptor *>(desc) + ->calculate( + workspace, workspace_size, output, input, weight, bias, + stream); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +} + +__INFINI_C infiniStatus_t infiniopDestroyMatmulAllReduceDescriptor( + infiniopMatmulAllReduceDescriptor_t desc) { + if (desc == nullptr) { + return INFINI_STATUS_SUCCESS; + } + switch (desc->device_type) { +#ifdef ENABLE_ASCEND_API + case INFINI_DEVICE_ASCEND: + delete reinterpret_cast< + const op::matmul_all_reduce::ascend::Descriptor *>(desc); + return INFINI_STATUS_SUCCESS; +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +}