aclnnTransposeBatchMatMul

Applicable Products

Product	Supported
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: Performs matrix multiplication computation on tensor x1 and tensor x2. Only three-dimensional tensors are supported. Tensors support transposition, and the transpose sequence changes based on the input sequence. permX1 represents the transpose sequence of tensor x1, and permX2 represents the transpose sequence of tensor x2. The sequence value 0 corresponds to the batch dimension, and the remaining two dimensions are used for matrix multiplication.
Example:
- When the shape of x1 is (B, M, K), the shape of x2 is (B, K, N), scale is None, and batchSplitFactor equals 1, the shape of the computed output out is (M, B, N).
- When the shape of x1 is (B, M, K), the shape of x2 is (B, K, N), scale is not None, and batchSplitFactor equals 1, the shape of the computed output out is (M, 1, B * N).
- When the shape of x1 is (B, M, K), the shape of x2 is (B, K, N), scale is None, and batchSplitFactor is greater than 1, the shape of the compute output out is (batchSplitFactor, M, B * N / batchSplitFactor).

Function Prototype

Each operator is divided into a two-phase API. You must first call the "aclnnTransposeBatchMatMulGetWorkspaceSize" API to obtain the input parameters and calculate the required workspace size based on the process, and then call the "aclnnTransposeBatchMatMul" API to perform the computation.

aclnnStatus aclnnTransposeBatchMatMulGetWorkspaceSize(
    const aclTensor    *x1,
    const aclTensor    *x2,
    const aclTensor    *bias,
    const aclTensor    *scale,
    const aclIntArray  *permX1,
    const aclIntArray  *permX2,
    const aclIntArray  *permY,
    int8_t             cubeMathType,
    const int32_t      batchSplitFactor,
    aclTensor          *out,
    uint64_t           *workspaceSize,
    aclOpExecutor      **executor)

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

aclnnTransposeBatchMatMulGetWorkSpaceSize

Parameters:

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
x1	Input	Represents the first matrix of the matrix multiplication.	The data type must satisfy the data type deduction rules with x2 (see Deduction Relationship and Constraints). Supports the following data types: BFLOAT16, FLOAT16, and FLOAT32. Does not support data type deduction when the data type of x1 is BFLOAT16 and the data type of x2 is FLOAT16. Does not support data type deduction when the data type of x1 is BFLOAT16 and the data type of x2 is FLOAT32.	BFLOAT16, FLOAT16, FLOAT32	ND	3	√
x2	Input	Represents the second matrix of the matrix multiplication.	The data type must comply with the data type deduction rules with x1 (see Deduction Relationship and Constraints). The reduce dimension of x2 must have the same size as the reduce dimension of x1. Supported data types: BFLOAT16, FLOAT16, FLOAT32. Does not support data type deduction when x1 and x2 are BFLOAT16 and FLOAT16 respectively. Does not support data type deduction when x1 and x2 are BFLOAT16 and FLOAT32 respectively.	BFLOAT16, FLOAT16, FLOAT32	ND	3	√
bias	Input	Represents the bias matrix for matrix multiplication.	Reserved parameter, currently not supported.	BFLOAT16, FLOAT16, FLOAT32	-	-	-
scale	Optional input.	Represents the quantization coefficient of the output matrix, which can be enabled when the input is FLOAT16 and the output is INT8.	The shape supports only one dimension and must equal [b*n].	INT64, UINT64	ND	1	√
permX1	Input	Represents the transpose sequence of the first matrix in matrix multiplication, an aclIntArray on the host side.	Supports [0, 1, 2] and [1, 0, 2].	INT64	-	1	-
permX2	Input	Represents the transpose sequence of the second matrix in matrix multiplication, an aclIntArray on the host side.	-	INT64	-	1	-
permY	Input	Represents the transpose sequence of the matrix multiplication output matrix, an aclIntArray on the host side.	Supports [1, 0, 2].	INT64	-	1	-
cubeMathType	Input	Specifies the computation logic of the Cube unit. It is an integer on the host side.	If the input data types have a mutual derivation relationship, this parameter processes the derived data type by default. The enumeration values are as follows: 0: KEEP_DTYPE. Keeps the input data type for computation. 1: ALLOW_FP32_DOWN_PRECISION. Supports down-precision computation of the input data. When the input data type is FLOAT32, it is converted to HFLOAT32 for computation. When the data type is other types, the input type is kept for computation. 2: USE_FP16. Supports converting the input to FLOAT16 precision for computation. This option is not supported when the input data type is BFLOAT16. 3: USE_HF32. Supports down-precision computation of the input to the HFLOAT32 data type. When the input data type is FLOAT32, it is converted to HFLOAT32 for computation.	INT8	-	-	-
batchSplitFactor	Input	Specifies the split size of the B dimension in the matrix multiplication output matrix. Host-side integer.	The value range is [1, B] and must be divisible by B. When scale is not empty, batchSplitFactor must equal 1.	INT32	-	-	-
out	Output	Represents the output matrix of the matrix multiplication, i.e., out in the formula.	The data type must be consistent with the data type derived from x1 and x2 (see Deduction Relationship and Constraints). When scale has a value, the output shape is (M, 1, B * N). When batchSplitFactor is greater than 1, the output shape of out is (batchSplitFactor, M, B * N / batchSplitFactor). Example 1: When M, K, N, B = 32, 512, 128, 16 and batchSplitFactor = 2, the output shape size of out is (2, 32, 1024). Example 2: When M, K, N, B = 32, 512, 128, 16 and batchSplitFactor = 4, the output shape size of out is (4, 32, 512).	BFLOAT16, FLOAT16, FLOAT32, INT8	ND	3	-

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 out is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data type of x1, x2, or out is not within the supported range.
		The second dimension of x1 does not equal the first dimension of x2.
		The dimension size of x1 or x2 does not equal 3.
		The data type of scale is not within the supported range.
		The value of batchSplitFactor is not within the supported range.

aclnnTransposeBatchMatMul

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 by the first-phase API aclnnTransposeBatchMatMulGetWorkSpaceSize.
executor	Input	Operator executor, which 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

Deterministic computation description: aclnnTransposeBatchMatMul defaults to a deterministic implementation.
Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products:
- The value range of B is [1, 65536), and the value range of N is [1, 65536).
- When the input shape of x1 is (B, M, K), K <= 65535; when the input shape of x1 is (M, B, K), B * K <= 65535.
- The second dimension of x2 or the third dimension of x2 cannot be divisible by 16.
- permX2 only supports input [0, 1, 2].
- When scale is not empty, batchSplitFactor must equal 1, the product of B and N must be less than 65536, and only type inference with FLOAT16 input and INT8 output is supported.
Ascend 950PR/Ascend 950DT:
- permX2 supports input [0, 1, 2] and [0, 2, 1].
- When scale is not empty, batchSplitFactor must equal 1, and only type inference with FLOAT16 input and INT8 output is supported.

Example

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

#include <iostream>
#include <vector>
#include <cmath>
#include "acl/acl.h"
#include "aclnnop/aclnn_transpose_batch_mat_mul.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)

// Convert the FP16 uint16_t representation to a float representation.
float Fp16ToFloat(uint16_t h) {
  int s = (h >> 15) & 0x1;              // sign
  int e = (h >> 10) & 0x1F;             // exponent
  int f =  h        & 0x3FF;            // fraction
  if (e == 0) {
    // Zero or Denormal
    if (f == 0) {
      return s ? -0.0f : 0.0f;
    }
    // Denormals
    float sig = f / 1024.0f;
    float result = sig * pow(2, -24);
    return s ? -result : result;
  } else if (e == 31) {
      // Infinity or NaN
      return f == 0 ? (s ? -INFINITY : INFINITY) : NAN;
  }
  // Normalized
  float result = (1.0f + f / 1024.0f) * pow(2, e - 15);
  return s ? -result : result;
}

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

int Init(int32_t deviceId, aclrtStream* stream) {
  // Boilerplate: resource initialization.
  auto ret = aclInit(nullptr);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
  ret = aclrtSetDevice(deviceId);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
  ret = aclrtCreateStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
  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);
  // Call aclrtMalloc to allocate memory on the device side.
  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 data from the host side to the device 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);

  // 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];
  }

  // 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 main() {
  // 1. (Boilerplate) Initialize the device/stream. See the ACL API manual.
  // Enter the deviceId based on the actual device.
  int32_t deviceId = 0;
  aclrtStream stream;
  auto ret = Init(deviceId, &stream);
  // Handle the check based on your requirements.
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

  // 2. Construct the input and output. Customize the construction based on the API.
  int32_t M = 32;
  int32_t K = 512;
  int32_t N = 128;
  int32_t Batch = 16;
  std::vector<int64_t> x1Shape = {M, Batch, K};
  std::vector<int64_t> x2Shape = {Batch, K, N};
  std::vector<int64_t> outShape = {M, Batch, N};
  std::vector<int64_t> permX1Series = {1, 0, 2};
  std::vector<int64_t> permX2Series = {0, 1, 2};
  std::vector<int64_t> permYSeries = {1, 0, 2};
  void* x1DeviceAddr = nullptr;
  void* x2DeviceAddr = nullptr;
  void* scaleDeviceAddr = nullptr;
  void* outDeviceAddr = nullptr;
  aclTensor* x1 = nullptr;
  aclTensor* x2 = nullptr;
  aclTensor* scale = nullptr;
  aclTensor* out = nullptr;
  std::vector<uint16_t> x1HostData(GetShapeSize(x1Shape),0x3C00);
  std::vector<uint16_t> x2HostData(GetShapeSize(x2Shape),0x3C00);
  std::vector<uint16_t> outHostData(GetShapeSize(outShape),0);
  int8_t cubeMathType = 1;
  int8_t batchSplitFactor = 1;

  // Create the x1 aclTensor.
  ret = CreateAclTensor(x1HostData, x1Shape, &x1DeviceAddr, aclDataType::ACL_FLOAT16, &x1);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the x2 aclTensor.
  ret = CreateAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_FLOAT16, &x2);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the out aclTensor.
  ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT16, &out);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  aclIntArray *permX1 = aclCreateIntArray(permX1Series.data(), permX1Series.size());
  aclIntArray *permX2 = aclCreateIntArray(permX2Series.data(), permX2Series.size());
  aclIntArray *permY = aclCreateIntArray(permYSeries.data(), permYSeries.size());
  uint64_t workspaceSize = 0;
  aclOpExecutor* executor = nullptr;

  // Example of the aclnnTransposeBatchMatMul API call.
  // 3. Call the CANN operator library API. Replace it with the specific API name.
  // Call the first-phase API of aclnnTransposeBatchMatMul.
  ret = aclnnTransposeBatchMatMulGetWorkspaceSize(x1, x2, (const aclTensor*)nullptr, (const aclTensor*)nullptr,
                                                  permX1, permX2, permY, cubeMathType, batchSplitFactor, out,
                                                  &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnTransposeBatchMatMulGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  // Apply for device memory based on the workspaceSize computed by the first-phase API.
  void* workspaceAddr = nullptr;
  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 of aclnnTransposeBatchMatMul.
  ret = aclnnTransposeBatchMatMul(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnTransposeBatchMatMul failed. ERROR: %d\n", ret); return ret);

  // 4. (Boilerplate) Synchronize and wait for task execution to complete.
  ret = aclrtSynchronizeStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

  // 5. Obtain the output value and copy the result from the Device side memory to the Host side. Modify based on the specific API definition.
  auto size = GetShapeSize(outShape);
  std::vector<uint16_t> resultData(size, 0);
  ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr,
                    size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
  for (int64_t i = 0; i < size; i++) {
    float fp16Float = Fp16ToFloat(resultData[i]);
    LOG_PRINT("result[%ld] is: %f\n", i, fp16Float);
  }

  // 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
  aclDestroyTensor(x1);
  aclDestroyTensor(x2);
  aclDestroyTensor(out);

  // 7. Release device resources. Modify based on the specific API definition.
  aclrtFree(x1DeviceAddr);
  aclrtFree(outDeviceAddr);
  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();
  return 0;
}