#include "Utility.h"
#include "AsyncUtility.h"
#include "Dialect/TritonAMDGPU/IR/Dialect.h"
#include "TritonAMDGPUToLLVM/GCNAsmFormat.h"
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/PatternMatch.h"
#include "triton/Conversion/TritonGPUToLLVM/Utility.h"
#include "triton/Dialect/Triton/IR/Dialect.h"
#include "triton/Dialect/Triton/IR/Types.h"
#include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h"
namespace tt = mlir::triton;
using mlir::triton::ModuleAxisInfoAnalysis;
using mlir::triton::amdgpu::ISAFamily;
using mlir::triton::gpu::appendOrGetExternFuncOp;

namespace mlir::LLVM::AMD {
namespace {

enum class ShflKind : uint32_t {
  bfly = 0,
  up = 1,
  idx = 2,
};

Value emitDpp(Location loc, RewriterBase &rewriter, Value old, Value src,
              DppCtrl dppCtrl, uint32_t rowMask = 0xf, uint32_t bankMask = 0xf,
              bool boundCtrl = true) {
  return ROCDL::DPPUpdateOp::create(
      rewriter, loc, src.getType(), old, src,
      rewriter.getI32IntegerAttr(static_cast<uint32_t>(dppCtrl)),
      rewriter.getI32IntegerAttr(rowMask), rewriter.getI32IntegerAttr(bankMask),
      rewriter.getBoolAttr(boundCtrl));
}

Value emitPermlaneX16Xor(Location loc, RewriterBase &rewriter, Value val,
                         uint32_t rowMask) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  // PermlaneX16 reads the opposite 16-lane half of a 32-lane wave. It uses a
  // 16-nibble selector to choose the source lane within that half.
  // We use it to perform a shuffleXor with mask `rowMask ^ 16`.
  assert(rowMask < 16 && "Expected a cross-row shuffleXor");
  auto buildSelectorMask = [&](unsigned startLane) {
    uint32_t sel = 0;
    for (unsigned lane = 0; lane < 8; ++lane)
      sel |= ((startLane + lane) ^ rowMask) << (lane * 4);
    return sel;
  };
  Value loSel = b.i32_val(buildSelectorMask(0));
  Value hiSel = b.i32_val(buildSelectorMask(8));
  return ROCDL::PermlaneX16Op::create(rewriter, loc, val.getType(), val, val,
                                      loSel, hiSel, true, false);
}

Value emitPermlaneSwapXor(Location loc, RewriterBase &rewriter, Value val,
                          uint32_t laneMask) {
  assert((laneMask == 16 || laneMask == 32) && "Expected a power of 2 mask");
  assert(val.getType().isInteger(32) && "permlane_swap operates on i32 values");
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  MLIRContext *ctx = rewriter.getContext();
  const char *intrinsic = laneMask == 32 ? "llvm.amdgcn.permlane32.swap"
                                         : "llvm.amdgcn.permlane16.swap";
  Type retTy = struct_ty({i32_ty, i32_ty});
  Value f = b.false_val();
  Value perm = LLVM::createLLVMIntrinsicCallOp(rewriter, loc, intrinsic, retTy,
                                               ValueRange{val, val, f, f})
                   ->getResult(0);
  Value v0 = b.extract_val(i32_ty, perm, 0);
  Value v1 = b.extract_val(i32_ty, perm, 1);
  Value isOriginalVal = b.icmp_eq(v0, val);
  return b.select(isOriginalVal, v1, v0);
}

Value shuffleCommonImpl(Location loc, RewriterBase &rewriter,
                        ISAFamily isaFamily, Value val, Value i, int strideInt,
                        ShflKind mode, Value clamp) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  unsigned bits = val.getType().getIntOrFloatBitWidth();

  // On AMD, the ds_swizzle_b32 and ds_permute_b32 instructions work on
  // 32bit/dwords so we need promote to 32 here.
  auto valType = val.getType();
  if (!valType.isInteger(32) && bits <= 32) {
    if (!valType.isIntOrIndex())
      val = b.bitcast(val, int_ty(bits));
    if (bits < 32)
      val = b.sext(i32_ty, val);

    val = shuffleCommonImpl(loc, rewriter, isaFamily, val, i, strideInt, mode,
                            clamp);

    if (bits < 32)
      val = b.trunc(int_ty(bits), val);
    if (!valType.isIntOrIndex())
      val = b.bitcast(val, valType);
    return val;
  }

  if (bits == 64) {
    Type vecTy = vec_ty(f32_ty, 2);
    Value vec = b.bitcast(val, vecTy);
    Value val0 = b.extract_element(f32_ty, vec, b.i32_val(0));
    Value val1 = b.extract_element(f32_ty, vec, b.i32_val(1));
    val0 = shuffleCommonImpl(loc, rewriter, isaFamily, val0, i, strideInt, mode,
                             clamp);
    val1 = shuffleCommonImpl(loc, rewriter, isaFamily, val1, i, strideInt, mode,
                             clamp);
    vec = b.undef(vecTy);
    vec = b.insert_element(vecTy, vec, val0, b.i32_val(0));
    vec = b.insert_element(vecTy, vec, val1, b.i32_val(1));
    return b.bitcast(vec, val.getType());
  }

  auto mod = rewriter.getBlock()->getParent()->getParentOfType<ModuleOp>();
  Value threadId = getThreadId(rewriter, loc);

  unsigned iWarpSize = triton::gpu::TritonGPUDialect::getThreadsPerWarp(mod);
  Value warpSize = b.i32_val(iWarpSize);
  Value laneId = b.urem(threadId, warpSize);
  auto bpermute = [&](Value lane) {
    // Multiple lineId by 4. (More on permute instruction semantics:
    // https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/instinct-mi200-cdna2-instruction-set-architecture.pdf#page=180
    Value byteOffset = b.i32_val(2);
    Value permuteAddr = b.shl(lane, byteOffset);
    return ROCDL::DsBpermuteOp::create(rewriter, loc, valType, permuteAddr,
                                       val);
  };

  switch (mode) {
  case ShflKind::bfly: {
    // We have the following tools to decompose shuffleXor into DPP primitives.
    // On CDNA:
    //  - xor by 15  : reflect the lanes within each group of 16 lanes,
    //  - xor by 8   : right rotate by 8 within each group of 16 lanes,
    //  - xor by 7   : reflect the lanes within each group of 8 lanes,
    //  - xor by 1-3 : perform a permutation within each group of 4 lanes.
    // On RDNA:
    //  - xor by 1-15 : row_xmask does exactly this.
    //
    // On RDNA, we also have permlanex16 which can perform xor by 16-31.
    // On CDNA, one can see that any shuffleXor with `mask` in [1, 15] can
    // be implemented using at most 2 DPP instructions by mapping the upper bits
    // of `mask` to the `xor by 7` and `xor by 8` instructions.
    uint32_t mask = strideInt;
    if (mask == 0)
      return val;
    assert(mask < iWarpSize &&
           "shuffle_xor expects a mask within the wave size");

    auto makeDppCtrl = [](DppCtrl base, uint32_t arg) {
      return static_cast<DppCtrl>(llvm::to_underlying(base) + arg);
    };
    auto makeQuadPermCtrl = [](uint32_t quadMask) {
      DppCtrl ctrl = DppCtrl::QUAD_PERM_FIRST;
      uint32_t ctrlBits = llvm::to_underlying(ctrl);
      for (unsigned lane = 0; lane < 4; ++lane)
        ctrlBits |= (lane ^ quadMask) << (lane * 2);
      return static_cast<DppCtrl>(ctrlBits);
    };
    bool usePermlaneSwap = isaFamily == ISAFamily::CDNA4 && (mask & 0x30);

    if (triton::amdgpu::isRDNA(isaFamily) || isaFamily == ISAFamily::GFX1250) {
      if (mask < 16)
        return emitDpp(loc, rewriter, val, val,
                       makeDppCtrl(DppCtrl::ROW_XMASK0, mask));
      else if (mask < 32)
        return emitPermlaneX16Xor(loc, rewriter, val, mask & 0xf);
    } else if ((triton::amdgpu::isCDNA(isaFamily) ||
                isaFamily == ISAFamily::GCN5_1) &&
               (mask < 16 || usePermlaneSwap)) {
      Value result = val;
      uint32_t highBitsDppBasis = 0;

      if (mask & 4)
        highBitsDppBasis ^= 7;
      if (mask & 8)
        highBitsDppBasis ^= 8;

      uint32_t quadMask = (mask & 0xf) ^ highBitsDppBasis;

      if (highBitsDppBasis) {
        DppCtrl highBitsDppCtrl;
        switch (highBitsDppBasis) {
        case 0x7:
          highBitsDppCtrl = DppCtrl::ROW_HALF_MIRROR;
          break;
        case 0x8:
          highBitsDppCtrl = makeDppCtrl(DppCtrl::ROW_ROR0, 8);
          break;
        case 0xf:
          highBitsDppCtrl = DppCtrl::ROW_MIRROR;
          break;
        }
        result = emitDpp(loc, rewriter, result, result, highBitsDppCtrl);
      }

      if (quadMask) {
        result =
            emitDpp(loc, rewriter, result, result, makeQuadPermCtrl(quadMask));
      }

      if (usePermlaneSwap) {
        if (mask & 16)
          result = emitPermlaneSwapXor(loc, rewriter, result, 16);
        if (mask & 32)
          result = emitPermlaneSwapXor(loc, rewriter, result, 32);
      }
      return result;
    } else {
      return bpermute(b.xor_(laneId, b.i32_val(mask)));
    }
  }
  case ShflKind::up: {
    Value mask = b.icmp_slt(laneId, i);
    Value delta = b.sub(laneId, i);
    Value index = b.select(mask, laneId, delta);
    return bpermute(index);
  }
  case ShflKind::idx:
    return bpermute(i);
  default:
    assert(false && "Unsupported ShflKind");
    break;
  }
  return Value();
}

Value shuffleCommon(Location loc, RewriterBase &rewriter, ISAFamily isaFamily,
                    Value val, Value i, int strideInt, ShflKind mode,
                    Value clamp) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  // To shuffle pointers, convert them to i64.
  Type valTy = val.getType();
  if (isa<LLVM::LLVMPointerType>(valTy))
    val = b.ptrtoint(i64_ty, val);
  Value result = shuffleCommonImpl(loc, rewriter, isaFamily, val, i, strideInt,
                                   mode, clamp);
  if (isa<LLVM::LLVMPointerType>(valTy))
    result = b.inttoptr(valTy, result);
  return result;
}

} // namespace

Value shuffleXor(Location loc, RewriterBase &rewriter, Value val, int i,
                 ISAFamily isaFamily) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  return shuffleCommon(loc, rewriter, isaFamily, val, b.i32_val(i), i,
                       ShflKind::bfly, b.i32_val(0x1f));
}

Value shuffleUp(Location loc, RewriterBase &rewriter, Value val, int i,
                ISAFamily isaFamily) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  return shuffleCommon(loc, rewriter, isaFamily, val, b.i32_val(i), i,
                       ShflKind::up, b.i32_val(0x0));
}

Value shuffleIdx(Location loc, RewriterBase &rewriter, Value val, int i,
                 ISAFamily isaFamily) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  return shuffleIdx(loc, rewriter, val, b.i32_val(i), isaFamily);
}

Value shuffleIdx(Location loc, RewriterBase &rewriter, Value val, Value i,
                 ISAFamily isaFamily) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  return shuffleCommon(loc, rewriter, isaFamily, val, i, 0, ShflKind::idx,
                       b.i32_val(0x1f));
}

Value permute(Location loc, RewriterBase &rewriter, Value x, Value y,
              Value selector) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  Value prmt_mask = selector;
  // convert from nybble mask to byte mask:
  prmt_mask =
      b.or_(b.and_(prmt_mask, b.i32_val(0x000000ff)),
            b.shl(b.and_(prmt_mask, b.i32_val(0x0000ff00)), b.i32_val(8)));
  prmt_mask =
      b.or_(b.and_(prmt_mask, b.i32_val(0x000f000f)),
            b.shl(b.and_(prmt_mask, b.i32_val(0x00f000f0)), b.i32_val(4)));
  Value args[] = {x, y, prmt_mask};
  auto op = createLLVMIntrinsicCallOp(rewriter, loc, "llvm.amdgcn.perm", i32_ty,
                                      args);
  return op.getResult(0);
}

// Utility function that returns flags <volatile, nontemporal> for a predicated
// Load or Store
// ---------------------------------
// Op   | cm  | volatile | NT
// -----+-----+---------------------
// Load | .ca |   F      | F
//      | .cg |   F      | T
//      | .cs |   F      | T
//      | .cv |   T      | X
// -----+-----+----------+---------
// Store| .wb |   F      | F
//      | .cg |   F      | F
//      | .cs |   F      | T
//      | .wt |   T      | X
// -----+-----+----------+---------
std::pair<bool, bool>
getCacheModifierFlagsForLoadStore(const triton::CacheModifier &cm,
                                  MemoryOp op) {
  switch (op) {
  case MemoryOp::Load: {
    switch (cm) {
    case triton::CacheModifier::CA:
      return std::make_pair(false, false);
    case triton::CacheModifier::CG:
      return std::make_pair(false, true);
    case triton::CacheModifier::CS:
      return std::make_pair(false, true);
    case triton::CacheModifier::CV:
      return std::make_pair(true, true);
    default:
      return std::make_pair(false, false);
    }
  }
  case MemoryOp::Store: {
    switch (cm) {
    case triton::CacheModifier::CG:
      return std::make_pair(false, false);
    case triton::CacheModifier::CS:
      return std::make_pair(false, true);
    case triton::CacheModifier::WT:
      return std::make_pair(true, true);
    default:
      return std::make_pair(false, false);
    }
  }
  }
  return std::make_pair(false, false);
}

Value llGetPid(Location loc, RewriterBase &rewriter, ModuleOp moduleOp,
               ProgramIDDim axis) {
  assert(moduleOp);

  int numCTAs = triton::gpu::TritonGPUDialect::getNumCTAs(moduleOp);
  if (numCTAs == 1) {
    // For single CTA the block id is the program id
    Value blockId = ::mlir::gpu::BlockIdOp::create(rewriter, loc,
                                                   mlir::gpu::Dimension(axis));
    return arith::IndexCastOp::create(rewriter, loc, i32_ty, blockId);
  }
  // For multiple CTAs the cluster id is the program id
  Type resType = rewriter.getI32Type();
  Value clusterIdx = nullptr;
  switch (axis) {
  case ProgramIDDim::X: {
    clusterIdx = ROCDL::ClusterIdXOp::create(rewriter, loc, resType);
    break;
  }
  case ProgramIDDim::Y: {
    clusterIdx = ROCDL::ClusterIdYOp::create(rewriter, loc, resType);
    break;
  }
  case ProgramIDDim::Z: {
    clusterIdx = ROCDL::ClusterIdZOp::create(rewriter, loc, resType);
    break;
  }
  }
  return clusterIdx;
}

// For multicast memory operations (e.g., cluster.load.async.to.lds), we need a
// bitmask indicating which CTAs in the CGA/cluster will access the same memory
// addresses. This allows the hardware to efficiently broadcast data to multiple
// CTAs. The linear layout's free variables in the block dimension tell us which
// CTAs form a "communication group" (i.e., access the same data):
//   - Free bit at position k: CTAs whose IDs differ only in bit k access
//     the same data and should be in the same multicast group.
//   - Fixed bits (non-free): Distinguish between different groups that
//     access different data.
// The multicast mask has bit i set if CTA i is in the same communication
// group as the current CTA. The free bits determine a groupMask whereas the
// non-free bits determine the group offset:
//   ctaMask = groupMask << groupOffset
// where:
//   - groupMask: Covers all 2^k CTAs in the group (k = number of free bits)
//   - groupOffset: Starting position of this group, determined by fixed bits
// As an example suppose we have 8 CTAs and freeVarMask = 0b101 (bits 0,2 free).
// This creates 2 groups of 4 CTAs each:
//   - Group 0: CTAs {0,1,4,5} (fixed bits = 0b000)
//   - Group 1: CTAs {2,3,6,7} (fixed bits = 0b010)
// For CTA 5 (0b101): groupOffset = 0b101 & 0b010 = 0 => ctaMask = 0b00110011
// For CTA 7 (0b111): groupOffset = 0b111 & 0b010 = 2 => ctaMask = 0b11001100
// If the group does not fit into `maxMaskPopcount` mask bits it is split into
// equally sized subgroups, see below.
Value emitCtaMulticastMask(RewriterBase &rewriter, Location loc, Value groupId,
                           const LinearLayout &regLayout,
                           unsigned maxMaskPopcount) {
  TritonLLVMOpBuilder b(loc, rewriter);

  auto kBlock = StringAttr::get(rewriter.getContext(), "block");
  auto freeVarMask = regLayout.getFreeVariableMasks()[kBlock];

  // If there are no free bits we do not share any data with other CTAs
  if (freeVarMask == 0) {
    return Value();
  }

  assert(maxMaskPopcount > 0 && "expected a non-zero multicast mask limit");

  // A communication group spans 2^(number of free bits) CTAs, so we can only
  // keep floor(log2(limit)) free bits. Dropping the highest free bits turns
  // them into subgroup selectors.
  // Example: freeVarMask = 0b1101 describes CTAs {0,1,4,5,8,9,12,13}; with a
  // limit of 5 we keep bits 0 and 2 so bit 3 selects {0,1,4,5} or {8,9,12,13}.
  int maxFreeBits = llvm::Log2_32(maxMaskPopcount);
  int multicastFreeVarMask = freeVarMask;
  while (llvm::popcount<uint32_t>(multicastFreeVarMask) > maxFreeBits)
    multicastFreeVarMask ^= llvm::bit_floor<uint32_t>(multicastFreeVarMask);

  if (multicastFreeVarMask != freeVarMask) {
    unsigned numSharingCTAs = 1u << llvm::popcount<uint32_t>(freeVarMask);
    unsigned numMulticastCTAs = 1u << maxFreeBits;
    emitRemark(loc) << "Multicast group contains " << numSharingCTAs
                    << " workgroups, exceeding the hardware limit of "
                    << maxMaskPopcount << " set mask bits; splitting it into "
                    << "subgroups of " << numMulticastCTAs
                    << " workgroups. A cluster layout sharing data "
                    << "among at most " << numMulticastCTAs
                    << " workgroups may perform better.";
  }

  // Construct the groupMask with 1s at all positions representing CTAs in the
  // (sub)group. We start with 0b1 and iterate over the retained free bits. For
  // every retained bit at position k, we copy the current pattern 2^k positions
  // higher.
  // Example for multicastFreeVarMask = 0b101, x = non determined yet:
  //   Initial:              groupMask = 0bxxxxxxx1 (positions {0})
  //   Bit 0 (retained):     groupMask = 0bxxxxxx11 (positions {0,1})
  //   Bit 1 (not retained): groupMask = 0bxxxx0011 (positions {0,1})
  //   Bit 2 (retained):     groupMask = 0b00110011 (positions {0,1,4,5})
  int groupMask = 1;
  for (int log2 = 0; log2 < regLayout.getInDimSizeLog2(kBlock); log2++) {
    if (!(multicastFreeVarMask & (1 << log2)))
      continue;
    groupMask = groupMask | (groupMask << (1 << log2));
  }
  // If all bits are retained we broadcast to all CTAs so return the group mask.
  if (multicastFreeVarMask == regLayout.getInDimSize(kBlock) - 1) {
    return b.i32_val(groupMask);
  }
  // The bits not retained (non-free bits plus subgroup selectors) set in the
  // ctaId determine the group offset. For every such bit set at position k, we
  // shift the groupMask by 2^k positions. This can be conveniently computed by
  // masking the ctaId with the inverse of multicastFreeVarMask. Note that this
  // never carries into a retained position because the shift amount and the
  // groupMask cover disjoint bits.
  // Example1: multicastFreeVarMask = 0b101
  //   ~multicastFreeVarMask = 0b010
  //   shiftAmount           = 0b101 & 0b010 = 0b000 (no shift needed)
  //   blockMask             = 0b110011 << 0 = 0b00110011
  // Example2: multicastFreeVarMask = 0b101, ctaId = 0b111 (cta 7)
  //   ~multicastFreeVarMask = 0b010
  //   shiftAmount           = 0b111 & 0b010 = 0b010 (shift by 2)
  //   blockMask             = 0b110011 << 2 = 0b11001100
  Value shiftAmount = b.and_(groupId, b.i32_val(~multicastFreeVarMask));
  Value ctaMask = b.shl(b.i32_val(groupMask), shiftAmount);
  return ctaMask;
}

Value llLoad(RewriterBase &rewriter, Location loc, Value ptr, Type elemTy,
             Value pred, Value falseVal, Value multicastMask,
             triton::CacheModifier cm, bool forceNoAliasAsyncLoads) {
  return triton::amdgpu::MaskedLoadOp::create(rewriter, loc, elemTy, ptr, pred,
                                              falseVal, multicastMask, cm,
                                              forceNoAliasAsyncLoads)
      .getResult();
}

void llStore(RewriterBase &rewriter, Location loc, Value ptr, Value val,
             Value pred, triton::CacheModifier cm,
             bool forceNoAliasAsyncLoads) {
  triton::amdgpu::MaskedStoreOp::create(rewriter, loc, ptr, val, pred, cm,
                                        forceNoAliasAsyncLoads);
}

// Create the auxiliary/cachepolicy value of ROCDL::RawPtrBufferLoad/StoreOp
//   CDNA3 and CDNA4: bit 0 = sc0, bit 1 = nt, bit 3 = swz, bit 4 = sc1
// Vector Memory instructions (Flat, Global, Scratch, and Buffer) have 3
// bits to control scope and cacheability:
// - SC[1:0] System Cache level: 0=wave, 1=group, 2=device, 3=system
// - NT Non-Temporal: 0=expect temporal reuse; 1=do not expect temporal reuse
//
// -------+-----+-----+-----+----+--
// Op     | cm  | SC1 | SC0 | NT |
// -------+-----+-----+-----+----+--
// Load   | .ca |  0  |  0  | 0  |
//        | .cg |  0  |  1  | 1  |
//        | .cs |  0  |  1  | 1  |
//        | .cv |  1  |  1  | x  |
// -------+-----+-----+-----+----+--
// Store  | .wb |  0  |  0  | 0  |
//        | .cg |  0  |  0  | 0  |
//        | .cs |  0  |  1  | 1  |
//        | .wt |  1  |  1  | x  |
// -------+-----+-----+-----+----+--
// Atomic | N/A |  0  |  1  | x  | Setting sc0 returns the pre-op value
//        | N/A |  1  |  0  | x  | Setting sc1 performs a system-scope atomic
// -------+-----+-----+-----+----+--
static int32_t
getCtrlBitsForCacheModifierOn_CDNA3_CDNA4(triton::CacheModifier cm,
                                          bool isLoad) {
  const int sc0Bit = 0b1, ntBit = 0b10, sc1Bit = 0b10000;
  int32_t aux = 0;
  switch (cm) {
  case triton::CacheModifier::CA:
    aux = 0;
    break;
  case triton::CacheModifier::CG:
    if (isLoad)
      aux |= sc0Bit | ntBit;
    break;
  case triton::CacheModifier::CS:
    aux |= sc0Bit | ntBit;
    break;
  case triton::CacheModifier::CV:
    assert(isLoad);
    aux |= sc0Bit | sc1Bit;
    break;
  case triton::CacheModifier::WB:
    assert(!isLoad);
    aux = 0;
    break;
  case triton::CacheModifier::WT:
    assert(!isLoad);
    aux |= sc0Bit | sc1Bit;
    break;
  default:
    aux = 0;
  }
  return aux;
}

// Create the auxiliary/cachepolicy value for RDNA3 (and 3.5) buffer ops
//
// LLVM CPol bit layout (from SIDefines.h):
//   bit 0 = GLC  (Graphics L1 cache, GL1)
//   bit 1 = SLC  (Graphics L2 cache, GL2)
//   bit 2 = DLC  (Memory-Attached Last-Level cache, MALL)
//
// Hardware semantics on RDNA3:
//   GLC=1  forces MISS_EVICT at GL1 (bypass / invalidate L1 line).
//   SLC=1  sets GL2 to STREAM (evict-first) policy, limiting L2 pollution.
//   DLC=1  non-temporal hint for MALL: data is not expected to be reused,
//          so skip MALL allocation. Ignored if MALL is not present.
//   DLC=0  temporal (default): allow normal MALL caching.
//
// Cache modifier definitions (PTX/Triton semantics):
//   .ca  Cache at All levels (L1+L2+MALL). Default load policy, LRU eviction.
//   .cg  Cache at Global level only (L2 and below, bypass L1).
//   .cs  Cache Streaming, likely accessed once. Evict-first to limit pollution.
//   .cv  Cache Volatile. Invalidates matching L2 line and re-fetches.
//   .wb  Write-Back at all cache levels (default store policy).
//   .wt  Write-Through, writes data directly to system memory.
//
// -------+-----+-----+-----+-----+-----------
// Op     | cm  | DLC | SLC | GLC | aux value
// -------+-----+-----+-----+-----+-----------
// Load   | .ca |  0  |  0  |  0  |  0
//        | .cg |  0  |  0  |  1  |  1
//        | .cs |  1  |  1  |  1  |  7
//        | .cv |  1  |  1  |  1  |  7
// -------+-----+-----+-----+-----+-----------
// Store  | .wb |  0  |  0  |  0  |  0
//        | .cg |  0  |  0  |  0  |  0
//        | .cs |  1  |  1  |  1  |  7
//        | .wt |  1  |  1  |  1  |  7
// -------+-----+-----+-----+-----+-----------
//
// Condensed hardware behavior:
//   DLC=0, SLC=0, GLC=0  -> temporal at all levels (MALL+GL2+GL1) (.ca/.wb)
//   DLC=0, SLC=0, GLC=1  -> bypass GL1, cache in GL2              (.cg load)
//   DLC=0, SLC=0, GLC=0  -> default (no special flags)            (.cg store)
//   DLC=1, SLC=1, GLC=1  -> non-temporal: bypass GL1, stream GL2,
//                           skip MALL                             (.cs/.cv/.wt)
static int32_t getCtrlBitsForCacheModifierOnRDNA3(triton::CacheModifier cm,
                                                  bool isLoad) {
  const int glcBit = 0b1, slcBit = 0b10, dlcBit = 0b100;
  int32_t aux = 0;
  switch (cm) {
  case triton::CacheModifier::CA:
    aux = 0;
    break;
  case triton::CacheModifier::CG:
    if (isLoad)
      aux |= glcBit;
    break;
  case triton::CacheModifier::CS:
    aux |= glcBit | slcBit | dlcBit;
    break;
  case triton::CacheModifier::CV:
    assert(isLoad);
    aux |= glcBit | slcBit | dlcBit;
    break;
  case triton::CacheModifier::WB:
    assert(!isLoad);
    aux = 0;
    break;
  case triton::CacheModifier::WT:
    assert(!isLoad);
    aux |= glcBit | slcBit | dlcBit;
    break;
  default:
    aux = 0;
  }
  return aux;
}

// Create the auxiliary/cache policy value for GFX12 family
// ROCDL::RawPtrBufferLoad/StoreOp Vector Memory instructions (Flat, Global,
// Scratch, and Buffer). These instructions have 2 control bits for scope and 3
// control bits for temporal hint:
// - SCOPE[1:0] System Cache level:
//    0 CU  Coherent among all CU/WG threads in L1 cache
//    1 SE  Coherent among all clients (threads) sharing a SE-cache(L2)
//    2 DEV Coherent among all threads on the same device
//    3 SYS Coherent in system
// - TH[2:0] Temporal Hint for load:
//    0 RT    regular temporal for both near and far caches
//    1 NT    non-temporal (re-use not expected) for both near and far caches
//    2 HT    High-priority temporal for both near and far caches
//    3 LU    Last-use (non-temporal AND discard dirty if it hits)
//    4 NT_RT non-temporal for near cache(s) and regular for far caches
//    5 RT_NT regular for near cache(s) and non-temporal for far caches
//    6 NT_HT non-temporal for near cache(s) and high-priority for far caches
//    7       reserved
// - TH[2:0] Temporal Hint for store:
//    0 RT    regular temporal for both near and far caches
//    1 NT    non-temporal (re-use not expected) for both near and far caches
//    2 HT    High-priority temporal for both near and far caches
//    3 WB    Same as "HT", but also overrides wr-rinse in far cache
//    4 NT_RT non-temporal for near cache(s) and regular for far caches
//    5 RT_NT regular for near cache(s) and non-temporal for far caches
//    6 NT_HT non-temporal for near cache(s) and HT for far caches
//    7 NT_WB non-temporal for near cache(s) and WB for far cache
//
// See detailed bit mapping of control bits of CPol enum
// in llvm source code: llvm/lib/Target/AMDGPU/SIDefines.h
//
// Mapping between
// -------+-----+-------+----+-
// Op     | cm  | SCOPE | TH |
// -------+-----+-------+----+-
// Load   | .ca |  CU   | RT |
//        | .cg |  DEV  | RT |
//        | .cs |  CU   | NT |
//        | .cv |  SYS  | LU | on gfx1250, this combination bypasses all caches
// -------+-----+-------+----+-
// Store  | .wb |  CU   | RT |
//        | .cg |  DEV  | RT |
//        | .cs |  CU   | NT |
//        | .wt |  SYS  | RT | behavior for gfx12
//        | .wt |  SYS  | WB | behavior for gfx1250, bypasses all caches
// -------+-----+-------+----+-
static int32_t getCtrlBitsForCacheModifierOn_GFX12(triton::CacheModifier cm,
                                                   bool isLoad,
                                                   bool cacheBypassAvailable) {
  const int scopeShift = 3;
  const int scopeCU = 0 << scopeShift;
  const int scopeDev = 2 << scopeShift;
  const int scopeSys = 3 << scopeShift;
  const int THRegular = 0;
  const int THNonTemp = 1;
  const int THLastUse = 3;
  const int THWriteBack = 3;
  int aux = -1;
  switch (cm) {
  case triton::CacheModifier::CA:
    aux = scopeCU | THRegular;
    break;
  case triton::CacheModifier::CG:
    aux = scopeDev | THRegular;
    break;
  case triton::CacheModifier::CS:
    aux = scopeCU | THNonTemp;
    break;
  case triton::CacheModifier::CV:
    assert(isLoad);
    aux = scopeSys | THLastUse;
    break;
  case triton::CacheModifier::WB:
    assert(!isLoad);
    aux = scopeCU | THRegular;
    break;
  case triton::CacheModifier::WT:
    assert(!isLoad);
    aux = scopeSys | (cacheBypassAvailable ? THWriteBack : THRegular);
    break;
  default:
    aux = 0;
  }
  return aux;
}

static int32_t getDefaultCtrlBitsForCacheModifier(triton::CacheModifier cm) {
  return 0;
}

// Cache modifiers changes how data is managed in the GPU's cache hierarchy:
// .ca: cache at all levels with LRU policy
// .cg: cache at L2, can use .ca or .cs
// .cs: cache streaming, use data once
// .cv: don't cache and fetch again
// .wb: write-back, writes back data at all cache levels
// .wt: write-through, write data directly to system memory
int32_t getCtrlBitsForCacheModifierOnTarget(
    triton::CacheModifier cm, bool isLoad,
    const mlir::triton::AMD::TargetInfo &targetInfo) {
  switch (targetInfo.getISAFamily()) {
  case ISAFamily::CDNA3:
  case ISAFamily::CDNA4:
    return getCtrlBitsForCacheModifierOn_CDNA3_CDNA4(cm, isLoad);
  case ISAFamily::RDNA3:
  case ISAFamily::RDNA4m:
    return getCtrlBitsForCacheModifierOnRDNA3(cm, isLoad);
  case ISAFamily::RDNA4:
    return getCtrlBitsForCacheModifierOn_GFX12(cm, isLoad, /*$ bypass*/ false);
  case ISAFamily::GFX1250:
    return getCtrlBitsForCacheModifierOn_GFX12(cm, isLoad, /*$ bypass*/ true);
  default:
    return getDefaultCtrlBitsForCacheModifier(cm);
  }
}

Type getPointerTypeWithShape(Value basePtr, Value offset) {
  Type basePtrType = basePtr.getType();
  auto offsetType = cast<RankedTensorType>(offset.getType());
  return offsetType.cloneWith(std::nullopt, basePtrType);
}

unsigned getContiguity(Value ptr, ModuleAxisInfoAnalysis &axisAnalysisPass) {
  auto tensorTy = dyn_cast<RankedTensorType>(ptr.getType());
  if (!tensorTy)
    return 1;
  return axisAnalysisPass.getContiguity(ptr);
}

unsigned getContiguity(Value ptr, Value offset,
                       ModuleAxisInfoAnalysis &axisAnalysisPass) {

  Type type = getPointerTypeWithShape(ptr, offset);
  RankedTensorType tensorTy = cast<RankedTensorType>(type);

  // To compute the contiguity of the scalar/warp-uniform ptr and offset pair we
  // need to look at the contiguity of the offsets and the alignment of the ptr
  auto elemNumBits = triton::getPointeeBitWidth(tensorTy);
  auto contiguity = axisAnalysisPass.getContiguity(offset, elemNumBits);

  // To get the alignment of the scalar ptr we need to look at the divisibility
  auto *axisInfo = axisAnalysisPass.getAxisInfo(ptr);
  auto maxMultipleBytes = axisInfo->getDivisibility(0);
  auto elemNumBytes = std::max<unsigned>(elemNumBits / 8, 1);
  auto align = std::max<unsigned>(maxMultipleBytes / elemNumBytes, 1);

  // FIXME (Alex): this should not be needed anymore because it's done inside
  // getContiguity, but we have an order issues with LL, so we keep this
  // until the LL order issue is fixed
  auto order = triton::gpu::getOrder(tensorTy);
  auto contigPerThread = triton::gpu::getContigPerThread(tensorTy);
  assert(order[0] < contigPerThread.size() &&
         "Unexpected contigPerThread size");
  contiguity = std::min(contiguity, contigPerThread[order[0]]);

  // Final contiguity is a min of the offset contiguity and pointer alignment
  return std::min(align, contiguity);
}

unsigned getVectorSize(Value ptr, ModuleAxisInfoAnalysis &axisAnalysisPass) {
  auto tensorTy = dyn_cast<RankedTensorType>(ptr.getType());
  if (!tensorTy)
    return 1;
  auto contiguity = getContiguity(ptr, axisAnalysisPass);
  auto pointeeBitWidth = triton::getPointeeBitWidth(tensorTy);
  return std::min<unsigned>(128 / pointeeBitWidth, contiguity);
}

unsigned getVectorSize(Value ptr, Value offset,
                       ModuleAxisInfoAnalysis &axisAnalysisPass) {
  auto contiguity = getContiguity(ptr, offset, axisAnalysisPass);
  auto pointeeBitWidth = triton::getPointeeBitWidth(ptr.getType());
  return std::min<unsigned>(128 / pointeeBitWidth, contiguity);
}

Type scaleDotElemTypeToMLIRType(MLIRContext *ctx, triton::ScaleDotElemType t) {
  switch (t) {
  case triton::ScaleDotElemType::FP16:
    return Float16Type::get(ctx);
  case triton::ScaleDotElemType::BF16:
    return BFloat16Type::get(ctx);
  case triton::ScaleDotElemType::E4M3:
    return Float8E4M3FNType::get(ctx);
  case triton::ScaleDotElemType::E5M2:
    return Float8E5M2Type::get(ctx);
  case triton::ScaleDotElemType::E3M2:
    return Float6E3M2FNType::get(ctx);
  case triton::ScaleDotElemType::E2M3:
    return Float6E2M3FNType::get(ctx);
  case triton::ScaleDotElemType::E2M1:
    return Float4E2M1FNType::get(ctx);
  default:
    llvm_unreachable("unsupported ScaleDotElemType!");
  }
}

bool canCoalesceWriteIntoSharedMemory(MLIRContext *ctx,
                                      const LinearLayout &srcToSharedLayout,
                                      unsigned threadsPerWarp) {
  auto kReg = StringAttr::get(ctx, "register");
  StringAttr kLane = StringAttr::get(ctx, "lane");
  auto kOffset = StringAttr::get(ctx, "offset");

  auto contig = srcToSharedLayout.getNumConsecutiveInOut();

  // Create a coalesced/identity layout and see if it divides srcToShared
  auto coalescedLayout =
      LinearLayout::identity1D(contig, kReg, kOffset) *
      LinearLayout::identity1D(threadsPerWarp, kLane, kOffset);

  return divideLeft(srcToSharedLayout, coalescedLayout).has_value();
}

// On architectures supporting scattering into LDS we are only constraint by the
// minimal vector size. On architectures not support scattering, e.g. gfx9,
// direct to LDS loads do not support per lane shared offsets. We need to ensure
// that each warp writes coalesced into shared memory. This means we cannot
// exceed the supported load width because splitting them would cause strided
// (non coalesced) writes. Additionally:
// 1. For *non* swizzled shared encodings we check if they result in coalesced
//    writes and can then lower them directly to the intrinsics.
// 2. For swizzled shared encodings we need to transfer the swizzling to the
//    source pointers. For now this is done by swizzling the pointers
//    between the lane of a warp via permute. This only works if the swizzle
//    pattern does not exchange elements between warps which holds for all
//    our swizzle patterns. There is still a check performed to not silently
//    produce wrong results if we invalidate the condition in the future
bool canLoadDirectToLDS(const triton::AMD::TargetInfo &targetInfo,
                        RankedTensorType srcTy, Attribute dstEnc,
                        ArrayRef<int64_t> dstAllocShape, unsigned &vectorSize) {
  // For padded encodings restrict vec by the min interval
  auto paddedEnc = dyn_cast<triton::gpu::PaddedSharedEncodingAttr>(dstEnc);
  if (paddedEnc) {
    unsigned maxAllowedVecSize = paddedEnc.getMinInterval();
    // Without scattering support, padding can only be inserted at warp
    // boundaries. This means minInterval must be a multiple of (vectorSize *
    // warpSize) which becomes vectorSize <= minInterval / warpSize.
    if (!targetInfo.supportsDirectToLdsScatter())
      maxAllowedVecSize = paddedEnc.getMinInterval() / targetInfo.getWarpSize();

    vectorSize = std::min(vectorSize, maxAllowedVecSize);
  }

  // Check that vectorSize is not smaller than the minimal supported vector size
  int elemBitWidth = tt::getPointeeBitWidth(srcTy);
  if (fitToValidDirectToLdsVecSize(vectorSize, elemBitWidth, targetInfo) == 0) {
    LDBG("unsupported global load to LDS vectorSize (" << vectorSize << ")");
    return false;
  }

  // Following checks are specific to architectures not supporting scattering
  if (targetInfo.supportsDirectToLdsScatter())
    return true;

  // Compute the blocked -> shared linear layout to check preconditions
  LinearLayout srcLayout = triton::gpu::toLinearLayout(srcTy);
  LinearLayout sharedLayout;
  if (paddedEnc) {
    sharedLayout = paddedEnc.getLinearComponent();
  } else {
    sharedLayout = triton::gpu::toLinearLayout(dstAllocShape, dstEnc);

    // Use a non swizzled layout since we apply swizzling to the src pointers
    auto swizzledEnc =
        dyn_cast<triton::gpu::SwizzledSharedEncodingAttr>(dstEnc);
    if (swizzledEnc && swizzledEnc.getMaxPhase() != 1) {
      auto flatSharedEnc = tt::gpu::SwizzledSharedEncodingAttr::get(
          srcTy.getContext(), swizzledEnc.getVec(), 1, 1,
          swizzledEnc.getOrder(), swizzledEnc.getCGALayout());
      sharedLayout = tt::gpu::toLinearLayout(dstAllocShape, flatSharedEnc);
    }
  }
  LinearLayout srcToSharedLayout = srcLayout.invertAndCompose(sharedLayout);

  auto contig = srcToSharedLayout.getNumConsecutiveInOut();
  // The reg->shared layout can only write `contig` consecutive elements
  // coalesced into LDS. A load narrower than `contig` would leave gaps between
  // lanes and produce strided writes, so we cannot lower it here.
  if (vectorSize < contig) {
    LDBG("Load vectorization (" << vectorSize << ") smaller than contiguity ("
                                << contig << ") resulting in strided writes");
    return false;
  }
  // If the requested load is wider than what the layout can write coalesced,
  // reduce it so that each load instruction writes exactly one coalesced chunk.
  // The lowering then emits multiple (narrower) direct-to-LDS loads.
  vectorSize = contig;

  // The reduced width must still be a supported direct-to-LDS load width.
  if (!targetInfo.supportsDirectToLdsLoadBitWidth(vectorSize * elemBitWidth)) {
    LDBG("coalesced load width (" << vectorSize
                                  << ") is not a supported direct-to-LDS load");
    return false;
  }

  if (!canCoalesceWriteIntoSharedMemory(srcTy.getContext(), srcToSharedLayout,
                                        targetInfo.getWarpSize())) {
    LDBG("Does not write coalesced into LDS");
    return false;
  }

  return true;
}

// For a region iter arg of an scf.for, return the yield operand that feeds it
// on the back-edge. Returns {} if the parent isn't scf.for or the arg is the
// induction variable.
static Value resolveLoopBackedge(BlockArgument bbArg) {
  auto forOp = dyn_cast<scf::ForOp>(bbArg.getOwner()->getParentOp());
  if (!forOp)
    return {};
  auto yieldOp = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
  int iterArgIdx = bbArg.getArgNumber() - forOp.getNumInductionVars();
  if (iterArgIdx < 0 || iterArgIdx >= (int)yieldOp.getNumOperands())
    return {};
  return yieldOp.getOperand(iterArgIdx);
}

bool isChainDotHead(tt::DotOpInterface dotOp, unsigned opIdx) {
  auto isInSameRegion = [&dotOp](Operation *op) {
    return op->getParentRegion() == dotOp->getParentRegion();
  };
  ForwardSliceOptions fwdOpt;
  fwdOpt.filter = isInSameRegion;
  SetVector<mlir::Operation *> fwdSlices;
  getForwardSlice(dotOp, &fwdSlices, fwdOpt);
  for (Operation *op : fwdSlices) {
    if (auto dOp = dyn_cast<tt::DotOpInterface>(op)) {
      assert(dOp != dotOp);
      Operation *dotOperand = (opIdx == 0) ? dOp.getA().getDefiningOp()
                                           : dOp.getB().getDefiningOp();
      if (dotOperand && fwdSlices.contains(dotOperand)) {
        return true;
      }
    }
  }

  // Cross-iteration: for each op in the forward slice (including dotOp itself)
  // that is consumed by a scf.yield, resolve to the corresponding iter arg and
  // check that iter arg's forward slice for a dot.
  fwdSlices.insert(dotOp);
  for (Operation *op : fwdSlices) {
    for (Value result : op->getResults()) {
      for (OpOperand &use : result.getUses()) {
        auto yieldOp = dyn_cast<scf::YieldOp>(use.getOwner());
        if (!yieldOp)
          continue;
        auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
        if (!forOp)
          continue;
        BlockArgument nextIterArg =
            forOp.getRegionIterArg(use.getOperandNumber());
        SetVector<Operation *> argFwdSlices;
        getForwardSlice(nextIterArg, &argFwdSlices, fwdOpt);
        for (Operation *argOp : argFwdSlices) {
          auto dOp = dyn_cast<tt::DotOpInterface>(argOp);
          if (!dOp || dOp == dotOp)
            continue;
          Value dotOperand = (opIdx == 0) ? dOp.getA() : dOp.getB();
          if (dotOperand == nextIterArg ||
              (dotOperand.getDefiningOp() &&
               argFwdSlices.contains(dotOperand.getDefiningOp())))
            return true;
        }
      }
    }
  }
  return false;
}

bool isChainDotTail(tt::DotOpInterface dotOp) {
  auto isInSameRegion = [&dotOp](Operation *op) {
    return op->getParentRegion() == dotOp->getParentRegion();
  };
  BackwardSliceOptions bwdOpt;
  bwdOpt.omitBlockArguments = true;
  bwdOpt.filter = isInSameRegion;
  SetVector<Operation *> bwdSlices;
  Operation *opA = dotOp.getA().getDefiningOp();
  if (opA) {
    (void)getBackwardSlice(opA, &bwdSlices, bwdOpt);
    if (llvm::find_if(bwdSlices, [](Operation *op) {
          return isa<tt::DotOpInterface>(op);
        }) != bwdSlices.end())
      return true;
  }

  // Cross-iteration: if operand A (or its backward slice) touches a block
  // arg, resolve via the yield back-edge and check the backward slice of
  // the yielded value for a dot (mirrors the intra-iteration check above).
  SmallVector<BlockArgument, 4> bbArgs;
  if (auto bbArg = dyn_cast<BlockArgument>(dotOp.getA())) {
    bbArgs.push_back(bbArg);
  } else if (opA) {
    bwdSlices.insert(opA);
    for (Operation *sliceOp : bwdSlices)
      for (Value operand : sliceOp->getOperands())
        if (auto bbArg = dyn_cast<BlockArgument>(operand))
          bbArgs.push_back(bbArg);
  }

  for (BlockArgument bbArg : bbArgs) {
    Value yieldVal = resolveLoopBackedge(bbArg);
    if (!yieldVal)
      continue;
    Operation *yieldDef = yieldVal.getDefiningOp();
    if (!yieldDef)
      continue;
    SetVector<Operation *> yieldBwdSlices;
    (void)getBackwardSlice(yieldDef, &yieldBwdSlices, bwdOpt);
    yieldBwdSlices.insert(yieldDef);
    if (llvm::find_if(yieldBwdSlices, [&dotOp](Operation *op) {
          return isa<tt::DotOpInterface>(op) && op != dotOp;
        }) != yieldBwdSlices.end())
      return true;
  }

  return false;
}

Value convertF8ToF32_SW(RewriterBase &rewriter, Location loc, Value fp8Val,
                        bool isE4M3FN) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  if (isE4M3FN) {
    // Position the 7 magnitude bits at f32 bits [26:20] and multiply by
    // 2^120 to correct the exponent bias (f32 bias 127 - E4M3 bias 7 = 120).
    Value absVal = b.and_(fp8Val, b.i8_val(0x7F));
    Value vi32 = b.shl(b.zext(i32_ty, absVal), b.i32_val(20));
    Value rawF32 = b.fmul(b.bitcast(vi32, f32_ty), b.f32_val(0x1p120f));
    Value signBit =
        b.shl(b.and_(b.zext(i32_ty, fp8Val), b.i32_val(0x80)), b.i32_val(24));
    return b.bitcast(b.or_(b.bitcast(rawF32, i32_ty), signBit), f32_ty);
  }
  // E5M2 -> f16 (same exponent format, simple shift) -> fpext to f32.
  Value i16Val = b.shl(b.zext(i16_ty, fp8Val), b.i16_val(8));
  return b.fpext(f32_ty, b.bitcast(i16Val, f16_ty));
}

SmallVector<Value> upcast8xMxfp4_SW(RewriterBase &rewriter, Operation *op,
                                    bool toFp16, Value packedVec,
                                    ISAFamily isaFamily, Value scale) {
  assert((isa<triton::gpu::Fp4ToFpOp, triton::amdgpu::ScaledUpcastFp4Op>(op)) &&
         "Expected Fp4ToFpOp or ScaledUpcastFp4Op");
  Location loc = op->getLoc();
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  auto permU32FnTy =
      LLVM::LLVMFunctionType::get(i32_ty, {i32_ty, i32_ty, i32_ty});
  LLVM::LLVMFuncOp funcOp =
      appendOrGetExternFuncOp(rewriter, op, "llvm.amdgcn.perm", permU32FnTy);

  // Start with 8 mxfp4 elements in a single i32 register
  // | e7e6 | e5e4 | e3e2 | e1e0 |
  Value input = b.bitcast(packedVec, i32_ty);

  // fp4 to bf16 for cdna3: fp4->fp8->fp32
  if (isaFamily == ISAFamily::CDNA3 && !toFp16) {
    // Step 1: extract EM bits for elements 0,2,4,6 and 1,3,5,7 respectively.
    // e2m1_6420_idx = | 0[0e6EM] | 0[0e4EM] | 0[0e2EM] | 0[0e0EM] |
    Value e2m1_6420_idx = b.and_(input, b.i32_val(0x07070707));
    // e2m1_7531_idx = | [0e7EM]0 | [0e5EM]0 | [0e3EM]0 | [0e1EM]0 |
    Value e2m1_7531_idx = b.and_(input, b.i32_val(0x70707070));
    e2m1_7531_idx = b.lshr(e2m1_7531_idx, b.i32_val(4));

    // Step 2: convert fp4 to fp8 using LUT
    Value resLutLo = b.i32_val(0xc4c0b800);
    Value resLutHi = b.i32_val(0xd4d0ccc8);
    Value res_6420 = LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                                            {resLutHi, resLutLo, e2m1_6420_idx})
                         .getResult();
    Value res_7531 = LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                                            {resLutHi, resLutLo, e2m1_7531_idx})
                         .getResult();

    // Step 3: extract sign bits
    Value s_6420 = b.or_(b.shl(input, b.i32_val(4)), b.i32_val(0x7f7f7f7f));
    Value s_7531 = b.or_(input, b.i32_val(0x7f7f7f7f));

    // Step 4:  assemble 4 packed fp8 values w/ sign
    res_6420 = b.and_(res_6420, s_6420);
    res_7531 = b.and_(res_7531, s_7531);

    // Step 5: convert fp8 to fp32
    Value res_20 = ROCDL::CvtPkF32Fp8Op::create(
        rewriter, loc, i64_ty, res_6420, rewriter.getIntegerAttr(i1_ty, 0));
    Value res_64 = ROCDL::CvtPkF32Fp8Op::create(
        rewriter, loc, i64_ty, res_6420, rewriter.getIntegerAttr(i1_ty, 1));
    Value res_31 = ROCDL::CvtPkF32Fp8Op::create(
        rewriter, loc, i64_ty, res_7531, rewriter.getIntegerAttr(i1_ty, 0));
    Value res_75 = ROCDL::CvtPkF32Fp8Op::create(
        rewriter, loc, i64_ty, res_7531, rewriter.getIntegerAttr(i1_ty, 1));
    SmallVector<Value> pkVals{res_20, res_64, res_31, res_75};
    if (scale) {
      // pack 2 values together to help llvm backend codegen
      Value scaleF32 =
          b.bitcast(b.shl(b.zext(i32_ty, scale), b.i32_val(23)), f32_ty);
      Type v2f32 = vec_ty(f32_ty, 2);
      Value pkScale = b.undef(v2f32);
      pkScale = b.insert_element(pkScale, scaleF32, b.i32_val(0));
      pkScale = b.insert_element(pkScale, scaleF32, b.i32_val(1));
      Type v2i32 = vec_ty(i32_ty, 2);
      for (unsigned i = 0; i < 4; i++) {
        Value pkScaled = b.fmul(pkScale, b.bitcast(pkVals[i], v2f32));
        pkVals[i] = (b.bitcast(pkScaled, v2i32));
      }
    } else {
      // bitcast to v2i32
      for (unsigned i = 0; i < 4; i++) {
        pkVals[i] = b.bitcast(pkVals[i], vec_ty(i32_ty, 2));
      }
    }
    Value e0 = b.extract_element(pkVals[0], b.i32_val(0));
    Value e1 = b.extract_element(pkVals[2], b.i32_val(0));
    Value e2 = b.extract_element(pkVals[0], b.i32_val(1));
    Value e3 = b.extract_element(pkVals[2], b.i32_val(1));
    Value e4 = b.extract_element(pkVals[1], b.i32_val(0));
    Value e5 = b.extract_element(pkVals[3], b.i32_val(0));
    Value e6 = b.extract_element(pkVals[1], b.i32_val(1));
    Value e7 = b.extract_element(pkVals[3], b.i32_val(1));
    SmallVector<Value, 8> f32Vals{e0, e1, e2, e3, e4, e5, e6, e7};
    Value sel = b.i32_val(0x07060302);
    SmallVector<Value> results;
    for (unsigned i = 0; i < 8; i += 2) {
      // v2f32->v2bf16: {e1.f32[31:16], e0.f32[31:16]}
      Value res = LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                                         {f32Vals[i + 1], f32Vals[i], sel})
                      .getResult();
      Type v2bf16 = vec_ty(bf16_ty, 2);
      res = b.bitcast(res, v2bf16);
      results.push_back(b.extract_element(res, b.i32_val(0)));
      results.push_back(b.extract_element(res, b.i32_val(1)));
    }
    return results;
  }

  // MXFP4 has 4 bits, S.EE.M, for Sign, Exponent, and Mantissa respectively.
  // For a specific S, we have a total of 8 bit patterns. We can encode all
  // these 8 resultant bf16/fp16 bit patterns in a lookup table (LUT). It
  // happens that llvm.amdgcn.perm supports selecting 4 bytes from 8 input bytes
  // using a 4-byte selector. So the overall idea is to use llvm.amdgcn.perm to
  // implement such a LUT; though we need to select the two bytes for the
  // resultant bf16/fp16 bit patterns separately. For the byte containing S, we
  // also need to handle the S and E bits separately.

  // FP4 has 4 bits: S.EE.M. Bf16/fp16 bit patterns for positive values:
  //
  // FP4    | BF16   | FP16   | Value
  // ------ | ------ | ------ | -----
  // 0.00.0 | 0x0000 | 0x0000 | + 0.0
  // 0.00.1 | 0x3f00 | 0x3800 | + 0.5
  // 0.01.0 | 0x3f80 | 0x3c00 | + 1.0
  // 0.01.1 | 0x3fc0 | 0x3e00 | + 1.5
  // 0.10.0 | 0x4000 | 0x4000 | + 2.0
  // 0.10.1 | 0x4040 | 0x4200 | + 3.0
  // 0.11.0 | 0x4080 | 0x4400 | + 4.0
  // 0.11.1 | 0x40c0 | 0x4600 | + 6.0
  //
  // Encode Byte #0 (M) for BF16/FP16 in a LUT.
  Value resB0LutLo = toFp16 ? b.i32_val(0) : b.i32_val(0xc0800000);
  Value resB0LutHi = toFp16 ? b.i32_val(0) : b.i32_val(0xc0804000);
  // Encode Byte #1 (EM, non-S part) for BF16/FP16 in a LUT.
  Value resB1LutLoNoS = toFp16 ? b.i32_val(0x3e3c3800) : b.i32_val(0x3f3f3f00);
  Value resB1LutHiNoS = toFp16 ? b.i32_val(0x46444240) : b.i32_val(0x40404040);

  // Step 1: extract EM bits for elements 0,2,4,6 and 1,3,5,7 respectively.
  // e2m1_6420_idx = | 0[0e6EM] | 0[0e4EM] | 0[0e2EM] | 0[0e0EM] |
  Value e2m1_6420_idx = b.and_(input, b.i32_val(0x07070707));
  // e2m1_7531_idx = | [0e7EM]0 | [0e5EM]0 | [0e3EM]0 | [0e1EM]0 |
  Value e2m1_7531_idx = b.and_(input, b.i32_val(0x70707070));
  // e2m1_7531_idx = | 0[0e7EM] | 0[0e5EM] | 0[0e3EM] | 0[0e1EM] |
  e2m1_7531_idx = b.lshr(e2m1_7531_idx, b.i32_val(4));

  // Step 2: extract S bit for elements 0,2,4,6 and 1,3,5,7
  // s_6420 = | 0[e6S000] | 0[e4S000] | 0[e2S000] | 0[e0S000] |
  Value s_6420 = b.and_(input, b.i32_val(0x08080808));
  // s_6420 = | [e6S000]0 | [e4S000]0 | [e2S000]0 | [e0S000]0 |
  s_6420 = b.shl(s_6420, b.i32_val(4));
  // s_7531 = | [e7S000]0 | [e5S000]0 | [e3S000]0 | [e1S000]0 |
  Value s_7531 = b.and_(input, b.i32_val(0x80808080));

  // Step 3: Upcast elements 0,2,4,6 to 4 16-bit elements
  // Select Byte #0. It's always 0 if upcasting to fp16.
  // resB0_6420 = | e6B0 | e4B0 | e2B0 | e0B0 |
  Value resB0_6420 = b.i32_val(0);
  if (!toFp16) {
    resB0_6420 = LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                                        {resB0LutHi, resB0LutLo, e2m1_6420_idx})
                     .getResult();
  }
  // Select Byte #1
  Value resB1NoS_6420 =
      LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                             {resB1LutHiNoS, resB1LutLoNoS, e2m1_6420_idx})
          .getResult();
  // resB1_6420 = | e6B1 | e4B1 | e2B1 | e0B1 |
  Value resB1_6420 = b.or_(resB1NoS_6420, s_6420);
  // Construct 16-bit values of e0 and e2
  // res_20 = | e2B1 | e2B0 | e0B1 | e0B0 | = | e2_f16 | e0_f16 |
  Value res_20 =
      LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                             {resB1_6420, resB0_6420, b.i32_val(0x05010400)})
          .getResult();
  // Construct 16-bit values of e4 and e6
  // res_64 = | e6B1 | e6B0 | e4B1 | e4B0 | = | e6_f16 | e4_f16 |
  Value res_64 =
      LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                             {resB1_6420, resB0_6420, b.i32_val(0x07030602)})
          .getResult();

  // Step 4: Upcast elements 1,3,5,7 to 4 16-bit elements
  // This is a copy of step 3 on different group of elements
  // Select Byte #0. It's always 0 if upcasting to fp16.
  // resB0_7531 = | e7B0 | e5B0 | e3B0 | e1B0 |
  Value resB0_7531 = b.i32_val(0);
  if (!toFp16) {
    resB0_7531 = LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                                        {resB0LutHi, resB0LutLo, e2m1_7531_idx})
                     .getResult();
  }
  // Select Byte #1
  Value resB1NoS_7531 =
      LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                             {resB1LutHiNoS, resB1LutLoNoS, e2m1_7531_idx})
          .getResult();
  // resB1_7531 = | e7B1 | e5B1 | e3B1 | e1B1 |
  Value resB1_7531 = b.or_(resB1NoS_7531, s_7531);
  // Construct 16-bit values of e1 and e3
  // res_31 = | e3B1 | e3B0 | e1B1 | e1B0 | = | e3_f16 | e1_f16 |
  Value res_31 =
      LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                             {resB1_7531, resB0_7531, b.i32_val(0x05010400)})
          .getResult();
  // Construct 16-bit values of e5 and e7
  // res_75 = | e7B1 | e7B0 | e5B1 | e5B0 | = | e7_f16 | e5_f16 |
  Value res_75 =
      LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                             {resB1_7531, resB0_7531, b.i32_val(0x07030602)})
          .getResult();

  // Step 5: Reorder 16-bit elements to be 0,1,2,3,4,5,6,7
  // res_10 = | e1_f16 | e0_f16 |
  Value res_10 = LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                                        {res_31, res_20, b.i32_val(0x05040100)})
                     .getResult();
  // res_32 = | e3_f16 | e2_f16 |
  Value res_32 = LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                                        {res_31, res_20, b.i32_val(0x07060302)})
                     .getResult();
  // res_54 = | e5_f16 | e4_f16 |
  Value res_54 = LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                                        {res_75, res_64, b.i32_val(0x05040100)})
                     .getResult();
  // res_76 = | e7_f16 | e6_f16 |
  Value res_76 = LLVM::createLLVMCallOp(rewriter, loc, funcOp,
                                        {res_75, res_64, b.i32_val(0x07060302)})
                     .getResult();

  SmallVector<Value, 4> pkVals{res_10, res_32, res_54, res_76};
  SmallVector<Value> results;
  Type elmTy = toFp16 ? f16_ty : bf16_ty;
  for (int j = 0; j < 4; j++) {
    Value elements = b.bitcast(pkVals[j], vec_ty(elmTy, 2));
    results.push_back(b.extract_element(elements, b.i32_val(0)));
    results.push_back(b.extract_element(elements, b.i32_val(1)));
  }
  return results;
}

} // namespace mlir::LLVM::AMD
