ops-tune-workbench-plugin 0.0.0-stage → 0.0.1-beta.2

package/skills/tool-ops-tune/ops-family/references/conversion/templates/dav3510/transpose_with_gather.template

@@ -0,0 +1,688 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file transpose_with_gather.h
+ * \brief transpose_with_gather
+ */
+#ifndef KERNEL_TRANSPOSE_WITH_GATHER_H_
+#define KERNEL_TRANSPOSE_WITH_GATHER_H_
+
+#include <type_traits>
+
+#include "reg_compute/kernel_reg_compute_datacopy_intf.h"
+#include "op_kernel/platform_util.h"
+#include "transpose_base.h"
+
+namespace Transpose {
+using namespace AscendC;
+using AscendC::Reg::CreateMask;
+//using AscendC::Reg::DataCopyGather;
+using AscendC::Reg::MaskReg;
+using AscendC::Reg::RegTensor;
+using AscendC::Reg::UpdateMask;
+
+constexpr int8_t UB_MAX_DIM_NUM = 6;
+constexpr int8_t NUM_TWO = 2;
+constexpr int8_t NUM_THREE = 3;
+
+template <typename T>
+class TransposeWithGather {
+public:
+    __aicore__ inline TransposeWithGather(){};
+    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, const GatherTransposeTilingData* tilingData, TPipe* pipe);
+    __aicore__ inline void Process();
+
+private:
+    __aicore__ inline void GetCoreLoopRange();
+    __aicore__ inline void InitAxes();
+    __aicore__ inline int32_t CalcUbAxesInOffset(int8_t begIdx);
+    __aicore__ inline int32_t CalcUbAxesOutOffset(int8_t begIdx);
+    __aicore__ inline int64_t CalcBlkAxesOffset(int8_t begIdx);
+    __aicore__ inline void GenIndex4OneDim(int32_t goffset);
+    __aicore__ inline void GenIndex4TwoDim(int32_t goffset);
+    __aicore__ inline void GenIndex4ThreeDim(int32_t goffset);
+    __aicore__ inline void GenGatherIndex(int32_t goffset);
+    __aicore__ inline void GenGatherIndex4AllPhase();
+    __aicore__ inline void CalcBlkAddr(int64_t idx);
+    __aicore__ inline void SetCopyInParams();
+    __aicore__ inline void CopyDataIn();
+    __aicore__ inline void GetOutLoopAxes();
+    __aicore__ inline void SetCopyOutParams();
+    __aicore__ inline void CopyDataOut();
+    __aicore__ inline void UpdateUbAxes(int64_t blkLpIdx);
+    __aicore__ inline void GatherData();
+
+private:
+    int64_t blockIdx_ = 0;
+    const GatherTransposeTilingData* td_ = nullptr;
+    TQue<QuePosition::VECIN, 1> xInQue_;
+    TQue<QuePosition::VECOUT, 1> xOutQue_;
+    TBuf<QuePosition::VECCALC> idxBuf_;
+    GlobalTensor<T> xGM_;
+    GlobalTensor<T> yGM_;
+
+    int64_t blkBeg_ = 0;
+    int64_t blkEnd_ = 0;
+    int64_t blkInAddr_ = 0;
+    int64_t blkOutAddr_ = 0;
+    int64_t blkInCutROffset_ = 0;
+    int64_t blkOutCutROffset_ = 0;
+    int64_t blkInCutAxisSize_ = 0;
+    int64_t blkOutCutAxisSize_ = 0;
+    int32_t inUbAxes_[UB_MAX_DIM_NUM] = {1, 1, 1, 1, 1, 1};
+    int32_t outUbAxes_[UB_MAX_DIM_NUM] = {1, 1, 1, 1, 1, 1};
+    /*  0: out_cut=false, in_cut=false
+     *  1: out_cut=false, in_cut=true
+     *  2: out_cut=true,  in_cut=false
+     *  3: out_cut=true,  in_cut=true
+     */
+    int32_t gIndexId_ = 0; // control index offset of gather
+
+    using RangeType_ = std::conditional_t<sizeof(T) <= sizeof(int16_t), int16_t, int32_t>;
+    using IdxType_ = std::conditional_t<sizeof(T) <= sizeof(int16_t), uint16_t, uint32_t>;
+    using CastType_ = std::conditional_t<sizeof(T) == 1,
+                                         std::conditional_t<std::is_same_v<T, uint8_t>, uint16_t, int16_t>, T>;
+    uint32_t vlSize_ = static_cast<uint32_t>(Ops::Base::GetVRegSize() / sizeof(CastType_));
+    uint32_t idxVLSize_ = static_cast<uint32_t>(Ops::Base::GetVRegSize() / sizeof(RangeType_));
+    LocalTensor<RangeType_> idxLocal_;
+    int32_t gElemPerBlock_ = static_cast<int32_t>(BLOCK_SIZE_BYTE / sizeof(RangeType_));
+    int32_t elemPerBlock_ = static_cast<int32_t>(BLOCK_SIZE_BYTE / sizeof(T));
+    int32_t gIdxOffset_ = 0;
+
+    LoopModeParams lpModeInParams_ = {1, 1, 0, 0, 0, 0};
+    LoopModeParams lpModeOutParams_ = {1, 1, 0, 0, 0, 0};
+    DataCopyExtParams copyInParams_;
+    DataCopyExtParams copyOutParams_;
+    uint32_t outUbAxis0_ = 1;
+    uint32_t outUbAxis1_ = 1;
+    uint32_t outUbAxis2_ = 1;
+    uint32_t outUbAxis0InROffset_ = 0;
+    uint32_t outUbAxis1InROffset_ = 0;
+    uint32_t outUbAxis2InROffset_ = 0;
+    int32_t inUbTailFactor_ = 0;
+    int32_t outUbTailFactor_ = 0;
+};
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::Init(GM_ADDR x, GM_ADDR y, const GatherTransposeTilingData* tilingData,
+                                                    TPipe* pipe)
+{
+    blockIdx_ = GetBlockIdx();
+    td_ = tilingData;
+
+    pipe->InitBuffer(xInQue_, BUFFER_NUM, td_->dataTensorSize);
+    pipe->InitBuffer(xOutQue_, BUFFER_NUM, td_->dataTensorSize);
+    pipe->InitBuffer(idxBuf_, td_->indexTensorSize);
+    idxLocal_ = idxBuf_.Get<RangeType_>();
+
+    xGM_.SetGlobalBuffer((__gm__ T*)x);
+    yGM_.SetGlobalBuffer((__gm__ T*)y);
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::Process()
+{
+    if (blockIdx_ >= td_->usedCoreCnt) {
+        return;
+    }
+    GetCoreLoopRange();
+    InitAxes();
+    GenGatherIndex4AllPhase();
+    for (int64_t blkLpIdx = blkBeg_; blkLpIdx < blkEnd_; ++blkLpIdx) {
+        CalcBlkAddr(blkLpIdx);
+        UpdateUbAxes(blkLpIdx);
+        CopyDataIn();
+        GetOutLoopAxes();
+        GatherData();
+        CopyDataOut();
+    }
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::GetCoreLoopRange()
+{
+    blkBeg_ = blockIdx_ * td_->blkFactor;
+    if (blockIdx_ != td_->usedCoreCnt - 1) {
+        blkEnd_ = blkBeg_ + td_->blkFactor;
+    } else {
+        blkEnd_ = blkBeg_ + td_->blkTailFactor;
+    }
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::InitAxes()
+{
+    for (int8_t i = 0; i < td_->ubAxesCnt; ++i) {
+        inUbAxes_[i] = td_->inUbAxes[i];
+        outUbAxes_[i] = td_->outUbAxes[i];
+    }
+    if (td_->blkInUbCutPos != -1) {
+        blkInCutROffset_ = CalcBlkAxesOffset(td_->blkInUbCutPos);
+        blkInCutAxisSize_ = td_->blkAxes[td_->blkInUbCutPos];
+    }
+    if (td_->blkOutUbCutPos != -1) {
+        blkOutCutROffset_ = CalcBlkAxesOffset(td_->blkOutUbCutPos);
+        blkOutCutAxisSize_ = td_->blkAxes[td_->blkOutUbCutPos];
+    }
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::UpdateUbAxes(int64_t blkLpIdx)
+{
+    bool isInSwitch = false;
+    bool isOutSwitch = false;
+
+    if (td_->blkInUbCutPos != -1) {
+        if (blkLpIdx / blkInCutROffset_ % blkInCutAxisSize_ == blkInCutAxisSize_ - 1) {
+            inUbAxes_[td_->inUbInCutPos] = inUbTailFactor_;
+            outUbAxes_[td_->outUbInCutPos] = inUbTailFactor_;
+            isInSwitch = true;
+        } else {
+            inUbAxes_[td_->inUbInCutPos] = td_->inUbCutAxisFactor;
+            outUbAxes_[td_->outUbInCutPos] = td_->inUbCutAxisFactor;
+            isInSwitch = false;
+        }
+    }
+    if (td_->blkOutUbCutPos != -1) {
+        if (blkLpIdx / blkOutCutROffset_ % blkOutCutAxisSize_ == blkOutCutAxisSize_ - 1) {
+            inUbAxes_[td_->inUbOutCutPos] = outUbTailFactor_;
+            outUbAxes_[td_->outUbOutCutPos] = outUbTailFactor_;
+            isOutSwitch = true;
+        } else {
+            inUbAxes_[td_->inUbOutCutPos] = td_->outUbCutAxisFactor;
+            outUbAxes_[td_->outUbOutCutPos] = td_->outUbCutAxisFactor;
+            isOutSwitch = false;
+        }
+    }
+
+    if (!isOutSwitch) {
+        if (!isInSwitch) {
+            gIndexId_ = 0;
+        } else {
+            gIndexId_ = 1;
+        }
+    } else {
+        if (!isInSwitch) {
+            gIndexId_ = NUM_TWO;
+        } else {
+            gIndexId_ = NUM_THREE;
+        }
+    }
+}
+
+template <typename T>
+__aicore__ inline int32_t TransposeWithGather<T>::CalcUbAxesInOffset(int8_t begIdx)
+{
+    int32_t res = 1;
+
+    // DataCopyPad blockCount and blockLen are consecutive
+    if (begIdx + 1 >= td_->inUbInCutPos) {
+        for (int8_t i = begIdx + 1; i < td_->ubAxesCnt; ++i) {
+            res *= inUbAxes_[i];
+        }
+    } else {
+        for (int8_t i = td_->inUbInCutPos - 1; i < td_->ubAxesCnt; ++i) {
+            res *= inUbAxes_[i];
+        }
+        // DataCopyPad loop2, loop1 are block align
+        res = Ops::Base::CeilAlign(res, elemPerBlock_);
+        for (int8_t j = begIdx + 1; j < td_->inUbInCutPos - 1; ++j) {
+            res *= inUbAxes_[j];
+        }
+    }
+    return res;
+}
+
+template <typename T>
+__aicore__ inline int32_t TransposeWithGather<T>::CalcUbAxesOutOffset(int8_t begIdx)
+{
+    int32_t res = 1;
+    for (int8_t i = begIdx + 1; i < td_->ubAxesCnt; ++i) {
+        res *= outUbAxes_[i];
+    }
+    return res;
+}
+
+template <typename T>
+__aicore__ inline int64_t TransposeWithGather<T>::CalcBlkAxesOffset(int8_t begIdx)
+{
+    int32_t res = 1;
+    for (int8_t i = begIdx + 1; i < td_->blkAxesCnt; ++i) {
+        res *= td_->blkAxes[i];
+    }
+    return res;
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::GatherData()
+{
+    uint32_t outLen = static_cast<uint32_t>(CalcUbAxesOutOffset(td_->outUbOutCutPos - 1));
+    uint32_t outLenAlign = Ops::Base::CeilAlign(outLen, static_cast<uint32_t>(elemPerBlock_));
+    uint16_t burstLpCnt = Ops::Base::CeilDiv(outLen, vlSize_);
+    uint32_t maskValue = outLen;
+    if constexpr (sizeof(T) == sizeof(int8_t)) {
+        maskValue *= sizeof(int16_t);
+    }
+
+    auto xInLocal = xInQue_.DeQue<T>();
+    auto xOutLocal = xOutQue_.AllocTensor<T>();
+    __local_mem__ T* xInAddr = (__local_mem__ T*)xInLocal.GetPhyAddr();
+    __local_mem__ T* xOutAddr = (__local_mem__ T*)xOutLocal.GetPhyAddr();
+    __local_mem__ RangeType_* idxAddr = (__local_mem__ RangeType_*)idxLocal_.GetPhyAddr();
+    idxAddr += static_cast<uint32_t>(gIndexId_ * gIdxOffset_);
+
+    __VEC_SCOPE__
+    {
+        RegTensor<T> xReg;
+        RegTensor<RangeType_> idxReg;
+        RegTensor<RangeType_> idxOriReg;
+        MaskReg maskIdx = CreateMask<RangeType_, MicroAPI::MaskPattern::ALL>();
+        MaskReg mask;
+
+        for (uint16_t lpIdx = 0; lpIdx < burstLpCnt; ++lpIdx) {
+            mask = UpdateMask<T>(maskValue);
+            MicroAPI::DataCopy(idxOriReg, idxAddr + lpIdx * vlSize_);
+            uint32_t outIdx = 0;
+            for (uint16_t axis2Idx = 0; axis2Idx < static_cast<uint16_t>(outUbAxis2_); ++axis2Idx) {
+                RangeType_ axis2Update = static_cast<RangeType_>(axis2Idx * outUbAxis2InROffset_);
+                for (uint16_t axis1Idx = 0; axis1Idx < static_cast<uint16_t>(outUbAxis1_); ++axis1Idx) {
+                    RangeType_ axis1Update = static_cast<RangeType_>(axis1Idx * outUbAxis1InROffset_);
+                    for (uint16_t axis0Idx = 0; axis0Idx < static_cast<uint16_t>(outUbAxis0_); ++axis0Idx) {
+                        RangeType_ axis0Update = static_cast<RangeType_>(axis0Idx * outUbAxis0InROffset_);
+                        RangeType_ idxUpdate = axis2Update + axis1Update + axis0Update;
+                        MicroAPI::Adds(idxReg, idxOriReg, idxUpdate, maskIdx);
+			DataCopyGatherImpl((MicroAPI::RegTensor<CastType_>&)xReg, xInAddr,
+                                       (MicroAPI::RegTensor<IdxType_>&)idxReg, mask);
+			//vgather2(xReg, xInAddr, idxReg, mask);
+                        if constexpr (sizeof(T) != sizeof(int8_t)) {
+                            MicroAPI::DataCopy(xOutAddr + outIdx * outLenAlign + lpIdx * vlSize_, xReg, mask);
+                        } else {
+                            __local_mem__ CastType_* xOutAddrB16 = reinterpret_cast<__local_mem__ CastType_*>(
+                                xOutAddr + outIdx * outLenAlign + lpIdx * vlSize_);
+                            MicroAPI::DataCopy<CastType_, MicroAPI::StoreDist::DIST_PACK_B16>(
+                                xOutAddrB16, (MicroAPI::RegTensor<CastType_>&)xReg, mask);
+                        }
+                        ++outIdx;
+                    }
+                }
+            }
+        }
+    }
+    xInQue_.FreeTensor(xInLocal);
+    xOutQue_.EnQue(xOutLocal);
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::GenIndex4OneDim(int32_t goffset)
+{
+    int8_t lastPerm = td_->ubPerm[td_->ubAxesCnt - 1];
+    uint32_t lastDimSize = static_cast<uint32_t>(outUbAxes_[td_->ubAxesCnt - 1]);
+    RangeType_ lastDimInOffset = static_cast<RangeType_>(CalcUbAxesInOffset(lastPerm));
+
+    __local_mem__ RangeType_* idxAddr = (__local_mem__ RangeType_*)idxLocal_.GetPhyAddr();
+    idxAddr += goffset;
+    uint16_t loopCnt = static_cast<uint16_t>(Ops::Base::CeilDiv(lastDimSize, idxVLSize_));
+
+    __VEC_SCOPE__
+    {
+        RegTensor<RangeType_> srcReg;
+        RegTensor<RangeType_> dstReg;
+        MaskReg mask;
+        MicroAPI::Arange(srcReg, 0);
+        for (uint16_t lpIdx = 0; lpIdx < loopCnt; ++lpIdx) {
+            mask = UpdateMask<RangeType_>(lastDimSize);
+            MicroAPI::Muls(dstReg, srcReg, lastDimInOffset, mask);
+            MicroAPI::DataCopy(idxAddr + lpIdx * idxVLSize_, dstReg, mask);
+            MicroAPI::Adds(srcReg, srcReg, idxVLSize_, mask);
+        }
+    }
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::GenIndex4TwoDim(int32_t goffset)
+{
+    int8_t lastPerm = td_->ubPerm[td_->ubAxesCnt - 1];
+    uint32_t lastDimSize = static_cast<uint32_t>(outUbAxes_[td_->ubAxesCnt - 1]);
+    RangeType_ lastDimInOffset = static_cast<RangeType_>(CalcUbAxesInOffset(lastPerm));
+    int8_t last2ndPerm = td_->ubPerm[td_->ubAxesCnt - NUM_TWO];
+    uint32_t last2ndDimSize = static_cast<uint32_t>(outUbAxes_[td_->ubAxesCnt - NUM_TWO]);
+    RangeType_ last2ndDimInOffset = static_cast<RangeType_>(CalcUbAxesInOffset(last2ndPerm));
+    auto totalDimSize = lastDimSize * last2ndDimSize;
+    RangeType_ lastDimInc = static_cast<RangeType_>(idxVLSize_ % lastDimSize);
+    RangeType_ last2ndDimInc = static_cast<RangeType_>(idxVLSize_ / lastDimSize);
+
+    __local_mem__ RangeType_* idxAddr = (__local_mem__ RangeType_*)idxLocal_.GetPhyAddr();
+    idxAddr += goffset;
+    uint16_t loopCnt = 0;
+    uint32_t leftDimSize = 0;
+    if (totalDimSize > idxVLSize_) {
+        leftDimSize = totalDimSize - idxVLSize_;
+        loopCnt = static_cast<uint16_t>(Ops::Base::CeilDiv(leftDimSize, idxVLSize_));
+    }
+
+    /*   input: (a, b, c, d)
+     *  output: (d, c, b, a)
+     *  generate index for a, b
+     */
+    __VEC_SCOPE__
+    {
+        RegTensor<RangeType_> idxReg;
+        RegTensor<RangeType_> dim0Reg;
+        RegTensor<RangeType_> tmpReg;
+        RegTensor<RangeType_> dim1Reg;
+        RegTensor<RangeType_> dstReg;
+        MaskReg mask = CreateMask<RangeType_, MicroAPI::MaskPattern::ALL>();
+        // vec_a: VL % a
+        MicroAPI::Arange(idxReg, 0);
+        MicroAPI::Duplicate(dim0Reg, lastDimSize);
+        MicroAPI::Div(tmpReg, idxReg, dim0Reg, mask);
+        MicroAPI::Copy(dim1Reg, tmpReg); // vec_b: VL / a
+        MicroAPI::Mul(tmpReg, tmpReg, dim0Reg, mask);
+        MicroAPI::Sub(dim0Reg, idxReg, tmpReg, mask);
+        // index: vec_a * a_in_offset + vec_b * b_in_offset
+        MicroAPI::Muls(tmpReg, dim0Reg, lastDimInOffset, mask);
+        MicroAPI::Muls(dstReg, dim1Reg, last2ndDimInOffset, mask);
+        MicroAPI::Add(dstReg, dstReg, tmpReg, mask);
+        MicroAPI::DataCopy(idxAddr, dstReg, mask);
+
+        MaskReg lpMask;
+        MaskReg selMask;
+        RegTensor<RangeType_> zeroReg;
+        RegTensor<RangeType_> oneReg;
+        RegTensor<RangeType_> cmpReg;
+        MicroAPI::Duplicate(zeroReg, 0);
+        MicroAPI::Duplicate(oneReg, 1);
+        for (uint16_t lpIdx = 0; lpIdx < loopCnt; ++lpIdx) {
+            lpMask = UpdateMask<RangeType_>(leftDimSize);
+            /*   vec_a += a_inc
+             *   cmp_a = vec_a >= a
+             *   vec_a = vec_a - cmp_a * a
+             */
+            MicroAPI::Adds(dim0Reg, dim0Reg, lastDimInc, lpMask);
+            MicroAPI::CompareScalar<RangeType_, CMPMODE::GE>(selMask, dim0Reg, lastDimSize, lpMask);
+            MicroAPI::Select(cmpReg, oneReg, zeroReg, selMask);
+            MicroAPI::Muls(tmpReg, cmpReg, lastDimSize, lpMask);
+            MicroAPI::Sub(dim0Reg, dim0Reg, tmpReg, lpMask);
+            // vec_b += (b_inc + cmp_a)
+            MicroAPI::Adds(dim1Reg, dim1Reg, last2ndDimInc, lpMask);
+            MicroAPI::Add(dim1Reg, dim1Reg, cmpReg, lpMask);
+            // index: vec_a * a_in_offset + vec_b * b_in_offset
+            MicroAPI::Muls(tmpReg, dim0Reg, lastDimInOffset, lpMask);
+            MicroAPI::Muls(dstReg, dim1Reg, last2ndDimInOffset, lpMask);
+            MicroAPI::Add(dstReg, dstReg, tmpReg, lpMask);
+            MicroAPI::DataCopy(idxAddr + (lpIdx + 1) * idxVLSize_, dstReg, lpMask);
+        }
+    }
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::GenIndex4ThreeDim(int32_t goffset)
+{
+    int8_t lastPerm = td_->ubPerm[td_->ubAxesCnt - 1];
+    uint32_t lastDimSize = static_cast<uint32_t>(outUbAxes_[td_->ubAxesCnt - 1]);
+    RangeType_ lastDimInOffset = static_cast<RangeType_>(CalcUbAxesInOffset(lastPerm));
+    int8_t last2ndPerm = td_->ubPerm[td_->ubAxesCnt - NUM_TWO];
+    uint32_t last2ndDimSize = static_cast<uint32_t>(outUbAxes_[td_->ubAxesCnt - NUM_TWO]);
+    RangeType_ last2ndDimInOffset = static_cast<RangeType_>(CalcUbAxesInOffset(last2ndPerm));
+    int8_t last3rdPerm = td_->ubPerm[td_->ubAxesCnt - NUM_THREE];
+    uint32_t last3rdDimSize = static_cast<uint32_t>(outUbAxes_[td_->ubAxesCnt - NUM_THREE]);
+    RangeType_ last3rdDimInOffset = static_cast<RangeType_>(CalcUbAxesInOffset(last3rdPerm));
+    auto totalDimSize = lastDimSize * last2ndDimSize * last3rdDimSize;
+    RangeType_ lastDimInc = static_cast<RangeType_>(idxVLSize_ % lastDimSize);
+    RangeType_ last2ndDimInc = static_cast<RangeType_>(idxVLSize_ / lastDimSize % last2ndDimSize);
+    RangeType_ last3rdDimInc = static_cast<RangeType_>(idxVLSize_ / (lastDimSize * last2ndDimSize));
+
+    __local_mem__ RangeType_* idxAddr = (__local_mem__ RangeType_*)idxLocal_.GetPhyAddr();
+    idxAddr += goffset;
+    uint16_t loopCnt = 0;
+    uint32_t leftDimSize = 0;
+    if (totalDimSize > idxVLSize_) {
+        leftDimSize = totalDimSize - idxVLSize_;
+        loopCnt = static_cast<uint16_t>(Ops::Base::CeilDiv(leftDimSize, idxVLSize_));
+    }
+
+    /*   input: (a, b, c, d)
+     *  output: (d, c, b, a)
+     *  generate index for a, b, c
+     */
+    __VEC_SCOPE__
+    {
+        RegTensor<RangeType_> idxReg;
+        RegTensor<RangeType_> dim0Reg;
+        RegTensor<RangeType_> tmpReg;
+        RegTensor<RangeType_> dim1Reg;
+        RegTensor<RangeType_> tmp1Reg;
+        RegTensor<RangeType_> dim2Reg;
+        RegTensor<RangeType_> dstReg;
+        MaskReg mask = CreateMask<RangeType_, MicroAPI::MaskPattern::ALL>();
+        // vec_a: VL % a
+        MicroAPI::Arange(idxReg, 0);
+        MicroAPI::Duplicate(dim0Reg, lastDimSize);
+        MicroAPI::Copy(dim2Reg, dim0Reg); // backup a
+        MicroAPI::Div(tmpReg, idxReg, dim0Reg, mask);
+        MicroAPI::Copy(dim1Reg, tmpReg); // backup VL / a
+        MicroAPI::Mul(tmpReg, tmpReg, dim0Reg, mask);
+        MicroAPI::Sub(dim0Reg, idxReg, tmpReg, mask);
+        // vec_b: VL / a % b
+        MicroAPI::Duplicate(tmp1Reg, last2ndDimSize);
+        MicroAPI::Mul(dim2Reg, dim2Reg, tmp1Reg, mask); // backup b
+        MicroAPI::Div(tmpReg, dim1Reg, tmp1Reg, mask);
+        MicroAPI::Mul(tmpReg, tmpReg, tmp1Reg, mask);
+        MicroAPI::Sub(dim1Reg, dim1Reg, tmpReg, mask);
+        // vec_c: VL / (a * b)
+        MicroAPI::Div(dim2Reg, idxReg, dim2Reg, mask);
+        // index: vec_a * a_in_offset + vec_b * b_in_offset + vec_c * c_in_offset
+        MicroAPI::Muls(tmpReg, dim0Reg, lastDimInOffset, mask);
+        MicroAPI::Muls(tmp1Reg, dim1Reg, last2ndDimInOffset, mask);
+        MicroAPI::Muls(dstReg, dim2Reg, last3rdDimInOffset, mask);
+        MicroAPI::Add(dstReg, dstReg, tmpReg, mask);
+        MicroAPI::Add(dstReg, dstReg, tmp1Reg, mask);
+        MicroAPI::DataCopy(idxAddr, dstReg, mask);
+
+        MaskReg lpMask;
+        MaskReg selMask;
+        RegTensor<RangeType_> zeroReg;
+        RegTensor<RangeType_> oneReg;
+        RegTensor<RangeType_> cmpReg;
+        MicroAPI::Duplicate(zeroReg, 0);
+        MicroAPI::Duplicate(oneReg, 1);
+        for (uint16_t lpIdx = 0; lpIdx < loopCnt; ++lpIdx) {
+            lpMask = UpdateMask<RangeType_>(leftDimSize);
+            /*   vec_a += a_inc
+             *   cmp_a = vec_a >= a
+             *   vec_a = vec_a - cmp_a * a
+             */
+            MicroAPI::Adds(dim0Reg, dim0Reg, lastDimInc, lpMask);
+            MicroAPI::CompareScalar<RangeType_, CMPMODE::GE>(selMask, dim0Reg, lastDimSize, lpMask);
+            MicroAPI::Select(cmpReg, oneReg, zeroReg, selMask);
+            MicroAPI::Muls(tmpReg, cmpReg, lastDimSize, lpMask);
+            MicroAPI::Sub(dim0Reg, dim0Reg, tmpReg, lpMask);
+            /*   vec_b += (b_inc + cmp_a)
+             *   cmp_b = vec_b >= b
+             *   vec_b = vec_b - cmp_b * b
+             */
+            MicroAPI::Adds(cmpReg, cmpReg, last2ndDimInc, lpMask);
+            MicroAPI::Add(dim1Reg, dim1Reg, cmpReg, lpMask);
+            MicroAPI::CompareScalar<RangeType_, CMPMODE::GE>(selMask, dim1Reg, last2ndDimSize, lpMask);
+            MicroAPI::Select(cmpReg, oneReg, zeroReg, selMask);
+            MicroAPI::Muls(tmpReg, cmpReg, last2ndDimSize, lpMask);
+            MicroAPI::Sub(dim1Reg, dim1Reg, tmpReg, lpMask);
+            // vec_c += (c_inc + cmp_b)
+            MicroAPI::Adds(dim2Reg, dim2Reg, last3rdDimInc, lpMask);
+            MicroAPI::Add(dim2Reg, dim2Reg, cmpReg, lpMask);
+            // index: vec_a * a_in_offset + vec_b * b_in_offset + vec_c * c_in_offset
+            MicroAPI::Muls(tmpReg, dim0Reg, lastDimInOffset, lpMask);
+            MicroAPI::Muls(tmp1Reg, dim1Reg, last2ndDimInOffset, lpMask);
+            MicroAPI::Muls(dstReg, dim2Reg, last3rdDimInOffset, lpMask);
+            MicroAPI::Add(dstReg, dstReg, tmpReg, lpMask);
+            MicroAPI::Add(dstReg, dstReg, tmp1Reg, lpMask);
+            MicroAPI::DataCopy(idxAddr + (lpIdx + 1) * idxVLSize_, dstReg, lpMask);
+        }
+    }
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::GenGatherIndex(int32_t goffset)
+{
+    int8_t dimNum = td_->ubAxesCnt - td_->outUbOutCutPos;
+    // max borrowed axes count is 3
+    if (dimNum == 1) {
+        GenIndex4OneDim(goffset);
+    } else if (dimNum == NUM_TWO) {
+        GenIndex4TwoDim(goffset);
+    } else {
+        GenIndex4ThreeDim(goffset);
+    }
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::GenGatherIndex4AllPhase()
+{
+    GenGatherIndex(0);
+    gIdxOffset_ = Ops::Base::CeilAlign(CalcUbAxesOutOffset(td_->outUbOutCutPos - 1), gElemPerBlock_);
+    inUbTailFactor_ = td_->inUbCutAxisSize % td_->inUbCutAxisFactor;
+    outUbTailFactor_ = td_->outUbCutAxisSize % td_->outUbCutAxisFactor;
+    if (td_->blkOutUbCutPos == -1 && td_->blkInUbCutPos != -1) {
+        inUbAxes_[td_->inUbInCutPos] = inUbTailFactor_;
+        outUbAxes_[td_->outUbInCutPos] = inUbTailFactor_;
+        GenGatherIndex(gIdxOffset_);
+    } else if (td_->blkOutUbCutPos != -1 && td_->blkInUbCutPos == -1) {
+        inUbAxes_[td_->inUbOutCutPos] = outUbTailFactor_;
+        outUbAxes_[td_->outUbOutCutPos] = outUbTailFactor_;
+        GenGatherIndex(gIdxOffset_ * NUM_TWO);
+    } else if (td_->blkOutUbCutPos != -1 && td_->blkInUbCutPos != -1) {
+        inUbAxes_[td_->inUbInCutPos] = inUbTailFactor_;
+        outUbAxes_[td_->outUbInCutPos] = inUbTailFactor_;
+        GenGatherIndex(gIdxOffset_);
+        inUbAxes_[td_->inUbOutCutPos] = outUbTailFactor_;
+        outUbAxes_[td_->outUbOutCutPos] = outUbTailFactor_;
+        inUbAxes_[td_->inUbInCutPos] = td_->inUbCutAxisFactor;
+        outUbAxes_[td_->outUbInCutPos] = td_->inUbCutAxisFactor;
+        GenGatherIndex(gIdxOffset_ * NUM_TWO);
+        inUbAxes_[td_->inUbInCutPos] = inUbTailFactor_;
+        outUbAxes_[td_->outUbInCutPos] = inUbTailFactor_;
+        GenGatherIndex(gIdxOffset_ * NUM_THREE);
+    }
+    // recover ub axes for next process
+    inUbAxes_[td_->inUbOutCutPos] = td_->outUbCutAxisFactor;
+    inUbAxes_[td_->inUbInCutPos] = td_->inUbCutAxisFactor;
+    outUbAxes_[td_->outUbInCutPos] = td_->inUbCutAxisFactor;
+    outUbAxes_[td_->outUbOutCutPos] = td_->outUbCutAxisFactor;
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::CalcBlkAddr(int64_t idx)
+{
+    int64_t tmpIdx = 0;
+    blkInAddr_ = 0;
+    blkOutAddr_ = 0;
+    for (int8_t i = 0; i < td_->blkAxesCnt; ++i) {
+        tmpIdx = idx / CalcBlkAxesOffset(i) % td_->blkAxes[i];
+        blkInAddr_ += tmpIdx * td_->blkAxesInAOffset[i];
+        blkOutAddr_ += tmpIdx * td_->blkAxesOutAOffset[i];
+    }
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::SetCopyInParams()
+{
+    int32_t inUbAxis0 = 1;
+    int32_t inUbAxis1 = 1;
+    int32_t inUbAxis2 = 1;
+    if (td_->inUbInCutPos == 1) {
+        inUbAxis0 = inUbAxes_[0];
+    } else if (td_->inUbInCutPos == NUM_TWO) {
+        inUbAxis0 = inUbAxes_[1];
+        inUbAxis1 = inUbAxes_[0];
+    } else if (td_->inUbInCutPos == NUM_THREE) {
+        inUbAxis0 = inUbAxes_[NUM_TWO];
+        inUbAxis1 = inUbAxes_[1];
+        inUbAxis2 = inUbAxes_[0];
+    }
+
+    uint32_t inCube = static_cast<size_t>(CalcUbAxesInOffset(td_->inUbInCutPos - 1)) * sizeof(T);
+    copyInParams_.blockCount = inUbAxis0;
+    copyInParams_.blockLen = inCube;
+    copyInParams_.srcStride = td_->axis0InSrcStride * sizeof(T) - inCube;
+    copyInParams_.dstStride = 0;
+    lpModeInParams_.loop1Size = inUbAxis1;
+    lpModeInParams_.loop2Size = inUbAxis2;
+    lpModeInParams_.loop1SrcStride = td_->axis1InSrcStride * sizeof(T);
+    lpModeInParams_.loop2SrcStride = td_->axis2InSrcStride * sizeof(T);
+    lpModeInParams_.loop1DstStride = Ops::Base::CeilAlign(
+        static_cast<int64_t>(copyInParams_.blockCount * copyInParams_.blockLen), BLOCK_SIZE_BYTE);
+    lpModeInParams_.loop2DstStride = lpModeInParams_.loop1Size * lpModeInParams_.loop1DstStride;
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::CopyDataIn()
+{
+    SetCopyInParams();
+    DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
+    auto xInLocal = xInQue_.AllocTensor<T>();
+    SetLoopModePara(lpModeInParams_, DataCopyMVType::OUT_TO_UB);
+    DataCopyPad<T, PaddingMode::Compact>(xInLocal, xGM_[blkInAddr_], copyInParams_, padParams);
+    ResetLoopModePara(DataCopyMVType::OUT_TO_UB);
+    xInQue_.EnQue(xInLocal);
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::GetOutLoopAxes()
+{
+    if (td_->outUbOutCutPos == 1) {
+        outUbAxis0_ = outUbAxes_[0];
+        outUbAxis0InROffset_ = CalcUbAxesInOffset(td_->ubPerm[0]);
+    } else if (td_->outUbOutCutPos == NUM_TWO) {
+        outUbAxis0_ = outUbAxes_[1];
+        outUbAxis1_ = outUbAxes_[0];
+        outUbAxis0InROffset_ = CalcUbAxesInOffset(td_->ubPerm[1]);
+        outUbAxis1InROffset_ = CalcUbAxesInOffset(td_->ubPerm[0]);
+    } else if (td_->outUbOutCutPos == NUM_THREE) {
+        outUbAxis0_ = outUbAxes_[NUM_TWO];
+        outUbAxis1_ = outUbAxes_[1];
+        outUbAxis2_ = outUbAxes_[0];
+        outUbAxis0InROffset_ = CalcUbAxesInOffset(td_->ubPerm[NUM_TWO]);
+        outUbAxis1InROffset_ = CalcUbAxesInOffset(td_->ubPerm[1]);
+        outUbAxis2InROffset_ = CalcUbAxesInOffset(td_->ubPerm[0]);
+    }
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::SetCopyOutParams()
+{
+    uint32_t outCube = static_cast<size_t>(CalcUbAxesOutOffset(td_->outUbOutCutPos - 1)) * sizeof(T);
+    uint32_t outCubeAlign = Ops::Base::CeilAlign(outCube, static_cast<uint32_t>(BLOCK_SIZE_BYTE));
+    copyOutParams_.blockCount = outUbAxis0_;
+    copyOutParams_.blockLen = outCube;
+    copyOutParams_.srcStride = 0;
+    copyOutParams_.dstStride = td_->axis0OutDstStride * sizeof(T) - outCube;
+    lpModeOutParams_.loop1Size = outUbAxis1_;
+    lpModeOutParams_.loop2Size = outUbAxis2_;
+    lpModeOutParams_.loop1SrcStride = outUbAxis0_ * outCubeAlign;
+    lpModeOutParams_.loop2SrcStride = outUbAxis1_ * lpModeOutParams_.loop1SrcStride;
+    lpModeOutParams_.loop1DstStride = td_->axis1OutDstStride * sizeof(T);
+    lpModeOutParams_.loop2DstStride = td_->axis2OutDstStride * sizeof(T);
+}
+
+template <typename T>
+__aicore__ inline void TransposeWithGather<T>::CopyDataOut()
+{
+    SetCopyOutParams();
+    auto xOutLocal = xOutQue_.DeQue<T>();
+    SetLoopModePara(lpModeOutParams_, DataCopyMVType::UB_TO_OUT);
+    DataCopyPad(yGM_[blkOutAddr_], xOutLocal, copyOutParams_);
+    ResetLoopModePara(DataCopyMVType::UB_TO_OUT);
+    xOutQue_.FreeTensor(xOutLocal);
+}
+
+} // namespace Transpose
+
+#endif // TRANSPOSE_CUT_TWO_AXIS