aclnnAlltoAllAllGatherBatchMatMul

📄 View Source Code

Applicable Products

Function

• Description: Fuses AllToAll and AllGather collective communication with BatchMatMul computation for parallel execution.

• Formula:

  The computation logic is as follows, where y1, y2, and y3 are outputs:

  $$x1 = AllToAll(x)$$ $$y2 = AllGather(x1)$$ $$y3 = BatchMatMul(y2, weight, bias)$$ $$y1 = activation_function(y3)$$

Function Prototype

Each operator uses a two-phase API. You must first call the "aclnnAlltoAllAllGatherBatchMatMulGetWorkspaceSize" API to obtain the workspace size required for computation and the executor that includes the operator computation process, and then call the "aclnnAlltoAllAllGatherBatchMatMul" API to perform the computation.

aclnnStatus aclnnAlltoAllAllGatherBatchMatMulGetWorkspaceSize(
    const aclTensor* x,
    const aclTensor* weight,
    const aclTensor* biasOptional,
    const char*      groupEp,
    const char*      groupTp,
    int64_t          epWorldSize,
    int64_t          tpWorldSize,
    int64_t          xShardType,
    int64_t          actType,
    aclTensor*       y1Out,
    aclTensor*       y2OutOptional,
    aclTensor*       y3OutOptional,
    uint64_t*        workspaceSize,
    aclOpExecutor**  executor)

aclnnStatus aclnnAlltoAllAllGatherBatchMatMul(
    void*           workspace,
    uint64_t        workspaceSize,
    aclOpExecutor*  executor,
    aclrtStream     stream)

aclnnAlltoAllAllGatherBatchMatMulGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
  x (aclTensor*)	Input	Serves as the left matrix for BatchMatMul computation after communication.	This input undergoes AllToAll and AllGather collective communication and must be 3-dimensional.	FLOAT16, BFLOAT16	ND	3	×
  weight(aclTensor*)	Input	Right matrix for BatchMatMul computation.	The data type remains consistent with x and must be 3-dimensional.	FLOAT16, BFLOAT16	ND	3	×
  biasOptional(aclTensor*)	Input	Bias for the BatchMatMul computation.	When x is FLOAT16, biasOptional must be FLOAT16. When x is BFLOAT16, biasOptional must be FLOAT32, supporting two or three dimensions. A null pointer can be passed in.	FLOAT16, FLOAT32	ND	2, 3	×
  groupEp(char*)	Input	Name of the ep communication domain, which is the communication domain for expert parallelism.	The string length must be greater than 0 and less than 128.	STRING	-	-	-
  groupTp(char*)	Input	Name of the TP communication domain, which is the communication domain for tensor parallelism.	The string length must be greater than 0 and less than 128.	STRING	-	-	-
  epWorldSize (int64_t)	Input.	Size of the ep communication domain.	Supports 2/4/8/16/32.	INT64	-	-	-
  tpWorldSize (int64_t)	Input	tp communication domain size.	Supports 2, 4, 8, 16, or 32.	INT64	-	-	-
  xShardType (int64_t)	Input	Indicates the dimension for AllGather on the tp domain.	0 indicates performing AllGather on the H dimension (i.e., the second dimension of x; x has three dimensions: the zeroth, first, and second dimensions) across the tp domain. 1 indicates performing AllGather on the C dimension (i.e., the first dimension of x) across the tp domain.	INT64	-	-	-
  actType (int64_t)	Input	Activation function type.	Supports input of 0/1/2/3/4, where 0 indicates no activation function. The mapping is [0: None, 1: GELU, 2: SiLU, 3: ReLU, 4: FastGELU].	INT64	-	-	-
  y1Out(aclTensor*)	Output	Final computation result. If an activation function is present, it is the output of the activation function; otherwise, it is the output of BatchMatMul.	Supports 3 dimensions. The data type remains consistent with the input x.	FLOAT16, BFLOAT16	ND	3	×
  y2OutOptional(aclTensor*)	Optional Output	Output of AllGather, which may be required for backward propagation.	Supports 3 dimensions. The data type must remain consistent with input x. A null pointer indicates that this output is not required.	FLOAT16, BFLOAT16	ND	3	×
  y3OutOptional(aclTensor*)	Optional output	Output of BatchMatMul when an activation function is present.	Supports 3 dimensions. The data type remains consistent with the input x. A null pointer indicates that this output is not required.	FLOAT16, BFLOAT16	ND	3	×
  workspaceSize (uint64_t*)	Output	Workspace size.	Returns the workspace size to be applied on the Device side.	UINT64	-	-	-
  executor(aclOpExecutor**)	Output	Returns the op executor, which includes the operator computation flow.	-	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:

<table style="undefined;table-layout: fixed; width: 1149px"><colgroup>
<col style="width: 282px">
<col style="width: 120px">
<col style="width: 747px">
</colgroup>
<thead>
<tr>
<th>Return Value</th>
<th>Error Code</th>
<th>Description</th>
</tr></thead>
<tbody>
<tr>
<td>ACLNN_ERR_PARAM_NULLPTR</td>
<td>161001</td>
<td>1. The passed x, weight, groupEp, groupTp, or y1Out is a null pointer.</td>
</tr>
<tr>
<td class="merged-cell" rowspan="6">ACLNN_ERR_PARAM_INVALID</td>
<td class="merged-cell" rowspan="6">161002</td>
<td>1. The string length of groupEp or groupTp is invalid.</td>
</tr>
<tr>
<td>2. An unsupported data type is input.</td>
</tr>
<tr>
<td>3. The attribute value is invalid.</td>
</tr>
<tr>
<td>4. The aclTensor dimension is invalid.</td>
</tr>
<tr>
<td>5. The aclTensor shape is invalid.</td>
</tr>
<tr>
<td>6. The optional output scenario is invalid — when the y3OutOptional optional output is required, actType must not be 0.</td>
</tr>
</tbody></table>

aclnnAlltoAllAllGatherBatchMatMul

• Parameters

  Parameter	Input/Output	Description
  workspace	Input	Memory address of the workspace applied on the Device side.
  workspaceSize	Input	Size of the workspace applied on the Device side, obtained through the first-phase API aclnnAlltoAllAllGatherBatchMatMulGetWorkspaceSize.
  Executor	Input	The operator executor, which includes the operator computation flow.
  Stream	Input	Specifies the stream for task execution.

• Return Value

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

Constraints

• Deterministic computation:

  • aclnnAlltoAllAllGatherBatchMatMul defaults to a deterministic implementation.

• Due to the requirements of collective communication and BatchMatMul computation, the input and output shapes must satisfy the following mathematical relationships (where ep = epWorldSize, tp = tpWorldSize):

  • AllGather along the H axis scenario, i.e., the xShardType is 0 scenario:

    • x: (E, C, H/tp)

    • weight: (E/ep, H, M/tp)

    • biasOptional: (E/ep, 1, M/tp) for three dimensions, or (E/ep, M/tp) for two dimensions when the pointer is not null.

    • y1Out: (E/ep, ep*C, M/tp)

    • y2OutOptional: (E/ep, ep*C, H)

    • y3OutOptional: (E/ep, ep*C, M/tp)

  • AllGather along the C axis scenario, that is, the scenario where xShardType is 1:

    • x: (E, C/tp, H)

    • weight: (E/ep, H, M/tp)

    • biasOptional: When the pointer is not null, (E/ep, 1, M/tp) for a 3D tensor, and (E/ep, M/tp) for a 2D tensor.

    • y1Out: (E/ep, ep*tp*C/tp, M/tp)

    • y2OutOptional: (E/ep, ep*tp*C/tp, H)

    • y3OutOptional: (E/ep, ep*tp*C/tp, M/tp)

• Data relationship description:

  • For example, x.size(0) equals E, and weight.size(0) equals E/ep, which indicates that x.size(0) = ep*weight.size(0), and x.size(0) is an integer multiple of ep. Other relationships are similar.

  • The value range of E is [2, 512], and E is an integer multiple of ep.

  • The value range of H is [1, 65535]. When xShardType is 0, H is an integer multiple of tp.

  • The value range of M/tp is [1, 65535].

  • The value range of E/ep is [1, 32].

  • Both ep and tp support only 2, 4, 8, 16, and 32.

  • The names of groupEp and groupTp cannot be the same.

  • C must be greater than 0, with the upper limit being the operator's device memory limit. When xShardType is 1, C is an integer multiple of tp.

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

  • Cross-supernode operations are not supported. Only intra-supernode operations are supported.

Example

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

• Atlas A3 training products/Atlas A3 inference products:

  #include <thread>
  #include <iostream>
  #include <string>
  #include <vector>
  #include "acl/acl.h"
  #include "hccl/hccl.h"
  #include "aclnnop/aclnn_all_to_all_all_gather_batch_matmul.h"
  
  #define CHECK_RET(cond, return_expr) \
      do {                             \
          if (!(cond)) {               \
              return_expr;             \
          }                            \
      } while (0)
  
  #define LOG_PRINT(message, ...)         \
      do {                                \
          printf(message, ##__VA_ARGS__); \
      } while(0)
  
  constexpr int EP_WORLD_SIZE = 4;
  constexpr int TP_WORLD_SIZE = 2;
  constexpr int DEV_NUM = EP_WORLD_SIZE * TP_WORLD_SIZE;
  
  int64_t GetShapeSize(const std::vector<int64_t> &shape)
  {
      int64_t shape_size = 1;
      for (auto i : shape) {
          shape_size *= i;
      }
      return shape_size;
  }
  
  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);
      auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc failed. ret: %d\n", ret); return ret);
      ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMemcpy failed. ret: %d\n", ret); return ret);
      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];
      }
      *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
          shape.data(), shape.size(), *deviceAddr);
      return 0;
  }
  
  struct Args {
      int rankId;
      HcclComm hcclEpComm;
      HcclComm hcclTpComm;
      aclrtStream stream;
      aclrtContext context;
  };
  
  int LaunchOneThreadAlltoAllAllGatherBmm(Args &args)
  {
      int ret = aclrtSetCurrentContext(args.context);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSetCurrentContext failed. ret: %d\n", ret); return ret);
      char hcomEpName[128] = {0};
      ret = HcclGetCommName(args.hcclEpComm, hcomEpName);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetEpCommName failed. ret: %d\n", ret); return -1);
      char hcomTpName[128] = {0};
      ret = HcclGetCommName(args.hcclTpComm, hcomTpName);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetTpCommName failed. ret: %d\n", ret); return -1);
      LOG_PRINT("[INFO] rank = %d, hcomEpName = %s, hcomTpName = %s, stream = %p, context = %p\n", args.rankId,
      hcomEpName, hcomTpName, args.stream, args.context);
  
      int64_t E = 2 * EP_WORLD_SIZE;
      int64_t C = 2 * TP_WORLD_SIZE;
      int64_t H = 6 * TP_WORLD_SIZE;
      int64_t M = 6 * TP_WORLD_SIZE;
      int64_t xShardType = 1; // Switch to 0 to enable the gather H-axis scenario.
      int64_t actType = 1;
  
      std::vector<int64_t> xShape;
      std::vector<int64_t> weightShape;
      std::vector<int64_t> biasShape;
      std::vector<int64_t> y1OutShape;
      std::vector<int64_t> y2OutShape;
      std::vector<int64_t> y3OutShape;
  
      if (xShardType == 1) {
          xShape = {E, C / TP_WORLD_SIZE, H};
          weightShape = {E / EP_WORLD_SIZE, H, M / TP_WORLD_SIZE};
          biasShape = {E / EP_WORLD_SIZE, 1, M / TP_WORLD_SIZE};
          y1OutShape = {E / EP_WORLD_SIZE, EP_WORLD_SIZE * TP_WORLD_SIZE * C / TP_WORLD_SIZE, M / TP_WORLD_SIZE};
          y2OutShape = {E / EP_WORLD_SIZE, EP_WORLD_SIZE * TP_WORLD_SIZE * C / TP_WORLD_SIZE, H};
          y3OutShape = {E / EP_WORLD_SIZE, EP_WORLD_SIZE * TP_WORLD_SIZE * C / TP_WORLD_SIZE, M / TP_WORLD_SIZE};
      } else if (xShardType == 0) {
          xShape = {E, C, H / TP_WORLD_SIZE};
          weightShape = {E / EP_WORLD_SIZE, H, M / TP_WORLD_SIZE};
          biasShape = {E / EP_WORLD_SIZE, 1, M / TP_WORLD_SIZE};
          y1OutShape = {E / EP_WORLD_SIZE, EP_WORLD_SIZE * C, M / TP_WORLD_SIZE};
          y2OutShape = {E / EP_WORLD_SIZE, EP_WORLD_SIZE * C, H};
          y3OutShape = {E / EP_WORLD_SIZE, EP_WORLD_SIZE * C, M / TP_WORLD_SIZE};
      } else {
          LOG_PRINT("[ERROR] unsupported xShardType = %ld.\n", xShardType);
          return -1;
      }
  
      void *xDeviceAddr = nullptr;
      void *weightDeviceAddr = nullptr;
      void *biasDeviceAddr = nullptr;
      void *y1OutDeviceAddr = nullptr;
      void *y2OutDeviceAddr = nullptr;
      void *y3OutDeviceAddr = nullptr;
      aclTensor *x = nullptr;
      aclTensor *weight = nullptr;
      aclTensor *bias = nullptr;
      aclTensor *y1Out = nullptr;
      aclTensor *y2Out = nullptr;
      aclTensor *y3Out = nullptr;
  
      uint64_t workspaceSize = 0;
      aclOpExecutor *executor = nullptr;
      void *workspaceAddr = nullptr;
  
      long long xShapeSize = GetShapeSize(xShape);
      long long weightShapeSize = GetShapeSize(weightShape);
      long long biasShapeSize = GetShapeSize(biasShape);
      long long y1OutShapeSize = GetShapeSize(y1OutShape);
      long long y2OutShapeSize = GetShapeSize(y2OutShape);
      long long y3OutShapeSize = GetShapeSize(y3OutShape);
  
      std::vector<int16_t> xHostData(xShapeSize, 1);
      std::vector<int16_t> weightHostData(weightShapeSize, 2);
      std::vector<int16_t> biasHostData(biasShapeSize, 3);
      std::vector<int16_t> y1OutHostData(y1OutShapeSize, 0);
      std::vector<int16_t> y2OutHostData(y2OutShapeSize, 0);
      std::vector<int16_t> y3OutHostData(y3OutShapeSize, 0);
  
      ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT16, &x);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(weightHostData, weightShape, &weightDeviceAddr, aclDataType::ACL_FLOAT16, &weight);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(y1OutHostData, y1OutShape, &y1OutDeviceAddr, aclDataType::ACL_FLOAT16, &y1Out);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(y2OutHostData, y2OutShape, &y2OutDeviceAddr, aclDataType::ACL_FLOAT16, &y2Out);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(y3OutHostData, y3OutShape, &y3OutDeviceAddr, aclDataType::ACL_FLOAT16, &y3Out);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
  
      // Call the first-phase API.
      ret = aclnnAlltoAllAllGatherBatchMatMulGetWorkspaceSize(x, weight, bias, hcomEpName, hcomTpName, EP_WORLD_SIZE,
          TP_WORLD_SIZE, xShardType, actType, y1Out, y2Out, y3Out, &workspaceSize, &executor);
      CHECK_RET(ret == ACL_SUCCESS,
          LOG_PRINT("[ERROR] aclnnAlltoAllAllGatherBatchMatMulGetWorkspaceSize failed. ret = %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("[ERROR] aclrtMalloc workspace failed. ret = %d \n", ret); return ret);
      }
      // Call the second-phase API.
      ret = aclnnAlltoAllAllGatherBatchMatMul(workspaceAddr, workspaceSize, executor, args.stream);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnAlltoAllAllGatherBatchMatMul failed. ret = %d \n", ret);
          return ret);
      // (Boilerplate) Synchronously wait for task execution to complete.
      ret = aclrtSynchronizeStreamWithTimeout(args.stream, 10000);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSynchronizeStreamWithTimeout failed. ret = %d \n", ret);
          return ret);
      LOG_PRINT("[INFO] device_%d aclnnAlltoAllAllGatherBatchMatMul execute successfully.\n", args.rankId);
      HcclCommDestroy(args.hcclEpComm);
      HcclCommDestroy(args.hcclTpComm);
      // Release device resources. Modify based on the specific API definition.
      if (x != nullptr) {
          aclDestroyTensor(x);
      }
      if (weight != nullptr) {
          aclDestroyTensor(weight);
      }
      if (bias != nullptr) {
          aclDestroyTensor(bias);
      }
      if (y1Out != nullptr) {
          aclDestroyTensor(y1Out);
      }
      if (y2Out != nullptr) {
          aclDestroyTensor(y2Out);
      }
      if (y3Out != nullptr) {
          aclDestroyTensor(y3Out);
      }
      if (xDeviceAddr != nullptr) {
          aclrtFree(xDeviceAddr);
      }
      if (weightDeviceAddr != nullptr) {
          aclrtFree(weightDeviceAddr);
      }
      if (biasDeviceAddr != nullptr) {
          aclrtFree(biasDeviceAddr);
      }
      if (y1OutDeviceAddr != nullptr) {
          aclrtFree(y1OutDeviceAddr);
      }
      if (y2OutDeviceAddr != nullptr) {
          aclrtFree(y2OutDeviceAddr);
      }
      if (y3OutDeviceAddr != nullptr) {
          aclrtFree(y3OutDeviceAddr);
      }
      if (workspaceSize > 0) {
          aclrtFree(workspaceAddr);
      }
      ret = aclrtDestroyStream(args.stream);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtDestroyStream failed. ret = %d \n", ret); return ret);
      ret = aclrtResetDevice(args.rankId);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtResetDevice failed. ret = %d \n", ret); return ret);
      return 0;
  }
  
  int main(int argc, char *argv[])
  {
      // This sample is implemented based on Atlas A3 and must run on Atlas A3.
      int ret = aclInit(nullptr);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclInit failed. ret = %d \n", ret); return ret);
      aclrtStream stream[DEV_NUM];
      aclrtContext context[DEV_NUM];
      for (uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
      ret = aclrtSetDevice(rankId);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSetDevice failed. ret = %d \n", ret); return ret);
      ret = aclrtCreateContext(&context[rankId], rankId);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateContext failed. ret = %d \n", ret); return ret);
      ret = aclrtCreateStream(&stream[rankId]);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateStream failed. ret = %d \n", ret); return ret);
      }
  
      int32_t devicesEp[DEV_NUM];
      int32_t devicesTp[DEV_NUM];
  
      // Initialize the ep domain. ep = 4  {0,2,4,6} {1,3,5,7}
      HcclComm commsEp[DEV_NUM];
      for (int i = 0; i < TP_WORLD_SIZE; i++) {
          for (int j =0; j < EP_WORLD_SIZE; j++) {
              devicesEp[j + i * EP_WORLD_SIZE] = i + j * TP_WORLD_SIZE;
          }
          ret = HcclCommInitAll(EP_WORLD_SIZE, &devicesEp[i * EP_WORLD_SIZE], &commsEp[i * EP_WORLD_SIZE]);
          CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclCommInitAll ep world %d failed. ret = %d \n", i, ret);
              return ret);
      }
  
      // Initialize the tp domain. tp = 4  {0,1},{2,3},{4,5},{6,7}
      HcclComm commsTp[DEV_NUM];
      for (int i = 0; i < EP_WORLD_SIZE; i++) {
          for (int j =0; j < TP_WORLD_SIZE; j++) {
              devicesTp[j + i * TP_WORLD_SIZE] = j + i * TP_WORLD_SIZE;
          }
          ret = HcclCommInitAll(TP_WORLD_SIZE, &devicesTp[i * TP_WORLD_SIZE], &commsTp[i * TP_WORLD_SIZE]);
          CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclCommInitAll tp world %d failed. ret = %d \n", i, ret);
              return ret);
      }
  
      Args args[DEV_NUM];
      // Start multithreading.
      std::vector<std::unique_ptr<std::thread>> threads(DEV_NUM);
      for (uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
          args[rankId].rankId = rankId;
          args[rankId].hcclEpComm = commsEp[rankId % TP_WORLD_SIZE * EP_WORLD_SIZE + rankId / TP_WORLD_SIZE];
          // Initialize TP in order, but within the same domain, device IDs are paired as [0,4], [1,5], ..., [3,7].
          args[rankId].hcclTpComm = commsTp[rankId];
          std::cout << "test devices id " << rankId << " = " << devicesTp[rankId] << std::endl;
          if (rankId == (DEV_NUM - 1)) {
              args[rankId].hcclTpComm = commsTp[rankId];
          }
          args[rankId].stream = stream[rankId];
          args[rankId].context = context[rankId];
          threads[rankId].reset(new std::thread(&LaunchOneThreadAlltoAllAllGatherBmm, std::ref(args[rankId])));
      }
      for (uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
          threads[rankId]->join();
      }
  
      aclFinalize();
      return 0;
  }