Feat/pr 1485 rebased

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,

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"

include/infinicore/ops/linear_allreduce.hpp

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+#include <infiniccl.h>
+#include <optional>
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(
+    LinearAllReduce,
+    Tensor,
+    const Tensor &,
+    const Tensor &,
+    const std::optional<Tensor> &,
+    infinicclComm_t);
+
+Tensor linear_allreduce(
+    Tensor input,
+    Tensor weight,
+    std::optional<Tensor> bias,
+    infinicclComm_t communicator);
+
+Tensor linear_allreduce_packed(
+    Tensor input,
+    Tensor packed_weight,
+    std::optional<Tensor> bias,
+    infinicclComm_t communicator);
+
+} // namespace infinicore::op

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"

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

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,

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,

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,

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,

src/infiniccl/infiniccl_impl.h

@@ -3,6 +3,8 @@
 
 #include "infiniccl.h"
 
+#include <cstdio>
+
 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<int>(comm->device_type),
+        comm->comm);
+    if (size < 0 || static_cast<size_t>(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,                                       \

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,

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,

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,

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<Tensor> &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<Tensor> &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<Tensor> 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<Tensor> bias,
+    infinicclComm_t communicator) {
+    return linear_allreduce_impl(
+        input, weight, bias, communicator, false);
+}
+
+Tensor linear_allreduce_packed(
+    Tensor input,
+    Tensor packed_weight,
+    std::optional<Tensor> bias,
+    infinicclComm_t communicator) {
+    return linear_allreduce_impl(
+        input, packed_weight, bias, communicator, true);
+}
+
+} // namespace infinicore::op