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 new file mode 100644 index 000000000..82aea3f18 --- /dev/null +++ b/include/infinicore/ops/linear_allreduce.hpp @@ -0,0 +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, + infinicclComm_t communicator); + +Tensor linear_allreduce_packed( + Tensor input, + Tensor packed_weight, + std::optional bias, + 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 new file mode 100644 index 000000000..c89c91b6e --- /dev/null +++ b/src/infinicore/ops/linear_allreduce/linear_allreduce.cc @@ -0,0 +1,89 @@ +#include "infinicore/ops/linear_allreduce.hpp" +#include "infinicore/ops/distributed/allreduce.hpp" +#include "infinicore/ops/linear.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +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); + } + INFINICORE_GRAPH_OP_DISPATCH( + output->device().getType(), output, input, weight, bias, communicator); +} + +void LinearAllReduce::execute( + Tensor output, + const Tensor &input, + const Tensor &weight, + const std::optional &bias, + infinicclComm_t communicator) { + 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; + } + + 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_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; + } +}