Documentation
Download

aclnnMatmulAlltoAll

Applicable Products

ProductSupported
Ascend 950PR/Ascend 950DT√
Atlas A3 training products/Atlas A3 inference products√
Atlas A2 training products/Atlas A2 inference products√
Atlas 200I/500 A2 inference products×
Atlas inference products×
Atlas training products×

Function

  • Description: Fuses Matmul computation, Permute (to ensure contiguous addresses after communication), and AlltoAll communication. Computation precedes communication.

  • Formula: Assume the shape of x1 is (BS, H1), the shape of x2 is (H1, H2), and rankSize is the number of NPU cards.

    computeOut=x1@x2+biaspermutedOut=computeOut.view(BS,rankSize,H2/rankSize).permute(1,0,2)output=AlltoAll(permutedOut).view(rankSize∗BS,H2/rankSize)computeOut = x1 @ x2 + bias \\ permutedOut = computeOut.view(BS, rankSize, H2/rankSize).permute(1, 0, 2) \\ output = AlltoAll(permutedOut).view(rankSize*BS, H2/rankSize)

Function Prototype

Each operator is divided into a two-phase API. The "aclnnMatmulAlltoAllGetWorkspaceSize" API must be called first to obtain the workspace size required for computation and the executor that encapsulates the operator computation flow. Then, the "aclnnMatmulAlltoAll" API is called to perform the computation.

aclnnStatus aclnnMatmulAlltoAllGetWorkspaceSize(
  const aclTensor*   x1,
  const aclTensor*   x2,
  const aclTensor*   biasOptional,
  const aclIntArray* alltoAllAxesOptional,
  const char*        group,
  bool               transposeX1,
  bool               transposeX2,
  const aclTensor*   output,
  uint64_t*          workspaceSize,
  aclOpExecutor**    executor)
aclnnStatus aclnnMatmulAlltoAll(
  void*          workspace,
  uint64_t       workspaceSize,
  aclOpExecutor* executor,
  aclrtStream    stream)

aclnnMatmulAlltoAllGetWorkspaceSize

  • Parameters:

    Parameter Input/Output Description Instruction Data Type Data Format Dimension (Shape) Non-Contiguous Tensor
    x1 Input Left matrix input of the fused operator, corresponding to x1 in the formula. Serves as the left matrix input for the MatMul computation. FLOAT16, BFLOAT16 ND 2D, with shape (BS, H1). x
    x2 Input The right matrix input of the fusion operator, which is also the right matrix of the MatMul computation. Directly used as the right matrix input for MatMul computation. FLOAT16, BFLOAT16 ND 2D, with shape (H1, H2). √
    biasOptional Optional input. Bias accumulated after the matrix multiplication, corresponding to bias in the formula. Supports passing a null pointer. The data type restrictions vary by device model. For details, see Constraints. FLOAT16, BFLOAT16, FLOAT32 ND 1D, with shape (H2) x
    alltoAllAxesOptional Input Optional input. Direction of data exchange for AlltoAll and Permute. Supports empty or [-1, -2]. When empty, defaults to [-1, -2], indicating that the input is transformed from (BS, H2) to (BS * rankSize, H2 / rankSize). aclIntArray* (element type: INT64) - 1D, with shape (2). -
    group Input Specifies the string that identifies the column group, i.e., the communication domain name. Obtain commName as this parameter through the Hccl API HcclGetCommName. The string length must be in the range (0, 128). STRING - - -
    transposeX1 Input Indicates whether the left matrix is transposed. Setting this parameter to True is not supported currently. bool - - -
    transposeX2 Input Indicates whether the right matrix has been transposed. Specifies that the right matrix shape is (H2, H1) when set to True. bool - - -
    output Output Final computation result. The data type is the same as that of input x1. FLOAT16, BFLOAT16 ND 2D, with shape (BS*rankSize, H2/rankSize). x
    workspaceSize Output Returns the workspace size to be applied for on the Device side. - UINT64 - - -
    executor Output Returns the op executor, which contains the computation flow of the operator. - aclOpExecutor* - - -

  • Return Value

    aclnnStatus: return code. For details, see aclnn Return Code.

    The first-phase API performs input parameter validation and returns an error in the following scenarios:

    Return Value Error Code Description
    ACLNN_ERR_PARAM_NULLPTR 161001 The required input and output tensor parameters are null pointers.
    ACLNN_ERR_PARAM_INVALID 161002 The data type of the input or output is not within the supported range.
    The input tensor is an empty tensor.
    The alltoAllAxesOptional is invalid.
    transposeX1 is true.
    The communication domain length is invalid.
    The input or output tensor dimensions are invalid.
    The input or output format is a private format.

aclnnMatmulAlltoAll

  • Parameters:

    Parameter Input/Output Description
    workspace Input Workspace memory address applied for on the Device side.
    workspaceSize Input Workspace size applied for on the Device side, obtained from the first-phase API aclnnMatmulAlltoAllGetWorkspaceSize.
    executor Input Operator executor that contains the operator computation process.
    stream Input Specifies the stream for task execution.

  • Return value

    aclnnStatus: return code. For details, see aclnn Return Code.

Constraints

  • Defaults to a deterministic computation.

  • Number of Cards (rankSize), which varies by device model:

    • Atlas A2 training products/Atlas A2 inference products: Supports 2, 4, or 8 cards.

    • Atlas A3 training products/Atlas A3 inference products: Supports 2, 4, 8, or 16 cards.

    • Ascend 950PR/Ascend 950DT: Supports 2, 4, 8, or 16 cards.

  • The variable H2 used in the shape in Parameters must be divisible by the number of NPU cards.

  • The range of H1 supports only [1, 65535].

  • The value of BS*rankSize and H2 must not exceed 2147483647 (INT32_MAX). The value of BS must not be less than 0, and the value of H2 must not be less than 2.

  • The support for empty tensors varies by device model:

    • Atlas A2 training products/Atlas A2 inference products: No empty tensors are supported.

    • Atlas A3 training products/Atlas A3 inference products: Only empty tensors where the first dimension (BS) of input x1 is 0 are supported. Other empty tensors are not supported.

    • Ascend 950PR/Ascend 950DT: Only empty tensors where the first dimension (BS) of input x1 is 0 are supported. Other empty tensors are not supported.

  • The support for non-contiguous tensors varies by device model:

    • Atlas A2 training products/Atlas A2 inference products: No non-contiguous tensors are supported.

    • Atlas A3 training products/Atlas A3 inference products: No non-contiguous tensors are supported.

    • Ascend 950PR/Ascend 950DT: Only x2 is supported as a non-contiguous tensor. Other non-contiguous tensors are not supported.

  • The data types of the x1 and x2 compute inputs must be consistent with the data type of the output compute output. The x1, x2, and output passed in must not be null pointers.

  • The data type of biasOptional has different restrictions depending on the device model:

    • Atlas A2 training products/Atlas A2 inference products: When the data type of the output compute output is FLOAT16, the data type of the biasOptional compute input supports FLOAT16. When the data type of the output compute output is BFLOAT16, the data type of the biasOptional compute input supports FLOAT32.

    • Atlas A3 training products/Atlas A3 inference products: When the data type of the x1/x2 compute inputs is FLOAT16, the data type of the biasOptional compute input supports FLOAT16. When the data type of the x1/x2 compute inputs is BFLOAT16, the data type of the biasOptional compute input supports FLOAT32.

    • Ascend 950PR/Ascend 950DT: When the data type of x1/x2 Compute Input is FLOAT16, the data type of biasOptional Compute Input supports FLOAT16 and FLOAT32. When the data type of x1/x2 Compute Input is BFLOAT16, the data type of biasOptional Compute Input supports BFLOAT16 and FLOAT32.

  • The communication-computation fusion operator does not support concurrent calls, and different communication-computation fusion operators also do not support concurrent calls.

  • Cross-super-node communication is not supported. Only intra-super-node communication is supported.

Example

The following provides example code for reference only. For details about compilation and running, see Compile and Run Samples.

Note: This sample code calls some HCCL collective communication library APIs: HcclGetCommName, HcclCommInitAll, and HcclCommDestroy. See HCCL API (C).

  • Atlas A2 training products/Atlas A2 inference products:

    #include <thread>
#include <iostream>
#include <string>
#include <cstring>
#include <vector>
#include <acl/acl.h>
#include <hccl/hccl.h>
#include "aclnn/opdev/fp16_t.h"
#include "aclnnop/aclnn_matmul_allto_all.h"

int ndev = 2;

#define CHECK_RET(cond, return_expr) \
do {                               \
    if (!(cond)) {                   \
    return_expr;                   \
    }                                \
} while (0)

#define LOG_PRINT(message, ...)     \
do {                              \
    printf(message, ##__VA_ARGS__); \
} while (0)

int64_t GetShapeSize(const std::vector<int64_t> &shape) {
    int64_t shapeSize = 1;
    for (auto i: shape) {
        shapeSize *= i;
    }
    return shapeSize;
}

template<typename T>
int CreateAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr,
                    aclDataType dataType, aclTensor **tensor) {
    auto size = GetShapeSize(shape) * sizeof(T);
    // Call aclrtMalloc to apply for device-side memory.
    auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
    // Call aclrtMemcpy to copy host-side data to the device-side memory.
    ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
    // Calculate the strides of a contiguous tensor.
    std::vector<int64_t> strides(shape.size(), 1);
    for (int64_t i = shape.size() - 2; i >= 0; i--) {
        strides[i] = shape[i + 1] * strides[i + 1];
    }
    // Call the aclCreateTensor API to create an aclTensor.
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                            shape.data(), shape.size(), *deviceAddr);
    return 0;
}

struct Args {
    uint32_t rankId;
    HcclComm hcclComm;
    aclrtStream stream;
    aclrtContext context;
};

int launchOneThreadMatmulAlltoAll(Args &args) {
    int ret;
    ret = aclrtSetCurrentContext(args.context);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetCurrentContext failed. ERROR: %d\n", ret); return ret);
    char hcom_name[128];
    ret = HcclGetCommName(args.hcclComm, hcom_name);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetCommName failed. ret = %d \n", ret); return -1);
    LOG_PRINT("[INFO] rank %d hcom: %s stream: %p, context : %p\n", args.rankId, hcom_name, args.stream,
            args.context);

    std::vector<int64_t> x1Shape = {32, 64};
    std::vector<int64_t> x2Shape = {64, 128};
    std::vector<int64_t> biasShape = {128};
    std::vector<int64_t> outShape = {64, 64};
    void *x1DeviceAddr = nullptr;
    void *x2DeviceAddr = nullptr;
    void *biasDeviceAddr = nullptr;
    void *outDeviceAddr = nullptr;
    aclTensor *x1 = nullptr;
    aclTensor *x2 = nullptr;
    aclTensor *bias = nullptr;
    aclTensor *out = nullptr;

    int64_t a2aAxes[2] = {-1, -2};
    aclIntArray* alltoAllAxesOptional = aclCreateIntArray(a2aAxes, static_cast<uint64_t>(2));
    uint64_t workspaceSize = 0;
    aclOpExecutor *executor;
    void *workspaceAddr = nullptr;

    long long x1ShapeSize = GetShapeSize(x1Shape);
    long long x2ShapeSize = GetShapeSize(x2Shape);
    long long biasShapeSize = GetShapeSize(biasShape);
    long long outShapeSize = GetShapeSize(outShape);
    std::vector<op::fp16_t> x1HostData(x1ShapeSize, 1);
    std::vector<op::fp16_t> x2HostData(x2ShapeSize, 1);
    std::vector<op::fp16_t> biasHostData(biasShapeSize, 1);
    std::vector<op::fp16_t> outHostData(outShapeSize, 0);
    // Create a tensor.
    ret = CreateAclTensor(x1HostData, x1Shape, &x1DeviceAddr, aclDataType::ACL_FLOAT16, &x1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_FLOAT16, &x2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(biasHostData, biasShape, &biasDeviceAddr, aclDataType::ACL_FLOAT16, &bias);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT16, &out);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Call the first-phase API.
    ret = aclnnMatmulAlltoAllGetWorkspaceSize(x1, x2, bias, alltoAllAxesOptional, hcom_name, false, false,
                                            out, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS,
            LOG_PRINT("aclnnMatmulAlltoAllGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Apply for device memory based on the workspaceSize calculated by the first-phase API.
    if (workspaceSize > 0) {
        ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
    }
    // Call the second-phase API.
    ret = aclnnMatmulAlltoAll(workspaceAddr, workspaceSize, executor, args.stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMatmulAlltoAll failed. ERROR: %d\n", ret); return ret);
    //(Boilerplate) Synchronously wait for the task execution to complete.
    ret = aclrtSynchronizeStreamWithTimeout(args.stream, 10000);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
    LOG_PRINT("device%d aclnnMatmulAlltoAll execute success \n", args.rankId);
    // Release device resources. Modify based on the specific API definition.
    if (x1 != nullptr) {
        aclDestroyTensor(x1);
    }
    if (x2 != nullptr) {
        aclDestroyTensor(x2);
    }
    if (bias != nullptr) {
        aclDestroyTensor(bias);
    }
    if (out != nullptr) {
        aclDestroyTensor(out);
    }
    if (x1DeviceAddr != nullptr) {
        aclrtFree(x1DeviceAddr);
    }
    if (x2DeviceAddr != nullptr) {
        aclrtFree(x2DeviceAddr);
    }
    if (biasDeviceAddr != nullptr) {
        aclrtFree(biasDeviceAddr);
    }
    if (outDeviceAddr != nullptr) {
        aclrtFree(outDeviceAddr);
    }
    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(args.stream);
    HcclCommDestroy(args.hcclComm);
    aclrtDestroyContext(args.context);
    aclrtResetDevice(args.rankId);
    return 0;
}

int main(int argc, char *argv[])
{
    // This sample is implemented based on Atlas A2 and must run on Atlas A2.
    int ret;
    int32_t devices[ndev];
    for (int i = 0; i < ndev; i++) {
        devices[i] = i;
    }
    HcclComm comms[128];
    ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
    // Initialize the collective communication domain.
    for (int i = 0; i < ndev; i++) {
        ret = aclrtSetDevice(devices[i]);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
    }
    ret = HcclCommInitAll(ndev, devices, comms);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("HcclCommInitAll failed. ERROR: %d\n", ret); return ret);
    Args args[ndev];
    aclrtStream stream[ndev];
    aclrtContext context[ndev];
    for (uint32_t rankId = 0; rankId < ndev; rankId++) {
        ret = aclrtSetDevice(rankId);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
        ret = aclrtCreateContext(&context[rankId], rankId);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateContext failed. ERROR: %d\n", ret); return ret);
        ret = aclrtCreateStream(&stream[rankId]);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
    }
    // Start multiple threads.
    std::vector<std::unique_ptr<std::thread>> threads(ndev);
    for (uint32_t rankId = 0; rankId < ndev; rankId++) {
        args[rankId].rankId = rankId;
        args[rankId].hcclComm = comms[rankId];
        args[rankId].stream = stream[rankId];
        args[rankId].context = context[rankId];
        threads[rankId].reset(new(std::nothrow) std::thread(&launchOneThreadMatmulAlltoAll, std::ref(args[rankId])));
    }
    for (uint32_t rankId = 0; rankId < ndev; rankId++) {
        threads[rankId]->join();
    }
    aclFinalize();
    return 0;
}

  • Ascend 950PR/Ascend 950DT:

    #include <thread>
#include <iostream>
#include <string>
#include <cstring>
#include <vector>
#include <acl/acl.h>
#include <hccl/hccl.h>
#include "aclnnop/aclnn_matmul_allto_all.h"

int ndev = 2;

#define CHECK_RET(cond, return_expr) \
do {                               \
    if (!(cond)) {                   \
    return_expr;                   \
    }                                \
} while (0)

#define LOG_PRINT(message, ...)     \
do {                              \
    printf(message, ##__VA_ARGS__); \
} while (0)

int64_t GetShapeSize(const std::vector<int64_t> &shape) {
    int64_t shapeSize = 1;
    for (auto i: shape) {
        shapeSize *= i;
    }
    return shapeSize;
}

template<typename T>
int CreateAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr,
                    aclDataType dataType, aclTensor **tensor) {
    auto size = GetShapeSize(shape) * sizeof(T);
    // Call aclrtMalloc to apply for device memory.
    auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
    // Copy data from the host to the device memory by calling aclrtMemcpy.
    ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
    // Compute the strides of a contiguous tensor.
    std::vector<int64_t> strides(shape.size(), 1);
    for (int64_t i = shape.size() - 2; i >= 0; i--) {
        strides[i] = shape[i + 1] * strides[i + 1];
    }
    // Create an aclTensor by calling aclCreateTensor.
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                            shape.data(), shape.size(), *deviceAddr);
    return 0;
}

struct Args {
    uint32_t rankId;
    HcclComm hcclComm;
    aclrtStream stream;
    aclrtContext context;
};

int launchOneThreadMatmulAlltoAll(Args &args) {
    int ret;
    ret = aclrtSetCurrentContext(args.context);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetCurrentContext failed. ERROR: %d\n", ret); return ret);
    char hcom_name[128];
    ret = HcclGetCommName(args.hcclComm, hcom_name);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetCommName failed. ret = %d \n", ret); return -1);
    LOG_PRINT("[INFO] rank %d hcom: %s stream: %p, context : %p\n", args.rankId, hcom_name, args.stream,
            args.context);

    std::vector<int64_t> x1Shape = {32, 64};
    std::vector<int64_t> x2Shape = {64, 128};
    std::vector<int64_t> biasShape = {128};
    std::vector<int64_t> outShape = {64, 64};
    void *x1DeviceAddr = nullptr;
    void *x2DeviceAddr = nullptr;
    void *biasDeviceAddr = nullptr;
    void *outDeviceAddr = nullptr;
    aclTensor *x1 = nullptr;
    aclTensor *x2 = nullptr;
    aclTensor *bias = nullptr;
    aclTensor *out = nullptr;

    int64_t a2aAxes[2] = {-1, -2};
    aclIntArray* alltoAllAxesOptional = aclCreateIntArray(a2aAxes, static_cast<uint64_t>(2));
    uint64_t workspaceSize = 0;
    aclOpExecutor *executor;
    void *workspaceAddr = nullptr;

    long long x1ShapeSize = GetShapeSize(x1Shape);
    long long x2ShapeSize = GetShapeSize(x2Shape);
    long long biasShapeSize = GetShapeSize(biasShape);
    long long outShapeSize = GetShapeSize(outShape);
    std::vector<int16_t> x1HostData(x1ShapeSize, 1);
    std::vector<int16_t> x2HostData(x2ShapeSize, 1);
    std::vector<int16_t> biasHostData(biasShapeSize, 1);
    std::vector<int16_t> outHostData(outShapeSize, 0);
    // Create a tensor.
    ret = CreateAclTensor(x1HostData, x1Shape, &x1DeviceAddr, aclDataType::ACL_FLOAT16, &x1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_FLOAT16, &x2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(biasHostData, biasShape, &biasDeviceAddr, aclDataType::ACL_FLOAT16, &bias);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT16, &out);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Call the first-phase API.
    ret = aclnnMatmulAlltoAllGetWorkspaceSize(x1, x2, bias, alltoAllAxesOptional, hcom_name, false, false,
                                            out, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS,
            LOG_PRINT("aclnnMatmulAlltoAllGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Apply for device memory based on the workspaceSize calculated by the first-phase API.
    if (workspaceSize > 0) {
        ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
    }
    // Call the second-phase API.
    ret = aclnnMatmulAlltoAll(workspaceAddr, workspaceSize, executor, args.stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMatmulAlltoAll failed. ERROR: %d\n", ret); return ret);
    //Synchronously wait for the task execution to complete. (boilerplate)
    ret = aclrtSynchronizeStreamWithTimeout(args.stream, 10000);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
    LOG_PRINT("device%d aclnnMatmulAlltoAll execute success \n", args.rankId);
    // Release device resources. Modify based on the specific API definition.
    if (x1 != nullptr) {
        aclDestroyTensor(x1);
    }
    if (x2 != nullptr) {
        aclDestroyTensor(x2);
    }
    if (bias != nullptr) {
        aclDestroyTensor(bias);
    }
    if (out != nullptr) {
        aclDestroyTensor(out);
    }
    if (x1DeviceAddr != nullptr) {
        aclrtFree(x1DeviceAddr);
    }
    if (x2DeviceAddr != nullptr) {
        aclrtFree(x2DeviceAddr);
    }
    if (biasDeviceAddr != nullptr) {
        aclrtFree(biasDeviceAddr);
    }
    if (outDeviceAddr != nullptr) {
        aclrtFree(outDeviceAddr);
    }
    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(args.stream);
    HcclCommDestroy(args.hcclComm);
    aclrtDestroyContext(args.context);
    aclrtResetDevice(args.rankId);
    return 0;
}

int main(int argc, char *argv[])
{
    // This sample is implemented based on Atlas A5 and must run on Atlas A5.
    int ret;
    int32_t devices[ndev];
    for (int i = 0; i < ndev; i++) {
        devices[i] = i;
    }
    HcclComm comms[128];
    ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
    // Initialize the collective communication domain.
    for (int i = 0; i < ndev; i++) {
        ret = aclrtSetDevice(devices[i]);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
    }
    ret = HcclCommInitAll(ndev, devices, comms);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("HcclCommInitAll failed. ERROR: %d\n", ret); return ret);
    Args args[ndev];
    aclrtStream stream[ndev];
    aclrtContext context[ndev];
    for (uint32_t rankId = 0; rankId < ndev; rankId++) {
        ret = aclrtSetDevice(rankId);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
        ret = aclrtCreateContext(&context[rankId], rankId);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateContext failed. ERROR: %d\n", ret); return ret);
        ret = aclrtCreateStream(&stream[rankId]);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
    }
    // Start multithreading.
    std::vector<std::unique_ptr<std::thread>> threads(ndev);
    for (uint32_t rankId = 0; rankId < ndev; rankId++) {
        args[rankId].rankId = rankId;
        args[rankId].hcclComm = comms[rankId];
        args[rankId].stream = stream[rankId];
        args[rankId].context = context[rankId];
        threads[rankId].reset(new(std::nothrow) std::thread(&launchOneThreadMatmulAlltoAll, std::ref(args[rankId])));
    }
    for (uint32_t rankId = 0; rankId < ndev; rankId++) {
        threads[rankId]->join();
    }
    aclFinalize();
    return 0;
}
This site uses cookies to improve your browsing experience. By continuing to browse this site, you agree to our use of cookies.View Details