aclnnQuantMatmulAllReduce

📄 View Source Code

Applicable Products

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

Function

• Description: Performs MatMul computation on the quantized inputs x1 and x2, then performs Dequant computation, then performs an Add operation with x3, and finally performs an AllReduce computation. Supports per-tensor and per-channel quantization modes.

• Formula:

  $$output= AllReduce(dequantScale*(x1_{int8}@x2_{int8} + bias_{int32}) + x3)$$

Function Prototype

Each operator is divided into a two-phase API. You must first call the aclnnQuantMatmulAllReduceGetWorkspaceSize API to obtain the workspace size required for computation and the executor that contains the operator computation process, and then call the aclnnQuantMatmulAllReduce API to perform the computation.

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

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

aclnnQuantMatmulAllReduceGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
  x1	Input	Left matrix for the MatMul computation, i.e., x1 in the formula.	The current version supports only 2D or 3D input. Supports the non-transpose scenario.	INT8	ND	2-3	No
  x2	Input	Right matrix for the MatMul computation, i.e., x2 in the formula.	Supports transposed and non-transposed scenarios. In ND format, supports non-contiguous tensors when the last two axes are transposed. Other non-contiguous tensors are not supported.	INT8	ND, FRACTAL_NZ	2	×
  bias	Input	Corresponds to the bias offset in the formula.	The current version supports only one-dimensional input.	INT32	ND	0-1	√
  x3	Input	The add computation after the MatMul computation, i.e., x3 in the formula.	The shape is consistent with the shape after the MatMul computation.	FLOAT16, BFLOAT16	ND	2	√
  dequantScale	Input	Dequantization coefficient after the MatMul computation, i.e., dequantScale in the formula.	Shape is (1) in the per-tensor scenario, and (n) or (1, n) in the per-channel scenario.	INT64, UINT64, BFLOAT16	ND	2	√
  group	Input	Communication domain name.	Obtained through the HCCL API "extern HcclResult HcclGetCommName(HcclComm comm, char* commName);", where commName is the group.	String	-	-	-
  reduceOp	Input	Type of the reduce operation.	The current version only supports input "sum".	String	-	-	-
  commTurn	Input	Number of communication data splits, i.e., total data volume divided by single communication volume.	The current version only supports input of 0.	INT64	-	-	-
  streamMode	Input	Enumeration of stream modes.	The current version only supports enumeration value 1.	INT64	-	-	-
  output	Output	Result of the MatMul computation and AllReduce communication, i.e., output in the formula.	The number of dimensions of output is the same as that of x1.	FLOAT16, BFLOAT16	ND	2-3	√
  workspaceSize	Output	Returns the workspace size to be requested on the Device side.	-	-	-	-	-
  executor	Output	Returns the op executor, which contains the operator computation flow.	-	-	-	-	-

  • Atlas A2 training products/Atlas A2 inference products:

    • The data format of x2 input supports ND (only two-dimensional input is supported in the current version) and FRACTAL_NZ format (only four-dimensional input is supported in the current version). When the data format of x2 is FRACTAL_NZ, use aclnnCalculateMatmulWeightSizeV2 and aclnnTransMatmulWeight to complete the conversion from ND data format to NZ data format. Non-contiguous tensors only support the transpose scenario.

  • Ascend 950PR/Ascend 950DT:

    • The data format of x2 input only supports ND (only two-dimensional input is supported in the current version).

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

aclnnQuantMatmulAllReduce

• Parameters

  Parameter	Input/Output	Description
  workspace	Input.	Memory address of the workspace applied for on the Device side.
  workspaceSize	Input.	Size of the workspace applied for on the Device side, obtained by the first-phase API aclnnQuantMatmulAllReduceGetWorkspaceSize.
  executor	Input	Operator executor that contains the operator computation flow.
  stream	Input	Specifies the stream for executing the task.

• Return Value

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

Constraints

• Deterministic computation:

  • Atlas A2 training products/Atlas A2 inference products: aclnnQuantMatmulAllReduce defaults to a non-deterministic implementation. Deterministic computation can be enabled by setting the HCCL_DETERMINISTIC environment variable to true.

  • Ascend 950PR/Ascend 950DT: aclnnQuantMatmulAllReduce defaults to a deterministic implementation.

• MC2 is disabled in incremental scenarios and enabled in full scenarios.

• The input x1 can be two-dimensional or three-dimensional, with shape (b, s, k) or (m, k). x2 must be two-dimensional, with shape (k, n). The k axis must satisfy the input parameter requirements of the mm operator, meaning the k axes must be equal.

• The size of m does not exceed 2147483647. The size of the last dimension of x1 and x2 does not exceed 65535. The last dimension of x1 refers to k, and the last dimension of x2 refers to k during transpose or n during non-transpose. If bias is non-null, its shape is (n). If x3 is non-null, its shape is the same as output.

• When the shape of input x1 is (b, s, k), the shape of output is (b, s, n). When the shape of input x1 is (m, k), the shape of output is (m, n).

• The passed x1, x2, dequantScale, or output must not be null pointers.

• The data types and data formats of x1, x2, dequantScale, output, bias (in non-null scenarios), and x3 (in non-null scenarios) must fall within the supported range.

• If the output type is FLOAT16, the data type of dequantScale is INT64 or UINT64; if the output type is BFLOAT16, the data type of dequantScale is BFLOAT16 and the data type of x3 is BFLOAT16.

• Only HCCS link all-mesh networking is supported.

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

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

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

• Empty tensor support:

  • Empty tensors are not supported.

Example

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

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

• Atlas A2 training products/Atlas A2 inference products, Ascend 950PR/Ascend 950DT:

  #include <iostream>
  #include <vector>
  #include <thread>
  #include "hccl/hccl.h"
  #include "aclnn/opdev/fp16_t.h"
  #include "aclnnop/aclnn_trans_matmul_weight.h"
  #include "aclnnop/aclnn_quant_matmul_all_reduce.h"
  
  int ndev = 2;
  
  #define ACL_CHECK(ret)                                                                                     \
      do {                                                                                                   \
          auto retcode = ret;                                                                                \
          if (retcode != ACL_SUCCESS) {                                                                      \
              printf("[ERROR] acl interface return err %s:%d, retcode: %d \n", __FILE__, __LINE__, retcode); \
              return retcode;                                                                                \
          }                                                                                                  \
      } while (0)
  
  #define CHECK_RET(cond, return_expr) \
  do {                               \
      if (!(cond)) {                   \
      return_expr;                   \
      }                                \
  } while (0)
  
  #define LOG_PRINT(message, ...)     \
  do {                              \
      printf(message, ##__VA_ARGS__); \
  } while (0)
  
  struct Args {
      uint32_t rankId;
      HcclComm hcclComm;
      aclrtStream stream;
      aclrtContext context;
      std::string format;
  };
  
  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 CreateWeightNzAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr,
                              aclDataType dataType, aclTensor **tensor, Args &args) {
      auto size = GetShapeSize(shape) * sizeof(T);
      const aclIntArray *mat2Size = aclCreateIntArray(shape.data(), shape.size());
      auto ret = aclnnCalculateMatmulWeightSizeV2(mat2Size, ACL_INT8, &size);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCalculateMatmulWeightSizeV2 failed. ERROR: %d\n", ret); return    ret);
  
      // Call aclrtMalloc to apply for device memory.
      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 host-side data to device-side 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);
  
      // 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];
      }
      // Create an aclTensor by calling the aclCreateTensor API.
      *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                              shape.data(), shape.size(), *deviceAddr);
  
      uint64_t transWorkspaceSize;
      aclOpExecutor *executor;
      void *transWorkspaceAddr = nullptr;
      ret = aclnnTransMatmulWeightGetWorkspaceSize(*tensor, &transWorkspaceSize, &executor);
      CHECK_RET(ret == ACL_SUCCESS && transWorkspaceSize > 0,
              printf("[ERROR] aclnnTransMatmulWeightGetWorkspaceSize failed. ret = %d \n", ret); return ret);
      ACL_CHECK(aclrtMalloc(&transWorkspaceAddr, transWorkspaceSize, ACL_MEM_MALLOC_HUGE_FIRST));
      ret = aclnnTransMatmulWeight(transWorkspaceAddr, transWorkspaceSize, executor, args.stream);
      CHECK_RET(ret == ACL_SUCCESS, printf("[ERROR] aclnnTransMatmulWeight failed. ret = %d \n", ret);return ret);
      ACL_CHECK(aclrtSynchronizeStreamWithTimeout(args.stream, 20000));
  
      return 0;
  }
  
  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);
      // Allocate device-side memory by calling aclrtMalloc.
      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 host-side data to device-side 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);
      // Calculate the strides of the 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;
  }
  
  int launchOneThreadQuantMatmulAllReduce(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> dequantScaleShape = {128};
      std::vector<int64_t> x3Shape = {32, 128};
      std::vector<int64_t> outShape = {32, 128};
      void *x1DeviceAddr = nullptr;
      void *x2DeviceAddr = nullptr;
      void *biasDeviceAddr = nullptr;
      void *dequantScaleDeviceAddr = nullptr;
      void *x3DeviceAddr = nullptr;
      void *outDeviceAddr = nullptr;
      aclTensor *x1 = nullptr;
      aclTensor *x2 = nullptr;
      aclTensor *bias = nullptr;
      aclTensor *dequantScale = nullptr;
      aclTensor *x3 = nullptr;
      aclTensor *out = nullptr;
  
      int64_t commTurn = 0;
      int64_t streamMode = 1;
      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 dequantScaleShapeSize = GetShapeSize(dequantScaleShape);
      long long x3ShapeSize = GetShapeSize(x3Shape);
      long long outShapeSize = GetShapeSize(outShape);
  
      std::vector<int8_t> x1HostData(x1ShapeSize, 1);
      std::vector<int8_t> x2HostData(x2ShapeSize, 1);
      std::vector<int32_t> biasHostData(biasShapeSize, 1);
      std::vector<uint64_t> dequantScaleHostData(dequantScaleShapeSize, 1);
      std::vector<op::fp16_t> x3HostData(x3ShapeSize, 1);
      std::vector<op::fp16_t> outHostData(outShapeSize, 0);
      // Create the tensor.
      ret = CreateAclTensor(x1HostData, x1Shape, &x1DeviceAddr, aclDataType::ACL_INT8, &x1);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      if (args.format == "NZ") {
          ret = CreateWeightNzAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_INT8, &x2, args);
      } else {
          ret = CreateAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_INT8, &x2);
      }
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(biasHostData, biasShape, &biasDeviceAddr, aclDataType::ACL_INT32, &bias);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(dequantScaleHostData, dequantScaleShape, &dequantScaleDeviceAddr,
                          aclDataType::ACL_UINT64, &dequantScale);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(x3HostData, x3Shape, &x3DeviceAddr, aclDataType::ACL_FLOAT16, &x3);
      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 = aclnnQuantMatmulAllReduceGetWorkspaceSize(x1, x2, bias, x3, dequantScale, hcom_name, "sum", commTurn, streamMode, out,
                                          &workspaceSize, &executor);
      CHECK_RET(ret == ACL_SUCCESS,
              LOG_PRINT("aclnnQuantMatmulAllReduceGetWorkspaceSize 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 = aclnnQuantMatmulAllReduce(workspaceAddr, workspaceSize, executor, args.stream);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnQuantMatmulAllReduce failed. ERROR: %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("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
      LOG_PRINT("device%d aclnnQuantMatmulAllReduce 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 (dequantScale != nullptr) {
          aclDestroyTensor(dequantScale);
      }
      if (x3 != nullptr) {
          aclDestroyTensor(x3);
      }
      if (out != nullptr) {
          aclDestroyTensor(out);
      }
      if (x1DeviceAddr != nullptr) {
          aclrtFree(x1DeviceAddr);
      }
      if (x2DeviceAddr != nullptr) {
          aclrtFree(x2DeviceAddr);
      }
      if (biasDeviceAddr != nullptr) {
          aclrtFree(biasDeviceAddr);
      }
      if (dequantScaleDeviceAddr != nullptr) {
          aclrtFree(dequantScaleDeviceAddr);
      }
      if (x3DeviceAddr != nullptr) {
          aclrtFree(x3DeviceAddr);
      }
      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[]) {
      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 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(&launchOneThreadQuantMatmulAllReduce, std::ref(args[rankId])));
      }
      for (uint32_t rankId = 0; rankId < ndev; rankId++) {
          threads[rankId]->join();
      }
      aclFinalize();
      return 0;
  }