#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
#include "mlir/Support/LogicalResult.h"
#include "triton/Analysis/Utility.h"
#include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h"
#include "triton/Conversion/TritonGPUToLLVM/Utility.h"

#include <optional>

#include "triton/Analysis/Allocation.h"
#include "triton/Dialect/Triton/IR/Types.h"
#include "triton/Dialect/Triton/IR/Utility.h"
#include "triton/Dialect/TritonGPU/IR/Attributes.h"
#include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h"
#include "triton/Dialect/TritonGPU/Transforms/Utility.h"
#include "triton/Tools/GenericSwizzling.h"
#include "triton/Tools/LayoutUtils.h"

#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h"

namespace {

using namespace mlir;
using namespace mlir::triton::gpu;
using TranspositionInfo = DecomposedWarpConversion::TranspositionInfo;

struct ConvertLayoutOpConversion
    : public ConvertOpToLLVMPattern<ConvertLayoutOp> {
  const TargetInfoBase &targetInfo;

  explicit ConvertLayoutOpConversion(LLVMTypeConverter &typeConverter,
                                     const TargetInfoBase &targetInfo,
                                     PatternBenefit benefit = 1)
      : ConvertOpToLLVMPattern(typeConverter, benefit), targetInfo(targetInfo) {
  }

  LogicalResult
  matchAndRewrite(ConvertLayoutOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    MLIRContext *ctx = op.getContext();

    auto srcTy = op.getSrc().getType();
    auto dstTy = op.getType();

    auto kBlock = str_attr("block");
    auto kWarp = str_attr("warp");
    auto kLane = str_attr("lane");
    auto kRegister = str_attr("register");

    auto srcLayout = toLinearLayout(srcTy).removeZeroBasesAlongDim(kRegister);
    auto dstLayout = toLinearLayout(dstTy).removeZeroBasesAlongDim(kRegister);
    LinearLayout conversion = minimalCvtLayout(srcLayout, dstLayout);

    auto dims = conversion.getInDimNames();
    auto srcEnc = cast<RankedTensorType>(srcTy).getEncoding();
    auto dstEnc = cast<RankedTensorType>(dstTy).getEncoding();
    if ((isGenericLinearEncoding(srcEnc) || isGenericLinearEncoding(dstEnc)) &&
        llvm::range_size(dims) > 1)
      return op.emitError("ConvertLayoutOp  supports GenericLinearEncoding "
                          " only when the conversion is transfer between "
                          "values in the same thread.");
    assert(to_vector(conversion.getInDimNames()) ==
           to_vector(conversion.getOutDimNames()));
    if (llvm::is_contained(dims, kBlock) || llvm::is_contained(dims, kWarp)) {
      // Transfer between values in the same CTA, or across CTAs. We move values
      // through (distributed) shared memory.
      transferSwizzlingLocalMem(op, adaptor.getSrc(), srcLayout, dstLayout,
                                rewriter);
      return success();
    } else if (llvm::is_contained(dims, kLane)) {
      // Case 3. Transfer between values in the same warp, in which case we try
      //         to move values using warp shuffles, though if the pattern is
      //         expensive enough we fall back to using shared memory
      if (cvtNeedsWarpShuffle(op))
        return transferWithinWarp(op, srcLayout, dstLayout, adaptor, rewriter);

      transferSwizzlingLocalMem(op, adaptor.getSrc(), srcLayout, dstLayout,
                                rewriter);
      return success();
    } else if (llvm::is_contained(dims, kRegister)) {
      // Case 4. Transfer between values in the same thread, in which case we
      //         simply reorder the elements of adaptor.getSrc().
      return transferWithinThread(op, conversion, adaptor, rewriter);
    } else {
      // Cast 5. The two layouts are equivalent. We should probably remove
      // these in RemoveLayoutConversion.
      assert(adaptor.getSrc().getType() ==
             getTypeConverter()->convertType(dstTy));
      rewriter.replaceOp(op, adaptor.getSrc());
      return success();
    }
  }

  LogicalResult
  transferWithinThread(ConvertLayoutOp op, const LinearLayout &conversion,
                       OpAdaptor adaptor,
                       ConversionPatternRewriter &rewriter) const {
    MLIRContext *ctx = op.getContext();
    auto loc = op.getLoc();
    auto kRegister = str_attr("register");
    auto dstTy = op.getType();
    auto dims = to_vector(conversion.getInDimNames());
    assert(dims.size() == 1 && dims.front() == kRegister &&
           "expected a register-only conversion");

    auto inVals = unpackUniqueTensorElements(loc, adaptor.getSrc(), rewriter);
    SmallVector<Value> outVals(conversion.getInDimSize(kRegister));
    for (int i = 0; i < outVals.size(); i++) {
      auto srcIdx = conversion.apply({{kRegister, i}}).begin()->second;
      assert(srcIdx < inVals.size());
      outVals[i] = inVals[srcIdx];
    }
    Value result = packUniqueTensorElements(loc, getTypeConverter(), outVals,
                                            rewriter, dstTy);
    rewriter.replaceOp(op, result);
    return success();
  }

  SmallVector<Value> transferSwizzlingLocalMemImpl(
      Location loc, ConversionPatternRewriter &rewriter,
      const LinearLayout &srcLayout, const LinearLayout &dstLayout,
      ArrayRef<Value> inVals, Type llvmElemTy, Value smemBase,
      Operation *sourceOp) const {
    auto *ctx = rewriter.getContext();
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    // We handle transformations recursively as they all need a preprocessing
    // and a postprocessing step.

    // Handle pointer types as 64-bit integers
    if (isa<LLVM::LLVMPointerType>(llvmElemTy)) {
      auto llvmElemTyPtr = i64_ty;
      auto newInVals = llvm::to_vector(llvm::map_range(inVals, [&](Value v) {
        return b.ptrtoint(llvmElemTyPtr, v).getResult();
      }));
      auto outVals = transferSwizzlingLocalMemImpl(
          loc, rewriter, srcLayout, dstLayout, newInVals, llvmElemTyPtr,
          smemBase, sourceOp);
      for (auto &v : outVals) {
        v = b.inttoptr(llvmElemTy, v);
      }
      return outVals;
    }

    // Handle sub-byte elements like i1
    if (llvmElemTy.getIntOrFloatBitWidth() < 8) {
      // Upcast to i8
      auto i8ElemTy = i8_ty;
      auto newInVals = llvm::to_vector(llvm::map_range(
          inVals, [&](Value v) { return b.zext(i8ElemTy, v).getResult(); }));
      auto outVals = transferSwizzlingLocalMemImpl(
          loc, rewriter, srcLayout, dstLayout, newInVals, i8ElemTy, smemBase,
          sourceOp);
      for (auto &v : outVals) {
        v = b.trunc(llvmElemTy, v);
      }
      return outVals;
    }

    // At this point we have a type that's at least 8-bit
    // and we don't have broadcasting in the registers
    auto bitwidth = llvmElemTy.getIntOrFloatBitWidth();
    int numBanks = targetInfo.getSharedMemoryBanks();
    int32_t vecBitwidth =
        triton::gpu::getVecBitwidthLdSt(srcLayout, dstLayout, bitwidth);
    auto [dstTile, srcTile] = targetInfo.getSharedLdStTiles(vecBitwidth);
    auto smem = optimalSwizzlingLdSt(srcLayout, dstLayout, bitwidth, numBanks,
                                     srcTile, dstTile);

    // Extract reps from smem
    auto kReg = str_attr("register");
    auto kWarp = str_attr("warp");
    auto kBlock = str_attr("block");
    auto kReps = str_attr("reps");
    auto nReps = smem.getInDimSize(kReps);
    auto reps = LinearLayout::identity1D(nReps, kReg, kReps);

    auto totalStoreCvt = srcLayout.invertAndCompose(smem);
    auto totalLoadCvt = invertAndComposeBlockLocal(smem, dstLayout);

    // The permutation exists by construction of the reps dimension in
    // optimalSwizzling
    auto permStore =
        regPermForDivide(totalStoreCvt, reps, /*left=*/false).value();
    totalStoreCvt = permStore.apply(totalStoreCvt);
    auto permutedInVals = permStore.apply(inVals);
    auto permLoad =
        regPermForDivide(totalLoadCvt, reps, /*left=*/false).value();
    totalLoadCvt = permLoad.apply(totalLoadCvt);

    // Remove the reps and flatten into offset
    auto storeCvt = *divideRight(totalStoreCvt, reps);
    auto loadCvt = *divideRight(totalLoadCvt, reps);
    auto kOffset = str_attr("offset");
    auto nBlock = storeCvt.getInDimSize(kBlock);
    storeCvt = storeCvt.reshapeOuts(
        {{kOffset, storeCvt.getTotalOutDimSize() / nBlock}, {kBlock, nBlock}});
    loadCvt = loadCvt.reshapeOuts(
        {{kOffset, loadCvt.getTotalOutDimSize() / nBlock}, {kBlock, nBlock}});

    auto tileSize = storeCvt.getInDimSize(kReg);

    assert(permutedInVals.size() == tileSize * nReps);
    SmallVector<Value> outVals;
    auto affineOffset = b.i32_val(0);
    auto maskSpanAffineOffset = 0;

    bool isWarpSync = mlir::isCvtDimSync(srcLayout, dstLayout, kWarp);
    bool isBlockSync = mlir::isCvtDimSync(srcLayout, dstLayout, kBlock);
    auto emitBarrier = [&]() {
      if (isWarpSync) {
        targetInfo.warpSync(loc, rewriter);
      } else if (isBlockSync) {
        targetInfo.barrier(loc, rewriter, triton::gpu::AddrSpace::Local);
      } else {
        targetInfo.clusterBarrier(loc, rewriter, sourceOp);
      }
    };

    auto [laneId, warpId] = getLaneAndWarpId(rewriter, loc);
    for (int i = 0; i < nReps; ++i) {
      if (i > 0)
        emitBarrier();
      auto tileInVals =
          ArrayRef<Value>(permutedInVals).slice(i * tileSize, tileSize);
      // Store
      lowerLdSt(loc, ctx, storeCvt, tileInVals, llvmElemTy, smemBase,
                /*paddingShifts=*/{}, affineOffset, maskSpanAffineOffset,
                /*affineBlockOffset=*/Value(), /*maskSpanAffineBlock=*/0,
                laneId, warpId, rewriter, targetInfo, /*maybeMaxVecElems=*/{},
                makeSharedStoreEmitter(targetInfo, b.true_val()));
      emitBarrier();
      // Load
      auto tileOutVals = lowerLdSt(
          loc, ctx, loadCvt, {}, llvmElemTy, smemBase, /*paddingShifts=*/{},
          affineOffset, maskSpanAffineOffset, /*affineBlockOffset=*/Value(),
          /*maskSpanAffineBlock=*/0, laneId, warpId, rewriter, targetInfo,
          /*maybeMaxVecElems=*/{}, makeSharedLoadEmitter(targetInfo));
      llvm::append_range(outVals, tileOutVals);
    }

    // Undo the permLoad used to divideRight
    outVals = permLoad.inverse().apply(outVals);
    return outVals;
  }

  void transferSwizzlingLocalMem(ConvertLayoutOp op, Value src,
                                 const LinearLayout &srcLayout,
                                 const LinearLayout &dstLayout,
                                 ConversionPatternRewriter &rewriter) const {
    auto loc = op.getLoc();
    auto srcTy = op.getSrc().getType();
    auto dstTy = op.getType();

    auto llvmElemTy = getTypeConverter()->convertType(srcTy.getElementType());
    auto smemBase =
        LLVM::getSharedMemoryBase(loc, rewriter, targetInfo, op.getOperation());
    auto inVals = unpackUniqueTensorElements(loc, src, rewriter);
    auto outVals = transferSwizzlingLocalMemImpl(
        loc, rewriter, srcLayout, dstLayout, inVals, llvmElemTy, smemBase, op);

    Value result = packUniqueTensorElements(loc, getTypeConverter(), outVals,
                                            rewriter, dstTy);
    rewriter.replaceOp(op, result);
  }

  // Use warp shuffles to implement a layout conversion where data only needs to
  // be moved within warps.
  LogicalResult transferWithinWarp(ConvertLayoutOp op,
                                   const LinearLayout &srcLayout,
                                   const LinearLayout &dstLayout,
                                   OpAdaptor adaptor,
                                   ConversionPatternRewriter &rewriter) const {
    auto loc = op.getLoc();
    auto *ctx = op.getContext();
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    auto srcTy = op.getSrc().getType();
    auto dstTy = op.getType();
    auto kReg = str_attr("register");
    auto kLane = str_attr("lane");
    auto elemTy = getTypeConverter()->convertType(srcTy.getElementType());
    int bitwidth = getIntOrFloatOrPtrBitWidth(elemTy);

    auto factors =
        getWarpLayoutConvertDecomposition(srcLayout, dstLayout, bitwidth);
    auto &[pReg, pLane, mixedTranspositions, nPack] = factors;
    int m = mixedTranspositions.size();
    bool pLaneIsTrivial = squareSublayoutIsIdentity(pLane, kLane);
    assert((m > 0 || !pLaneIsTrivial) && "Shuffles not needed for conversion");

    // The desired layout conversion can be expressed as a permutation P of
    // hardware index bits for the `kLane` and `kReg` dimensions. The `factors`
    // of P describe a decomposition
    //
    //                 P = P_mixed \circ P_lane \circ P_reg,
    //
    // where P_reg and P_lane are permutations involving only register or only
    // lane index bits and P_mixed is a product of disjoint transpositions of
    // register index bits with lane index bits. Our goal is to implement P
    // using predicated selects and warp-shuffles. We have two tools for this:
    //  - An out-of-place `Ship` method which implements one mixed transposition
    //    at a time using 1.5 * R selects/permutes and .5 * R shuffles each.
    //  - An in-place `Swap` method which can simultaneously implement P_lane
    //    and multiple mixed transpositions at a time using 2 * m * R selects/
    //    permutes and either (1 - (1/2)^m) * R shuffles if `pLaneIsTrivial` and
    //    R shuffles otherwise.
    // Here, R denotes the number of 32-bit registers in use after packing (or
    // splitting, if applied to 64-bit types or pointers), and in the `Swap`
    // method, `m` denotes the number of mixed transpositions passed in.
    auto inVals = unpackUniqueTensorElements(loc, adaptor.getSrc(), rewriter);

    // If the target layout has a larger register dimension than the source
    // layout, then we broadcast along the register dimension to match size. The
    // removal of broadcasting above and introduction here is expected by the
    // `factors`.
    int regDim = inVals.size();
    int pRegDim = pReg.getInDimSize(kReg);
    if (pRegDim > regDim) {
      SmallVector<Value> original(inVals.begin(), inVals.end());
      inVals.clear();
      inVals.reserve(pRegDim);
      while (inVals.size() < pRegDim)
        inVals.append(original.begin(), original.end());
      regDim = pRegDim;
    }

    // Apply pReg.
    SmallVector<Value> newInVals(regDim);
    for (const auto &[i, v] : llvm::enumerate(inVals))
      newInVals[pReg.apply({{kReg, i}})[0].second] = v;
    inVals = std::move(newInVals);

    // Pack registers if possible.
    int elemsPerVec = 1 << nPack;
    int bitsPerVecElem = 32 / elemsPerVec;
    if (elemsPerVec > 1) {
      SmallVector<Value> packedVals;
      packedVals.reserve(regDim / elemsPerVec);
      if (bitwidth == 8 && bitsPerVecElem == 16) {
        // TODO: Can remove `if` part of `if-else` once ptxas bugfix lands.
        for (int i = 0; i < regDim; i += elemsPerVec) {
          Value x0 = b.zext(i32_ty, b.bitcast(inVals[i], int_ty(bitwidth)));
          Value x1 = b.zext(i32_ty, b.bitcast(inVals[i + 1], int_ty(bitwidth)));
          x1 = b.shl(x1, b.i32_val(16));
          packedVals.emplace_back(b.or_(x0, x1));
        }
      } else {
        if (bitwidth < bitsPerVecElem) {
          for (Value &v : inVals) {
            if (elemTy != int_ty(bitwidth))
              v = b.bitcast(v, int_ty(bitwidth));
            v = b.zext(int_ty(bitsPerVecElem), v);
          }
        }
        for (int i = 0; i < regDim; i += elemsPerVec) {
          auto slice = ArrayRef<Value>(inVals).slice(i, elemsPerVec);
          Value v = packLLVector(loc, slice, rewriter);
          v = b.bitcast(v, i32_ty);
          packedVals.emplace_back(v);
        }
      }
      inVals = std::move(packedVals);
    }

    auto isShippable = [](const TranspositionInfo &t) {
      // The `Ship` method cannot mix elements from different registers in the
      // same lane, so we are restricted to cycles like (l0 r1), (l0 r2), and
      // (l0 r0 r1) which do not use both high and low register bits.
      return t.topPreSel == t.topPostSel ||
             (t.topPreSel == 0x5140 && t.topPostSel == 0x6240) ||
             (t.topPreSel == 0x6420 && t.topPostSel == 0x5410) ||
             (t.topPreSel == 0x3210 && t.topPostSel == 0x3120);
    };

    SmallVector<Value> outVals;
    if (m == 1 && pLaneIsTrivial && isShippable(mixedTranspositions[0])) {
      outVals = transferWithinWarpShipImpl(loc, rewriter, inVals, nPack,
                                           mixedTranspositions[0]);
    } else {
      outVals = transferWithinWarpSwapImpl(loc, rewriter, inVals, nPack, pLane,
                                           pLaneIsTrivial, mixedTranspositions);
    }

    // Unpack registers if needed.
    if (elemsPerVec > 1) {
      SmallVector<Value> unpackedVals;
      unpackedVals.reserve(regDim);
      auto packedTy =
          bitwidth < bitsPerVecElem ? int_ty(bitsPerVecElem) : elemTy;
      auto vecTy = vec_ty(packedTy, elemsPerVec);
      auto unpackVal = [&](Value v) {
        v = b.bitcast(v, vecTy);
        return unpackLLVector(loc, v, rewriter);
      };
      for (auto v : outVals) {
        auto unpacked = unpackVal(v);
        unpackedVals.append(unpacked.begin(), unpacked.end());
      }
      if (bitwidth < bitsPerVecElem) {
        for (Value &v : unpackedVals) {
          v = b.trunc(int_ty(bitwidth), v);
          if (elemTy != int_ty(bitwidth))
            v = b.bitcast(v, elemTy);
        }
      }
      outVals = std::move(unpackedVals);
    }

    // If `dstLayout` has a smaller `kReg` dimension than `srcLayout` after
    // broadcasting is removed, then drop the extra registers from `outVals`.
    outVals.resize(dstLayout.getInDimSize(kReg));

    Value result = packUniqueTensorElements(loc, getTypeConverter(), outVals,
                                            rewriter, dstTy);
    rewriter.replaceOp(op, result);
    return success();
  }

  SmallVector<Value> transferWithinWarpSwapImpl(
      Location loc, ConversionPatternRewriter &rewriter, ArrayRef<Value> inVals,
      int nPack, const LinearLayout &pLane, bool pLaneIsTrivial,
      ArrayRef<TranspositionInfo> mixedTranspositions) const {
    auto *ctx = rewriter.getContext();
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    auto kLane = str_attr("lane");

    SmallVector<Value> vals(inVals.begin(), inVals.end());
    int numRegs = inVals.size();
    // A single mixed transposition (r_i l_j) which swaps the i-th register
    // index bit and the j-th lane index bit of an element applies a tiled 2x2
    // block transpose with block size (1 << i) by (1 << j) to the data. This
    // can be realized as:
    //
    //             [ A B ] selp [ A D ] shfl [ A D ] selp [ A C ]
    //             [ C D ] ---> [ C B ] ---> [ B C ] ---> [ B D ].
    //
    // In linear-algebraic terms, this is the factorization over GF(2):
    //
    //   1. r_i ^= l_j (selp)                     selp    shfl    selp
    //   2. l_j ^= r_i (shfl)        [ 0 1 ]     [ 1 1 ] [ 1 0 ] [ 1 1 ]
    //   3. r_i ^= l_j (selp),       [ 1 0 ]  =  [ 0 1 ] [ 1 1 ] [ 0 1 ],
    //
    // where we pass in bits as column vectors [r_i, l_j].
    //
    // When the transpositions are all disjoint, we can group the three stages
    // of each transposition together. The two combined `selp` stages each use
    // `numRegs` selects per transposition, while the `shfl` stage only requires
    // code emission when at least one of the `r_i` bits is on, resulting in
    // `(1 - (1/2)^m) * numRegs` shuffles in total. If `pLane` is nontrivial,
    // then we can conjugate its effects through the first two stages and fuse
    // it with the second stage, resulting in `numRegs` shuffles instead.
    Value laneId = getLaneId(rewriter, loc);
    auto pLaneInv = pLane.invert();
    const auto &pLInvBases = pLaneInv.getBases().lookup(kLane);

    // Implement r_i ^= l_j using `numRegs` independent selects or permutes.
    auto applySwap = [&](TranspositionInfo t, bool preShuf) {
      int rIdx = t.transposition.first - nPack;
      int origLIdx = t.transposition.second;
      int lIdx = preShuf ? llvm::Log2_32(pLInvBases[origLIdx][0]) : origLIdx;
      uint16_t topSel = preShuf ? t.topPreSel : t.topPostSel;
      uint16_t botSel = preShuf ? t.botPreSel : t.botPostSel;

      SmallVector<Value> newVals(numRegs);
      Value lBitVal = b.and_(laneId, b.i32_val(1 << lIdx));
      Value lBitOff = b.icmp_eq(lBitVal, b.i32_val(0));

      int tileSize = 1 << (rIdx + 1);
      int numTiles = numRegs / tileSize;
      for (int tileIdx = 0; tileIdx < numTiles; ++tileIdx) {
        int baseIdx = tileIdx * tileSize;
        for (int i = 0; i < tileSize / 2; ++i) {
          int r0 = baseIdx + i;
          int r1 = r0 + (1 << rIdx);
          Value v0 = vals[r0];
          Value v1 = vals[r1];
          if (topSel == 0x3210 && botSel == 0x7654) {
            newVals[r0] = b.select(lBitOff, v0, v1);
            newVals[r1] = b.select(lBitOff, v1, v0);
          } else {
            Value sel00 = b.i32_val(topSel);
            Value sel01 = b.i32_val(preShuf ? botSel : (topSel ^ 0x4444));
            Value sel10 = b.i32_val(botSel);
            Value sel11 = b.i32_val(preShuf ? topSel : (botSel ^ 0x4444));
            Value sel1 = b.select(lBitOff, sel00, sel01);
            Value sel2 = b.select(lBitOff, sel10, sel11);
            newVals[r0] = targetInfo.permute(rewriter, loc, v0, v1, sel1);
            newVals[r1] = targetInfo.permute(rewriter, loc, v0, v1, sel2);
          }
        }
      }
      return newVals;
    };

    // Stage 1 (selp/prmt)
    for (const auto &t : mixedTranspositions)
      vals = applySwap(t, /*preShuf=*/true);
    // Stage 2 (shfl)
    Value laneIdPerm;
    if (!pLaneIsTrivial)
      laneIdPerm = triton::gpu::matrixVectorProd(b, pLaneInv, laneId);
    for (int r = 0; r < numRegs; ++r) {
      int mask = 0;
      for (const auto &t : mixedTranspositions) {
        int rIdx = t.transposition.first - nPack;
        int lIdx = t.transposition.second;
        if (r & (1 << rIdx)) {
          mask |= pLInvBases[lIdx][0];
        }
      }
      if (pLaneIsTrivial) {
        if (mask != 0)
          vals[r] = targetInfo.shuffleXor(rewriter, loc, vals[r], mask);
      } else {
        Value srcIdx = b.xor_(laneIdPerm, b.i32_val(mask));
        vals[r] = targetInfo.shuffleIdx(rewriter, loc, vals[r], srcIdx);
      }
    }
    // Stage 3 (selp/prmt)
    for (const auto &t : mixedTranspositions)
      vals = applySwap(t, /*preShuf=*/false);
    return vals;
  }

  SmallVector<Value>
  transferWithinWarpShipImpl(Location loc, ConversionPatternRewriter &rewriter,
                             ArrayRef<Value> inVals, int nPack,
                             TranspositionInfo t) const {
    // Implements the effects of a single mixed transposition as in
    // `transferWithinWarpSwapImpl`, but uses auxiliary registers to hold the
    // values to be shuffled, resulting in fewer emitted instructions.
    int numRegs = inVals.size();
    int rIdx = t.transposition.first - nPack;
    int lIdx = t.transposition.second;
    int tileSize = 1 << (rIdx + 1);
    int numTiles = numRegs / tileSize;

    auto b = TritonLLVMOpBuilder(loc, rewriter);
    Value laneId = getLaneId(rewriter, loc);
    Value lBitVal = b.and_(laneId, b.i32_val(1 << lIdx));
    Value lBitOff = b.icmp_eq(lBitVal, b.i32_val(0));
    SmallVector<Value> outVals(numRegs);

    auto postShipSel = [](uint16_t preSel, uint16_t postSel, bool rBitOn) {
      // When selecting non-shipped values, we need to compose with the
      // preshuffle masks since the values were not previously permuted.
      uint16_t ret = 0;
      for (size_t k = 0; k < 4; ++k) {
        uint16_t postNib = (postSel >> (4 * k)) & 0xF;
        uint16_t nib = postNib;
        if (postNib < 4) {
          uint16_t preNib = (preSel >> (4 * postNib)) & 0xF;
          nib = rBitOn ? (preNib - 4) : preNib;
        }
        ret |= nib << (4 * k);
      }
      return ret;
    };

    for (int tileIdx = 0; tileIdx < numTiles; ++tileIdx) {
      int baseIdx = tileIdx * tileSize;
      for (int i = 0; i < tileSize / 2; ++i) {
        int r0 = baseIdx + i;
        int r1 = r0 + (1 << rIdx);
        Value v0 = inVals[r0];
        Value v1 = inVals[r1];
        if (t.topPreSel == 0x3210 && t.topPostSel == 0x3210) {
          Value valToShip = b.select(lBitOff, v1, v0);
          Value shippedVal =
              targetInfo.shuffleXor(rewriter, loc, valToShip, (1 << lIdx));
          outVals[r0] = b.select(lBitOff, v0, shippedVal);
          outVals[r1] = b.select(lBitOff, shippedVal, v1);
        } else {
          Value shipSel =
              b.select(lBitOff, b.i32_val(t.botPreSel), b.i32_val(t.topPreSel));
          Value valToShip = targetInfo.permute(rewriter, loc, v0, v1, shipSel);
          Value shippedVal =
              targetInfo.shuffleXor(rewriter, loc, valToShip, (1 << lIdx));
          uint16_t sel00 =
              postShipSel(t.topPreSel, t.topPostSel, /*rBitOn=*/false);
          uint16_t sel01 =
              postShipSel(t.botPreSel, t.topPostSel ^ 0x4444, /*rBitOn=*/false);
          uint16_t sel10 =
              postShipSel(t.topPreSel, t.botPostSel, /*rBitOn=*/true);
          uint16_t sel11 =
              postShipSel(t.botPreSel, t.botPostSel ^ 0x4444, /*rBitOn=*/true);
          Value sel1 = b.select(lBitOff, b.i32_val(sel00), b.i32_val(sel01));
          Value sel2 = b.select(lBitOff, b.i32_val(sel10), b.i32_val(sel11));
          outVals[r0] = targetInfo.permute(rewriter, loc, v0, shippedVal, sel1);
          outVals[r1] = targetInfo.permute(rewriter, loc, v1, shippedVal, sel2);
        }
      }
    }
    return outVals;
  }
};

} // namespace

void mlir::triton::populateConvertLayoutOpToLLVMPatterns(
    LLVMTypeConverter &typeConverter, const TargetInfoBase &targetInfo,
    RewritePatternSet &patterns, PatternBenefit benefit) {
  patterns.add<ConvertLayoutOpConversion>(typeConverter, targetInfo, benefit);
}
