#include "triton/Conversion/TritonGPUToLLVM/Utility.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/IR/Attributes.h"
#include "mlir/Transforms/RegionUtils.h"
#include "triton/Analysis/Allocation.h"
#include "triton/Analysis/Utility.h"
#include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h"
#include "triton/Dialect/Triton/IR/Dialect.h"
#include "triton/Dialect/TritonGPU/IR/Attributes.h"
#include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h"
#include "triton/Dialect/TritonGPU/Transforms/Utility.h"
#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
#include "triton/Tools/GenericSwizzling.h"
#include "triton/Tools/LayoutUtils.h"
#include "triton/Tools/LinearLayout.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/MathExtras.h"

#include <functional>

#if defined(_MSC_VER) && !defined(__clang__)
// from https://gist.github.com/pps83/3210a2f980fd02bb2ba2e5a1fc4a2ef0
#include <intrin.h>

static int __builtin_clz(unsigned x) {
  unsigned long r;
  _BitScanReverse(&r, x);
  return static_cast<int>(r ^ 31);
}

static int __builtin_ctz(unsigned x) {
  unsigned long r;
  _BitScanForward(&r, x);
  return static_cast<int>(r);
}

#endif

namespace mlir {

namespace triton::gpu {

std::pair<SmallVector<LocalMemOpTile>, SmallVector<LocalMemOpTile>>
getSrcDstTiles(const TargetInfoBase &targetInfo, int bitwidth,
               bool crossCTALoads) {
  assert(bitwidth <= 128 && "bitwidth must be <= 128");
  assert(llvm::isPowerOf2_32(bitwidth) && "bitwidth must be a power of two");
  SmallVector<LocalMemOpTile> src;
  SmallVector<LocalMemOpTile> dst;

  // ld.shared/st.shared
  auto ldstshared = LocalMemOpTile{{}, {0, 1, 2}};
  src.push_back(ldstshared);
  dst.push_back(ldstshared);

  if (targetInfo.supportLdMatrix() || targetInfo.supportStMatrix()) {
    // ldmatrix/stmatrix
    if (bitwidth <= 32) {
      auto ldstmatrix = LocalMemOpTile{{0, 1}, {2, 3, 4}};
      if (targetInfo.supportStMatrix()) {
        src.push_back(ldstmatrix);
      }
      // We do cross-CTA reads but in-CTA writes
      // ldmatrix/stmatrix do not support cross-CTA transfers.
      if (!crossCTALoads && targetInfo.supportLdMatrix()) {
        dst.push_back(ldstmatrix);
      }
    }
    // ldmatrix.trans/stmatrix.trans
    if (bitwidth == 16) {
      auto ldstmatrixtrans = LocalMemOpTile{{2, 3, 4}, {0, 1}};
      if (targetInfo.supportStMatrix()) {
        src.push_back(ldstmatrixtrans);
      }
      if (!crossCTALoads && targetInfo.supportLdMatrix()) {
        dst.push_back(ldstmatrixtrans);
      }
    }
  }
  return {std::move(src), std::move(dst)};
}

Type getFunctionType(Type resultType, ValueRange operands) {
  SmallVector<Type> operandTypes(operands.getTypes());
  return LLVM::LLVMFunctionType::get(resultType, operandTypes);
}

LLVM::LLVMFuncOp appendOrGetExternFuncOp(RewriterBase &rewriter, Operation *op,
                                         StringRef funcName, Type funcType,
                                         StringRef libname /*= ""*/,
                                         StringRef libpath /*= ""*/) {
  using LLVM::LLVMFuncOp;

  auto funcAttr = StringAttr::get(op->getContext(), funcName);
  Operation *funcOp = SymbolTable::lookupNearestSymbolFrom(op, funcAttr);
  if (funcOp)
    return cast<LLVMFuncOp>(*funcOp);

  Operation *parent = op;
  if (!isa<LLVM::LLVMFuncOp>(op))
    parent = op->getParentOfType<LLVM::LLVMFuncOp>();
  OpBuilder b(parent);
  auto ret = LLVMFuncOp::create(b, op->getLoc(), funcName, funcType);
  ret.getOperation()->setAttr("libname",
                              StringAttr::get(op->getContext(), libname));
  ret.getOperation()->setAttr("libpath",
                              StringAttr::get(op->getContext(), libpath));
  return ret;
}

Value matrixVectorProd(TritonLLVMOpBuilder &b, const LinearLayout &A, Value x) {
  assert(A.getNumInDims() == 1);
  assert(A.getNumOutDims() == 1);
  auto flatten = [](const std::vector<std::vector<int32_t>> &matrix) {
    SmallVector<int32_t> ret;
    for (const auto &row : matrix) {
      ret.push_back(row[0]);
    }
    return ret;
  };
  auto nCol = A.getTotalInDimSizeLog2();
  auto nRow = A.getTotalOutDimSizeLog2();
  SmallVector<int32_t> matrix = flatten(A.getBases().begin()->second);
  assert(matrix.size() == nCol);

  // Row-wise popcount to detect rows that appear exactly once across columns.
  uint32_t rowsUnique = 0;
  {
    SmallVector<int> rowPopCnt(nRow, 0);
    for (int c = 0; c < nCol; ++c) {
      uint32_t colBits = matrix[c];
      for (int r = 0; r < nRow; ++r) {
        if (colBits & (1u << r))
          ++rowPopCnt[r];
      }
    }
    for (int r = 0; r < nRow; ++r) {
      if (rowPopCnt[r] == 1)
        rowsUnique |= 1u << r;
    }
  }

  // We iterate the matrix following the diagonals and build
  // (x & mask_i) << s_i terms. Prefer OR for diagonals whose rows are unique,
  // then XOR everything else. This tends to encourage mad.lo codegen.
  auto getMaskAndAllRowsUnique = [&](int i) -> std::pair<uint32_t, bool> {
    uint32_t mask = 0;
    int row = i < 0 ? -i : 0;
    int col = i < 0 ? 0 : i;
    bool allRowsUnique = true;
    while (row < nRow && col < nCol) {
      uint32_t bitValue = (matrix[col] >> row) & 1u;
      mask |= bitValue << col;
      allRowsUnique &= ((rowsUnique >> row) & 1u) == 1u;
      ++row;
      ++col;
    }
    return {mask, allRowsUnique};
  };

  uint32_t explicitCols = 0;

  {
    SmallVector<uint32_t> masks;
    for (int i = -nRow + 1; i < nCol; i++) {
      masks.push_back(std::get<0>(getMaskAndAllRowsUnique(i)));
    }
    bool reachedFixedPoint = false;
    while (!reachedFixedPoint) {
      reachedFixedPoint = true;
      for (uint32_t m : masks) {
        uint32_t c = m & ~explicitCols;
        if (llvm::isPowerOf2_32(c)) {
          // found a single-element diagonal
          explicitCols |= c;
          reachedFixedPoint = false;
        }
      }
    }
  }

  // handle any diagonals that have survived
  SmallVector<Value> ors;
  SmallVector<Value> xors;
  for (int i = -nRow + 1; i < nCol; i++) {
    auto [mask, allRowsUnique] = getMaskAndAllRowsUnique(i);
    mask &= ~explicitCols;
    if (mask == 0)
      continue;
    auto masked = b.and_(x, b.i32_val(mask));
    auto shifted = i >= 0 ? Value(b.lshr(masked, b.i32_val(i)))
                          : Value(b.shl(masked, b.i32_val(-i)));
    if (allRowsUnique) {
      ors.push_back(shifted);
    } else {
      xors.push_back(shifted);
    }
  }

  // handle any explicit columns:
  Value zero = b.i32_val(0);
  for (int i = 0; i < nCol; i++) {
    if ((explicitCols >> i) & 1) {
      int32_t basis = matrix[i];
      if (basis == 0)
        continue;
      Value term;
      Value bit = b.and_(x, b.i32_val(1 << i));
      uint32_t basisBits = basis;
      if (llvm::isPowerOf2_32(basisBits)) {
        unsigned row = llvm::countr_zero(basisBits);
        unsigned col = i;
        if (row == col)
          term = bit;
        else if (row > col)
          term = b.shl(bit, b.i32_val(row - col));
        else
          term = b.lshr(bit, b.i32_val(col - row));
      } else {
        Value bit_is_zero = b.icmp_eq(bit, zero);
        term = b.select(bit_is_zero, zero, b.i32_val(basis));
      }
      if ((rowsUnique & basis) == basis) {
        ors.push_back(term);
      } else {
        xors.push_back(term);
      }
    }
  }

  auto treeReduce = [&](SmallVector<Value> &terms,
                        std::function<Value(Value, Value)> op) -> Value {
    if (terms.empty())
      return b.i32_val(0);
    while (terms.size() > 1) {
      SmallVector<Value> next;
      for (size_t i = 0; i + 1 < terms.size(); i += 2)
        next.push_back(op(terms[i], terms[i + 1]));
      if (terms.size() % 2 == 1)
        next.push_back(terms.back());
      terms = std::move(next);
    }
    return terms[0];
  };

  auto orPart = treeReduce(
      ors, [&b](Value x, Value y) { return b.or_(x, y, /*disjoint=*/true); });
  auto xorPart =
      treeReduce(xors, [&b](Value x, Value y) { return b.xor_(x, y); });
  return b.or_(orPart, xorPart, /*disjoint=*/true);
}

Value maybeAnd(OpBuilder &builder, Location loc, Value a, Value b) {
  auto tb = TritonLLVMOpBuilder(loc, builder);
  if (a && b) {
    return tb.and_(a, b);
  }
  return a ? a : b;
}

Value emitRedundantThreadPredicate(
    const llvm::MapVector<StringAttr, int32_t> &freeVarMasks,
    ConversionPatternRewriter &rewriter, Location loc,
    const TargetInfoBase &targetInfo) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  auto ctx = rewriter.getContext();
  auto kLane = str_attr("lane");
  auto kWarp = str_attr("warp");
  auto kBlock = str_attr("block");

  Value zero = b.i32_val(0);
  auto [laneId, warpId] = getLaneAndWarpId(rewriter, loc);
  Value blockId = freeVarMasks.lookup(kBlock) == 0
                      ? zero
                      : targetInfo.getClusterCTAId(rewriter, loc);

  Value pred;
  auto dimNames = {kLane, kWarp, kBlock};
  auto dimIds = {laneId, warpId, blockId};
  for (auto [dimName, dimId] : llvm::zip(dimNames, dimIds)) {
    int32_t mask = freeVarMasks.lookup(dimName);
    if (mask != 0) {
      auto dimPred = b.icmp_eq(b.and_(dimId, b.i32_val(mask)), zero);
      pred = maybeAnd(rewriter, loc, pred, dimPred);
    }
  }
  return pred;
}

} // namespace triton::gpu

SmallVector<std::pair<StringAttr, Value>>
applyLinearLayout(Location loc, RewriterBase &rewriter,
                  const LinearLayout &layout,
                  ArrayRef<std::pair<StringAttr, Value>> indices) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  assert(layout.getNumInDims() == indices.size());
  assert(llvm::equal(layout.getInDimNames(), llvm::make_first_range(indices)));
  // Trivial layout
  if (layout.getNumOutDims() == 0) {
    return {};
  }

  // This function can emit a lot of MLIR code, which ultimately makes
  // compilation slow.  (We think this shouldn't be the case -- it's not *that*
  // much code -- but we're not clear on how to fix the slowness, which happens
  // in the bowels of MLIR.)
  //
  // As a result we go through some contortions to avoid emitting code where
  // possible.

  // Manually constant-fold the layout where possible.
  SmallVector<std::pair<StringAttr, int32_t>> constantIns;
  SmallVector<std::pair<StringAttr, Value>> nonConstantIns;
  for (auto [inDimName, idx] : indices) {
    APInt constant;
    if (matchPattern(idx, m_ConstantInt(&constant))) {
      constantIns.push_back({inDimName, constant.getSExtValue()});
    } else {
      constantIns.push_back({inDimName, 0});
      nonConstantIns.push_back({inDimName, idx});
    }
  }

  // Compute constant part of the output and wrap it as values
  Value zero = b.i32_val(0);
  SmallVector<std::pair<StringAttr, Value>> outIndices;
  for (auto [outDimName, constant] : layout.apply(constantIns)) {
    if (constant == 0)
      outIndices.push_back({outDimName, zero});
    else
      outIndices.push_back({outDimName, b.i32_val(constant)});
  }

  if (nonConstantIns.size() == 0) {
    return outIndices;
  }

  SmallVector<StringAttr> inDimNames;
  // Concatenate input
  Value x = b.i32_val(0);
  int shift = 0;
  for (auto [inDimName, idx] : nonConstantIns) {
    inDimNames.push_back(inDimName);
    x = b.or_(x, b.shl(idx, b.i32_val(shift)));
    shift += layout.getInDimSizeLog2(inDimName);
  }

  for (auto &[outDimName, outIdx] : outIndices) {
    // Apply flattened sublayout for this output
    auto matrix = layout.sublayout(inDimNames, outDimName).flattenIns();
    auto out = triton::gpu::matrixVectorProd(b, matrix, x);
    outIdx = b.xor_(outIdx, out);
  }

  return outIndices;
}

std::optional<int> getWarpGroupStartWarpId(Block *block) {
  using namespace triton::gpu;

  // Look for an enclosing `ttg.warp_specialize` op.
  while (block && block->getParentOp() &&
         !isa<WarpSpecializePartitionsOp>(block->getParentOp()))
    block = block->getParentOp()->getBlock();
  if (!block || !block->getParentOp())
    return {};

  auto partitions = cast<WarpSpecializePartitionsOp>(block->getParentOp());
  unsigned idx = block->getParent()->getRegionNumber();
  WarpSpecializeOp ws = partitions.getParentOp();
  std::optional<ArrayRef<int32_t>> startIds = ws.getWarpGroupStartIds();
  assert(startIds && "cannot get warp group ID before warp group allocation");
  int32_t warpStartId = (*startIds)[idx];
  return warpStartId;
}

std::optional<int> getWarpGroupStartThreadId(Block *block) {
  using namespace triton::gpu;

  std::optional<int> warpStartId = getWarpGroupStartWarpId(block);
  if (!warpStartId)
    return {};

  int threadsPerWarp = TritonGPUDialect::getThreadsPerWarp(
      block->getParentOp()->getParentOfType<ModuleOp>());
  return *warpStartId * threadsPerWarp;
}

Value getThreadId(OpBuilder &rewriter, Location loc) {
  Value tid =
      ::mlir::gpu::ThreadIdOp::create(rewriter, loc, ::mlir::gpu::Dimension::x);
  tid = arith::IndexCastOp::create(rewriter, loc, i32_ty, tid);

  Operation *lookupPt = &rewriter.getInsertionBlock()->front();
  int threadsPerWarp = triton::gpu::lookupThreadsPerWarp(rewriter);
  int numWarps = triton::gpu::lookupNumWarps(lookupPt);
  int upperBound = numWarps * threadsPerWarp;

  TritonLLVMOpBuilder b(loc, rewriter);

  // If this is being created inside a warp specialize op, compute the relative
  // thread ID within the warp group.
  if (std::optional<int> startId =
          getWarpGroupStartThreadId(rewriter.getInsertionBlock())) {
    tid = arith::SubIOp::create(rewriter, loc, tid, b.i32_val(*startId));
  }

  assert(llvm::isPowerOf2_32(upperBound));
  // help LLVM's known bits analysis:
  tid = b.and_(tid, b.i32_val(upperBound - 1));

  return tid;
}

std::pair<Value, Value> getLaneAndWarpId(OpBuilder &rewriter, Location loc) {
  TritonLLVMOpBuilder b(loc, rewriter);
  Value tid = getThreadId(rewriter, loc);
  int threadsPerWarp = triton::gpu::lookupThreadsPerWarp(rewriter);
  Value warpSizeVal = b.i32_val(threadsPerWarp);

  // If there is only one warp, the warp ID is always 0.
  Operation *lookupPt = &rewriter.getInsertionBlock()->front();
  Value laneId;
  Value warpId;
  if (triton::gpu::lookupNumWarps(lookupPt) == 1) {
    laneId = tid;
    warpId = b.i32_val(0);
  } else {
    laneId = b.urem(tid, warpSizeVal);
    warpId = mlir::triton::gpu::WarpIdOp::create(rewriter, loc,
                                                 /*omitUniformHint=*/true);
  }

  return {laneId, warpId};
}

Value getLaneId(OpBuilder &rewriter, Location loc) {
  return getLaneAndWarpId(rewriter, loc).first;
}

// Helper function: applies linear layout vectorized over register indices
static SmallVector<SmallVector<std::pair<StringAttr, Value>>>
applyLinearLayoutVec(Location loc, RewriterBase &rewriter,
                     const LinearLayout &layout,
                     ArrayRef<std::pair<StringAttr, Value>> indices,
                     ArrayRef<uint32_t> registers) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  MLIRContext *ctx = rewriter.getContext();

  StringAttr kRegister = str_attr("register");

  // Precompute the base (with register = 0)
  SmallVector<std::pair<StringAttr, Value>> indicesWithZeroReg;
  for (const auto &[attr, val] : indices) {
    if (attr == kRegister)
      indicesWithZeroReg.emplace_back(attr, b.i32_val(0));
    else
      indicesWithZeroReg.emplace_back(attr, val);
  }

  auto baseIndices =
      applyLinearLayout(loc, rewriter, layout, indicesWithZeroReg);

  SmallVector<SmallVector<std::pair<StringAttr, Value>>> ret;

  // Iterate over registers, applying XOR trick
  for (auto reg : registers) {
    SmallVector<std::pair<StringAttr, int32_t>> constRegIndices;
    for (const auto &[attr, val] : indices) {
      constRegIndices.emplace_back(attr, attr == kRegister ? reg : 0);
    }
    auto regIndices = layout.apply(constRegIndices);

    SmallVector<std::pair<StringAttr, Value>> combinedIndices;
    for (auto [base, regIdx] : llvm::zip(baseIndices, regIndices)) {
      assert(base.first == regIdx.first);
      Value combined = b.xor_(base.second, b.i32_val(regIdx.second));
      combinedIndices.emplace_back(base.first, combined);
    }

    ret.push_back(combinedIndices);
  }

  return ret;
}

// Refactored emitIndices function using applyLinearLayoutVec
SmallVector<SmallVector<Value>>
emitIndices(Location loc, RewriterBase &rewriter, const TargetInfoBase &target,
            const LinearLayout &ll, RankedTensorType type, bool withCTAOffset) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  MLIRContext *ctx = rewriter.getContext();
  auto shape = type.getShape();

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

  auto [laneId, warpId] = getLaneAndWarpId(rewriter, loc);
  Value blockId =
      withCTAOffset ? target.getClusterCTAId(rewriter, loc) : b.i32_val(0);

  SmallVector<std::pair<StringAttr, Value>> commonIndices = {
      {kRegister, b.i32_val(0)},
      {kLane, laneId},
      {kWarp, warpId},
      {kBlock, blockId}};

  // Vectorize over registers
  SmallVector<uint32_t> registerIndices;
  for (unsigned reg = 0; reg < ll.getInDimSize(kRegister); ++reg)
    registerIndices.push_back(reg);

  auto vecIndices =
      applyLinearLayoutVec(loc, rewriter, ll, commonIndices, registerIndices);

  unsigned rank = shape.size();
  SmallVector<SmallVector<Value>> ret;
  for (auto &indices : vecIndices) {
    SmallVector<Value> vals;
    assert(indices.size() == rank);
    for (auto &idx : indices)
      vals.push_back(idx.second);
    ret.push_back(vals);
  }

  return ret;
}

SmallVector<SmallVector<Value>>
emitIndices(Location loc, RewriterBase &rewriter, const TargetInfoBase &target,
            Attribute layout, RankedTensorType type, bool withCTAOffset) {
  LinearLayout ll = triton::gpu::toLinearLayout(type.getShape(), layout);
  return emitIndices(loc, rewriter, target, ll, type, withCTAOffset);
}

SmallVector<std::pair<Value, Value>> computeBlockLocalOffsets(
    Location loc, triton::gpu::MemDescType memDescTy,
    const LinearLayout &regLayout, ArrayRef<Value> idxValues, unsigned axis,
    RewriterBase &rewriter, const TargetInfoBase &targetInfo) {
  MLIRContext *ctx = memDescTy.getContext();
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  auto sharedLayout = triton::gpu::toLinearLayoutIgnoringPadding(memDescTy);
  auto allDims = standardOutDimNames(ctx, memDescTy.getRank());
  auto kRegister = str_attr("register");
  auto kLane = str_attr("lane");
  auto kWarp = str_attr("warp");
  auto kOffset = str_attr("offset");
  auto kBlock = str_attr("block");
  StringAttr axisDim = allDims[axis];
  LinearLayout orderedRegLayout = regLayout.transposeOuts(allDims);
  assert(idxValues.size() == regLayout.getInDimSize(kRegister));

  auto toI32 = [&](Value idx) -> Value {
    unsigned idxWidth = idx.getType().getIntOrFloatBitWidth();
    if (idxWidth > 32)
      return b.trunc(i32_ty, idx);
    if (idxWidth < 32)
      return b.zext(i32_ty, idx);
    return idx;
  };

  SmallVector<std::pair<Value, Value>> offsetAndBlock;
  offsetAndBlock.reserve(idxValues.size());

  auto [laneId, warpId] = getLaneAndWarpId(rewriter, loc);
  SmallVector<uint32_t> registerIndices;
  for (uint32_t i = 0; i < idxValues.size(); ++i)
    registerIndices.push_back(i);
  // Replace the statically-known coordinate along axis with an independent
  // input carrying the runtime index.
  SmallVector<StringAttr> nonIndexedDims = allDims;
  nonIndexedDims.erase(nonIndexedDims.begin() + axis);
  assert(!orderedRegLayout.hasInDim(axisDim));
  LinearLayout indexedLayout =
      orderedRegLayout.sublayout(
          llvm::to_vector(orderedRegLayout.getInDimNames()), nonIndexedDims) *
      LinearLayout::identity1D(sharedLayout.getOutDimSize(axisDim), axisDim,
                               axisDim);
  indexedLayout = indexedLayout.transposeOuts(allDims);
  LinearLayout cvt = invertAndComposeBlockLocal(sharedLayout, indexedLayout);
  bool crossCTA = !cvt.isIdentityOnOutDim(kBlock);
  bool blockAffectsOffset = !cvt.sublayoutIsZero({kBlock}, {kOffset});

  Value blockId = b.i32_val(0);
  if (crossCTA || blockAffectsOffset)
    blockId = targetInfo.getClusterCTAId(rewriter, loc);

  auto bases = applyLinearLayoutVec(loc, rewriter, cvt,
                                    {{kRegister, b.i32_val(0)},
                                     {kLane, laneId},
                                     {kWarp, warpId},
                                     {kBlock, blockId},
                                     {axisDim, b.i32_val(0)}},
                                    registerIndices);
  LinearLayout indexCvt = cvt.sublayout({axisDim}, {kOffset, kBlock});
  for (auto [base, idxVal] : llvm::zip(bases, idxValues)) {
    auto index =
        applyLinearLayout(loc, rewriter, indexCvt, {{axisDim, toI32(idxVal)}});
    Value offset = b.xor_(base[0].second, index[0].second);
    Value block = b.xor_(base[1].second, index[1].second);
    offsetAndBlock.push_back({offset, crossCTA ? block : Value()});
  }
  return offsetAndBlock;
}

SmallVector<LocalSharedMemoryAddress>
materializeLocalAddrs(Location loc, triton::gpu::MemDescType memDescTy,
                      SharedMemoryObject smemObj, Type llvmElemTy,
                      ArrayRef<std::pair<Value, Value>> offsetAndBlock,
                      RewriterBase &rewriter) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);

  // Get the subslice affine offset and target CTA.
  auto [affineOffset, affineBlockOffset] =
      smemObj.getShmemOffsetAndBlock(loc, rewriter, memDescTy);
  bool hasAffineBlock =
      SharedMemoryObject::getMaskSpanOffsetsAndBlocks(memDescTy).second != 0;
  auto bitwidth = getIntOrFloatOrPtrBitWidth(llvmElemTy);
  SmallVector<LocalSharedMemoryAddress> addrs;
  addrs.reserve(offsetAndBlock.size());
  for (auto [offset, blockId] : offsetAndBlock) {
    // For subslices, the physical offset and target CTA are computed as:
    //   physical_offset = L⁻¹(coords) ⊕ L⁻¹(subslice_logical_offset)
    //   target_cta = block(L⁻¹(coords)) ⊕
    //                block(L⁻¹(subslice_logical_offset))
    //
    // We use XOR for consistency with lowerLdSt. MemDescSubsliceOp::verify()
    // enforces:
    // 1. Subslice offsets must be multiples of the tile size
    // 2. Subslice offsets must map to power-of-2 physical offsets
    //
    // These constraints ensure the bit ranges of L⁻¹(coords) and
    // L⁻¹(subslice_offset) are disjoint, so XOR and addition are equivalent.
    offset = b.xor_(offset, affineOffset);
    if (hasAffineBlock) {
      if (!blockId)
        blockId = b.i32_val(0);
      blockId = b.xor_(blockId, affineBlockOffset);
    }

    // Add padding offset for padded layouts (non-linear component)
    Value ptr;
    if (triton::gpu::isPaddedEncoding(memDescTy.getEncoding())) {
      // Convert offset to bytes for padding calculation
      Value offsetBytes = b.mul(offset, b.i32_val(bitwidth / 8));
      auto shifts = getPaddedSharedShifts(memDescTy.getEncoding(), bitwidth,
                                          /*offsetInBytes=*/true);
      // GEP in bytes: base + offset*elemSize + padOffset
      Value totalOffset = applyPadding(loc, rewriter, offsetBytes, shifts);
      ptr = b.gep(smemObj.getBase().getType(), i8_ty, smemObj.getBase(),
                  totalOffset);
    } else {
      ptr = b.gep(smemObj.getBase().getType(), llvmElemTy, smemObj.getBase(),
                  offset);
    }
    addrs.push_back({ptr, blockId});
  }
  return addrs;
}

FailureOr<LocalAtomicScatterRMWInfo> prepareLocalAtomicScatterRMW(
    triton::gpu::LocalAtomicScatterRMWOp op, Value dst, Value indices,
    Value inputValues, Value mask, ConversionPatternRewriter &rewriter,
    const TargetInfoBase &targetInfo, const LLVMTypeConverter *typeConverter) {
  auto loc = op.getLoc();
  auto *ctx = op.getContext();
  auto valuesTy = cast<RankedTensorType>(op.getValues().getType());
  auto memDescTy = cast<triton::gpu::MemDescType>(op.getDst().getType());
  if (isa<triton::gpu::PartitionedSharedEncodingAttr>(
          memDescTy.getEncoding())) {
    return failure();
  }

  auto llvmElemTy = typeConverter->convertType(memDescTy.getElementType());
  auto smemObj =
      LLVM::getSharedMemoryObjectFromStruct(loc, dst, llvmElemTy, rewriter);
  SmallVector<Value> idxValues =
      unpackUniqueTensorElements(loc, indices, rewriter);
  SmallVector<Value> values =
      unpackUniqueTensorElements(loc, inputValues, rewriter);
  SmallVector<Value> maskValues;
  if (mask)
    maskValues = unpackUniqueTensorElements(loc, mask, rewriter);

  LinearLayout regLayout = triton::gpu::toLinearLayout(valuesTy);
  auto freeVarMasks = regLayout.getFreeVariableMasks();
  auto removeBroadcast = actionRemoveBroadcastedRegs(regLayout);
  Value threadPred = triton::gpu::emitRedundantThreadPredicate(
      freeVarMasks, rewriter, loc, targetInfo);
  LinearLayout activeRegLayout =
      regLayout.removeZeroBasesAlongDim(str_attr("register"));
  auto offsetAndBlock =
      computeBlockLocalOffsets(loc, memDescTy, activeRegLayout, idxValues,
                               op.getAxis(), rewriter, targetInfo);
  SmallVector<LocalSharedMemoryAddress> addrs = materializeLocalAddrs(
      loc, memDescTy, smemObj, llvmElemTy, offsetAndBlock, rewriter);

  return LocalAtomicScatterRMWInfo{valuesTy,        llvmElemTy, regLayout,
                                   removeBroadcast, threadPred, values,
                                   maskValues,      addrs};
}

SmallVector<std::pair<unsigned, unsigned>>
getPaddedSharedShifts(Attribute enc, unsigned bitwidth, bool offsetInBytes) {
  auto padded = triton::gpu::getPaddedEncoding(enc);
  if (!padded)
    return {};

  SmallVector<std::pair<unsigned, unsigned>> shifts;
  assert(bitwidth >= 8 && (bitwidth % 8 == 0) &&
         "bitwidth must be a positive multiple of 8 for padding");
  uint64_t offScale = offsetInBytes ? (bitwidth / 8) : 1;
  for (auto [interval, padding] :
       llvm::zip_equal(padded.getIntervals(), padded.getPaddings())) {
    uint64_t intervalScaled = static_cast<uint64_t>(interval) * offScale;
    uint64_t paddingScaled = static_cast<uint64_t>(padding) * offScale;
    unsigned i = llvm::Log2_64(intervalScaled);
    unsigned p = llvm::Log2_64(paddingScaled);
    assert(i < 32 && p < 32 && "shift amount must be < 32 for i32 offsets");
    shifts.push_back({i, p});
  }
  return shifts;
}

Value applyPadding(Location loc, RewriterBase &rewriter, Value baseOffset,
                   ArrayRef<std::pair<unsigned, unsigned>> shifts) {
  if (shifts.empty())
    return baseOffset;
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  Value pad = b.i32_val(0);
  for (auto [i, p] : shifts)
    pad = b.add(pad, b.shl(b.lshr(baseOffset, b.i32_val(i)), b.i32_val(p)));
  return b.add(baseOffset, pad);
}

uint32_t applyPadding(uint32_t baseOffset,
                      ArrayRef<std::pair<unsigned, unsigned>> shifts) {
  uint64_t pad = 0;
  for (auto [i, p] : shifts)
    pad += (static_cast<uint64_t>(baseOffset) >> i) << p;
  uint64_t out = baseOffset + pad;
  assert(out <= std::numeric_limits<uint32_t>::max() &&
         "padded offset must be within 32-bit range");
  return static_cast<uint32_t>(out);
}

LowerLdStCallback makeSharedStoreEmitter(const TargetInfoBase &targetInfo,
                                         Value pred) {
  return [&targetInfo, pred](RewriterBase &rewriter, Location loc,
                             ArrayRef<Value> vals, Value shmemAddr, int idx,
                             VectorType vecTy,
                             Value ctaId) -> SmallVector<Value> {
    auto length = vecTy.getNumElements();
    Value valsVec =
        packLLVector(loc, ArrayRef<Value>(vals).slice(idx, length), rewriter);
    targetInfo.storeDShared(rewriter, loc, shmemAddr, ctaId, valsVec, pred);
    return {};
  };
}

LowerLdStCallback makeSharedLoadEmitter(const TargetInfoBase &targetInfo,
                                        Operation *localLoadOp) {
  return [&targetInfo, localLoadOp](RewriterBase &rewriter, Location loc,
                                    ArrayRef<Value> vals, Value shmemAddr, int,
                                    VectorType vecTy,
                                    Value ctaId) -> SmallVector<Value> {
    assert(vals.empty());
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    Value valsVec =
        targetInfo.loadDShared(rewriter, loc, shmemAddr, ctaId, vecTy,
                               /*pred=*/b.true_val(), localLoadOp);
    return unpackLLVector(loc, valsVec, rewriter);
  };
}

SmallVector<Value>
lowerLdSt(Location loc, MLIRContext *ctx, LinearLayout cvt,
          ArrayRef<Value> valsArray, // Input for store, output for load
          Type llvmElemTy, ArrayRef<Value> smemBases,
          ArrayRef<std::pair<unsigned, unsigned>> paddingShifts,
          Value affineOffset, uint64_t maskSpanAffineOffset,
          Value affineBlockOffset, uint64_t maskSpanAffineBlock, Value laneId,
          Value warpId, RewriterBase &rewriter,
          const TargetInfoBase &targetInfo, std::optional<int> maybeMaxVecElems,
          LowerLdStCallback lowerInst) {
  assert(!smemBases.empty() && "smemBases cannot be empty");
  auto vals = to_vector(valsArray);
  bool isStore = !vals.empty();
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  auto smemPtrTy = ptr_ty(ctx, targetInfo.getSharedAddressSpace());
  auto kReg = str_attr("register");
  auto kLane = str_attr("lane");
  auto kWarp = str_attr("warp");
  auto kBlock = str_attr("block");
  auto kOffset = str_attr("offset");
  auto kPartition = str_attr("partition");
  auto bitwidth = getIntOrFloatOrPtrBitWidth(llvmElemTy);
  bool crossCTA = !cvt.isIdentityOnOutDim(kBlock);

  // Either we have multiple bases with a matching partition dimension,
  // or we have a single base.
  assert((smemBases.size() == 1 ||
          (cvt.hasOutDim(kPartition) &&
           cvt.getOutDimSize(kPartition) == smemBases.size())) &&
         "smemBases size must match partition dimension size");
  bool isPartitioned = smemBases.size() > 1;

  // Extract the partition sublayout before stripping it for vectorization.
  auto inDimNames = to_vector(cvt.getInDimNames());
  LinearLayout partitionLayout;
  Value basesVec;
  if (isPartitioned) {
    partitionLayout = cvt.sublayout(inDimNames, {kPartition});
    basesVec = LLVM::buildBasePtrVector(loc, rewriter, smemBases);
  }

  // Strip kPartition output for vectorization analysis.
  SmallVector<StringAttr> outDims = to_vector(cvt.getOutDimNames());
  llvm::erase(outDims, kPartition);
  cvt = cvt.sublayout(inDimNames, outDims);

  auto [elemsPerVec, permutation] =
      largestVectorisation(ctx, cvt, bitwidth, maybeMaxVecElems);

  cvt = permutation.apply(cvt);
  if (isStore) {
    vals = permutation.apply(vals);
  }

  auto tile = LinearLayout::identity1D(elemsPerVec, kReg, kOffset);
  auto quot = divideLeft(cvt, tile);
  assert(quot.has_value() && "cvt must be divisible by tile");
  LinearLayout reps = zerosLike(tile) * *quot;
  assert(reps.hasInDim(kBlock));
  LinearLayout addrLayout =
      LinearLayout({{kLane, reps.getBases().lookup(kLane)},
                    {kWarp, reps.getBases().lookup(kWarp)},
                    {kBlock, reps.getBases().lookup(kBlock)}},
                   reps.getOutDims(), false);
  auto [nAdditive, permStrides] = actionAdditiveStrides(
      reps, addrLayout, maskSpanAffineOffset, maskSpanAffineBlock, elemsPerVec);
  reps = permStrides.apply(reps);

  if (isPartitioned) {
    partitionLayout = permutation.apply(partitionLayout);
    partitionLayout = permStrides.apply(partitionLayout);
  }
  if (isStore) {
    vals = permStrides.apply(vals);
  }

  // PTX expects the address increments to be done in bytes
  // If we don't perform the computations in i8, the compiler would
  // have to divide the computation by bitwdith / 8 and then lift this
  // shl, which often it's not able to do.
  auto i8Tile =
      LinearLayout::zeros1D(bitwidth / 8, kReg, kOffset, bitwidth / 8);
  auto i8AddrLayout = i8Tile * addrLayout;

  bool layoutUsesBlockId = !reps.isTrivialOver({kBlock});
  bool useBlockId = layoutUsesBlockId || maskSpanAffineBlock != 0;
  Value blockId = b.i32_val(0);
  if (layoutUsesBlockId) {
    blockId = targetInfo.getClusterCTAId(rewriter, loc);
  }

  auto baseI8AndCTA = applyLinearLayout(loc, rewriter, i8AddrLayout,
                                        {{kReg, b.i32_val(0)},
                                         {kLane, laneId},
                                         {kWarp, warpId},
                                         {kBlock, blockId}});
  auto regBaseI8 = baseI8AndCTA[0].second;
  Value targetCtaId;
  if (crossCTA) {
    targetCtaId = baseI8AndCTA[1].second;
    if (maskSpanAffineBlock != 0)
      targetCtaId = b.xor_(targetCtaId, affineBlockOffset);
  }

  // It's fine that we don't compute the offset in bytes as affineOffset
  // will be folded into a constant
  auto affineOffsetI8 = b.mul(affineOffset, b.i32_val(bitwidth / 8));
  bool hasPadding = !paddingShifts.empty();
  Value paddedAffineOffsetI8 = b.i32_val(0);
  if (hasPadding && maskSpanAffineOffset != 0) {
    // `maskSpanAffineOffset != 0` indicates the affine offsets come from
    // MemDescSubsliceOp, whose verifier guarantees that the affine offsets are
    // bitwise disjoint from other offset contributors. Padding can thus be
    // applied separately. This helps LLVM reuse base pointers.
    paddedAffineOffsetI8 =
        applyPadding(loc, rewriter, affineOffsetI8, paddingShifts);
  } else {
    regBaseI8 = b.xor_(regBaseI8, affineOffsetI8);
  }

  // Pre-compute the base pointers for partitioned tensors.
  SmallVector<Value> partitionBases;
  if (isPartitioned) {
    // Loop-invariant partition contribution (register = 0).
    Value dynamicPartition = applyLinearLayout(loc, rewriter, partitionLayout,
                                               {{kReg, b.i32_val(0)},
                                                {kLane, laneId},
                                                {kWarp, warpId},
                                                {kBlock, blockId}})[0]
                                 .second;
    unsigned numPartitions = smemBases.size();
    partitionBases.resize(numPartitions);
    for (unsigned c = 0; c < numPartitions; ++c) {
      // partitionBases[c] holds the base for register partitions that map to
      // `c`, so reordering by `dynamicPartition ^ c` folds the loop-invariant
      // part in once.
      Value idx = b.xor_(dynamicPartition, b.i32_val(c));
      partitionBases[c] = b.extract_element(basesVec, idx);
    }
  }

  SmallVector<Value> outVals;
  auto vecTy = vec_ty(llvmElemTy, elemsPerVec);
  for (int i = 0; i < cvt.getInDimSize(kReg); i += nAdditive) {
    auto idxAndBlock =
        reps.apply({{kReg, i}, {kLane, 0}, {kWarp, 0}, {kBlock, 0}});
    auto regIdxI8 = idxAndBlock[0].second * (bitwidth / 8);
    Value offset = b.xor_(regBaseI8, b.i32_val(regIdxI8));
    if (hasPadding) {
      offset = applyPadding(loc, rewriter, offset, paddingShifts);
      if (maskSpanAffineOffset != 0)
        offset = b.add(offset, paddedAffineOffsetI8);
    }
    Value ctaOffset = b.i32_val(0);
    if (crossCTA) {
      ctaOffset = b.xor_(targetCtaId, b.i32_val(idxAndBlock[1].second));
    }
    for (int j = 0; j < nAdditive; j += elemsPerVec) {
      // all these constants will go as immediate values to LDS/STS
      auto idxAndBlockAdd =
          reps.apply({{kReg, j}, {kLane, 0}, {kWarp, 0}, {kBlock, 0}});
      auto regIdxAddI8 = idxAndBlockAdd[0].second * (bitwidth / 8);
      // `actionAdditiveStrides` forces `regIdxAddI8` and `offset` to be bitwise
      // disjoint, so we can calculate their padding contributions separately.
      regIdxAddI8 = applyPadding(regIdxAddI8, paddingShifts);
      Value innerOffset = b.add(offset, b.i32_val(regIdxAddI8));

      // Select the appropriate base pointer for partitioned tensors
      Value smemBase = smemBases[0];
      if (isPartitioned) {
        // The register-only partition contribution is a compile-time constant,
        // so it indexes the pre-reordered array as a literal.
        unsigned regPartition =
            partitionLayout
                .apply({{kReg, i + j}, {kLane, 0}, {kWarp, 0}, {kBlock, 0}})[0]
                .second;
        smemBase = partitionBases[regPartition];
      }

      Value innerCtaOffset;
      if (crossCTA) {
        innerCtaOffset = b.add(ctaOffset, b.i32_val(idxAndBlockAdd[1].second));
      }
      auto vecAddr = b.gep(smemPtrTy, i8_ty, smemBase, innerOffset,
                           LLVM::GEPNoWrapFlags::inbounds);
      llvm::append_range(outVals, lowerInst(rewriter, loc, vals, vecAddr, i + j,
                                            vecTy, innerCtaOffset));
    }
  }

  // Permute the values back if we are loading
  if (!isStore) {
    auto invPermStrides = permStrides.inverse();
    outVals = invPermStrides.apply(outVals);
    auto invPerm = permutation.inverse();
    outVals = invPerm.apply(outVals);
  }
  return outVals;
}

SmallVector<Value>
lowerLocalLdSt(Location loc, MLIRContext *ctx,
               LinearLayout cvt, // Map from registers to offset[, partition]
               ArrayRef<Value> valsArray, // Input for store, empty for load
               Type llvmElemTy, triton::gpu::MemDescType srcTy,
               SharedMemoryObject smemObj, RewriterBase &rewriter,
               const TargetInfoBase &targetInfo, LowerLdStCallback lowerInst) {

  auto isStore = !valsArray.empty();
  // Remove broadcasting in the registers
  auto removeBroadcastSrc = actionRemoveBroadcastedRegs(cvt);
  if (!removeBroadcastSrc.isIdentity()) {
    auto prmtCvt = removeBroadcastSrc.apply(cvt);
    auto inVals = to_vector(valsArray);
    if (isStore) {
      inVals = removeBroadcastSrc.apply(inVals);
    }
    auto outVals = lowerLocalLdSt(loc, ctx, prmtCvt, inVals, llvmElemTy, srcTy,
                                  smemObj, rewriter, targetInfo, lowerInst);
    if (!isStore) {
      outVals = broadcastAs(outVals, cvt);
    }
    return outVals;
  }
  auto [affineOffset, affineBlockOffset] =
      smemObj.getShmemOffsetAndBlock(loc, rewriter, srcTy);
  auto [maskSpanAffineOffset, maskSpanAffineBlock] =
      smemObj.getMaskSpanOffsetsAndBlocks(srcTy);

  // Extract padding info from padded encoding (standalone or inside
  // partitioned)
  std::optional<int> maybeMaxVecElems;
  SmallVector<std::pair<unsigned, unsigned>> paddingShifts;
  if (triton::gpu::isPaddedEncoding(srcTy.getEncoding())) {
    maybeMaxVecElems = triton::gpu::getMinInterval(srcTy.getEncoding());
    auto bitwidth = getIntOrFloatOrPtrBitWidth(llvmElemTy);
    paddingShifts = getPaddedSharedShifts(srcTy.getEncoding(), bitwidth,
                                          /*offsetInBytes=*/true);
  }

  // Get shared memory bases from the SharedMemoryObject.
  // For non-partitioned tensors, this returns a single base.
  // For partitioned tensors, this returns all bases (one per partition).
  SmallVector<Value> smemBases(smemObj.getBases().begin(),
                               smemObj.getBases().end());
  auto [laneId, warpId] = getLaneAndWarpId(rewriter, loc);
  return lowerLdSt(loc, ctx, cvt, valsArray, llvmElemTy, smemBases,
                   paddingShifts, affineOffset, maskSpanAffineOffset,
                   affineBlockOffset, maskSpanAffineBlock, laneId, warpId,
                   rewriter, targetInfo, maybeMaxVecElems, lowerInst);
}

SmallVector<Value> unpackLLElements(Location loc, Value llvmStruct,
                                    RewriterBase &rewriter) {
  assert(bool(llvmStruct) && "can not unpack null values");
  if (llvmStruct.getType().isIntOrIndexOrFloat() ||
      isa<triton::PointerType>(llvmStruct.getType()) ||
      isa<LLVM::LLVMPointerType>(llvmStruct.getType()))
    return {llvmStruct};
  ArrayRef<Type> types =
      cast<LLVM::LLVMStructType>(llvmStruct.getType()).getBody();
  SmallVector<Value> results(types.size());
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  for (unsigned i = 0; i < types.size(); ++i) {
    Type type = types[i];
    results[i] = b.extract_val(type, llvmStruct, i);
  }
  return results;
}

SmallVector<Value> unpackUniqueTensorElements(Location loc, Value llvmStruct,
                                              RewriterBase &rewriter) {
  return unpackLLElements(loc, llvmStruct, rewriter);
}

SmallVector<Value> unpackTensorElements(Location loc, Value llvmStruct,
                                        RewriterBase &rewriter,
                                        Type originalType) {
  if (auto tensorTy = dyn_cast<RankedTensorType>(originalType))
    return broadcastAs(unpackUniqueTensorElements(loc, llvmStruct, rewriter),
                       triton::gpu::toLinearLayout(tensorTy));
  return unpackLLElements(loc, llvmStruct, rewriter);
}

Value packLLElements(Location loc, const LLVMTypeConverter *typeConverter,
                     ValueRange resultVals, RewriterBase &rewriter, Type type) {
  auto structType =
      dyn_cast<LLVM::LLVMStructType>(typeConverter->convertType(type));
  if (!structType) {
    assert(resultVals.size() == 1);
    return *resultVals.begin();
  }

  auto elementTypes = structType.getBody();
  if (elementTypes.size() != resultVals.size()) {
    emitError(loc) << " size mismatch when packing elements for LLVM struct"
                   << " expected " << elementTypes.size() << " but got "
                   << resultVals.size();
    llvm::report_fatal_error(
        "size mismatch when packing elements for LLVM struct");
  }
  Value llvmStruct = LLVM::UndefOp::create(rewriter, loc, structType);
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  for (auto [i, value] : llvm::enumerate(resultVals)) {
    assert(value && "unexpected null value");
    if (value.getType() != elementTypes[i]) {
      LDBG("type " << type << " structType " << structType);
      LDBG("value " << value);
      emitError(loc) << "invalid element type in packLLElements. Expected "
                     << elementTypes[i] << " but got " << value.getType();
      llvm::report_fatal_error(
          "element type mismatch when packing elements for LLVM struct");
    }
    llvmStruct = b.insert_val(structType, llvmStruct, value, i);
  }
  return llvmStruct;
}

Value packUniqueTensorElements(Location loc,
                               const LLVMTypeConverter *typeConverter,
                               ValueRange resultVals, RewriterBase &rewriter,
                               Type type) {
  return packLLElements(loc, typeConverter, resultVals, rewriter, type);
}

Value packTensorElements(Location loc, const LLVMTypeConverter *typeConverter,
                         ValueRange resultVals, RewriterBase &rewriter,
                         Type type) {
  if (auto tensorTy = dyn_cast<RankedTensorType>(type)) {
    auto uniqueResultVals =
        actionRemoveBroadcastedRegs(triton::gpu::toLinearLayout(tensorTy))
            .apply(resultVals);
    return packUniqueTensorElements(loc, typeConverter, uniqueResultVals,
                                    rewriter, type);
  }
  return packLLElements(loc, typeConverter, resultVals, rewriter, type);
}

SmallVector<Value> unpackLLVector(Location loc, Value llvmVec,
                                  RewriterBase &rewriter) {
  assert(bool(llvmVec) && "cannot unpack null value");
  if (llvmVec.getType().isIntOrIndexOrFloat() ||
      isa<triton::PointerType>(llvmVec.getType()) ||
      isa<LLVM::LLVMPointerType>(llvmVec.getType()))
    return {llvmVec};

  auto b = TritonLLVMOpBuilder(loc, rewriter);
  SmallVector<Value> results;
  for (int i = 0; i < cast<VectorType>(llvmVec.getType()).getNumElements();
       i++) {
    results.push_back(b.extract_element(llvmVec, b.i32_val(i)));
  }
  return results;
}

Value packLLVector(Location loc, ValueRange vals, RewriterBase &rewriter) {
  assert(vals.size() > 0);
  auto vecType = vec_ty(vals[0].getType(), vals.size());
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  Value vec = b.undef(vecType);
  for (int i = 0; i < vals.size(); i++) {
    vec = b.insert_element(vec, vals[i], b.i32_val(i));
  }
  return vec;
}

std::optional<LLVM::AtomicBinOp> matchAtomicOp(RMWOp atomicOp) {
  switch (atomicOp) {
  case RMWOp::AND:
    return LLVM::AtomicBinOp::_and;
  case RMWOp::OR:
    return LLVM::AtomicBinOp::_or;
  case RMWOp::XOR:
    return LLVM::AtomicBinOp::_xor;
  case RMWOp::ADD:
    return LLVM::AtomicBinOp::add;
  case RMWOp::FADD:
    return LLVM::AtomicBinOp::fadd;
  case RMWOp::MAX:
    return LLVM::AtomicBinOp::max;
  case RMWOp::MIN:
    return LLVM::AtomicBinOp::min;
  case RMWOp::UMAX:
    return LLVM::AtomicBinOp::umax;
  case RMWOp::UMIN:
    return LLVM::AtomicBinOp::umin;
  case RMWOp::XCHG:
    return LLVM::AtomicBinOp::xchg;
  default:
    return {};
  }
}

std::optional<LLVM::AtomicOrdering> getMemoryOrdering(MemSemantic memOrdering) {
  switch (memOrdering) {
  case MemSemantic::RELAXED:
    return LLVM::AtomicOrdering::monotonic;
  case MemSemantic::ACQUIRE:
    return LLVM::AtomicOrdering::acquire;
  case MemSemantic::RELEASE:
    return LLVM::AtomicOrdering::release;
  case MemSemantic::ACQUIRE_RELEASE:
    return LLVM::AtomicOrdering::acq_rel;
  default:
    return {};
  }
}

void insertAtomicOrderingBarriers(Operation *op, MemSemantic memOrdering,
                                  bool emitBarrierAfter, RewriterBase &rewriter,
                                  const TargetInfoBase &targetInfo) {
  auto emitBarrier = [&] {
    if (triton::gpu::lookupNumCTAs(op) == 1)
      targetInfo.barrier(op->getLoc(), rewriter, triton::gpu::AddrSpace::Local);
    else
      targetInfo.clusterBarrier(op->getLoc(), rewriter, op);
  };

  if (memOrdering == MemSemantic::RELEASE ||
      memOrdering == MemSemantic::ACQUIRE_RELEASE) {
    OpBuilder::InsertionGuard guard(rewriter);
    rewriter.setInsertionPoint(op);
    emitBarrier();
  }
  if (emitBarrierAfter && (memOrdering == MemSemantic::ACQUIRE ||
                           memOrdering == MemSemantic::ACQUIRE_RELEASE)) {
    OpBuilder::InsertionGuard guard(rewriter);
    rewriter.setInsertionPointAfter(op);
    emitBarrier();
  }
}

Value broadcastScalarAtomicResult(Operation *op, Type valueElemTy,
                                  Value resultVal,
                                  ConversionPatternRewriter &rewriter,
                                  TritonLLVMOpBuilder &b, Value threadPred,
                                  const TargetInfoBase &targetInfo) {
  if (!op->hasAttr("allocation.offset"))
    return resultVal;

  auto loc = op->getLoc();
  Value smemBase = LLVM::getSharedMemoryBase(loc, rewriter, targetInfo, op);
  targetInfo.storeShared(rewriter, loc, smemBase, resultVal, threadPred);
  if (triton::gpu::lookupNumCTAs(op) == 1) {
    targetInfo.barrier(loc, rewriter, triton::gpu::AddrSpace::Local);
    return targetInfo.loadShared(rewriter, loc, smemBase, valueElemTy,
                                 b.true_val());
  }

  targetInfo.clusterBarrier(loc, rewriter, op);
  return targetInfo.loadDShared(rewriter, loc, smemBase, b.i32_val(0),
                                valueElemTy, b.true_val());
}

llvm::MapVector<StringAttr, int32_t> getAllFreeVarMasks(MLIRContext *ctx) {
  // Mask where all elements are redundant
  auto kReg = str_attr("reg");
  auto kLane = str_attr("lane");
  auto kWarp = str_attr("warp");
  auto kBlock = str_attr("block");

  int32_t fullMask = -1;
  llvm::MapVector<StringAttr, int32_t> ret;
  for (auto dimName : {kReg, kLane, kWarp, kBlock}) {
    ret[dimName] = fullMask;
  }
  return ret;
}

llvm::MapVector<StringAttr, int32_t> getFreeVariableMasks(Type type) {
  auto ctx = type.getContext();
  auto tensorTy = dyn_cast<RankedTensorType>(type);
  if (!tensorTy) {
    return getAllFreeVarMasks(ctx);
  }
  auto ll = triton::gpu::toLinearLayout(tensorTy);
  return ll.getFreeVariableMasks();
}

SmallVector<SmallVector<unsigned>> emitOffsetForLayout(Attribute layout,
                                                       RankedTensorType type) {
  MLIRContext *ctx = layout.getContext();
  auto shape = type.getShape();
  unsigned rank = shape.size();

  auto ll = triton::gpu::toLinearLayout(type);

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

  SmallVector<SmallVector<unsigned>> offsets;
  for (int i = 0; i < ll.getInDimSize(str_attr("register")); i++) {
    auto idxs = ll.apply({{kRegister, i}, {kLane, 0}, {kWarp, 0}, {kBlock, 0}});
    assert(idxs.size() == rank);
    for (unsigned k = 0; k < rank; ++k) {
      assert(idxs[k].first == str_attr("dim" + std::to_string(k)));
    }
    offsets.push_back(
        llvm::to_vector_of<unsigned>(llvm::make_second_range(idxs)));
  }
  return offsets;
}

namespace LLVM {
using namespace mlir::triton;
using mlir::triton::gpu::getOrder;

Value createConstantI1(Location loc, OpBuilder &rewriter, bool v) {
  auto i1ty = rewriter.getIntegerType(1);
  return LLVM::ConstantOp::create(rewriter, loc, i1ty,
                                  IntegerAttr::get(i1ty, v));
}

Value createConstantI32(Location loc, OpBuilder &rewriter, int32_t v) {
  auto i32ty = rewriter.getIntegerType(32);
  return LLVM::ConstantOp::create(rewriter, loc, i32ty,
                                  IntegerAttr::get(i32ty, v));
}

Value createConstantI64(Location loc, OpBuilder &rewriter, int64_t v) {
  auto i64ty = rewriter.getIntegerType(64);
  return LLVM::ConstantOp::create(rewriter, loc, i64ty,
                                  IntegerAttr::get(i64ty, v));
}

Value createConstantF16(Location loc, OpBuilder &rewriter, float v) {
  auto type = type::f16Ty(rewriter.getContext());
  return LLVM::ConstantOp::create(rewriter, loc, type,
                                  rewriter.getF16FloatAttr(v));
}

Value createConstantBF16(Location loc, OpBuilder &rewriter, float v) {
  APFloat apf(v);
  bool ignored;
  apf.convert(APFloat::BFloat(), APFloat::rmNearestTiesToEven, &ignored);
  auto type = type::bf16Ty(rewriter.getContext());
  auto attr = FloatAttr::get(type, apf);
  return LLVM::ConstantOp::create(rewriter, loc, type, attr);
}

Value createConstantF32(Location loc, OpBuilder &rewriter, float v) {
  auto type = type::f32Ty(rewriter.getContext());
  return LLVM::ConstantOp::create(rewriter, loc, type,
                                  rewriter.getF32FloatAttr(v));
}

Value createConstantF64(Location loc, OpBuilder &rewriter, double v) {
  auto type = type::f64Ty(rewriter.getContext());
  return LLVM::ConstantOp::create(rewriter, loc, type,
                                  rewriter.getF64FloatAttr(v));
}

Value createNaNConstant(Location loc, OpBuilder &rewriter, Type type) {
  if (!isa<FloatType>(type)) {
    llvm::report_fatal_error("Creating NaN constant for non-float type!");
  }
  return LLVM::ConstantOp::create(
      rewriter, loc, type,
      APFloat::getNaN(cast<FloatType>(type).getFloatSemantics()));
}

// Create an index type constant.
Value createIndexConstant(OpBuilder &builder, Location loc,
                          const TypeConverter *converter, int64_t value) {
  Type ty = converter->convertType(builder.getIndexType());
  return LLVM::ConstantOp::create(builder, loc, ty,
                                  builder.getIntegerAttr(ty, value));
}

// Create an integer constant of \param width bits.
Value createLLVMIntegerConstant(OpBuilder &builder, Location loc, short width,
                                int64_t value) {
  Type ty = builder.getIntegerType(width);
  return LLVM::ConstantOp::create(builder, loc, ty,
                                  builder.getIntegerAttr(ty, value));
}

LLVM::CallOp createLLVMCallOp(OpBuilder &builder, Location loc,
                              LLVMFuncOp funcOp, ValueRange args) {
  auto op = LLVM::CallOp::create(builder, loc, funcOp, args);
  op.getProperties().setOpBundleSizes(builder.getDenseI32ArrayAttr({}));
  op.getProperties().setOperandSegmentSizes({static_cast<int>(args.size()), 0});
  return op;
}

LLVM::CallIntrinsicOp
createLLVMIntrinsicCallOp(OpBuilder &builder, Location loc, StringRef intrinsic,
                          TypeRange types, ValueRange args) {
  auto op = LLVM::CallIntrinsicOp::create(builder, loc, types, args);
  op.getProperties().setIntrin(builder.getStringAttr(intrinsic));
  op.getProperties().setOpBundleSizes(builder.getDenseI32ArrayAttr({}));
  op.getProperties().setOperandSegmentSizes({static_cast<int>(args.size()), 0});
  return op;
}

SharedMemoryObject::SharedMemoryObject(Value base, Type baseElemType,
                                       ArrayRef<Value> offsets)
    : bases({base}), baseElemType(baseElemType),
      offsets(offsets.begin(), offsets.end()) {}

SharedMemoryObject::SharedMemoryObject(ArrayRef<Value> bases, Type baseElemType,
                                       ArrayRef<Value> offsets)
    : bases(bases.begin(), bases.end()), baseElemType(baseElemType),
      offsets(offsets.begin(), offsets.end()) {
  assert(!bases.empty() && "SharedMemoryObject must have at least one base");
}

SharedMemoryObject::SharedMemoryObject(Value base, Type baseElemType,
                                       int64_t rank, Location loc,
                                       RewriterBase &rewriter)
    : bases({base}), baseElemType(baseElemType) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  offsets.append(rank, b.i32_val(0));
}

SharedMemoryObject::SharedMemoryObject(ArrayRef<Value> bases, Type baseElemType,
                                       int64_t rank, Location loc,
                                       RewriterBase &rewriter)
    : bases(bases.begin(), bases.end()), baseElemType(baseElemType) {
  assert(!bases.empty() && "SharedMemoryObject must have at least one base");
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  offsets.append(rank, b.i32_val(0));
}

SmallVector<Value> SharedMemoryObject::getElems() const {
  SmallVector<Value> elems;
  // All bases come first, then offsets
  elems.append(bases.begin(), bases.end());
  elems.append(offsets.begin(), offsets.end());
  return elems;
}

SmallVector<Type> SharedMemoryObject::getTypes() const {
  SmallVector<Type> types;
  // All base pointer types first
  for (Value base : bases) {
    types.push_back(base.getType());
  }
  // Then all offset types (i32)
  types.append(offsets.size(), IntegerType::get(bases[0].getContext(), 32));
  return types;
}

std::pair<uint64_t, uint64_t> SharedMemoryObject::getMaskSpanOffsetsAndBlocks(
    triton::gpu::MemDescType srcTy) {
  auto ctx = srcTy.getContext();
  auto encoding = srcTy.getEncoding();
  auto shape = triton::gpu::dropPipeliningDim(srcTy.getShape(), encoding);
  auto allocShape =
      triton::gpu::dropPipeliningDim(srcTy.getAllocShape(), encoding);

  // Early exist when there is no subview
  if (allocShape == shape) {
    return {0, 0};
  }
  auto totalLl = triton::gpu::toLinearLayoutIgnoringPadding(
      allocShape, srcTy.getEncoding());
  // Map from dimNames to offset, block
  auto invLl = totalLl.pseudoinvert();
  SmallVector<std::pair<StringAttr, int32_t>> logicalOffsets;
  for (auto dim : standardOutDimNames(ctx, shape.size())) {
    logicalOffsets.push_back({dim, 0});
  }

  uint64_t offsetMask = 0;
  uint64_t blockMask = 0;
  for (auto [dim, shapes] : llvm::enumerate(llvm::zip(shape, allocShape))) {
    auto [shape, allocShape] = shapes;
    for (int j = llvm::Log2_32(shape); j < llvm::Log2_32(allocShape); ++j) {
      logicalOffsets[dim].second = 1 << j;
      auto offsetAndBlock = invLl.apply(logicalOffsets);
      offsetMask |= offsetAndBlock[0].second;
      blockMask |= offsetAndBlock[1].second;
    }
    // Reset the offset for the next dimension
    logicalOffsets[dim].second = 0;
  }
  return {offsetMask, blockMask};
}

uint64_t
SharedMemoryObject::getMaskSpanOffsets(triton::gpu::MemDescType srcTy) {
  return getMaskSpanOffsetsAndBlocks(srcTy).first;
}

std::pair<Value, Value> SharedMemoryObject::getShmemOffsetAndBlock(
    Location loc, RewriterBase &rewriter,
    triton::gpu::MemDescType srcTy) const {
  auto ctx = srcTy.getContext();
  auto b = TritonLLVMOpBuilder(loc, rewriter);

  // If it did not have a memdesc_subslice we don't need to compute the offset
  // as it is zero
  if (!isAffineSharedMemoryAccess(srcTy)) {
    return {b.i32_val(0), b.i32_val(0)};
  }

  // We return the offset without the padding. The padding will be added in the
  // lowering
  LinearLayout ll = triton::gpu::toLinearLayoutIgnoringPadding(srcTy);

  auto layoutOffsets =
      triton::gpu::dropPipeliningDim(ArrayRef(offsets), srcTy.getEncoding());
  auto dimNames = standardOutDimNames(ctx, layoutOffsets.size());
  SmallVector<std::pair<StringAttr, Value>> logicalOffsets;
  for (auto [dim, offset] : llvm::zip(dimNames, layoutOffsets)) {
    logicalOffsets.push_back({dim, offset});
  }

  auto offsetAndBlock =
      applyLinearLayout(loc, rewriter, ll.pseudoinvert(), logicalOffsets);
  return {offsetAndBlock[0].second, offsetAndBlock[1].second};
}

Value SharedMemoryObject::getShmemOffset(Location loc, RewriterBase &rewriter,
                                         triton::gpu::MemDescType srcTy) const {
  return getShmemOffsetAndBlock(loc, rewriter, srcTy).first;
}

Value SharedMemoryObject::getShmemAffineBase(
    Location loc, RewriterBase &rewriter,
    triton::gpu::MemDescType srcTy) const {
  assert(bases.size() == 1 &&
         "getShmemAffineBase does not support partitioned shared memory");
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  Value offset = getShmemOffset(loc, rewriter, srcTy);
  return b.gep(bases[0].getType(), baseElemType, bases[0], offset);
}

Value getStructFromSharedMemoryObject(Location loc,
                                      const SharedMemoryObject &smemObj,
                                      RewriterBase &rewriter) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  auto elems = smemObj.getElems();
  auto types = smemObj.getTypes();
  auto structTy =
      LLVM::LLVMStructType::getLiteral(rewriter.getContext(), types);
  // pack into struct
  Value llvmStruct = LLVM::UndefOp::create(rewriter, loc, structTy);
  for (const auto &v : llvm::enumerate(elems)) {
    assert(v.value() && "can not insert null values");
    llvmStruct = b.insert_val(structTy, llvmStruct, v.value(), v.index());
  }
  return llvmStruct;
}

SharedMemoryObject getSharedMemoryObjectFromStruct(Location loc,
                                                   Value llvmStruct,
                                                   Type elemTy,
                                                   RewriterBase &rewriter) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  ArrayRef<Type> types =
      cast<LLVM::LLVMStructType>(llvmStruct.getType()).getBody();
  SmallVector<Value> elems(types.size());
  for (unsigned i = 0; i < types.size(); ++i) {
    Type type = types[i];
    elems[i] = b.extract_val(type, llvmStruct, i);
  }

  // Determine how many bases there are by counting pointer types at the start.
  // The struct layout is: [base0, base1, ..., baseN-1, offset0, offset1, ...]
  // All bases are pointer types, all offsets are i32.
  size_t numBases = llvm::count_if(
      types, [](Type t) { return isa<LLVM::LLVMPointerType>(t); });
  assert(numBases > 0 &&
         "SharedMemoryObject struct must have at least one base");

  return {/*bases=*/{elems.begin(), elems.begin() + numBases},
          /*baseElemType=*/elemTy,
          /*offsets=*/{elems.begin() + numBases, elems.end()}};
}

// Build a vector containing multiple base pointers for dynamic indexing.
Value buildBasePtrVector(Location loc, RewriterBase &rewriter,
                         ArrayRef<Value> smemBases) {
  assert(smemBases.size() > 1 && "Need multiple bases to build a vector");
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  auto ptrTy = smemBases[0].getType();
  auto vecTy = VectorType::get({static_cast<int64_t>(smemBases.size())}, ptrTy);
  Value basesVec = b.undef(vecTy);
  for (size_t i = 0; i < smemBases.size(); ++i) {
    basesVec = b.insert_element(basesVec, smemBases[i], b.i32_val(i));
  }
  return basesVec;
}

Value getStackPointer(RewriterBase &rewriter, FunctionOpInterface funcOp) {
  // See NOTE: [Additional Function Arguments]
  if (!isKernel(funcOp)) {
    return funcOp.getArgument(funcOp.getNumArguments() + kSharedMemoryOffset);
  }

  auto mod = funcOp->getParentOfType<ModuleOp>();
  auto globalBase = dyn_cast<LLVM::GlobalOp>(mod.lookupSymbol("global_smem"));
  assert(globalBase);
  return LLVM::AddressOfOp::create(rewriter, funcOp.getLoc(), globalBase);
}

static Value getScratchPtrWithOffset(Location loc, RewriterBase &rewriter,
                                     Value gmemBase, Value allocOffset) {
  if (!allocOffset)
    return gmemBase;

  auto ptrTy = mlir::LLVM::LLVMPointerType::get(rewriter.getContext(), 1);
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  return b.gep(ptrTy, i8_ty, gmemBase, allocOffset);
}

static Value getScratchPtrImpl(Location loc, RewriterBase &rewriter,
                               const TargetInfoBase &targetInfo,
                               FunctionOpInterface funcOp, Value allocOffset,
                               int32_t bufferArgOffset, StringRef allocSizeAttr,
                               bool addOffsetIfNoAllocSizeAttr,
                               bool currentCTA = true) {
  // See NOTE: [Additional Function Arguments]
  auto gmemBase =
      funcOp.getArgument(funcOp.getNumArguments() + bufferArgOffset);

  if (!isKernel(funcOp))
    return getScratchPtrWithOffset(loc, rewriter, gmemBase, allocOffset);

  ModuleOp mod = funcOp.getOperation()->getParentOfType<ModuleOp>();
  auto allocSizeAttrVal =
      mod.getOperation()->getAttrOfType<mlir::IntegerAttr>(allocSizeAttr);
  if (!allocSizeAttrVal || allocSizeAttrVal.getValue().isZero()) {
    return addOffsetIfNoAllocSizeAttr
               ? getScratchPtrWithOffset(loc, rewriter, gmemBase, allocOffset)
               : gmemBase;
  }

  Value gridIdx[3];
  Value gridDim[2];
  for (int k = 0; k < 3; ++k)
    gridIdx[k] = GetProgramIdOp::create(rewriter, loc, k);
  for (int k = 0; k < 2; ++k)
    gridDim[k] = GetNumProgramsOp::create(rewriter, loc, k);

  auto b = TritonLLVMOpBuilder(loc, rewriter);
  auto zextToI64 = [&](Value value) {
    if (value.getType() == i64_ty)
      return value;
    return b.zext(i64_ty, value).getResult();
  };
  Value linearId = zextToI64(gridIdx[2]);
  for (int k = 0; k < 2; ++k)
    linearId = b.add(zextToI64(gridIdx[1 - k]),
                     b.mul(linearId, zextToI64(gridDim[1 - k])));

  auto numCTAs = triton::gpu::TritonGPUDialect::getNumCTAs(mod);
  if (numCTAs > 1) {
    linearId = b.mul(linearId, b.i64_val(numCTAs));
    // currentCTA sets whether to rebase the linearId to the CTA id or
    // just keep the pointer to the whole tensor
    if (currentCTA)
      linearId =
          b.add(linearId, zextToI64(targetInfo.getClusterCTAId(rewriter, loc)));
  }

  auto allocSize = allocSizeAttrVal.getValue().getZExtValue();
  Value offset = b.mul(linearId, b.i64_val(allocSize));
  if (allocOffset)
    offset = b.add(offset, zextToI64(allocOffset));

  auto *ctx = rewriter.getContext();
  return b.gep(mlir::LLVM::LLVMPointerType::get(ctx, 1), i8_ty, gmemBase,
               offset);
}

Value getGlobalScratchPtr(Location loc, RewriterBase &rewriter,
                          const TargetInfoBase &targetInfo,
                          FunctionOpInterface funcOp, Value allocOffset) {
  return getScratchPtrImpl(loc, rewriter, targetInfo, funcOp, allocOffset,
                           kGlobalScratchBufferOffset,
                           "ttg.global_scratch_memory_size",
                           /*addOffsetIfNoAllocSizeAttr=*/false);
}

Value getProfileScratchPtr(Location loc, RewriterBase &rewriter,
                           const TargetInfoBase &targetInfo,
                           FunctionOpInterface funcOp, Value allocOffset,
                           bool currentCTA) {
  return getScratchPtrImpl(loc, rewriter, targetInfo, funcOp, allocOffset,
                           kProfileScratchBufferOffset,
                           "ttg.profile_scratch_memory_size",
                           /*addOffsetIfNoAllocSizeAttr=*/true, currentCTA);
}

Value getSharedMemoryBase(Location loc, RewriterBase &rewriter,
                          const TargetInfoBase &target, Operation *op) {
  auto ptrTy = LLVM::LLVMPointerType::get(rewriter.getContext(),
                                          target.getSharedAddressSpace());
  auto func = op->template getParentOfType<FunctionOpInterface>();
  if (!func)
    func = cast<FunctionOpInterface>(op);

  assert(op->hasAttr("allocation.offset"));
  auto offsetAttr = op->getAttr("allocation.offset");
  assert(isa<IntegerAttr>(offsetAttr) &&
         "getSharedMemoryBase() called on partitioned tensor. "
         "Use getSharedMemoryBases() instead.");
  size_t offset = cast<IntegerAttr>(offsetAttr).getValue().getZExtValue();

  auto b = TritonLLVMOpBuilder(loc, rewriter);
  Value offVal = b.i32_val(offset);
  Value base =
      b.gep(ptrTy, i8_ty, LLVM::getStackPointer(rewriter, func), offVal);
  return base;
}

SmallVector<int64_t> getPartitionOffsets(Operation *op) {
  SmallVector<int64_t> offsets;
  if (!op->hasAttr("allocation.offset"))
    return offsets;

  auto offsetAttr = op->getAttr("allocation.offset");
  if (auto intAttr = dyn_cast<IntegerAttr>(offsetAttr)) {
    // Single offset (non-partitioned tensor)
    offsets.push_back(intAttr.getValue().getZExtValue());
  } else if (auto arrayAttr = dyn_cast<ArrayAttr>(offsetAttr)) {
    // Multiple offsets (partitioned tensor)
    for (auto elem : arrayAttr) {
      offsets.push_back(cast<IntegerAttr>(elem).getValue().getZExtValue());
    }
  }
  return offsets;
}

SmallVector<Value> getSharedMemoryBases(Location loc, RewriterBase &rewriter,
                                        const TargetInfoBase &target,
                                        Operation *op) {
  SmallVector<Value> bases;
  auto ptrTy = LLVM::LLVMPointerType::get(rewriter.getContext(),
                                          target.getSharedAddressSpace());
  auto func = op->template getParentOfType<FunctionOpInterface>();
  if (!func)
    func = cast<FunctionOpInterface>(op);

  auto offsets = getPartitionOffsets(op);
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  Value stackPtr = LLVM::getStackPointer(rewriter, func);

  for (int64_t offset : offsets) {
    Value offVal = b.i32_val(offset);
    Value base = b.gep(ptrTy, i8_ty, stackPtr, offVal);
    bases.push_back(base);
  }
  return bases;
}

namespace {

// Extract the bits of `a` that are set in `mask`
Value pext_i32(RewriterBase &rewriter, Location loc, Value a, uint32_t mask) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  assert(a.getType() == i32_ty && "a must be i32");
  // Handle width = 32 to avoid doing 1 << 32
  if (mask == 0xFFFFFFFF)
    return a;

  // Implements the blocked algorithm from
  // https://forums.developer.nvidia.com/t/pdep-and-pext-functionality-for-cuda/270973
  uint32_t mskConst = mask;
  uint32_t extcnt = 0;
  Value result = b.i32_val(0);
  while (mskConst) {
    uint32_t oldmsk = mskConst;
    uint32_t bitgrplsb = mskConst & (-mskConst);
    mskConst &= bitgrplsb + mskConst;
    uint32_t bitgrp = mskConst ^ oldmsk;
    uint32_t lsbpos = 31 - __builtin_clz(bitgrplsb);
    // like popcount for a number 0..01..1..0 but portable
    uint32_t grplen = __builtin_ctz(~(bitgrp >> lsbpos));
    uint32_t shift = lsbpos - extcnt;
    extcnt += grplen;
    result =
        b.or_(result, b.lshr(b.and_(b.i32_val(bitgrp), a), b.i32_val(shift)));
  }
  return result;
}

// Puts the bits of `a` that are set in `mask` into the bits of `result`
Value pdep_i32(RewriterBase &rewriter, Location loc, Value a, uint32_t mask) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  assert(a.getType() == i32_ty && "a must be i32");

  if (mask == 0)
    return b.i32_val(0);
  assert(mask < 64 && "mask must be less than 64");

  // Blocked algorithm (same grouping trick as the pext example).
  uint32_t mskConst = mask;
  uint32_t depcnt = 0; // how many source bits from `a` we've consumed
  Value result = b.i32_val(0);

  while (mskConst) {
    uint32_t oldmsk = mskConst;

    // Isolate lsb set bit, then clear the lowest contiguous run of 1s.
    uint32_t bitgrplsb = mskConst & (~mskConst + 1); // m & -m
    mskConst &= (bitgrplsb + mskConst);
    uint32_t bitgrp = mskConst ^ oldmsk; // the cleared run (contiguous 1s)

    // Group start position and length.
    uint32_t lsbpos = __builtin_ctz(bitgrplsb);
    uint32_t grplen = __builtin_ctz(~(bitgrp >> lsbpos));

    // Align the next grplen bits of `a` to the group's lsb, then mask to the
    // group.
    uint32_t shift =
        lsbpos - depcnt; // non-negative invariant for this traversal order
    depcnt += grplen;

    Value deposited = b.and_(b.shl(a, b.i32_val(shift)), b.i32_val(bitgrp));
    result = b.or_(result, deposited);
  }

  return result;
}

} // namespace

std::tuple<SmallVector<Value>, Value>
delinearize(RewriterBase &rewriter, Location loc,
            triton::gpu::DistributedEncodingTrait layout,
            ArrayRef<int64_t> shape, StringAttr dimName, Value linear) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  auto ll = triton::gpu::toLinearLayout(shape, layout);
  assert(ll.hasInDim(dimName));
  int32_t freeVarMask = ll.getFreeVariableMasks()[dimName];
  auto isRepresentative = b.true_val();
  if (freeVarMask != 0) {
    isRepresentative =
        b.icmp_eq(b.and_(b.i32_val(freeVarMask), linear), b.i32_val(0));
    // We remove the bits of linear that are set to one in freeVarMask
    int32_t nonFreeVarMask = ~freeVarMask & (ll.getInDimSize(dimName) - 1);
    linear = pext_i32(rewriter, loc, linear, nonFreeVarMask);
  }

  auto linearLayout = triton::gpu::LinearEncodingAttr::get(
      rewriter.getContext(), std::move(ll));
  auto orderDim = linearLayout.orderPerDim(dimName, linearLayout.getOrder());
  auto shapeDim = linearLayout.basesPerDim(dimName, /*skipBroadcast=*/true);
  auto multiDim = delinearize(rewriter, loc, linear, shapeDim, orderDim);

  return std::make_tuple(std::move(multiDim), isRepresentative);
}

// Convert an \param index to a multi-dim coordinate given \param shape and
// \param order.
SmallVector<Value> delinearize(RewriterBase &rewriter, Location loc,
                               Value linear, ArrayRef<unsigned> shape,
                               ArrayRef<unsigned> order) {
  unsigned rank = shape.size();
  assert(rank == order.size());
  auto reordered = applyPermutation(shape, order);
  SmallVector<Value> reorderedMultiDim(rank);
  if (auto constantOp = linear.getDefiningOp<arith::ConstantOp>()) {
    unsigned intVal = mlir::cast<IntegerAttr>(constantOp.getValue())
                          .getValue()
                          .getSExtValue();
    reorderedMultiDim = delinearize(rewriter, loc, intVal, reordered);
  } else {
    reorderedMultiDim = delinearize(rewriter, loc, linear, reordered);
  }
  SmallVector<Value> multiDim(rank);
  for (unsigned i = 0; i < rank; ++i) {
    multiDim[order[i]] = reorderedMultiDim[i];
  }
  return multiDim;
}

SmallVector<Value> delinearize(RewriterBase &rewriter, Location loc,
                               unsigned linear, ArrayRef<unsigned> shape) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  unsigned rank = shape.size();
  assert(rank > 0);
  SmallVector<Value> multiDim(rank);
  unsigned remained = linear;
  for (auto &&en : llvm::enumerate(shape)) {
    unsigned dimSize = en.value();
    multiDim[en.index()] = b.i32_val(remained % dimSize);
    remained = remained / dimSize;
  }
  return multiDim;
}

SmallVector<Value> delinearize(RewriterBase &rewriter, Location loc,
                               Value linear, ArrayRef<unsigned> shape) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  unsigned rank = shape.size();
  assert(rank > 0);
  SmallVector<Value> multiDim(rank);
  Value remained = linear;
  for (auto &&en : llvm::enumerate(shape)) {
    Value dimSize = b.i32_val(en.value());
    multiDim[en.index()] = b.urem(remained, dimSize);
    remained = b.udiv(remained, dimSize);
  }
  return multiDim;
}

SmallVector<unsigned> delinearize(unsigned linear, ArrayRef<unsigned> shape,
                                  ArrayRef<unsigned> order) {
  auto rank = shape.size();
  assert(order.size() == rank);
  SmallVector<unsigned> multiDim(rank);
  for (auto dim : order) {
    multiDim[dim] = linear % shape[dim];
    linear /= shape[dim];
  }
  assert(linear == 0);
  return multiDim;
}

Value linearize(RewriterBase &rewriter, Location loc, ArrayRef<Value> multiDim,
                ArrayRef<unsigned> shape, ArrayRef<unsigned> order) {
  return linearize(rewriter, loc, applyPermutation(multiDim, order),
                   applyPermutation(shape, order));
}

Value linearize(RewriterBase &rewriter, Location loc, ArrayRef<Value> multiDim,
                ArrayRef<unsigned> shape) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  auto rank = multiDim.size();
  Value linear = b.i32_val(0);
  if (rank > 0) {
    linear = multiDim.back();
    for (auto [dim, dimShape] :
         llvm::reverse(llvm::zip(multiDim.drop_back(), shape.drop_back()))) {
      Value dimSize = b.i32_val(dimShape);
      linear = b.add(b.mul(linear, dimSize), dim);
    }
  }
  return linear;
}

Value linearize(RewriterBase &rewriter, Location loc, ArrayRef<Value> multiDim,
                triton::gpu::LinearEncodingAttr encoding, StringAttr dimName) {
  auto orderDim = encoding.orderPerDim(dimName, encoding.getOrder());
  auto shapeDim = encoding.basesPerDim(dimName, /*skipBroadcast=*/true);
  auto linear = linearize(rewriter, loc, multiDim, shapeDim, orderDim);
  auto ll = encoding.getLinearLayout();
  int32_t freeVarMask = ll.getFreeVariableMasks().lookup(dimName);
  if (freeVarMask != 0) {
    int32_t nonFreeVarMask = ~freeVarMask & (ll.getInDimSize(dimName) - 1);
    linear = pdep_i32(rewriter, loc, linear, nonFreeVarMask);
  }
  return linear;
}

size_t linearize(ArrayRef<unsigned> multiDim, ArrayRef<unsigned> shape,
                 ArrayRef<unsigned> order) {
  size_t linear = 0;
  for (unsigned dim : llvm::reverse(order))
    linear = linear * shape[dim] + multiDim[dim];
  return linear;
}

LLVM::GlobalOp getOrInsertGlobalConstant(RewriterBase &rewriter,
                                         ModuleOp module, Type type,
                                         Attribute content, StringRef key) {
  for (auto op : module.getOps<LLVM::GlobalOp>()) {
    if (op.getConstant() && op.getLinkage() == LLVM::Linkage::Internal &&
        op.getType() == type && op.getValueAttr() == content)
      return op;
  }

  unsigned stringNumber = 0;
  SmallString<16> name;
  do {
    name.clear();
    (key + Twine(stringNumber++)).toStringRef(name);
  } while (module.lookupSymbol(name));
  RewriterBase::InsertionGuard guard(rewriter);
  rewriter.setInsertionPointToStart(module.getBody());
  return LLVM::GlobalOp::create(rewriter, UnknownLoc::get(type.getContext()),
                                type, /*isConstant=*/true,
                                LLVM::Linkage::Internal, name, content);
}

Value addStringToModule(Location loc, RewriterBase &rewriter, StringRef key,
                        StringRef content) {
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  auto moduleOp = rewriter.getBlock()->getParent()->getParentOfType<ModuleOp>();
  auto ctx = moduleOp.getContext();

  llvm::SmallString<64> contentStr(content);
  size_t contentSize = contentStr.size_in_bytes();
  auto globalType = LLVM::LLVMArrayType::get(i8_ty, contentSize);
  auto contentAttr = rewriter.getStringAttr(contentStr);
  LLVM::GlobalOp global = getOrInsertGlobalConstant(
      rewriter, moduleOp, globalType, contentAttr, key);

  Value zero = b.i32_val(0);
  Type globalPtrType = LLVM::LLVMPointerType::get(ctx, global.getAddrSpace());
  Value globalPtr = LLVM::AddressOfOp::create(
      rewriter, UnknownLoc::get(ctx), globalPtrType, global.getSymName());
  Value stringStart =
      b.gep(ptr_ty(ctx), i8_ty, globalPtr, SmallVector<Value>({zero}));
  return stringStart;
}

} // namespace LLVM

Value dot(RewriterBase &rewriter, Location loc, ArrayRef<Value> offsets,
          ArrayRef<Value> strides) {
  assert(offsets.size() == strides.size());
  auto b = TritonLLVMOpBuilder(loc, rewriter);
  Value ret = b.i32_val(0);
  for (auto [offset, stride] : llvm::zip(offsets, strides)) {
    ret = b.add(ret, b.mul(offset, stride));
  }
  return ret;
}

// Isolated a single warp specialize op from above.
static void
makeWarpGroupsIsolatedFromAbove(triton::gpu::WarpSpecializeOp wsOp) {
  SetVector<Value> captures;
  auto partOp = wsOp.getPartitionOp();
  getUsedValuesDefinedAbove(partOp.getPartitionRegions(), captures);
  for (Value capture : captures) {
    partOp->insertOperands(partOp.getNumOperands(), capture);
    for (Region &region : partOp.getPartitionRegions()) {
      BlockArgument arg =
          region.addArgument(capture.getType(), capture.getLoc());
      replaceAllUsesInRegionWith(capture, arg, region);
    }
  }
}

void makeAllWarpGroupsIsolatedFromAbove(Operation *op) {
  op->walk([](triton::gpu::WarpSpecializeOp wsOp) {
    makeWarpGroupsIsolatedFromAbove(wsOp);
  });
}

// TODO: Is there a better way to do this? This needs to be fixed upstream.
void fixUpLoopAnnotation(ModuleOp mod) {
  mod->walk([](Operation *op) {
    if (isa<LLVM::BrOp, LLVM::CondBrOp>(op)) {
      if (op->hasAttr("llvm.loop_annotation")) {
        auto loopMD = dyn_cast<LLVM::LoopAnnotationAttr>(
            op->getAttr("llvm.loop_annotation"));
        if (loopMD) {
          if (auto brOp = dyn_cast<LLVM::BrOp>(op)) {
            brOp.setLoopAnnotationAttr(loopMD);
          } else if (auto condBrOp = dyn_cast<LLVM::CondBrOp>(op)) {
            condBrOp.setLoopAnnotationAttr(loopMD);
          }
        }
      }
    }
  });
}

SmallVector<Value> inlineRegionImpl(RewriterBase &rewriter, Region &region,
                                    ArrayRef<Value> args,
                                    mlir::TypeID terminatorTypeId,
                                    Location loc) {
  // Inline regions with multiple blocks
  //
  //        Before                                   After
  //                                              ┌─────────┐
  //                                              │ op1     │
  //                    ┌──────────┐              │ cf.br   │
  //                    │region[0] │              └────┬────┘
  //                    │cf.cond_br├─┐            ┌────▼─────┐
  //                    └────┬─────┘ │            │region[0] │
  //                         │       │            │cf.cond_br├─┐
  // ┌───────┐          ┌────▼────┐  │            └────┬─────┘ │
  // │  op1  │  IP      │region[1]│  │            ┌────▼────┐  │
  // │       │◄───      │yield ...│  │            │region[1]│  │
  // │  op2  │          └─────────┘  │          ┌─┤cf.br    │  │
  // └───────┘                       │          │ └─────────┘  │
  //                    ┌─────────┐  │          │ ┌─────────┐  │
  //                    │region[2]│◄─┘          │ │region[2]│◄─┘
  //                    │yield    │             │ │cf.br    │
  //                    └─────────┘             │ └────┬────┘
  //                                            │ ┌────▼────┐
  //                                            └►│op2      │
  //                                              └─────────┘
  auto *curBlock = rewriter.getInsertionBlock();
  auto opPosition = rewriter.getInsertionPoint();
  auto *remainingOpsBlock = curBlock->splitBlock(opPosition);

  IRMapping regionMap;
  Region &parent = *curBlock->getParent();
  rewriter.cloneRegionBefore(region, parent, parent.end(), regionMap);
  rewriter.setInsertionPointToEnd(curBlock);
  LLVM::BrOp::create(rewriter, loc, args, regionMap.lookup(&region.front()));

  ValueRange terminatorOperands;
  for (Block &origBlock : region) {
    Block *newBlock = regionMap.lookup(&origBlock);
    rewriter.moveBlockBefore(newBlock, remainingOpsBlock);

    auto terminator = newBlock->getTerminator();
    if (terminator->getRegisteredInfo()->getTypeID() == terminatorTypeId) {
      terminatorOperands = terminator->getOperands();
      rewriter.setInsertionPointAfter(terminator);
      rewriter.replaceOpWithNewOp<LLVM::BrOp>(terminator, terminatorOperands,
                                              remainingOpsBlock);
    }
  }

  rewriter.setInsertionPointToStart(remainingOpsBlock);
  SmallVector<Value> vals;
  for (auto resultTy : terminatorOperands.getType()) {
    auto val = remainingOpsBlock->addArgument(resultTy, loc);
    vals.push_back(val);
  }
  return vals;
}

std::tuple<Block *, Block *, Block *> createIfBlock(RewriterBase &b,
                                                    Location loc, Value cnd) {
  Block *prevBlock = b.getInsertionBlock();
  Block *thenBlock = prevBlock->splitBlock(b.getInsertionPoint());
  Block *ifBlock = b.createBlock(thenBlock);
  b.setInsertionPointToEnd(ifBlock);
  LLVM::BrOp::create(b, loc, thenBlock);
  b.setInsertionPointToEnd(prevBlock);
  LLVM::CondBrOp::create(b, loc, cnd, ifBlock, thenBlock);
  b.setInsertionPointToStart(thenBlock);

  return {prevBlock, ifBlock, thenBlock};
}

void finalizeTensorAtomicResults(Operation *op, RankedTensorType tensorTy,
                                 ConversionPatternRewriter &rewriter,
                                 SmallVector<Value> &resultVals,
                                 Type valueElemTy, TritonLLVMOpBuilder &b,
                                 Value threadPred,
                                 const TargetInfoBase &targetInfo,
                                 const LLVMTypeConverter *typeConverter) {
  auto *ctx = rewriter.getContext();
  auto loc = op->getLoc();
  if (!op->hasAttr("allocation.offset")) {
    // No broadcasting, just pack the values into a struct
    Value resultStruct =
        packTensorElements(loc, typeConverter, resultVals, rewriter, tensorTy);
    rewriter.replaceOp(op, {resultStruct});
    return;
  }

  auto kOffset = str_attr("offset");
  auto kBlock = str_attr("block");
  auto kOwnerBlock = str_attr("owner_block");
  auto regLayout = triton::gpu::toLinearLayout(tensorTy);
  auto removeRegBroadcast = actionRemoveBroadcastedRegs(regLayout);
  auto uniqueRegLayout = removeRegBroadcast.apply(regLayout);
  auto uniqueResultVals = removeRegBroadcast.apply(resultVals);

  auto shapePerCTA =
      convertType<unsigned>(triton::gpu::getShapePerCTA(tensorTy));
  auto ctaShared =
      identityStandardND(kOffset, shapePerCTA, triton::gpu::getOrder(tensorTy));
  auto sharedLayout = triton::gpu::combineCtaCgaWithShape(
      ctaShared, triton::gpu::getCGALayout(tensorTy.getEncoding()),
      tensorTy.getShape());
  auto dimOffset = sharedLayout.getInDimSize(kOffset);
  auto dimBlock = sharedLayout.getInDimSize(kBlock);

  // Preserve CTA identity for stores so producers write to their local
  // allocation. Rename the shared block dimension for loads so the inverse
  // maps broadcast CTA bits to the canonical owner CTA instead.
  auto storeCvt = uniqueRegLayout.invertAndCompose(sharedLayout);
  assert(storeCvt.isTrivialOver({kBlock}));
  auto ownerShared =
      sharedLayout.reshapeIns({{kOffset, dimOffset}, {kOwnerBlock, dimBlock}});
  auto loadCvt = uniqueRegLayout.invertAndCompose(ownerShared);
  loadCvt = loadCvt.reshapeOuts({{kOffset, dimOffset}, {kBlock, dimBlock}});
  bool crossCTA = !loadCvt.isTrivialOver({kBlock});
  auto smemBase = LLVM::getSharedMemoryBase(loc, rewriter, targetInfo, op);

  auto [laneId, warpId] = getLaneAndWarpId(rewriter, loc);
  SmallVector<Value> smemBases = {smemBase};
  lowerLdSt(loc, ctx, storeCvt, uniqueResultVals, valueElemTy, smemBases,
            /*paddingShifts=*/{}, /*affineOffset=*/b.i32_val(0),
            /*maskSpanAffineOffset=*/0, /*affineBlockOffset=*/Value(),
            /*maskSpanAffineBlock=*/0, laneId, warpId, rewriter, targetInfo,
            /*maybeMaxVecElems=*/{},
            makeSharedStoreEmitter(targetInfo, threadPred));
  if (crossCTA)
    targetInfo.clusterBarrier(loc, rewriter, op);
  else
    targetInfo.barrier(loc, rewriter, triton::gpu::AddrSpace::Local);

  resultVals =
      lowerLdSt(loc, ctx, loadCvt, /*valsArray=*/{}, valueElemTy, smemBases,
                /*paddingShifts=*/{}, /*affineOffset=*/b.i32_val(0),
                /*maskSpanAffineOffset=*/0, /*affineBlockOffset=*/Value(),
                /*maskSpanAffineBlock=*/0, laneId, warpId, rewriter, targetInfo,
                /*maybeMaxVecElems=*/{}, makeSharedLoadEmitter(targetInfo));
  // Create the result struct and replace the operation
  Value resultStruct = packUniqueTensorElements(loc, typeConverter, resultVals,
                                                rewriter, tensorTy);
  rewriter.replaceOp(op, {resultStruct});
}

// Only retain those attributes that are not constructed by
// `LLVMFuncOp::build`. If `filterArgAttrs` is set, also filter out argument
// attributes.
void filterFuncAttributes(triton::FuncOp op, bool filterArgAttrs,
                          SmallVectorImpl<NamedAttribute> &result) {

  for (const auto &attr : op->getAttrs()) {
    if (attr.getName() == SymbolTable::getSymbolAttrName() ||
        attr.getName() == op.getFunctionTypeAttrName() ||
        attr.getName() == "std.varargs" ||
        attr.getName() == triton::gpu::AttrNumWarpsName ||
        (filterArgAttrs && attr.getName() == op.getArgAttrsAttrName()))
      continue;
    result.push_back(attr);
  }
}

triton::FuncOp amendFuncOp(triton::FuncOp funcOp,
                           ConversionPatternRewriter &rewriter,
                           const TargetInfoBase &targetInfo) {
  // Push back two new arguments that indicate the current pointer to shared
  // memory and global scratch memory.
  auto loc = funcOp.getLoc();
  auto ctx = funcOp->getContext();
  auto sharedPtrTy =
      LLVM::LLVMPointerType::get(ctx, targetInfo.getSharedAddressSpace());
  auto globalPtrTy = LLVM::LLVMPointerType::get(ctx, 1);
  auto profilePtrTy = LLVM::LLVMPointerType::get(ctx, 1);

  // 1. Modify the function type to add the new arguments.
  auto funcTy = funcOp.getFunctionType();
  auto amendedInputTy = llvm::to_vector<4>(funcTy.getInputs());
  bool isKernel = triton::isKernel(funcOp);
  if (isKernel && targetInfo.isCuda()) {
    for (auto i : llvm::seq(amendedInputTy.size())) {
      if (isa<triton::TensorDescInterface>(amendedInputTy[i])) {
        funcOp.setArgAttr(i, "tt.nv_tma_desc",
                          mlir::IntegerAttr::get(i32_ty, 1));
      }
    }
  }
  if (!isKernel) {
    amendedInputTy.push_back(sharedPtrTy);
  }
  amendedInputTy.push_back(globalPtrTy);
  amendedInputTy.push_back(profilePtrTy);
  auto amendedFuncTy =
      FunctionType::get(ctx, amendedInputTy, funcTy.getResults());
  // 2. Modify the argument attributes to add the new argument.
  SmallVector<NamedAttribute> amendedAttrs;
  filterFuncAttributes(funcOp, /*filterArgAttrs=*/true, amendedAttrs);
  if (auto argAttrs = funcOp.getAllArgAttrs()) {
    llvm::SmallVector<mlir::Attribute> amendedArgAttrs(argAttrs.begin(),
                                                       argAttrs.end());
    while (amendedArgAttrs.size() < amendedInputTy.size()) {
      amendedArgAttrs.emplace_back(DictionaryAttr::get(ctx));
    }
    amendedAttrs.push_back(rewriter.getNamedAttr(
        funcOp.getArgAttrsAttrName(), rewriter.getArrayAttr(amendedArgAttrs)));
  }

  // 3. Add the new arguments to the region
  auto amendedFuncOp = triton::FuncOp::create(
      rewriter, funcOp.getLoc(), funcOp.getName(), amendedFuncTy, amendedAttrs);
  auto &region = funcOp.getBody();
  if (!isKernel) {
    region.addArgument(sharedPtrTy, loc);
  }
  region.addArgument(globalPtrTy, loc);
  region.addArgument(profilePtrTy, loc);
  rewriter.inlineRegionBefore(region, amendedFuncOp.getBody(),
                              amendedFuncOp.end());
  return amendedFuncOp;
}

void handleArgPtrDatatype(triton::FuncOp funcOp, LLVM::LLVMFuncOp &llvmFuncOp) {
  // The convertion from triton::PointerType to LLVM::LLVMPointerType losts
  // the pointee datatype information.
  // This function add back the pointee datatype information to arg attribute.
  FunctionType fty = funcOp.getFunctionType();
  for (unsigned i = 0; i < fty.getNumInputs(); ++i) {
    auto argType = fty.getInput(i);
    if (auto argPtrType = dyn_cast<triton::PointerType>(argType)) {
      auto argDType = argPtrType.getPointeeType();
      llvmFuncOp.setArgAttr(i, "tt.pointee_type",
                            mlir::TypeAttr::get(argDType));
    }
  }
}

} // namespace mlir
