功能

改变输入Tensor的排列顺序，在多个维度上进行转置。

图1 TransposeOperation

定义

struct TransposeParam {
    SVector<int32_t> perm;
    uint8_t rsv[8] = {0};
};

参数列表

成员名称	类型	默认值	描述
perm	SVector<int32_t>	-	指示输入维度的重排结果。例如：输入tensor的shape是[a,b,c]，perm是[0,2,1]，表明交换第一维和第二维，输出tensor的shape是[a,c,b]。 perm需要保证输入正确，维度和输入x一致。
rsv[8]	uint8_t	{0}	预留参数。

成员名称

类型

默认值

描述

perm

SVector<int32_t>

指示输入维度的重排结果。

例如：输入tensor的shape是[a,b,c]，perm是[0,2,1]，表明交换第一维和第二维，输出tensor的shape是[a,c,b]。

perm需要保证输入正确，维度和输入x一致。

rsv[8]

uint8_t

{0}

预留参数。

输入

参数	维度	数据类型	格式
x	[dim_0, dim_1, ..., dim_n]	float16/bf16/int64/int8/int32	ND

输出

参数	维度	数据类型	格式
output	维度由参数确定	float16/bf16/int64/int8/int32 数据类型与x保持一致	ND

参数

维度

数据类型

格式

output

维度由参数确定

float16/bf16/int64/int8/int32

数据类型与x保持一致

规格约束

Atlas 200I/500 A2 推理产品上仅支持float16数据类型。

算子调用示例（C++）

前置条件和编译命令请参见算子调用示例。

场景：基础场景。

#include <iostream>
#include <vector>
#include <numeric>
#include "acl/acl.h"
#include "atb/operation.h"
#include "atb/types.h"
#include "atb/atb_infer.h"

#include "demo_util.h"

const uint32_t DIM1 = 2;
const uint32_t DIM2 = 3;

/**
 * @brief 准备atb::VariantPack中的所有输入tensor
 * @param contextPtr context指针
 * @param stream stream
 * @return atb::SVector<atb::Tensor> atb::VariantPack中的输入tensor
 * @note 需要传入所有host侧tensor
 */
atb::SVector<atb::Tensor> PrepareInTensor(atb::Context *contextPtr, aclrtStream stream)
{
    // 创建x tensor
    std::vector<float> xData(DIM1 * DIM2, 1.0);
    std::vector<int64_t> xShape = {DIM1, DIM2};
    atb::Tensor tensorX =
        CreateTensorFromVector(contextPtr, stream, xData, ACL_FLOAT16, aclFormat::ACL_FORMAT_ND, xShape);
    atb::SVector<atb::Tensor> inTensors = {tensorX};
    return inTensors;
}

/**
 * @brief 创建一个Reduce的Operation，并设置参数
 * @return atb::Operation * 返回一个Operation指针
 */
atb::Operation *PrepareOperation()
{
    atb::infer::TransposeParam transposeOpParam;
    atb::SVector<int32_t> perm = {1, 0};
    transposeOpParam.perm = perm;
    atb::Operation *transposeOp = nullptr;
    CHECK_STATUS(atb::CreateOperation(transposeOpParam, &transposeOp));
    return transposeOp;
}

int main(int argc, char **argv)
{
    // 设置卡号、创建context、设置stream
    CHECK_STATUS(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_STATUS(aclrtSetDevice(deviceId));
    atb::Context *context = nullptr;
    CHECK_STATUS(atb::CreateContext(&context));
    void *stream = nullptr;
    CHECK_STATUS(aclrtCreateStream(&stream));
    context->SetExecuteStream(stream);

    // Transpose示例
    atb::Operation *transposeOp = PrepareOperation();
    // 准备输入tensor
    atb::VariantPack transposeVariantPack;
    transposeVariantPack.inTensors = PrepareInTensor(context, stream);  // 放入输入tensor
    atb::Tensor tensorOut = CreateTensor(ACL_FLOAT16, aclFormat::ACL_FORMAT_ND, {DIM2, DIM1});
    transposeVariantPack.outTensors = {tensorOut};  // 放入输出tensor

    uint64_t workspaceSize = 0;
    // 计算workspace大小
    CHECK_STATUS(transposeOp->Setup(transposeVariantPack, workspaceSize, context));
    uint8_t *workspacePtr = nullptr;
    if (workspaceSize > 0) {
        CHECK_STATUS(aclrtMalloc((void **)(&workspacePtr), workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST));
    }
    // reduce执行
    transposeOp->Execute(transposeVariantPack, workspacePtr, workspaceSize, context);
    CHECK_STATUS(aclrtSynchronizeStream(stream));  // 流同步，等待device侧任务计算完成
    for (atb::Tensor &inTensor : transposeVariantPack.inTensors) {
        CHECK_STATUS(aclrtFree(inTensor.deviceData));
    }
    if (workspaceSize > 0) {
        CHECK_STATUS(aclrtFree(workspacePtr));
    }
    CHECK_STATUS(atb::DestroyOperation(transposeOp));  // operation，对象概念，先释放
    CHECK_STATUS(aclrtDestroyStream(stream));
    CHECK_STATUS(DestroyContext(context));  // context，全局资源，后释放
    std::cout << "Transpose demo success!" << std::endl;
    return 0;
}