aclnnMaxPool3dWithArgmax

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

Operator function:
- For the input channel of the input signal, performs 3D max pooling and outputs the pooled values out and indices.
- Description of input dims: N - batch, C - channel, D - depth, W - width, H - height.
- When D H W exceeds the int32 range, split the D axis on the model dimension.
Formula:
- Formula for each element in the output tensor:
  $out(N_i, C_j, d, h, w) = \max\limits_{{k\in[0,k_{D}-1],m\in[0,k_{H}-1],n\in[0,k_{W}-1]}}input(N_i,C_j,stride[0]\times d + k, stride[1]\times h + m, stride[2]\times w + n)$
- Shape derivation formula for the out tensor (default ceilMode=false, i.e., floor rounding):

[N, C, D_{out}, H_{out}, W_{out}]=[N,C,\lfloor{\frac{D_{in}+2 \times {padding[0] - dilation[0] \times(kernelSize[0] - 1) - 1}}{stride[0]}}\rfloor + 1,\\ \lfloor{\frac{H_{in}+2 \times {padding[1] - dilation[1] \times(kernelSize[1] - 1) - 1}}{stride[1]}}\rfloor + 1, \\ \lfloor{\frac{W_{in}+2 \times {padding[2] - dilation[2] \times(kernelSize[2] - 1) - 1}}{stride[2]}}\rfloor + 1]

Shape derivation formula for the out tensor (defaults to ceilMode=true, i.e., rounding up):

[N, C, D_{out}, H_{out}, W_{out}]=[N,C,\lceil{\frac{D_{in}+2 \times {padding[0] - dilation[0] \times(kernelSize[0] - 1) - 1}}{stride[0]}}\rceil + 1,\\ \lceil{\frac{H_{in}+2 \times {padding[1] - dilation[1] \times(kernelSize[1] - 1) - 1}}{stride[1]}}\rceil + 1, \\ \lceil{\frac{W_{in}+2 \times {padding[2] - dilation[2] \times(kernelSize[2] - 1) - 1}}{stride[2]}}\rceil + 1]

Function Prototype

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

aclnnStatus aclnnMaxPool3dWithArgmaxGetWorkspaceSize(
  const aclTensor   *self,
  const aclIntArray *kernelSize,
  const aclIntArray *stride,
  const aclIntArray *padding,
  const aclIntArray *dilation,
  bool               ceilMode,
  aclTensor         *out,
  aclTensor         *indices,
  uint64_t          *workspaceSize,
  aclOpExecutor     **executor)

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

aclnnMaxPool3dWithArgmaxGetWorkspaceSize

Parameters:

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Shape	Non-contiguous Tensor
self	Input	The tensor to be computed.	-	FLOAT32, FLOAT16, BFLOAT16	NCDHW, NDHWC, ND	4-5	√
kernelSize	Input	Window size of max pooling.	The array length must be 1 or 3, and all array elements must be greater than 0.	INT64	-	-	-
stride	Input	Stride of the sliding window.	The array length must be 0, 1, or 3, and all array elements must be greater than 0. When the array length is 0, the kernelSize value is used as strides internally.	INT64	-	-	-
padding	Input	Number of layers padded on each side.	The padded positions are filled with negative infinity. The array length must be 1 or 3, and each array element must be greater than or equal to 0 and less than or equal to kernelSize/2.	INT64	-	-	-
dilation	Input	Controls the stride of elements within the window.	The array length must be 1 or 3, and all array elements must be greater than 0.	INT64	-	-	-
ceilMode	Input	Rounding method used when computing the output shape.	If True, the output shape is computed using ceiling rounding; if False, floor rounding is used.	BOOL	-	-	-
out	Output	Output tensor, which is the result after pooling.	The data type is the same as self. The shape is derived from the formula above.	BFLOAT16, FLOAT16, FLOAT32	NCDHW, NDHWC, ND	4-5	√
indices	Output	A tensor consisting of the index positions of the maximum values.	The shape is consistent with out.	INT32, INT64	NCDHW, NDHWC, ND	4-5	√
workspaceSize	Output	Returns the workspace size to be requested on the Device side.	-	-	-	-	-
executor	Output	Returns the operator executor, which contains the operator computation flow.	-	-	-	-	-

Atlas A3 training products/Atlas A3 inference products, Atlas A2 training products/Atlas A2 inference products: The dilation element values only support 1; the indices data type does not support INT64. The input data format does not support NDHWC. depthheightwidth does not support values greater than max int32.

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 Code	Error Code	Description
ACLNN_ERR_PARAM_NULLPTR	161001	The input self or out is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data type of self is not within the supported range.
		The data format of self is not within the supported range.
		The shape of self does not meet the requirements.
		A dimension of the out shape derived from the formula is 0.
		The element values of kernelSize, stride, padding, or dilation do not meet the requirements.
		The length of kernelSize, stride, padding, or dilation does not meet the requirements.
		The platform is not supported.

aclnnMaxPool3dWithArgmax

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 aclnnMaxPool3dWithArgmaxGetWorkspaceSize.
executor	Input	Operator executor, which contains 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: aclnnMaxPool3dWithArgmax defaults to a deterministic implementation.
The kernelSize, stride, padding, dilation, and ceilMode parameters must ensure that no axis in the output shape is less than 1.
When ceilMode is True, if the sliding window is entirely on the right padding region, this output result will be ignored.

Example

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

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

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) {
    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 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 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.
    // Set deviceId based on your 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 interface.
    std::vector<int64_t> selfShape = {1, 1, 2, 2, 2};
    std::vector<int64_t> outShape = {1, 1, 1, 1, 1};
    std::vector<int64_t> indicesShape = {1, 1, 1, 1, 1};
    std::vector<int64_t> kernelSizeData = {2, 2, 2};
    std::vector<int64_t> strideData = {2, 2, 2};
    std::vector<int64_t> paddingData = {0, 0, 0};
    std::vector<int64_t> dilationData = {1, 1, 1};
    void* selfDeviceAddr = nullptr;
    void* outDeviceAddr = nullptr;
    void* indicesDeviceAddr = nullptr;
    aclTensor* self = nullptr;
    aclTensor* out = nullptr;
    aclTensor* indices = nullptr;
    std::vector<float> selfHostData = {1, 6, 2, 8, 4, 5, 7, 3};
    std::vector<float> outHostData = {0};
    std::vector<int32_t> indicesHostData = {0};

    // Create the self aclTensor.
    ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the out aclTensor.
    ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the indices aclTensor.
    ret = CreateAclTensor(indicesHostData, indicesShape, &indicesDeviceAddr, aclDataType::ACL_INT32, &indices);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // Create the input array.
    aclIntArray* kernelSize = aclCreateIntArray(kernelSizeData.data(), 3);
    aclIntArray* stride = aclCreateIntArray(strideData.data(), 3);
    aclIntArray* padding = aclCreateIntArray(paddingData.data(), 3);
    aclIntArray* dilation = aclCreateIntArray(dilationData.data(), 3);
    const bool ceilMode = false;

    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;

    // Example of calling the aclnnMaxPool3dWithArgmax API.
    // 3. Call the first-phase API of aclnnMaxPool3dWithArgmax.
    ret = aclnnMaxPool3dWithArgmaxGetWorkspaceSize(self, kernelSize, stride, padding, dilation, ceilMode, out, indices, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMaxPool3dWithArgmaxGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Allocate device memory based on the workspaceSize calculated 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 aclnnMaxPool3dWithArgmax.
    ret = aclnnMaxPool3dWithArgmax(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMaxPool3dWithArgmax failed. ERROR: %d\n", ret); return ret);

    // 4. 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<float> 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 output result from device to host failed. ERROR: %d\n", ret); return ret);
    for (int64_t i = 0; i < size; i++) {
        LOG_PRINT("output[%ld] is: %f\n", i, resultData[i]);
    }

    size = GetShapeSize(indicesShape);
    std::vector<int> indicesResultData(size, 0);
    ret = aclrtMemcpy(indicesResultData.data(), indicesResultData.size() * sizeof(indicesResultData[0]), indicesDeviceAddr,
                      size * sizeof(indicesResultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy indices result from device to host failed. ERROR: %d\n", ret); return ret);
    for (int64_t i = 0; i < size; i++) {
        LOG_PRINT("indices[%ld] is: %d\n", i, indicesResultData[i]);
    }

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

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