From 03f0fa98b12baebb51751af14f41357045cbd7c7 Mon Sep 17 00:00:00 2001 From: ShaneWu Date: Tue, 11 Aug 2026 11:10:43 +0800 Subject: [PATCH] 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