aclnnMatmulReduceScatter

📄 View Source Code

Applicable Products

NOTE When using this API, ensure that the driver firmware package and the CANN package are both of the matching 8.0.RC2 version or a matching later version. Otherwise, errors such as Bus Error will occur.

Function

• Description: Completes the mm + reduce_scatter_base computation.

• Formula:

  $$output=reduceScatter(x1@x2+bias)$$

Function Prototype

Each operator follows the two-phase API. You must first call the "aclnnMatmulReduceScatterGetWorkspaceSize" API to obtain the workspace size required for computation and the executor that encapsulates the operator computation process, and then call the "aclnnMatmulReduceScatter" API to perform the computation.

aclnnStatus aclnnMatmulReduceScatterGetWorkspaceSize(
    const aclTensor* x1,
    const aclTensor* x2,
    const aclTensor* bias,
    const char*      group,
    const char*      reduceOp,
    int64_t          commTurn,
    int64_t          streamMode,
    const aclTensor* output,
    uint64_t*        workspaceSize,
    aclOpExecutor**  executor)

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

aclnnMatmulReduceScatterGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
  x1	Input	Corresponds to x1 in the formula.	Supports empty Tensor. The data type must be consistent with x2. The current version only supports two-dimensional shape input and only supports the non-transpose scenario.	FLOAT16, BFLOAT16	ND	2	×
  x2	Input	Refers to x2 in the formula.	Supports empty tensors. The data type is consistent with x1. The current version only supports two-dimensional input, in both transpose and non-transpose scenarios. Only supports non-contiguous tensors in the two-axis transpose case. [Non-contiguous tensors] in other scenarios are not supported.	FLOAT16, BFLOAT16	ND	2	√ (applicable only to transpose scenarios)
  bias	Input.	Bias in the formula.	Supports passing a null pointer. The current version only supports one-dimensional input.	FLOAT16, BFLOAT16	ND	1	×
  group	Input.	Name of the communication domain.	Obtained through the Hccl API "extern HcclResult HcclGetCommName(HcclComm comm, char* commName);", where commName is the group.	-	-	-	-
  reduceOp	Input	Reduce operation type.	Currently, only "sum" is supported.	-	-	-	-
  commTurn	Input	Number of communication data splits, i.e., total data volume / single communication volume.	The current version only supports input 0.	-	-	-	-
  streamMode	Input	Stream mode enumeration.	Currently only supports the enumeration value 1.	-	-	-	-
  output	Output	Result of the AllGather communication and MatMul computation, i.e., the output in the formula.	Supports empty tensor. Consistent with the data type of x1.	FLOAT16, BFLOAT16	ND	2	×
  workspaceSize	Output	Returns the workspace size to be applied on the Device side.	-	-	-	-	-
  executor	Output	Returns the operator executor, which contains the operator computation process.	-	-	-	-	-

  • Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products:

    • bias: Currently does not support non-zero input scenarios.

  • Ascend 950PR/Ascend 950DT:

    • bias: Supports non-zero input scenarios.

• 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 passed x1, x2, or output is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The data type of x1, x2, bias, or output is not within the supported range.
  		streamMode is not within the valid range.
  		The shape of x1, x2, or output does not meet the constraint requirements.

aclnnMatmulReduceScatter

• Parameters

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

• Return value

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

Constraints

• Deterministic computation:

  • aclnnMatmulReduceScatter defaults to a non-deterministic implementation. Deterministic computation can be enabled by calling aclrtCtxSetSysParamOpt.

• The input x1 is 2-dimensional with a shape of (m, k), where m must be an integer multiple of the number of devices (rank_size).

• The input x2 must be 2-dimensional with a shape of (k, n). The axes must meet the input parameter requirements of the mm operator, with the k axes being equal, and the value range of the k axis is [256, 65535).

• For the empty tensor scenario supported by x1/x2, m and n can be empty, k cannot be empty, and the following conditions must be met:

  • m is empty, k is not empty, n is not empty;

  • m is not empty, k is not empty, n is empty;

  • m is empty, k is not empty, n is empty.

• The x2 matrix supports both transpose and non-transpose scenarios, while the x1 matrix only supports the non-transpose scenario.

• The data types of the x1 and x2 computation inputs must be consistent with the data type of the output computation output.

• bias does not currently support scenarios where the input is non-zero.

• The output is 2-dimensional, with shape (m/rank_size, n), where rank_size is the number of cards.

• Atlas A2 training products/Atlas A2 inference products: 2, 4, and 8 cards are supported, and only hccs link all mesh networking is supported.

• Atlas A3 training products/Atlas A3 inference products:

  • Supports 2, 4, 8, 16, and 32 cards, and only supports hccs link double ring networking.

  • The total collective communication data volume of reduceScatter(x1@x2+bias) cannot exceed 16*256 MB. The total collective communication data volume is calculated as: m * n * sizeof(output_dtype). Due to different shapes, the internal implementation of the operator may vary, and the actual supported total communication volume may be slightly less than this value.

• Ascend 950PR/Ascend 950DT: Supports 2, 4, 8, 16, 32, and 64 cards, and only supports hccs link all mesh networking.

• Atlas A2 training products/Atlas A2 inference products: The compute-communication fusion MC2 operators in a model only support the same communication domain.

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, Atlas A3 training products/Atlas A3 inference products:

  #include <thread>
  #include <iostream>
  #include <vector>
  #include <memory>
  #include <cstdio>
  #include "hccl/hccl.h"
  #include "aclnn/opdev/fp16_t.h"
  #include "aclnnop/aclnn_matmul_reduce_scatter.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 DEV_NUM = 2;
  
  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 hcclComm;
      aclrtStream stream;
      aclrtContext context;
    };
  
  int launchOneThread_MmReduceScatter(Args &args)
  {
      int ret = aclrtSetCurrentContext(args.context);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSetCurrentContext failed. ret = %d \n", ret); return ret);
  
      char hcomName[128] = {0};
      ret = HcclGetCommName(args.hcclComm, hcomName);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetCommName failed. ERROR: %d\n", ret); return -1);
      LOG_PRINT("[INFO] rank = %d, hcomName = %s, stream = %p\n", args.rankId, hcomName, args.stream);
      std::vector<int64_t> x1Shape = {1024, 256};
      std::vector<int64_t> x2Shape = {256, 512};
      std::vector<int64_t> biasShape = {512};
      std::vector<int64_t> outShape = {1024 / DEV_NUM, 512};
      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 commTurn = 0;
      int64_t streamMode = 1;
      uint64_t workspaceSize = 0;
      aclOpExecutor *executor = nullptr;
      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, 0);
      std::vector<op::fp16_t> x2HostData(x2ShapeSize, 0);
      std::vector<op::fp16_t> biasHostData(biasShapeSize, 0);
      std::vector<op::fp16_t> outHostData(outShapeSize, 0);
      // Create 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(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT16, &out);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
  
      // Call the first-phase API.
      ret = aclnnMatmulReduceScatterGetWorkspaceSize(
          x1, x2, bias, hcomName, "sum", commTurn, streamMode, out, &workspaceSize, &executor);
      CHECK_RET(ret == ACL_SUCCESS,
          LOG_PRINT("[ERROR] aclnnMatmulReduceScatterGetWorkspaceSize 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 = aclnnMatmulReduceScatter(workspaceAddr, workspaceSize, executor, args.stream);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMatmulReduceScatter 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 aclnnMatmulReduceScatter execute successfully.\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);
      }
      ret = HcclCommDestroy(args.hcclComm);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclCommDestroy failed. ret = %d \n", ret); return ret);
      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);
      ret = aclrtDestroyContext(args.context);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtDestroyContext 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("aclrtCreateContext failed. ERROR: %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 devices[DEV_NUM];
      for (int i = 0; i < DEV_NUM; i++) {
          devices[i] = i;
      }
      // Initialize the collective communication domain.
      HcclComm comms[DEV_NUM];
      ret = HcclCommInitAll(DEV_NUM, devices, comms);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclCommInitAll failed. ret = %d \n", ret); return ret);
  
      Args args[DEV_NUM];
      // Start multiple threads.
      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].hcclComm = comms[rankId];
          args[rankId].context = context[rankId];
          args[rankId].stream = stream[rankId];
          threads[rankId].reset(new(std::nothrow) std::thread(&launchOneThread_MmReduceScatter, std::ref(args[rankId])));
      }
      for (uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
          threads[rankId]->join();
      }
      aclFinalize();
      return 0;
  }