diff --git a/README.md b/README.md index eb468eb804cb943734d8dc4a26aa98baecf7c180..5893f56aeb7219dfb5379354e325eff00050c920 100644 --- a/README.md +++ b/README.md @@ -339,12 +339,12 @@ triton autotune性能配置说明参考本仓库的 docs\sources\getting-started | | gather | × | × | × | × | × | ✓ | ✓ | ✓ | × | | Atomic Ops | atomic_add | ✓ | ✓ | ✓ | × | × | ✓ | ✓ | ✓ | × | | | atomic_and | × | × | × | × | × | × | × | × | × | -| | atomic_cas | × | × | × | × | × | × | × | × | × | +| | atomic_cas | ✓ | ✓ | ✓ | × | ✓ | ✓ | ✓ | × | × | | | atomic_max | ✓ | ✓ | ✓ | × | × | ✓ | ✓ | ✓ | × | | | atomic_min | ✓ | ✓ | ✓ | × | × | ✓ | ✓ | ✓ | × | -| | atomic_or | × | × | × | × | × | × | × | × | × | -| | atomic_xchg | × | × | × | × | × | × | × | × | × | -| | atomic_xor | × | × | × | × | × | × | × | × | × | +| | atomic_or | ✓ | ✓ | ✓ | × | ✓ | × | × | × | × | +| | atomic_xchg | ✓ | ✓ | ✓ | × | ✓ | ✓ | ✓ | × | × | +| | atomic_xor | ✓ | ✓ | ✓ | × | ✓ | × | × | × | × | | Random Number Generation | randint4x | × | × | ✓ | × | × | × | × | × | × | | | randint | × | × | ✓ | × | × | × | × | × | × | | | rand | × | × | × | × | × | × | ✓ | × | × | @@ -379,6 +379,10 @@ triton autotune性能配置说明参考本仓库的 docs\sources\getting-started - atomic_min: 不支持标量（包括长度为1的tensor）访存 +- atomic_or: sem只支持默认值"acq_rel"模式，其他值均按默认值处理；scope只支持默认值"gpu"，其他值均按默认值处理 + +- atomic_xor: sem只支持默认值"acq_rel"模式，其他值均按默认值处理；scope只支持默认值"gpu"，其他值均按默认值处理 + - permute: 不支持不相邻轴转置，如`(0, 1, 2) -> (2, 1, 0)` - trans: 不支持不相邻轴转置，如`(0, 1, 2) -> (2, 1, 0)` diff --git a/ascend/backend/compiler.py b/ascend/backend/compiler.py index c2ee7ab53e9245a7bf711c66edf41334a7de04b6..df04dd51632f3cf37540584e053f38857162e4d1 100644 --- a/ascend/backend/compiler.py +++ b/ascend/backend/compiler.py @@ -199,6 +199,15 @@ def llir_to_cpuasm(llir: str, metadata, opt): return Path(dst_path).read_text() +def __get_metadata_attr_by_callback(lib, postfix: str, metadata, meta_key: str): + func_symbol = metadata["kernel_name"] + postfix + if hasattr(lib, func_symbol): + callback_func = getattr(lib, func_symbol) + callback_func.restype = ctypes.c_int64 + callback_func.argtypes = [] + metadata[meta_key] = callback_func() + + def linalg_to_bin_enable_npu_compile(linalg: str, metadata, opt): # Note: Compiled Kernel requires to estimate size of shared memory to occupy # Currently, NPU backend does not limit on shared memory @@ -249,13 +258,10 @@ def linalg_to_bin_enable_npu_compile(linalg: str, metadata, opt): ret = subprocess.run(cmd_list, capture_output=True, check=True) if Path(callback_path).is_file(): lib = ctypes.CDLL(callback_path) - callback_func = getattr( - lib, metadata["kernel_name"] + "_infer_workspace_shape_function" - ) - callback_func.restype = ctypes.c_int64 - callback_func.argtypes = [] - metadata["workspace_size"] = callback_func() - + __get_metadata_attr_by_callback(lib, "_infer_workspace_shape_function", metadata, "workspace_size") + __get_metadata_attr_by_callback(lib, "_infer_sync_block_lock_num_function", metadata, "lock_num") + __get_metadata_attr_by_callback(lib, "_infer_sync_block_lock_init_function", metadata, "lock_init_val") + return Path(bin_path).read_bytes() diff --git a/ascend/backend/driver.py b/ascend/backend/driver.py index 151161ff734075a698a2472d5b327130b071b84f..ab7901f9d78c81523931a6941e53bfe0448b6aff 100644 --- a/ascend/backend/driver.py +++ b/ascend/backend/driver.py @@ -66,13 +66,18 @@ class NPULauncher(object): debug_mode = metadata.debug workspace_size = int(metadata.workspace_size) \ if hasattr(metadata, 'workspace_size') else -1 + lock_init_value = int(metadata.lock_init_value) \ + if hasattr(metadata, 'lock_init_value') else 0 + lock_num = int(metadata.lock_num) \ + if hasattr(metadata, 'lock_num') else -1 constants = src.constants if hasattr(src, "constants") else dict() cst_key = lambda i: src.fn.arg_names.index(i) if isinstance(i, str) else i constants = {cst_key(key): value for key, value in constants.items()} signature = {cst_key(key): value for key, value in src.signature.items()} mix_mode = metadata.mix_mode wrapper_src = generate_npu_wrapper_src(constants, signature, \ - workspace_size, mix_mode) + workspace_size, mix_mode, \ + lock_num, lock_init_value) so_launcher_path = make_npu_launcher_stub(wrapper_src, debug_mode) # initialize launcher import importlib.util @@ -207,7 +212,7 @@ def make_npu_launcher_stub(src, debug=False): # the template is from triton-adapter HEAD. Wrapping the generated kernel binary into a python module -def generate_npu_wrapper_src(constants, signature, workspace_size, mix_mode): +def generate_npu_wrapper_src(constants, signature, workspace_size, mix_mode, lock_num, lock_ini_val): import os def _ty_to_cpp(ty): if ty[0] == '*': @@ -467,9 +472,24 @@ static void _launch(const char* kernelName, const void* func, rtStream_t stream, return {'ret' if enable_taskqueue else ''}; }} // stub argument for workspace + void *syncBlockLock = NULL; void *workspace_addr = NULL; - {f''' uint16_t ModuleId = 0; + {f''' + uint64_t syncBlockLockSize = {lock_num} * sizeof(int64_t); + ret = rtMalloc(reinterpret_cast(&syncBlockLock), + syncBlockLockSize, RT_MEMORY_HBM, 0); + if (ret != RT_ERROR_NONE) {{ + return {'ret' if enable_taskqueue else ''}; + }} + std::vector lockInitData({lock_num}, {lock_ini_val}); + ret = rtMemcpy(syncBlockLock, syncBlockLockSize, reinterpret_cast(lockInitData.data()), + syncBlockLockSize, RT_MEMCPY_HOST_TO_DEVICE); + if (ret != RT_ERROR_NONE) {{ + return {'ret' if enable_taskqueue else ''}; + }} + ''' if lock_num > 0 else ''} + {f''' uint64_t totalWorkSpaceSize = {workspace_size} * blockNum; ret = rtMalloc(reinterpret_cast(&workspace_addr), totalWorkSpaceSize, RT_MEMORY_HBM, ModuleId); @@ -479,12 +499,14 @@ static void _launch(const char* kernelName, const void* func, rtStream_t stream, ''' if workspace_size > 0 else ''} struct __attribute__((packed)) {{ void* ffts_addr __attribute__((aligned(8))); + void* syncBlockLock __attribute__((aligned(8))); void* workspace_addr __attribute__((aligned(8))); {' '.join(f'{_ty_to_cpp(ty)} arg{i} __attribute__((aligned({4 if ty[0] != "*" and ty[-2:] != "64" else 8})));' for i, ty in signature.items() if i not in constants)} {' '.join(f'{_ty_to_cpp(ty)} grid{mark} __attribute__((aligned(4)));' for mark, ty in grid_info.items())} {'void* DTData __attribute__((aligned(8)));' if enable_device_print else ''} }} args = {{ static_cast(ffts_addr), + static_cast(syncBlockLock), static_cast(workspace_addr), {', '.join(f'static_cast<{_ty_to_cpp(ty)}>(arg{i})' for i, ty in signature.items() if i not in constants)}, {', '.join(f'static_cast<{_ty_to_cpp(ty)}>(grid{mark})' for mark, ty in grid_info.items())} diff --git a/ascend/triton-adapter/include/TritonToLinalg/LoadStoreConverter.h b/ascend/triton-adapter/include/TritonToLinalg/LoadStoreConverter.h index 7bb438e74b6aa5722183ff7ea0fb1156550f947f..c1d2dbcddc46f12b1d32cf4536fc893b0b281f4f 100644 --- a/ascend/triton-adapter/include/TritonToLinalg/LoadStoreConverter.h +++ b/ascend/triton-adapter/include/TritonToLinalg/LoadStoreConverter.h @@ -123,6 +123,24 @@ class ScalarAtomicRMWCanonicalizer PatternRewriter &rewriter) const override; }; +class ScalarAtomicCASCanonicalizer + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(triton::AtomicCASOp op, + PatternRewriter &rewriter) const override; +}; + +class AtomicCASConverter : public OpConversionPattern { +public: + explicit AtomicCASConverter(MLIRContext *context) : + OpConversionPattern(context) {} + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::AtomicCASOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + class AtomicRMWConverter : public OpConversionPattern { private: Value createAtomicBinaryOps(OpBuilder &builder, Location loc, @@ -157,6 +175,8 @@ private: } else { binaryOp = builder.create(loc, lhs, rhs); } + } else if (rmwOp == triton::RMWOp::XCHG) { + binaryOp = rhs; } else { op.emitOpError("unsupported atomic RMW operation: "); llvm_unreachable( diff --git a/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp b/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp index 6f66cd7be44004409702d6b626f3d53e6903885a..9ba73de4408aae40c022364fdb8a6a226d0dca36 100644 --- a/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp +++ b/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp @@ -516,10 +516,150 @@ AtomicRMWConverter::matchAndRewrite(triton::AtomicRMWOp op, OpAdaptor adaptor, return success(); } +LogicalResult +AtomicCASConverter::matchAndRewrite(triton::AtomicCASOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + // If the result of AtomicCASOp is not used, we don't need to load the old + // data stored at the ptr + auto ptr = adaptor.getPtr(); + auto cmp = op.getCmp(); + auto val = op.getVal(); + auto loc = op.getLoc(); + + auto resType = dyn_cast(op.getResult().getType()); + if (!resType) { + return rewriter.notifyMatchFailure( + op, "atomicCASConverter: scalar will be handled by " + "ScalarAtomicCASCanonicalizer"); + } + + // 1. Simple case where no mask is used. + auto type = dyn_cast(ptr.getType()); + if (!type) { + // Seen when implicit broadcasting is done late in a chain of + // operations. The workaround is to broadcast the pointers early in the + // address calculation. A proper fix is complicated, but at least we can + // provide a better error message. + return rewriter.notifyMatchFailure( + op, "AtomicCASOp expects a memref, not a memref of pointers"); + } + + auto dstMemref = ptr; + // Well, linalg structure op wouldn't support mixed tensor/buffer semantics + // any more in latest LLVM(triton LLVM dependency has involed this), so we + // need to convert tensor to buffer early. + auto dstType = dstMemref.getType(); + Value inputMemref = + rewriter.create(loc, dstType, val); + + Value cmpMemref = + rewriter.create(loc, dstType, cmp); + + // 3. If needed, handle the return value of atomic op + // + // tt.atomicRMW op has two part of feature + // 1. load the old data at the ptr + // 2. atomically store the data on ub to the ptr + // at the same time it perform the action it has been assigned + // So we lower this op to load + atomically store + // + // The first part is not necessary when the returned value of atomic op + // is not used, it will be deleted cause it's meaningless + // Here, we preemptively determine whether it will be used + // and decide whether it is necessary to create the load process based on + // this assessment. + // + // logic of handling is copied + if (!op.getResult().use_empty()) { + auto tensorType = + RankedTensorType::get(type.getShape(), type.getElementType()); + auto alloc = rewriter.create( + loc, MemRefType::get(type.getShape(), type.getElementType())); + + // For the return value, don't need to care about mask for now + // this op don't support other, so we best not fill it + rewriter.create(loc, ptr, alloc); + Value tensor = rewriter.create( + loc, tensorType, alloc, true /* restrict */, true /* writable */); + rewriter.replaceOp(op, tensor); + } + + // create element-wise map + int64_t rank = type.getRank(); + SmallVector inputDims; + auto context = rewriter.getContext(); + + for (int i = 0; i < rank; i++) { + inputDims.push_back(getAffineDimExpr(i, context)); + } + + SmallVector indexingMaps; + // As mask has been erased for now + // the number of input must be 2 + // the input memref is also the output memref + // Thus, there are a total of four inputs and outputs. + // so here we have 4 map to create + for (int i = 0; i < 4; i++) { // 4: 3 input and 1 output + indexingMaps.push_back(AffineMap::get(rank, 0, inputDims, context)); + } + + auto linalgOp = rewriter.create( + loc, ValueRange{dstMemref, cmpMemref, inputMemref}, + mlir::ValueRange{dstMemref}, indexingMaps, + mlir::ConverterUtils::getNParallelLoopsAttrs(rank), + [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange blockArgs) { + Value lhs = blockArgs[0]; + Value rhs = blockArgs[1]; + Value setValue = blockArgs[2]; + Value cond; + if (mlir::isa(lhs.getType())) { + cond = nestedBuilder.create(nestedLoc, + arith::CmpFPredicate::UEQ, + lhs, rhs); + } else { + cond = nestedBuilder.create(nestedLoc, + arith::CmpIPredicate::eq, + lhs, rhs); + } + auto ifOp = nestedBuilder.create(nestedLoc, TypeRange{setValue.getType()}, cond, true); + { + OpBuilder::InsertionGuard guard(nestedBuilder); + nestedBuilder.setInsertionPointToEnd(&ifOp.getThenRegion().front()); + nestedBuilder.create(nestedLoc, setValue); + } + { + OpBuilder::InsertionGuard guard(nestedBuilder); + nestedBuilder.setInsertionPointToEnd(&ifOp.getElseRegion().front()); + nestedBuilder.create(nestedLoc, lhs); + } + nestedBuilder.setInsertionPointToEnd(nestedBuilder.getBlock()); + nestedBuilder.create(nestedLoc, ifOp.getResult(0)); + }); + + const StringRef genericAtomicRMW = "GenericAtomicRMW"; + const StringRef memSemantic = "MemSemantic"; + const StringRef memSyncScope = "MemSyncScope"; + auto attr = mlir::StringAttr::get(context, "cas"); + + linalgOp->setAttr(genericAtomicRMW, attr); + linalgOp->setAttr(memSemantic, + rewriter.getStringAttr(stringifyEnum(op.getSem()))); + linalgOp->setAttr(memSyncScope, + rewriter.getStringAttr(stringifyEnum(op.getScope()))); + + linalgOp->setAttr("Software", rewriter.getUnitAttr()); + + // if the result hasn't been replace by load + // we need to erase it here + if (op.getResult().use_empty()) { + rewriter.eraseOp(op); + } + return success(); +} + LogicalResult ScalarStoreCanonicalizer::matchAndRewrite(triton::StoreOp op, PatternRewriter &rewriter) const { - if (!op.getValue().getType().isIntOrIndexOrFloat()) { return rewriter.notifyMatchFailure( op, "ScalarStoreCanonicalizer handles scalar store scene!"); @@ -541,7 +681,6 @@ ScalarStoreCanonicalizer::matchAndRewrite(triton::StoreOp op, LogicalResult ScalarAtomicRMWCanonicalizer::matchAndRewrite(triton::AtomicRMWOp op, PatternRewriter &rewriter) const { - if (!op.getVal().getType().isIntOrIndexOrFloat()) { return rewriter.notifyMatchFailure( op, "ScalarAtomicRMWCanonicalizer handles scalar atomic rmw op scene!"); @@ -564,6 +703,31 @@ ScalarAtomicRMWCanonicalizer::matchAndRewrite(triton::AtomicRMWOp op, return success(); } +LogicalResult +ScalarAtomicCASCanonicalizer::matchAndRewrite(triton::AtomicCASOp op, + PatternRewriter &rewriter) const { + if (!op.getVal().getType().isIntOrIndexOrFloat() && !op.getCmp().getType().isIntOrIndexOrFloat()) { + return rewriter.notifyMatchFailure( + op, "ScalarAtomicCASCanonicalizer handles scalar atomic cas op scene!"); + } + + auto ptr = op.getPtr(); + auto ptrTy = RankedTensorType::get({(int64_t)1}, ptr.getType()); + auto ptrSplat = rewriter.create(op.getLoc(), ptrTy, ptr); + auto cmpTy = RankedTensorType::get({(int64_t)1}, op.getCmp().getType()); + auto cmpSplat = + rewriter.create(op.getLoc(), cmpTy, op.getCmp()); + auto valTy = RankedTensorType::get({(int64_t)1}, op.getVal().getType()); + auto valSplat = + rewriter.create(op.getLoc(), valTy, op.getVal()); + + auto newAtomicOp = rewriter.create( + op.getLoc(), valTy, ptrSplat, cmpSplat, valSplat, + op.getSem(), op.getScope()); + rewriter.replaceOp(op, newAtomicOp); + return success(); +} + // The atomic max op with float input will be devided into // two atomic max ops with integer input // One handles the part of the tensor greater than zero diff --git a/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp b/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp index a0a1fe6daa731d94de95e6ddfead73f8334210f8..c29f900a62bcc846da6e3c400b76afeabc5f5c92 100644 --- a/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp +++ b/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp @@ -374,11 +374,13 @@ void TritonToLinalgPass::populateTritonToLinalgCanonicalizationPatterns(RewriteP patterns.add(patterns.getContext()); patterns.add, LoadStoreConverter::LoadStoreCanonicalizer, - LoadStoreConverter::LoadStoreCanonicalizer>(patterns.getContext()); + LoadStoreConverter::LoadStoreCanonicalizer, + LoadStoreConverter::LoadStoreCanonicalizer>(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add< TTOpConverters::ScalarMathCanonicalizer, @@ -438,6 +440,7 @@ void TritonToLinalgPass::populateTritonToLinalgConversionPatterns( patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); @@ -604,7 +607,19 @@ void TritonToLinalgPass::runOnOperation() { continue; auto context = func.getContext(); - constexpr int64_t workspaceArgIdx = 0; + constexpr int64_t syncBlockLockArgIdx = 0; + NamedAttribute syncBlockLockArgAttr(StringAttr::get(context, "syncBlockLock"), + UnitAttr::get(context)); + MemRefType syncBlockLockArgType = + MemRefType::get(SmallVector(1, ShapedType::kDynamic), + IntegerType::get(context, 8)); + func.insertArgument(syncBlockLockArgIdx, // argIndex + syncBlockLockArgType, // argType + nullptr, func->getLoc()); // dicAttr + func->setAttr("SyncBlockLockArgIdx", + IntegerAttr::get(IntegerType::get(&getContext(), 64), 0)); // 64: 64位整型 + + constexpr int64_t workspaceArgIdx = 1; MemRefType workspaceArgType = MemRefType::get(SmallVector(1, ShapedType::kDynamic), IntegerType::get(context, 8)); @@ -615,7 +630,7 @@ void TritonToLinalgPass::runOnOperation() { /*argType*/ workspaceArgType, /*dicAttr*/ nullptr, func->getLoc()); func->setAttr("WorkspaceArgIdx", - IntegerAttr::get(IntegerType::get(&getContext(), 64), 0)); + IntegerAttr::get(IntegerType::get(&getContext(), 64), 1)); // 64: 64位整型 } // Fix the Location info diff --git a/ascend/triton-adapter/lib/TritonToLinalg/UseAnalysis.cpp b/ascend/triton-adapter/lib/TritonToLinalg/UseAnalysis.cpp index 4b096316d477785cd6be5f34f5eba8ba623ea714..0580b7cc35b2c79519e96664676e302295659d05 100644 --- a/ascend/triton-adapter/lib/TritonToLinalg/UseAnalysis.cpp +++ b/ascend/triton-adapter/lib/TritonToLinalg/UseAnalysis.cpp @@ -89,6 +89,12 @@ void triton::UseAnalysis::visitOperation(Operation *op, propagateUse(operands[2], UseType::MetaUse); } }) + .Case([&](auto atomicOp) { + propagateUse(operands[0], UseType::MetaUse); + propagateUse(operands[1], UseType::DataUse); + propagateUse(operands[2], UseType::DataUse); + auto value = atomicOp.getVal(); + }) .Case([&](auto dot) { propagateResults(operands[0], results); propagateResults(operands[1], results); @@ -212,6 +218,11 @@ LogicalResult triton::runUseAnalysis(triton::FuncOp &funcOp) { if (result == ptr || result == mask) metaUsers.insert(user); }) + .Case([&](auto atomicOp) { + auto ptr = atomicOp.getPtr(); + if (result == ptr) + metaUsers.insert(user); + }) .Case([&](auto dot) { auto opc = dot.getC(); triton::SplatOp splat; diff --git a/docs/sources/python-api/outline.md b/docs/sources/python-api/outline.md index 0ac738117eaa6ceb41353534ca1b54087831fc42..40a4ccc5109a46e9d03154bf30ef46e8ec8c51e6 100644 --- a/docs/sources/python-api/outline.md +++ b/docs/sources/python-api/outline.md @@ -90,12 +90,12 @@ | | gather | × | × | × | × | × | ✓ | ✓ | ✓ | × | | Atomic Ops | atomic_add | ✓ | ✓ | ✓ | × | × | ✓ | ✓ | ✓ | × | | | atomic_and | × | × | × | × | × | × | × | × | × | -| | atomic_cas | × | × | × | × | × | × | × | × | × | +| | atomic_cas | ✓ | ✓ | ✓ | × | ✓ | ✓ | ✓ | × | × | | | atomic_max | ✓ | ✓ | ✓ | × | × | ✓ | ✓ | ✓ | × | | | atomic_min | ✓ | ✓ | ✓ | × | × | ✓ | ✓ | ✓ | × | -| | atomic_or | × | × | × | × | × | × | × | × | × | -| | atomic_xchg | × | × | × | × | × | × | × | × | × | -| | atomic_xor | × | × | × | × | × | × | × | × | × | +| | atomic_or | ✓ | ✓ | ✓ | × | ✓ | × | × | × | × | +| | atomic_xchg | ✓ | ✓ | ✓ | × | ✓ | ✓ | ✓ | × | × | +| | atomic_xor | ✓ | ✓ | ✓ | × | ✓ | × | × | × | × | | Random Number Generation | randint4x | × | × | ✓ | × | × | × | × | × | × | | | randint | × | × | ✓ | × | × | × | × | × | × | | | rand | × | × | × | × | × | × | ✓ | × | × | @@ -130,6 +130,10 @@ - atomic_min: 不支持标量（包括长度为1的tensor）访存 +- atomic_or: sem只支持默认值"acq_rel"模式，其他值均按默认值处理；scope只支持默认值"gpu"，其他值均按默认值处理 + +- atomic_xor: sem只支持默认值"acq_rel"模式，其他值均按默认值处理；scope只支持默认值"gpu"，其他值均按默认值处理 + - permute: 不支持不相邻轴转置，如`(0, 1, 2) -> (2, 1, 0)` - trans: 不支持不相邻轴转置，如`(0, 1, 2) -> (2, 1, 0)`