#include "AsyncUtility.h"
#include "AtomicRMWOpsEmitter.h"
#include "BufferOpsEmitter.h"
#include "Dialect/TritonAMDGPU/IR/Dialect.h"
#include "PatternTritonGPUOpToLLVM.h"
#include "TDMUtility.h"
#include "TargetInfo.h"
#include "Utility.h"
#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Transforms/DialectConversion.h"
#include "triton/Conversion/TritonGPUToLLVM/Utility.h"
#include "triton/Dialect/Triton/IR/Types.h"
#include "triton/Dialect/TritonGPU/IR/Attributes.h"
#include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h"
#include "triton/Dialect/TritonGPU/Transforms/Utility.h"
#include "triton/Tools/LayoutUtils.h"

using namespace mlir;
using namespace mlir::triton::gpu;

using ::mlir::LLVM::getSharedMemoryBase;
using ::mlir::LLVM::AMD::getVectorSize;
using ::mlir::LLVM::AMD::llLoad;
using ::mlir::LLVM::AMD::llStore;
using ::mlir::triton::gpu::getTotalElemsPerThread;
using triton::amdgpu::ISAFamily;

namespace {

std::optional<const char *> getAMDGPUMemScopeStr(MemSyncScope scope) {
  switch (scope) {
  case MemSyncScope::GPU:
    return "agent";
  case MemSyncScope::CTA:
    return "workgroup";
  // The default AMDHSA LLVM Sync Scope is "system", so no string is
  // provided here
  case MemSyncScope::SYSTEM:
  default:
    return "";
  }
}

std::pair<bool, bool> getOrderingFlags(MemSemantic memOrdering) {
  bool emitReleaseFence = false;
  bool emitAcquireFence = false;
  switch (memOrdering) {
  case MemSemantic::RELAXED:
    // In this case, no memory fences are needed
    break;
  case MemSemantic::RELEASE:
    emitReleaseFence = true;
    break;
  case MemSemantic::ACQUIRE:
    emitAcquireFence = true;
    break;
  case MemSemantic::ACQUIRE_RELEASE:
    emitAcquireFence = true;
    emitReleaseFence = true;
  default:
    // default == acq_rel, so we emit the same barriers
    emitAcquireFence = true;
    emitReleaseFence = true;
  }
  return {emitAcquireFence, emitReleaseFence};
}

LogicalResult emitFence(Operation *op, ConversionPatternRewriter &rewriter,
                        Location loc, MemSemantic memOrdering,
                        MemSyncScope memScope, bool preAtomic) {
  // This function emits an LLVM::FenceOp which will get lowered by the
  // LLVM backend to the right scope and ordering instructions, as
  // described in the "atomicrmw" entries for "global" address-space,
  // in the "AMDHSA Memory Model Code Sequences GFX942"
  // table in https://llvm.org/docs/AMDGPUUsage.html#memory-model-gfx942
  //
  // Triton supports three scopes for atomic access
  // 1. System
  // 2. GPU (default) ('Agent' for AMDGPU)
  // 3. CTA ('Workgroup' for AMDGPU)
  //
  // and 4 orderings
  // 1. Relaxed
  // 2. Acquire
  // 3. Release
  // 4. AcquireRelease
  //
  // The following table shows the scope and ordering instructions that
  // are emitted by this function for each combination of scope and ordering
  // for buffer-atomic instructions.
  //
  // Note: In the following comments, "[buffer-atomic_0.. buffer-atomic_n]"
  // represents a sequence of buffer-atomic instructions that are lowered from
  // a single tl.atomic_*
  //
  // Unordered(Relaxed):
  //   agent/workgroup: Instr seq: [buffer-atomic_0.. buffer-atomic_n]
  //                    No scope/ordering instrs are required.
  //   system: //TODO:
  // Acquire:
  //   workgroup: Instr seq: [buffer-atomic_0.. buffer-atomic_n]
  //              All waves in the workgroup use same L1 and L2.
  //              No scope/ordering instrs are required.
  //   agent: Instr seq: [buffer-atomic_0.. buffer-atomic_n],
  //                     s_waitcnt vmcnt(0), buffer_inv sc1=1
  //          Waves across an agent may use different L1 and L2.
  //          Atomic ops bypass L1 and operate on L2.
  //          s_waitcnt vmcnt(0) ensures that the atomicrmw has completed
  //          before invalidating the cache. buffer_inv sc1=1 will a) L1:
  //          invalidate cache b) L2: Invalidate non-coherently modified lines
  //          if multiple L2s are configured, NOP otherwise. This buffer_inv
  //          ensures that following loads do not see stale global values.
  //   system: //TODO:
  //
  // Release:
  //   workgroup: Instr seq: [buffer-atomic_0.. buffer-atomic_n]
  //              All waves in the workgroup use same L1 and L2 so all
  //              previous global writes of a waver are visible to all other
  //              waves in the workgroup. LDS operations for all waves are
  //              executed in a total global ordering and are observed by all
  //              waves in the workgroup. So LDS stores issued before the
  //              release will be visible to LDS loads after the read of the
  //              released buffer-atomic. So, swait_cnt lgkmcnt is not
  //              required.
  //   agent: Instr seq: buffer_wbl2 sc1=1, s_waitcnt vmcnt(0),
  //                     [buffer-atomic_0.. buffer-atomic_n]
  //          buffer_wbl2 sc1=1 ensures that dirtly L2 lines are visible to
  //          CUs that don't use the same L2.
  //          From SIMemoryLegalizer.cpp SIGfx940CacheControl::insertRelease:
  //            "Inserting a "S_WAITCNT vmcnt(0)" before is not required
  //             because the hardware does not reorder memory operations by
  //             the same wave with respect to a following "BUFFER_WBL2".
  //             The "BUFFER_WBL2" is guaranteed to initiate writeback of
  //             any dirty cache lines of earlier writes by the same wave.
  //             A "S_WAITCNT vmcnt(0)" is needed after to ensure the writeback
  //             has completed.""
  //   system: //TODO:
  //
  // AcquireRelease:
  //   Instr seq: Release scope/order insts,
  //              [buffer-atomic_0..buffer-atomic_n],
  //              Acquire scope/order instrs.
  //
  // LLVM::FenceOp lowering will emit the required cache ops and s_waitcnt
  // vmcnt(0) instrs

  auto [emitAcquireFence, emitReleaseFence] = getOrderingFlags(memOrdering);
  if (MemSyncScope::SYSTEM == memScope)
    return rewriter.notifyMatchFailure(
        op, "System memory scope is not supported for Buffer Atomic Ops");
  auto scopeStr = getAMDGPUMemScopeStr(memScope);
  if (!scopeStr)
    return rewriter.notifyMatchFailure(
        op, "Unsupported memory scope for Buffer Atomic Ops");

  StringAttr scope = mlir::StringAttr::get(loc.getContext(), *scopeStr);

  if (emitReleaseFence && preAtomic) {
    LLVM::FenceOp::create(rewriter, loc, TypeRange{},
                          LLVM::AtomicOrdering::release, scope);
  }

  if (emitAcquireFence && !preAtomic) {
    LLVM::FenceOp::create(rewriter, loc, TypeRange{},
                          LLVM::AtomicOrdering::acquire, scope);
  }
  return success();
}

Value emitRedundantThreadPredicateNonNull(
    const llvm::MapVector<StringAttr, int32_t> &freeVarMasks,
    ConversionPatternRewriter &rewriter, Location loc,
    const TargetInfo &targetInfo) {
  auto res =
      emitRedundantThreadPredicate(freeVarMasks, rewriter, loc, targetInfo);
  if (!res) {
    TritonLLVMOpBuilder b(loc, rewriter);
    return b.i1_val(true);
  }
  return res;
}

std::pair<Block *, Block *> emitBranch(RewriterBase &rewriter, Location loc,
                                       Value cond) {
  Block *currentBlock = rewriter.getInsertionBlock();
  Block *after = currentBlock->splitBlock(rewriter.getInsertionPoint());
  Block *body = rewriter.createBlock(after);
  rewriter.setInsertionPointToEnd(currentBlock);
  LLVM::CondBrOp::create(rewriter, loc, cond, body, after);
  rewriter.setInsertionPointToStart(body);
  LLVM::BrOp::create(rewriter, loc, after);
  rewriter.setInsertionPointToStart(body);
  return {body, after};
}

bool isRedundantThreadPredWarpUniform(
    const llvm::MapVector<StringAttr, int32_t> &masks, MLIRContext *ctx) {
  return masks.lookup(StringAttr::get(ctx, "lane")) == 0 &&
         masks.lookup(StringAttr::get(ctx, "reg")) == 0;
}

Value selectLdsAddressForPredicate(TritonLLVMOpBuilder &b, Value pred,
                                   Value shmemAddr) {
  Value outOfRangeAddress =
      b.inttoptr(shmemAddr.getType(), b.i32_val(0x7FFFFFFF));
  return b.select(pred, shmemAddr, outOfRangeAddress);
}

// Contains some helper functions for both Load and Store conversions.
struct LoadStoreConversionBase {
  explicit LoadStoreConversionBase(const AMD::TargetInfo &targetInfo,
                                   ModuleAxisInfoAnalysis &axisAnalysisPass)
      : targetInfo(targetInfo), axisAnalysisPass(axisAnalysisPass) {}

  // Create a LLVM vector of type `vecTy` containing all zeros
  Value createZeroVector(OpBuilder &builder, Location loc,
                         VectorType vecTy) const {
    mlir::Attribute zeroAttr = builder.getZeroAttr(vecTy.getElementType());
    auto denseValue =
        DenseElementsAttr::get(cast<mlir::ShapedType>(vecTy), zeroAttr);
    Value zeroVal = LLVM::ConstantOp::create(builder, loc, vecTy, denseValue);
    return zeroVal;
  }

  // Given a vector of values `elems` and a starting point `start`, create a
  // LLVM vector of length `vec` whose elements are `elems[start, ...,
  // elems+vec-1]`
  Value packElementRangeIntoVector(RewriterBase &rewriter,
                                   const LLVMTypeConverter *typeConverter,
                                   Location loc, VectorType vecTy,
                                   ArrayRef<Value> elems, int64_t start) const {
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    int64_t vec = vecTy.getNumElements();
    // If we need to mask the loaded value with other elements
    Value v = b.undef(vecTy);
    for (size_t s = 0; s < vec; ++s) {
      Value otherElem = elems[start + s];
      Value indexVal =
          LLVM::createIndexConstant(rewriter, loc, typeConverter, s);
      v = b.insert_element(vecTy, v, otherElem, indexVal);
    }
    return v;
  }

  // Return a tensor of pointers with the same type of `basePtr` and the same
  // shape of `offset`
  Type getPointerTypeWithShape(Value basePtr, Value offset) const {
    Type basePtrType = basePtr.getType();
    auto offsetType = cast<RankedTensorType>(offset.getType());
    return offsetType.cloneWith(std::nullopt, basePtrType);
  }

  // Unpack the elements contained in a `llvmStruct` into a `SmallVector` of
  // `Value`s. While you do that, check also the alignment of the mask and
  // update the vector length `vec` accordingly
  SmallVector<Value>
  getMaskElemsAndUpdateVeclen(ConversionPatternRewriter &rewriter, Location loc,
                              Value llMask, Value mask, unsigned &vec) const {
    SmallVector<Value> maskElems;
    if (llMask) {
      vec = std::min<size_t>(vec, getMaskAlignment(mask));
      maskElems = unpackTensorElements(loc, llMask, rewriter, mask.getType());
    }
    return maskElems;
  }

  unsigned getMaskAlignment(Value mask) const {
    return axisAnalysisPass.getMaskAlignment(mask);
  }

protected:
  const AMD::TargetInfo &targetInfo;
  ModuleAxisInfoAnalysis &axisAnalysisPass;
};

// Contains some helper functions for direct to lds loads.
struct DirectToLdsLoadConversionBase : public LoadStoreConversionBase {
  explicit DirectToLdsLoadConversionBase(
      const AMD::TargetInfo &targetInfo,
      ModuleAxisInfoAnalysis &axisAnalysisPass)
      : LoadStoreConversionBase(targetInfo, axisAnalysisPass) {}

  // For each load emit the computation to get the lane id offset which holds
  // the source pointers/offsets we need to store to shared memory
  SmallVector<Value>
  emitSwizzledLaneOffsets(RewriterBase &rewriter, Operation *op,
                          RankedTensorType srcTy, MemDescType swizzledTy,
                          MemDescType flatTy, Value llDst, Type resElemTy,
                          unsigned vec) const {
    auto loc = op->getLoc();
    TritonLLVMOpBuilder b(loc, rewriter);

    // Create regToShared layout for the swizzled and flat encoding
    auto regLayout = triton::gpu::toLinearLayout(srcTy);

    auto sharedSwizz = triton::gpu::toLinearLayout(swizzledTy);
    auto sharedFlat = triton::gpu::toLinearLayout(flatTy);

    auto regToSharedSwizzled = regLayout.invertAndCompose(sharedSwizz);
    auto regToSharedFlat = regLayout.invertAndCompose(sharedFlat);

    MLIRContext *ctx = rewriter.getContext();
    StringAttr kBlock = str_attr("block");
    StringAttr kRegister = str_attr("register");
    StringAttr kLane = str_attr("lane");
    StringAttr kWarp = str_attr("warp");
    auto [laneId, warpId] = getLaneAndWarpId(rewriter, loc);
    Value blockId = b.i32_val(0);

    int numberOfLoads = regToSharedSwizzled.getInDimSize(kRegister) / vec;

    // For each load compute the difference between the flat and the swizzled
    // linear offsets into shared memory
    // TODO (alex): this is only correct as long as the lds view is a contiguous
    // block. So this can break if we slice along the 2 minor dimensions
    SmallVector<Value> swizzledOffsets;
    swizzledOffsets.reserve(numberOfLoads);
    auto vecVal = b.i32_val(vec);
    for (int i = 0; i < numberOfLoads; i++) {
      auto regId = b.i32_val(i * vec);

      std::array<std::pair<StringAttr, Value>, 4> indices{{
          {kRegister, regId},
          {kLane, laneId},
          {kWarp, warpId},
          {kBlock, blockId},
      }};

      Value swizzledOffset =
          applyLinearLayout(loc, rewriter, regToSharedSwizzled, indices)[0]
              .second;
      Value flatOffset =
          applyLinearLayout(loc, rewriter, regToSharedFlat, indices)[0].second;

      // Normalize the offset by vecTy to obtain the offset in lanes
      auto laneOffet = b.sdiv(b.sub(swizzledOffset, flatOffset), vecVal);
      swizzledOffsets.push_back(laneOffet);
    }
    return swizzledOffsets;
  }

  // Swizzle the mask (1bit) based on selectLane via ballot
  Value shuffleMask(RewriterBase &rewriter, TritonLLVMOpBuilder &b,
                    Location loc, const TargetInfoBase &targetInfo,
                    Value selectLane, Value mask) const {
    auto warpMask =
        targetInfo.ballot(rewriter, loc, rewriter.getI64Type(), mask);
    // Extract the selectLane bit
    auto bitMask = b.lshr(warpMask, b.zext(rewriter.getI64Type(), selectLane));
    return b.trunc(i1_ty, bitMask);
  }

  SmallVector<Value>
  zipAsyncCopyValues(RewriterBase &rewriter, Location loc, unsigned vec,
                     ArrayRef<Value> srcElems, Type srcTy,
                     ArrayRef<Value> maskElems, ArrayRef<Value> otherElems,
                     Type otherTy, ArrayRef<Value> swizzledLaneOffsets) const {
    TritonLLVMOpBuilder b(loc, rewriter);
    SmallVector<Value> loadVals;
    auto structTy = LLVM::LLVMStructType::getLiteral(
        rewriter.getContext(), ArrayRef<Type>{srcTy, i1_ty, otherTy, i32_ty});
    for (int i = 0; i < srcElems.size(); i++) {
      Value packedArr = LLVM::UndefOp::create(rewriter, loc, structTy);
      // src
      packedArr = b.insert_val(packedArr, srcElems[i], 0);
      // mask
      auto maskElem = maskElems.empty() ? b.true_val() : maskElems[i];
      packedArr = b.insert_val(packedArr, maskElem, 1);
      // other
      if (!otherElems.empty())
        packedArr = b.insert_val(packedArr, otherElems[i], 2);
      // swizzleOffset are per vec so we need to duplicate values vec times
      auto swizzleOffset = swizzledLaneOffsets.empty()
                               ? b.i32_val(0)
                               : swizzledLaneOffsets[i / vec];
      packedArr = b.insert_val(packedArr, swizzleOffset, 3);

      loadVals.push_back(packedArr);
    }
    return loadVals;
  }

  auto unzipAsyncCopyValues(RewriterBase &rewriter, Location loc, int startIdx,
                            ArrayRef<Value> values, Type srcTy, Type otherTy,
                            bool hasOther, unsigned vec) const {
    TritonLLVMOpBuilder b(loc, rewriter);
    auto structElem = values[startIdx];
    Value offsetElem = b.extract_val(srcTy, structElem, 0);
    Value maskElem = b.extract_val(i1_ty, structElem, 1);
    // Gather other elements
    SmallVector<Value> otherElems;
    if (hasOther) {
      for (int i = 0; i < vec; i++) {
        otherElems.push_back(b.extract_val(otherTy, values[startIdx + i], 2));
      }
    }

    Value swizzleLaneOffset = b.extract_val(i32_ty, structElem, 3);

    return std::make_tuple(offsetElem, maskElem, std::move(otherElems),
                           swizzleLaneOffset);
  }

  void applySwizzling(RewriterBase &rewriter, Location loc, Value &srcOrOffset,
                      Value &mask, Value laneId,
                      Value swizzleLaneOffset) const {
    TritonLLVMOpBuilder b(loc, rewriter);
    // laneId + swizzleOffset will always stay inside the warp [0,
    // threadsPerWarp) because we only swizzle inside a warp
    Value swizzledLaneId = b.add(laneId, swizzleLaneOffset);
    // Shuffle based on swizzleLaneId to apply the swizzling
    srcOrOffset =
        targetInfo.shuffleIdx(rewriter, loc, srcOrOffset, swizzledLaneId);

    if (mask) {
      mask = shuffleMask(rewriter, b, loc, targetInfo, swizzledLaneId, mask);
    }
  }

  // Unified helper for async copy between global and shared memory.
  // Works for both load (global→shared) and store (shared→global).
  // Parameters:
  //   globalTy: The global memory tensor type (src for load, dst for store)
  //   sharedTy: The shared memory descriptor type (dst for load, src for store)
  //   vals: Values to process (packed pointers/masks)
  //   llShared: LLVM value for shared memory struct
  //   isLoad: true for global→shared, false for shared→global
  //   isaFamily: ISA family (only used for load multicast)
  //   lowerInst: Callback to emit the actual load/store instruction
  LogicalResult lowerDirectLDSAsyncCopy(
      RewriterBase &rewriter, Location loc, RankedTensorType globalTy,
      MemDescType sharedTy, SmallVector<Value> vals, Value llShared,
      Type resElemTy, unsigned vec, bool isLoad,
      std::function<SmallVector<Value>(RewriterBase &, Location,
                                       ArrayRef<Value>, Value, int, VectorType,
                                       Value)>
          lowerInst) const {
    TritonLLVMOpBuilder b(loc, rewriter);
    auto *ctx = rewriter.getContext();

    // Build global to shared layout and remove broadcasted registers
    auto globalLayout = triton::gpu::toLinearLayout(globalTy);
    auto removeBroadcast = actionRemoveBroadcastedRegs(globalLayout);
    globalLayout = removeBroadcast.apply(globalLayout);
    vals = removeBroadcast.apply(vals);

    LinearLayout sharedLayout;
    if (auto paddedEnc = dyn_cast<triton::gpu::PaddedSharedEncodingAttr>(
            sharedTy.getEncoding())) {
      sharedLayout = paddedEnc.getLinearComponent();
    } else {
      sharedLayout = triton::gpu::toLinearLayout(sharedTy);
    }
    auto cvt = globalLayout.invertAndCompose(sharedLayout);
    if (!cvt.isTrivialOver({str_attr("block")})) {
      return emitError(loc, isLoad ? "direct to lds loads do not support "
                                     "non-trivial block dimension"
                                   : "direct from lds stores do not support "
                                     "non-trivial block dimension");
    }

    // Multicast is only supported for loads
    Value ctaMulticastMask;
    if (isLoad && targetInfo.supportsMultiCTALaunch()) {
      ctaMulticastMask = LLVM::AMD::emitCtaMulticastMask(
          rewriter, loc, targetInfo.getClusterCTAId(rewriter, loc),
          globalLayout, targetInfo.getMaxMulticastMaskPopcount());
    }

    auto smemObj = LLVM::getSharedMemoryObjectFromStruct(loc, llShared,
                                                         resElemTy, rewriter);
    auto affineOffset = smemObj.getShmemOffset(loc, rewriter, sharedTy);
    auto maskSpanAffineOffset =
        SharedMemoryObject::getMaskSpanOffsets(sharedTy);

    auto [laneId, warpId] = getLaneAndWarpId(rewriter, loc);

    auto paddingShifts = getPaddedSharedShifts(
        sharedTy.getEncoding(), sharedTy.getElementTypeBitWidth(),
        /*offsetInBytes=*/true);

    auto lowerInstForwardMulticastMask =
        [&](RewriterBase &rewriter, Location loc, ArrayRef<Value> vals,
            Value shmemAddr, int idx, VectorType vecTy, Value ctaId) {
          assert(!ctaId && "NYI");
          return lowerInst(rewriter, loc, vals, shmemAddr, idx, vecTy,
                           ctaMulticastMask);
        };

    // For loads on GFX9 (no scattering support), the address should be the
    // start address (scalar) of the warp
    if (isLoad && !targetInfo.supportsDirectToLdsScatter()) {
      laneId = b.i32_val(0);
    }

    lowerLdSt(loc, ctx, cvt, vals, resElemTy, smemObj.getBase(), paddingShifts,
              affineOffset, maskSpanAffineOffset, /*affineBlockOffset=*/Value(),
              /*maskSpanAffineBlock=*/0, laneId, warpId, rewriter, targetInfo,
              vec, lowerInstForwardMulticastMask);
    return success();
  }

  void emitOtherStore(RewriterBase &rewriter, Location loc,
                      const LLVMTypeConverter *typeConverter, VectorType vecTy,
                      Value mask, ArrayRef<Value> otherElems, Value shmemAddr,
                      Value laneId, bool requiresSrcPtrSwizzling,
                      Value swizzleLaneOffset) const {
    TritonLLVMOpBuilder b(loc, rewriter);
    Value storeVal = packElementRangeIntoVector(rewriter, typeConverter, loc,
                                                vecTy, otherElems, 0);
    Type ptrTy = shmemAddr.getType();
    Value ldsAddr = shmemAddr;
    // When scattering is unsupported, shmemAddr is the warp base address.
    // Use shmemAddr + lane_id [+ swizzleOffset] to compute each lane's address.
    if (!targetInfo.supportsDirectToLdsScatter()) {
      ldsAddr = b.gep(ptrTy, vecTy, shmemAddr, laneId);
      if (requiresSrcPtrSwizzling)
        ldsAddr = b.gep(ptrTy, vecTy, ldsAddr, swizzleLaneOffset);
    }

    // For architectures supporting scattering to LDS we mask without a branch
    // so we also mask the local writes by selecting OOB LDS address to avoid
    // adding a branch just for other writes
    if (targetInfo.supportsDirectToLdsScatter()) {
      auto invMask = b.icmp_ne(mask, b.true_val());
      ldsAddr = selectLdsAddressForPredicate(b, invMask, shmemAddr);
      // Set the mask to false since we mask via the LDS address. False mask
      // means that others values are written to LDS, see icmp_ne below
      mask = b.false_val();
    }

    llStore(rewriter, loc, ldsAddr, storeVal, b.icmp_ne(mask, b.true_val()),
            CacheModifier::NONE, targetInfo.requiresAliasInfoForAsyncOps());
  }
};

struct LoadOpConversion : public ConvertOpToLLVMPattern<triton::LoadOp>,
                          public LoadStoreConversionBase {
  LoadOpConversion(LLVMTypeConverter &converter,
                   const AMD::TargetInfo &targetInfo,
                   ModuleAxisInfoAnalysis &axisAnalysisPass,
                   PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        LoadStoreConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::LoadOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op->getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    // original values
    Value ptr = op.getPtr();
    Value mask = op.getMask();
    Value other = op.getOther();

    // adaptor values
    Value llPtr = adaptor.getPtr();
    Value llMask = adaptor.getMask();
    Value llOther = adaptor.getOther();

    // Determine the vectorization size
    Type valueTy = op.getType();
    Type valueElemTy =
        typeConverter->convertType(getElementTypeOrSelf(valueTy));
    unsigned vec = getVectorSize(ptr, axisAnalysisPass);
    unsigned numElems = getTotalElemsPerThread(ptr.getType());

    // Get the LLVM values for pointers
    auto ptrElems = unpackTensorElements(loc, llPtr, rewriter, ptr.getType());
    assert(ptrElems.size() == numElems);

    // Get the LLVM values for mask
    SmallVector<Value> maskElems =
        getMaskElemsAndUpdateVeclen(rewriter, loc, llMask, mask, vec);

    SmallVector<Value> otherElems;
    if (other)
      otherElems =
          unpackTensorElements(loc, llOther, rewriter, other.getType());

    Value multicastMask;
    if (targetInfo.supportsMultiCTALaunch()) {
      if (auto tensorTy = dyn_cast<RankedTensorType>(ptr.getType())) {
        Value clusterCTAId = targetInfo.getClusterCTAId(rewriter, loc);
        auto regLayout = triton::gpu::toLinearLayout(tensorTy);
        multicastMask = LLVM::AMD::emitCtaMulticastMask(
            rewriter, loc, clusterCTAId, regLayout,
            targetInfo.getMaxMulticastMaskPopcount());
      }
    }

    // vectorized iteration through all the pointer/mask/other elements
    const int valueElemNBits =
        std::max(8u, valueElemTy.getIntOrFloatBitWidth());
    const int numVecs = numElems / vec;

    auto cacheMod = op.getCache();
    SmallVector<Value> loadedVals;
    Type vecTy = LLVM::getVectorType(valueElemTy, vec);
    for (size_t vecStart = 0; vecStart < numElems; vecStart += vec) {
      const size_t maxWordWidth = std::max<size_t>(32, valueElemNBits);
      const size_t totalWidth = valueElemNBits * vec;
      const size_t width = std::min(totalWidth, maxWordWidth);
      const size_t nWords = std::max<size_t>(1, totalWidth / width);
      const size_t wordNElems = width / valueElemNBits;
      assert(wordNElems * nWords * numVecs == numElems);

      Value pred = mask ? maskElems[vecStart] : b.int_val(1, 1);
      Value ptr = ptrElems[vecStart];

      Value falseVal = createZeroVector(rewriter, loc, cast<VectorType>(vecTy));
      // If we need to mask the loaded value with other elements
      if (otherElems.size() != 0)
        falseVal = packElementRangeIntoVector(
            rewriter, this->getTypeConverter(), loc, cast<VectorType>(vecTy),
            otherElems, vecStart);

      Value loadVal = llLoad(rewriter, loc, ptr, vecTy, pred, falseVal,
                             multicastMask, cacheMod);
      for (size_t ii = 0; ii < vec; ++ii) {
        Value vecIdx = createIndexAttrConstant(
            rewriter, loc, getTypeConverter()->getIndexType(), ii);
        Value loaded = b.extract_element(valueElemTy, loadVal, vecIdx);
        loadedVals.push_back(loaded);
      }
    } // end vec

    Value resultStruct = packTensorElements(loc, getTypeConverter(), loadedVals,
                                            rewriter, valueTy);

    rewriter.replaceOp(op, {resultStruct});
    return success();
  }
};

struct BufferLoadOpConversion
    : public ConvertOpToLLVMPattern<triton::amdgpu::BufferLoadOp>,
      public LoadStoreConversionBase {
  BufferLoadOpConversion(LLVMTypeConverter &converter,
                         const AMD::TargetInfo &targetInfo,
                         ModuleAxisInfoAnalysis &axisAnalysisPass,
                         PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        LoadStoreConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::amdgpu::BufferLoadOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op->getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    LLVM::AMD::BufferEmitter bufferEmitter(rewriter, loc, targetInfo);

    // original values
    Value ptr = op.getPtr();
    Value offset = op.getOffsets();
    Value mask = op.getMask();
    auto cacheMod = op.getCache();

    // Converted values
    Value llPtr = adaptor.getPtr();
    Value llOffset = adaptor.getOffsets();
    Value llMask = adaptor.getMask();
    Value llOther = adaptor.getOther();
    Value llStride = adaptor.getStride();

    // Determine the vectorization size
    Type valueTy = op.getType();
    Type valueElemTy =
        typeConverter->convertType(getElementTypeOrSelf(valueTy));
    Type ptrType = getPointerTypeWithShape(ptr, offset);
    unsigned numElems = getTotalElemsPerThread(ptrType);
    unsigned vec = getVectorSize(ptr, offset, axisAnalysisPass);
    // If the op has a contiguity hint use it to increase the vector size.
    vec = std::max(vec, op.getContiguity());

    // Get the offset
    SmallVector<Value> offsetElems =
        unpackTensorElements(loc, llOffset, rewriter, offset.getType());
    assert(offsetElems.size() == numElems);

    // Get the mask
    SmallVector<Value> maskElems =
        getMaskElemsAndUpdateVeclen(rewriter, loc, llMask, mask, vec);

    // Get the `other` value (if any)
    SmallVector<Value> otherElems;
    if (llOther)
      otherElems =
          unpackTensorElements(loc, llOther, rewriter, op.getOther().getType());

    // Create the resource descriptor and then emit the buffer_load intrinsic(s)
    Value rsrcDesc = bufferEmitter.createResourceDescriptor(llPtr, llStride);
    SmallVector<Value> loadedVals;
    Type vecTy = LLVM::getVectorType(valueElemTy, vec);
    for (size_t vecStart = 0; vecStart < numElems; vecStart += vec) {
      Value pred = mask ? maskElems[vecStart] : b.int_val(1, 1);
      Value falseVal = createZeroVector(rewriter, loc, cast<VectorType>(vecTy));
      if (otherElems.size() != 0)
        falseVal = packElementRangeIntoVector(
            rewriter, this->getTypeConverter(), loc, cast<VectorType>(vecTy),
            otherElems, vecStart);
      Value loadVal = bufferEmitter.emitLoad(
          vecTy, rsrcDesc, offsetElems[vecStart], pred, falseVal, cacheMod);
      for (size_t ii = 0; ii < vec; ++ii) {
        Value vecIdx = createIndexAttrConstant(
            rewriter, loc, getTypeConverter()->getIndexType(), ii);
        Value loaded = b.extract_element(valueElemTy, loadVal, vecIdx);
        loadedVals.push_back(loaded);
      }
    } // end vec

    Value resultStruct = packTensorElements(loc, getTypeConverter(), loadedVals,
                                            rewriter, valueTy);

    rewriter.replaceOp(op, {resultStruct});
    return success();
  }
};

struct BufferLoadToLocalOpConversion
    : public ConvertOpToLLVMPattern<triton::amdgpu::BufferLoadToLocalOp>,
      public DirectToLdsLoadConversionBase {
  BufferLoadToLocalOpConversion(LLVMTypeConverter &converter,
                                const AMD::TargetInfo &targetInfo,
                                ModuleAxisInfoAnalysis &axisAnalysisPass,
                                PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        DirectToLdsLoadConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::amdgpu::BufferLoadToLocalOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op->getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    LLVM::AMD::BufferEmitter bufferEmitter(rewriter, loc, targetInfo);

    // Original values
    Value ptr = op.getPtr();
    Value offset = op.getOffsets();
    Value mask = op.getMask();

    // Converted values
    Value llPtr = adaptor.getPtr();
    Value llOffset = adaptor.getOffsets();
    Value llDst = adaptor.getDest();
    Value llMask = adaptor.getMask();
    Value llOther = adaptor.getOther();
    Value llStride = adaptor.getStride();

    RankedTensorType ptrType =
        cast<RankedTensorType>(getPointerTypeWithShape(ptr, offset));

    // We can load N elements at a time if:
    //  1. Every group of N source pointers are contiguous.  For example, if
    //     N=2, then the pointers should be [x, x+1, y, y+1, ...].
    //  2. The mask (if present) has "alignment" N, meaning that each group of N
    //     mask bits are the same.  For example if N=2, the mask must be
    //     [x, x, y, y, ...].
    unsigned vec = getVectorSize(ptr, offset, axisAnalysisPass);
    SmallVector<Value> maskElems =
        getMaskElemsAndUpdateVeclen(rewriter, loc, llMask, mask, vec);

    SmallVector<Value> offsetElems =
        unpackTensorElements(loc, llOffset, rewriter, offset.getType());
    SmallVector<Value> otherElems;
    if (llOther)
      otherElems =
          unpackTensorElements(loc, llOther, rewriter, op.getOther().getType());

    auto dstTy = op.getDest().getType();
    auto resElemTy = getTypeConverter()->convertType(dstTy.getElementType());
    auto dstEnc = dstTy.getEncoding();

    // If the op has a contiguity hint use it to increase the vector size.
    vec = std::max(vec, op.getContiguity());

    if (!LLVM::AMD::canLoadDirectToLDS(targetInfo, ptrType, dstEnc,
                                       dstTy.getAllocShape(), vec)) {
      return failure();
    }

    // For swizzled layouts we need to use the non swizzled layout to compute
    // the LDS addresses since we gather into LDS
    auto flatDstTy = dstTy;
    SmallVector<Value> swizzledLaneOffsets;

    auto maybeSwizzledEnc = dyn_cast<SwizzledSharedEncodingAttr>(dstEnc);
    bool requiresSrcPtrSwizzling = !targetInfo.supportsDirectToLdsScatter() &&
                                   maybeSwizzledEnc &&
                                   maybeSwizzledEnc.getMaxPhase() != 1;
    if (requiresSrcPtrSwizzling) {
      // TODO (alex): this is only correct as long as the lds view is a
      // contiguous block. So this can break if we slice along the 2 minor
      // dimensions.
      auto flatSharedEnc = SwizzledSharedEncodingAttr::get(
          op->getContext(), maybeSwizzledEnc.getVec(), 1, 1,
          maybeSwizzledEnc.getOrder(), maybeSwizzledEnc.getCGALayout());
      flatDstTy = MemDescType::get(dstTy.getShape(), dstTy.getElementType(),
                                   flatSharedEnc, dstTy.getMemorySpace());
      swizzledLaneOffsets = emitSwizzledLaneOffsets(
          rewriter, op, ptrType, dstTy, flatDstTy, llDst, resElemTy, vec);
    }

    auto offsetTy = offsetElems[0].getType();
    bool hasOther = !otherElems.empty();
    auto otherTy = hasOther ? otherElems[0].getType() : i1_ty;
    // Zip buffer_offset, mask, other, swizzleOffsets for lowerLdSt
    auto loadVals =
        zipAsyncCopyValues(rewriter, loc, vec, offsetElems, offsetTy, maskElems,
                           otherElems, otherTy, swizzledLaneOffsets);

    // Create the resource descriptor and then emit the buffer_loads to lds
    // based on the collected shared addresses and vector size
    Value rsrcDesc = bufferEmitter.createResourceDescriptor(llPtr, llStride);

    auto freeVarMasks = getFreeVariableMasks(ptrType);
    Value threadPred = emitRedundantThreadPredicateNonNull(
        freeVarMasks, rewriter, loc, targetInfo);
    bool isThreadPredWarpUniform =
        isRedundantThreadPredWarpUniform(freeVarMasks, rewriter.getContext());

    auto laneId = getLaneId(rewriter, loc);
    auto emitBufferLoadLds = [&](RewriterBase &rewriter, Location loc,
                                 ArrayRef<Value> loadVals, Value shmemAddr,
                                 int startIdx, VectorType vecTy,
                                 Value multicastMask) -> SmallVector<Value> {
      auto [offsetElem, maskElem, otherElems, swizzleLaneOffset] =
          unzipAsyncCopyValues(rewriter, loc, startIdx, loadVals, offsetTy,
                               otherTy, hasOther, vecTy.getNumElements());
      int vecBits = vecTy.getNumElements() * vecTy.getElementTypeBitWidth();
      assert(targetInfo.supportsDirectToLdsLoadBitWidth(vecBits));
      Value vecBytesVal = b.i32_val(vecBits / 8);

      Value maybeSwizzledMaskElem = maskElem;
      if (requiresSrcPtrSwizzling)
        applySwizzling(rewriter, loc, offsetElem, maybeSwizzledMaskElem, laneId,
                       swizzleLaneOffset);

      // Buffer-load-to-local supports zero-fill for per-lane masks by adjusting
      // the src offset to be OOB. Redundant-thread predication still needs a
      // branch when there are other values, otherwise inactive threads
      // zero-fill values loaded by active lanes from another warp.
      // Optimization: for warp-uniform thread predicates and no other values we
      // can avoid the branch by selecting an out-of-range *shared* address. The
      // HW will drop the load before fetching the data from global memory so we
      // will not overwrite values.
      if (isThreadPredWarpUniform && !hasOther) {
        Value predicatedAddress =
            selectLdsAddressForPredicate(b, threadPred, shmemAddr);
        auto bufferLoadToLds = bufferEmitter.emitLoadToLds(
            vecTy, vecBytesVal, rsrcDesc, offsetElem, predicatedAddress,
            maybeSwizzledMaskElem, op.getCache());
        if (targetInfo.requiresAliasInfoForAsyncOps())
          AMD::addAsyncCopyAliasScope(bufferLoadToLds);
      } else {
        Value pred =
            hasOther ? b.and_(threadPred, maybeSwizzledMaskElem) : threadPred;
        auto [loadBlock, afterLoadBlock] = emitBranch(rewriter, loc, pred);

        auto bufferLoadToLds = bufferEmitter.emitLoadToLds(
            vecTy, vecBytesVal, rsrcDesc, offsetElem, shmemAddr,
            hasOther ? b.true_val() : maybeSwizzledMaskElem, op.getCache());
        if (targetInfo.requiresAliasInfoForAsyncOps())
          AMD::addAsyncCopyAliasScope(bufferLoadToLds);

        if (hasOther) {
          emitOtherStore(rewriter, loc, this->getTypeConverter(), vecTy,
                         maskElem, otherElems, shmemAddr, laneId,
                         requiresSrcPtrSwizzling, swizzleLaneOffset);
        }

        rewriter.setInsertionPointToStart(afterLoadBlock);
      }

      return {};
    };

    auto res = lowerDirectLDSAsyncCopy(rewriter, loc, ptrType, flatDstTy,
                                       loadVals, llDst, resElemTy, vec,
                                       /*isLoad=*/true, emitBufferLoadLds);
    if (failed(res)) {
      return failure();
    }

    // Drop the result token.
    Value zero = LLVM::ConstantOp::create(rewriter, op.getLoc(),
                                          IntegerType::get(op.getContext(), 32),
                                          rewriter.getI32IntegerAttr(0));
    rewriter.replaceOp(op, zero);
    return success();
  }
};

struct AsyncCopyGlobalToLocalOpConversion
    : public ConvertOpToLLVMPattern<triton::gpu::AsyncCopyGlobalToLocalOp>,
      public DirectToLdsLoadConversionBase {
  AsyncCopyGlobalToLocalOpConversion(LLVMTypeConverter &converter,
                                     const AMD::TargetInfo &targetInfo,
                                     ModuleAxisInfoAnalysis &axisAnalysisPass,
                                     PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        DirectToLdsLoadConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::gpu::AsyncCopyGlobalToLocalOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);

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

    auto dstTy = op.getResult().getType();
    auto dstEnc = dstTy.getEncoding();
    auto resElemTy = getTypeConverter()->convertType(dstTy.getElementType());
    Value llDst = adaptor.getResult();

    // We can load N elements at a time if:
    //  1. Every group of N source pointers are contiguous.  For example, if
    //     N=2, then the pointers should be [x, x+1, y, y+1, ...].
    //  2. The mask (if present) has "alignment" N, meaning that each group of N
    //     mask bits are the same.  For example if N=2, the mask must be
    //     [x, x, y, y, ...].
    unsigned vec = getVectorSize(op.getSrc(), axisAnalysisPass);
    auto maskElements = getMaskElemsAndUpdateVeclen(
        rewriter, loc, adaptor.getMask(), op.getMask(), vec);

    auto srcElems = unpackTensorElements(loc, adaptor.getSrc(), rewriter,
                                         op.getSrc().getType());
    SmallVector<Value> otherElems;
    if (op.getOther())
      otherElems = unpackTensorElements(loc, adaptor.getOther(), rewriter,
                                        op.getOther().getType());

    // If the op has a contiguity hint use it to increase the vector size.
    vec = std::max(vec, op.getContiguity());

    if (!LLVM::AMD::canLoadDirectToLDS(targetInfo, srcTy, dstEnc,
                                       dstTy.getAllocShape(), vec)) {
      return failure();
    }

    // For swizzled layouts we need to use the non swizzled layout to compute
    // the LDS addresses since we gather into LDS
    auto flatDstTy = dstTy;
    SmallVector<Value> swizzledLaneOffsets;
    auto maybeSwizzledEnc = dyn_cast<SwizzledSharedEncodingAttr>(dstEnc);
    bool requiresSrcPtrSwizzling = !targetInfo.supportsDirectToLdsScatter() &&
                                   maybeSwizzledEnc &&
                                   maybeSwizzledEnc.getMaxPhase() != 1;
    if (requiresSrcPtrSwizzling) {
      auto flatSharedEnc = SwizzledSharedEncodingAttr::get(
          op->getContext(), maybeSwizzledEnc.getVec(), 1, 1,
          maybeSwizzledEnc.getOrder(), maybeSwizzledEnc.getCGALayout());
      flatDstTy = MemDescType::get(dstTy.getShape(), dstTy.getElementType(),
                                   flatSharedEnc, dstTy.getMemorySpace());
      swizzledLaneOffsets = emitSwizzledLaneOffsets(
          rewriter, op, srcTy, dstTy, flatDstTy, llDst, resElemTy, vec);
    }

    Type srcPtrTy = srcElems[0].getType();
    bool hasOther = !otherElems.empty();
    Type otherTy = hasOther ? otherElems[0].getType() : i1_ty;
    // Zip buffer_offset, mask, other, swizzleOffsets for lowerLdSt
    SmallVector<Value> loadVals =
        zipAsyncCopyValues(rewriter, loc, vec, srcElems, srcPtrTy, maskElements,
                           otherElems, otherTy, swizzledLaneOffsets);

    auto freeVarMasks = getFreeVariableMasks(srcTy);
    // We load redundant data on different CTAs so each CTA has a copy in its
    // shared memory; the multicast mask will be used by the hardware to
    // efficiently broadcast to different CTAs.
    freeVarMasks[rewriter.getStringAttr("block")] = 0;
    Value threadPred = emitRedundantThreadPredicateNonNull(
        freeVarMasks, rewriter, loc, targetInfo);
    bool hasMask = static_cast<bool>(op.getMask());
    bool threadPredIsWarpUniform =
        isRedundantThreadPredWarpUniform(freeVarMasks, rewriter.getContext());

    auto laneId = getLaneId(rewriter, loc);
    auto emitGlobalLoadLds = [&](RewriterBase &rewriter, Location loc,
                                 ArrayRef<Value> loadValues, Value shmemAddr,
                                 int startIdx, VectorType vecTy,
                                 Value multicastMask) -> SmallVector<Value> {
      auto [srcElem, maskElem, otherElems, swizzleLaneOffset] =
          unzipAsyncCopyValues(rewriter, loc, startIdx, loadValues, srcPtrTy,
                               otherTy, hasOther, vecTy.getNumElements());
      int vecBits = vecTy.getNumElements() * vecTy.getElementTypeBitWidth();
      assert(targetInfo.supportsDirectToLdsLoadBitWidth(vecBits));
      Value maybeSwizzledMaskElem = maskElem;

      if (requiresSrcPtrSwizzling)
        applySwizzling(rewriter, loc, srcElem, maybeSwizzledMaskElem, laneId,
                       swizzleLaneOffset);

      // Predicate load based on threadPred && swizzledMask
      auto cond = b.and_(threadPred, maybeSwizzledMaskElem);

      if (targetInfo.supportsDirectToLdsScatter() ||
          (threadPredIsWarpUniform && !hasMask)) {
        // For architectures supporting per lane LDS addresses or if the
        // predicate is warp-uniform, mask loads by setting the *shared* address
        // to out of range, the HW will drop the load before fetching the data
        // from global memory.
        Value predicatedAddress =
            selectLdsAddressForPredicate(b, cond, shmemAddr);

        emitAsyncLoad(rewriter, loc, targetInfo, vecBits, srcElem,
                      predicatedAddress, op.getCache(), multicastMask);
      } else {
        // For architectures not supporting per lane LDS addresses we need to
        // emit a branch.
        auto [loadBlock, afterLoadBlock] = emitBranch(rewriter, loc, cond);

        emitAsyncLoad(rewriter, loc, targetInfo, vecBits, srcElem, shmemAddr,
                      op.getCache(), multicastMask);

        rewriter.setInsertionPointToStart(afterLoadBlock);
      }

      if (hasOther) {
        emitOtherStore(rewriter, loc, this->getTypeConverter(), vecTy, maskElem,
                       otherElems, shmemAddr, laneId, requiresSrcPtrSwizzling,
                       swizzleLaneOffset);
      }

      return {};
    };

    auto res = lowerDirectLDSAsyncCopy(rewriter, loc, srcTy, flatDstTy,
                                       loadVals, llDst, resElemTy, vec,
                                       /*isLoad=*/true, emitGlobalLoadLds);
    if (failed(res)) {
      return failure();
    }

    // Drop the result token.
    Value zero = LLVM::ConstantOp::create(rewriter, op.getLoc(),
                                          IntegerType::get(op.getContext(), 32),
                                          rewriter.getI32IntegerAttr(0));
    rewriter.replaceOp(op, zero);
    return success();
  }

  void emitAsyncLoad(RewriterBase &rewriter, Location loc,
                     AMD::TargetInfo targetInfo, int vecBits, Value srcPtr,
                     Value shmemAddr, triton::CacheModifier cacheMod,
                     Value multicastMask) const {
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    int32_t cacheModifiers =
        mlir::LLVM::AMD::getCtrlBitsForCacheModifierOnTarget(
            cacheMod, /*isLoad=*/true, targetInfo);
    auto auxAttr = rewriter.getI32IntegerAttr(cacheModifiers);

    if (targetInfo.useAsyncMarks()) {
      // Use the async intrinsic so LLVM tracks these via asyncmark
      auto asyncLoadOp = ROCDL::GlobalLoadAsyncLDSOp::create(
          rewriter, loc, srcPtr, shmemAddr, vecBits / 8, /*offset=*/0u, auxAttr,
          nullptr, nullptr, nullptr);
      if (targetInfo.requiresAliasInfoForAsyncOps())
        AMD::addAsyncCopyAliasScope(asyncLoadOp);
    } else if (targetInfo.getISAFamily() == ISAFamily::GFX1250) {
      switch (vecBits) {
      case 32:
        if (multicastMask)
          ROCDL::ClusterLoadAsyncToLDSB32Op::create(
              rewriter, loc, srcPtr, shmemAddr, 0u, auxAttr, multicastMask,
              nullptr, nullptr, nullptr);
        else
          ROCDL::GlobalLoadAsyncToLDSB32Op::create(rewriter, loc, srcPtr,
                                                   shmemAddr, 0u, auxAttr,
                                                   nullptr, nullptr, nullptr);
        break;
      case 64:
        if (multicastMask)
          ROCDL::ClusterLoadAsyncToLDSB64Op::create(
              rewriter, loc, srcPtr, shmemAddr, 0u, auxAttr, multicastMask,
              nullptr, nullptr, nullptr);
        else
          ROCDL::GlobalLoadAsyncToLDSB64Op::create(rewriter, loc, srcPtr,
                                                   shmemAddr, 0u, auxAttr,
                                                   nullptr, nullptr, nullptr);
        break;
      case 128:
        if (multicastMask)
          ROCDL::ClusterLoadAsyncToLDSB128Op::create(
              rewriter, loc, srcPtr, shmemAddr, 0u, auxAttr, multicastMask,
              nullptr, nullptr, nullptr);
        else
          ROCDL::GlobalLoadAsyncToLDSB128Op::create(rewriter, loc, srcPtr,
                                                    shmemAddr, 0u, auxAttr,
                                                    nullptr, nullptr, nullptr);
        break;
      default:
        llvm_unreachable("Unsupported vec size for async load");
      }
    }
  }
};

struct AsyncCopyLocalToGlobalOpConversion
    : public ConvertOpToLLVMPattern<triton::amdgpu::AsyncCopyLocalToGlobalOp>,
      public DirectToLdsLoadConversionBase {
  AsyncCopyLocalToGlobalOpConversion(LLVMTypeConverter &converter,
                                     const AMD::TargetInfo &targetInfo,
                                     ModuleAxisInfoAnalysis &axisAnalysisPass,
                                     PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        DirectToLdsLoadConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::amdgpu::AsyncCopyLocalToGlobalOp op,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Only supported on GFX1250
    if (targetInfo.getISAFamily() != ISAFamily::GFX1250) {
      return rewriter.notifyMatchFailure(
          op, "async_copy_local_to_global only supported on GFX1250");
    }

    auto loc = op.getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);

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

    auto dstTy = op.getDst().getType();
    auto resElemTy = getTypeConverter()->convertType(srcTy.getElementType());
    Value llSrc = adaptor.getSrc();

    // We can store N elements at a time if:
    //  1. Every group of N destination pointers are contiguous.
    //  2. The mask (if present) has "alignment" N.
    unsigned vec = getVectorSize(op.getDst(), axisAnalysisPass);
    auto maskElements = getMaskElemsAndUpdateVeclen(
        rewriter, loc, adaptor.getMask(), op.getMask(), vec);

    auto dstElems = unpackTensorElements(loc, adaptor.getDst(), rewriter,
                                         op.getDst().getType());

    // If the op has a contiguity hint use it to increase the vector size.
    vec = std::max(vec, op.getContiguity());

    // For padded encodings restrict vec by the min interval
    auto srcEnc = srcTy.getEncoding();
    if (auto padEnc = dyn_cast<PaddedSharedEncodingAttr>(srcEnc)) {
      vec = std::min(vec, padEnc.getMinInterval());
    }

    Type dstPtrTy = dstElems[0].getType();
    // Zip dst_ptr, mask for lowerLdSt
    SmallVector<Value> storeVals = zipAsyncCopyValues(
        rewriter, loc, vec, dstElems, dstPtrTy, maskElements, {}, i1_ty, {});

    auto freeVarMasks = getFreeVariableMasks(dstTy);
    Value threadPred = emitRedundantThreadPredicateNonNull(
        freeVarMasks, rewriter, loc, targetInfo);

    auto emitGlobalStoreLds =
        [this, &op, &b, threadPred, dstPtrTy](
            RewriterBase &rewriter, Location loc, ArrayRef<Value> storeValues,
            Value shmemAddr, int startIdx, VectorType vecTy,
            Value /*multicastMask*/) -> SmallVector<Value> {
      auto [dstElem, maskElem, unused1, unused2] =
          unzipAsyncCopyValues(rewriter, loc, startIdx, storeValues, dstPtrTy,
                               i1_ty, false, vecTy.getNumElements());
      int vecBits = vecTy.getNumElements() * vecTy.getElementTypeBitWidth();

      // Predicate store based on threadPred && mask
      auto cond = b.and_(threadPred, maskElem);
      auto [storeBlock, afterStoreBlock] = emitBranch(rewriter, loc, cond);

      emitAsyncStore(rewriter, loc, targetInfo, vecBits, dstElem, shmemAddr,
                     op.getCache());

      rewriter.setInsertionPointToStart(afterStoreBlock);

      return {};
    };

    auto res = lowerDirectLDSAsyncCopy(rewriter, loc, dstTy, srcTy, storeVals,
                                       llSrc, resElemTy, vec,
                                       /*isLoad=*/false, emitGlobalStoreLds);
    if (failed(res)) {
      return failure();
    }

    // Drop the result token.
    Value zero = LLVM::ConstantOp::create(rewriter, op.getLoc(),
                                          IntegerType::get(op.getContext(), 32),
                                          rewriter.getI32IntegerAttr(0));
    rewriter.replaceOp(op, zero);
    return success();
  }

  void emitAsyncStore(RewriterBase &rewriter, Location loc,
                      AMD::TargetInfo targetInfo, int vecBits, Value dstPtr,
                      Value shmemAddr, triton::CacheModifier cacheMod) const {
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    int32_t cacheModifiers =
        mlir::LLVM::AMD::getCtrlBitsForCacheModifierOnTarget(
            cacheMod, /*isLoad=*/false, targetInfo);

    auto emitStore = [&](int bits, Value dst, Value shmem) {
      std::string intrinsic =
          "llvm.amdgcn.global.store.async.from.lds.b" + std::to_string(bits);
      LLVM::createLLVMIntrinsicCallOp(
          rewriter, loc, intrinsic, {},
          {dst, shmem, b.i32_val(0), b.i32_val(cacheModifiers)});
    };

    // If vecBits is not supported but vecBits/2 is, split into two
    // stores
    if (!targetInfo.supportsDirectFromLdsStoreBitWidth(vecBits) &&
        targetInfo.supportsDirectFromLdsStoreBitWidth(vecBits / 2)) {
      int halfVecBits = vecBits / 2;
      int halfVecBytes = halfVecBits / 8;
      // First half store
      emitStore(halfVecBits, dstPtr, shmemAddr);
      // Second half store (advance pointers by halfVecBytes)
      Value dstPtr2 = b.gep(ptr_ty(rewriter.getContext(), 1), i8_ty, dstPtr,
                            b.i32_val(halfVecBytes));
      Value shmemAddr2 = b.gep(ptr_ty(rewriter.getContext(), 3), i8_ty,
                               shmemAddr, b.i32_val(halfVecBytes));
      emitStore(halfVecBits, dstPtr2, shmemAddr2);
    } else {
      assert(targetInfo.supportsDirectFromLdsStoreBitWidth(vecBits));
      emitStore(vecBits, dstPtr, shmemAddr);
    }
  }
};

struct AsyncTDMCopyGlobalToLocalOpConversion
    : public ConvertOpToLLVMPattern<
          triton::amdgpu::AsyncTDMCopyGlobalToLocalOp>,
      public LoadStoreConversionBase {
  AsyncTDMCopyGlobalToLocalOpConversion(
      LLVMTypeConverter &converter, const AMD::TargetInfo &targetInfo,
      ModuleAxisInfoAnalysis &axisAnalysisPass, PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        LoadStoreConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::amdgpu::AsyncTDMCopyGlobalToLocalOp op,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);

    auto tensorDescTy = op.getDesc().getType();
    auto encoding = tensorDescTy.getSharedLayout();
    Type elementType =
        getTypeConverter()->convertType(tensorDescTy.getElementType());
    // Use descBlockTy to query shared layout because TDM lowering logic expects
    // the descriptor's dimensionality. For rank-reducing loads, destination
    // shared memory may have fewer dimensions than the descriptor block type.
    triton::LinearLayout sharedLayout =
        toLinearLayoutIgnoringPadding(tensorDescTy.getShape(), encoding);
    // Extract padding information if present
    unsigned padInterval = 0;
    unsigned padAmount = 0;
    if (auto padEnc = getPaddedEncoding(encoding)) {
      assert(padEnc.getIntervals().size() == 1 &&
             padEnc.getPaddings().size() == 1);
      padInterval = padEnc.getIntervals()[0];
      padAmount = padEnc.getPaddings()[0];
    }

    Value multicastMask;
    if (targetInfo.supportsMultiCTALaunch()) {
      multicastMask = LLVM::AMD::emitCtaMulticastMask(
          rewriter, loc, targetInfo.getClusterCTAId(rewriter, loc),
          sharedLayout, targetInfo.getMaxMulticastMaskPopcount());
    }

    SmallVector<Value> desc =
        mlir::LLVM::AMD::unpackTDMDescriptor(rewriter, loc, adaptor.getDesc());

    SmallVector<int64_t> blockShape = llvm::to_vector(tensorDescTy.getShape());

    // 2D tensors: 2 vector groups (group0: <4 x i32>, group1: <8 x i32>)
    // 3D-5D tensors: 4 vector groups (+group2, group3: <4 x i32>)
    assert((blockShape.size() <= 2 && desc.size() == 2) ||
           (blockShape.size() > 2 && desc.size() == 4));

    auto dstMemObj = LLVM::getSharedMemoryObjectFromStruct(
        loc, adaptor.getResult(), elementType, rewriter);
    // Get all base pointers (multiple for partitioned encoding)
    SmallVector<Value> dstPtrs = llvm::to_vector(dstMemObj.getBases());
    // Positioning lives in the descriptor; the copy only does per-warp
    // distribution, so the user offset is zero.
    SmallVector<Value> offset(blockShape.size(), b.i32_val(0));
    int numWarps = triton::gpu::lookupNumWarps(op);

    Value barrierPtr = nullptr;
    if (op.getBarrier()) {
      auto smemObj = LLVM::getSharedMemoryObjectFromStruct(
          loc, adaptor.getBarrier(),
          typeConverter->convertType(
              op.getBarrier().getType().getElementType()),
          rewriter);
      barrierPtr = smemObj.getBase();
    }

    auto ctaId = targetInfo.getClusterCTAId(rewriter, loc);

    auto shapePerCTA =
        triton::gpu::getShapePerCTA(encoding, tensorDescTy.getShape());

    std::optional<uint32_t> warpUsedHint;
    if (auto hintAttr = op.getWarpUsedHintAttr())
      warpUsedHint = static_cast<uint32_t>(hintAttr.getInt());

    auto cacheMod = op.getCache();
    auto auxBits = mlir::LLVM::AMD::getCtrlBitsForCacheModifierOnTarget(
        cacheMod, /*isLoad*/ true, targetInfo);

    // Placeholder: the copy inherits pred from the descriptor.
    Value pred = b.i32_val(1);
    mlir::LLVM::AMD::emitTDMLoadStore(
        rewriter, loc, getTypeConverter(), desc, shapePerCTA, numWarps,
        padInterval, padAmount, offset, dstPtrs, pred, multicastMask,
        elementType, barrierPtr, /*isLoad=*/true, sharedLayout, encoding, ctaId,
        auxBits, warpUsedHint, /*isPureForm=*/true);

    rewriter.eraseOp(op);
    return success();
  }
};

struct AsyncTDMFusedCopyGlobalToLocalOpConversion
    : public ConvertOpToLLVMPattern<
          triton::amdgpu::AsyncTDMFusedCopyGlobalToLocalOp>,
      public LoadStoreConversionBase {
  AsyncTDMFusedCopyGlobalToLocalOpConversion(
      LLVMTypeConverter &converter, const AMD::TargetInfo &targetInfo,
      ModuleAxisInfoAnalysis &axisAnalysisPass, PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        LoadStoreConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::amdgpu::AsyncTDMFusedCopyGlobalToLocalOp op,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    size_t numMembers = op.getDescs().size();
    int numWarps = triton::gpu::lookupNumWarps(op);
    Value ctaId = targetInfo.getClusterCTAId(rewriter, loc);

    SmallVector<mlir::LLVM::AMD::TDMFusedLoadMemberInfo, 4> members(numMembers);
    SmallVector<uint32_t, 4> memberHints;

    for (size_t i = 0; i < numMembers; ++i) {
      auto descTy = cast<triton::TensorDescType>(op.getDescs()[i].getType());
      auto enc = descTy.getSharedLayout();
      mlir::LLVM::AMD::TDMFusedLoadMemberInfo &m = members[i];

      m.elementType = getTypeConverter()->convertType(descTy.getElementType());
      m.sharedLayout = toLinearLayoutIgnoringPadding(descTy.getShape(), enc);
      if (auto padEnc = getPaddedEncoding(enc)) {
        assert(padEnc.getIntervals().size() == 1 &&
               padEnc.getPaddings().size() == 1);
        m.padInterval = padEnc.getIntervals()[0];
        m.padAmount = padEnc.getPaddings()[0];
      }
      if (targetInfo.supportsMultiCTALaunch())
        m.multicastMask = LLVM::AMD::emitCtaMulticastMask(
            rewriter, loc, ctaId, m.sharedLayout,
            targetInfo.getMaxMulticastMaskPopcount());

      m.sharedEncoding = enc;
      m.shapePerCTA =
          llvm::to_vector(triton::gpu::getShapePerCTA(enc, descTy.getShape()));
      m.desc = mlir::LLVM::AMD::unpackTDMDescriptor(rewriter, loc,
                                                    adaptor.getDescs()[i]);
      m.copyOffsets.append(descTy.getShape().size(), b.i32_val(0));

      auto dstMemObj = LLVM::getSharedMemoryObjectFromStruct(
          loc, adaptor.getDests()[i], m.elementType, rewriter);
      m.dstPtrs = llvm::to_vector(dstMemObj.getBases());
      m.pred = Value();
      memberHints.push_back(static_cast<uint32_t>(op.getWarpUsedHints()[i]));
    }

    auto auxBits = mlir::LLVM::AMD::getCtrlBitsForCacheModifierOnTarget(
        op.getCache(), /*isLoad*/ true, targetInfo);
    mlir::LLVM::AMD::emitTDMLoadFused(rewriter, loc, getTypeConverter(),
                                      members, numWarps, ctaId, auxBits,
                                      memberHints);

    rewriter.eraseOp(op);
    return success();
  }
};

struct AsyncTDMCopyLocalToGlobalOpConversion
    : public ConvertOpToLLVMPattern<
          triton::amdgpu::AsyncTDMCopyLocalToGlobalOp>,
      public LoadStoreConversionBase {
  AsyncTDMCopyLocalToGlobalOpConversion(
      LLVMTypeConverter &converter, const AMD::TargetInfo &targetInfo,
      ModuleAxisInfoAnalysis &axisAnalysisPass, PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        LoadStoreConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::amdgpu::AsyncTDMCopyLocalToGlobalOp op,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);

    auto tensorDescTy = op.getDesc().getType();
    auto smemTy = op.getSrc().getType();
    Type elementType = getTypeConverter()->convertType(smemTy.getElementType());

    SmallVector<Value> desc =
        mlir::LLVM::AMD::unpackTDMDescriptor(rewriter, loc, adaptor.getDesc());

    SmallVector<int64_t> blockShape = llvm::to_vector(tensorDescTy.getShape());

    // 2D tensors: 2 vector groups (group0: <4 x i32>, group1: <8 x i32>)
    // 3D-5D tensors: 4 vector groups (+group2, group3: <4 x i32>)
    assert((blockShape.size() <= 2 && desc.size() == 2) ||
           (blockShape.size() > 2 && desc.size() == 4));

    auto dstMemObj = LLVM::getSharedMemoryObjectFromStruct(
        loc, adaptor.getSrc(), elementType, rewriter);
    // Get all base pointers (multiple for partitioned encoding)
    SmallVector<Value> srcPtrs = llvm::to_vector(dstMemObj.getBases());
    // Positioning lives in the descriptor; the copy only does per-warp
    // distribution, so the user offset is zero.
    SmallVector<Value> offset(blockShape.size(), b.i32_val(0));
    int numWarps = triton::gpu::lookupNumWarps(op);

    Value barrierPtr = nullptr;
    if (op.getBarrier()) {
      auto smemObj = LLVM::getSharedMemoryObjectFromStruct(
          loc, adaptor.getBarrier(),
          typeConverter->convertType(
              op.getBarrier().getType().getElementType()),
          rewriter);
      barrierPtr = smemObj.getBase();
    }

    auto encoding = smemTy.getEncoding();
    auto paddedEnc = getPaddedEncoding(encoding);
    unsigned padInterval = 0;
    unsigned padAmount = 0;

    if (paddedEnc) {
      padInterval = paddedEnc.getIntervals()[0];
      padAmount = paddedEnc.getPaddings()[0];
    }

    triton::LinearLayout sharedLayout = toLinearLayoutIgnoringPadding(smemTy);

    auto ctaId = targetInfo.getClusterCTAId(rewriter, loc);

    auto shapePerCTA = triton::gpu::getShapePerCTA(smemTy);

    auto cacheMod = op.getCache();
    auto auxBits = mlir::LLVM::AMD::getCtrlBitsForCacheModifierOnTarget(
        cacheMod, /*isLoad*/ false, targetInfo);

    // Placeholder: the copy inherits pred from the descriptor.
    Value pred = arith::ConstantIntOp::create(rewriter, loc, 1, 32);
    mlir::LLVM::AMD::emitTDMLoadStore(
        rewriter, loc, getTypeConverter(), desc, shapePerCTA, numWarps,
        padInterval, padAmount, offset, srcPtrs, pred,
        /*multicastMask=*/{}, elementType, barrierPtr,
        /*isLoad=*/false, sharedLayout, encoding, ctaId, auxBits,
        /*warpUsedHint=*/std::nullopt, /*isPureForm=*/true);

    rewriter.eraseOp(op);
    return success();
  }
};

struct AsyncTDMScatterOpConversion
    : public ConvertOpToLLVMPattern<triton::amdgpu::AsyncTDMScatterOp>,
      public LoadStoreConversionBase {
  AsyncTDMScatterOpConversion(LLVMTypeConverter &converter,
                              const AMD::TargetInfo &targetInfo,
                              ModuleAxisInfoAnalysis &axisAnalysisPass,
                              PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        LoadStoreConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::amdgpu::AsyncTDMScatterOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);

    // Multi-CTA not supported for scatter
    if (lookupNumCTAs(op) > 1) {
      return op.emitError("TDM scatter does not support multi-CTA");
    }

    auto tensorDescTy = op.getDesc().getType();
    auto smemTy = op.getSrc().getType();
    Type elementType = getTypeConverter()->convertType(smemTy.getElementType());

    SmallVector<Value> desc =
        mlir::LLVM::AMD::unpackTDMDescriptor(rewriter, loc, adaptor.getDesc());

    SmallVector<int64_t> blockShape = llvm::to_vector(tensorDescTy.getShape());

    // Scatter only supports 2D tensors
    assert(blockShape.size() == 2 &&
           "TDM scatter mode only supports 2D tensors");

    auto srcMemObj = LLVM::getSharedMemoryObjectFromStruct(
        loc, adaptor.getSrc(), elementType, rewriter);
    Value srcPtr = srcMemObj.getBase();

    Value barrierPtr = nullptr;
    if (op.getBarrier()) {
      auto smemObj = LLVM::getSharedMemoryObjectFromStruct(
          loc, adaptor.getBarrier(),
          typeConverter->convertType(
              op.getBarrier().getType().getElementType()),
          rewriter);
      barrierPtr = smemObj.getBase();
    }

    // Get the destination row indices for scatter
    SmallVector<Value> dstRowIndices =
        unpackTensorElements(loc, adaptor.getDstRowIndices(), rewriter,
                             op.getDstRowIndices().getType());

    auto shapePerCTA = triton::gpu::getShapePerCTA(smemTy);

    auto dstRowIndicesType =
        cast<RankedTensorType>(op.getDstRowIndices().getType());

    // Extract padding information if present
    auto paddedEnc =
        llvm::dyn_cast<PaddedSharedEncodingAttr>(smemTy.getEncoding());
    unsigned padInterval = 0;
    unsigned padAmount = 0;
    if (paddedEnc) {
      assert(paddedEnc.getIntervals().size() == 1 &&
             paddedEnc.getPaddings().size() == 1);
      padInterval = paddedEnc.getIntervals()[0];
      padAmount = paddedEnc.getPaddings()[0];
    }

    // Create the CGA layout
    triton::LinearLayout sharedLayout;
    if (paddedEnc) {
      sharedLayout = paddedEnc.getLinearComponent();
    } else {
      sharedLayout = triton::gpu::toLinearLayout(smemTy);
    }
    auto kBlock = rewriter.getStringAttr("block");
    auto cgaLayout = sharedLayout.sublayout(
        {kBlock}, to_vector(sharedLayout.getOutDimNames()));
    auto ctaId = targetInfo.getClusterCTAId(rewriter, loc);
    int numWarps = triton::gpu::lookupNumWarps(op);

    if (failed(mlir::LLVM::AMD::emitTDMGatherScatter(
            rewriter, loc, getTypeConverter(), desc, shapePerCTA, padInterval,
            padAmount, srcPtr, /*multicastMask=*/{}, elementType, barrierPtr,
            cgaLayout, ctaId, dstRowIndices,
            /*isGather=*/false, numWarps, dstRowIndicesType)))
      return failure();

    rewriter.eraseOp(op);
    return success();
  }
};

struct AsyncTDMGatherOpConversion
    : public ConvertOpToLLVMPattern<triton::amdgpu::AsyncTDMGatherOp>,
      public LoadStoreConversionBase {
  AsyncTDMGatherOpConversion(LLVMTypeConverter &converter,
                             const AMD::TargetInfo &targetInfo,
                             ModuleAxisInfoAnalysis &axisAnalysisPass,
                             PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        LoadStoreConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::amdgpu::AsyncTDMGatherOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);

    auto tensorDescTy = op.getDesc().getType();
    auto smemTy = op.getDst().getType();
    auto paddedEnc =
        llvm::dyn_cast<PaddedSharedEncodingAttr>(smemTy.getEncoding());
    Type elementType = getTypeConverter()->convertType(smemTy.getElementType());

    SmallVector<Value> desc =
        mlir::LLVM::AMD::unpackTDMDescriptor(rewriter, loc, adaptor.getDesc());

    SmallVector<int64_t> blockShape = llvm::to_vector(tensorDescTy.getShape());

    // Gather only supports 2D tensors
    assert(blockShape.size() == 2 &&
           "TDM gather mode only supports 2D tensors");

    auto dstMemObj = LLVM::getSharedMemoryObjectFromStruct(
        loc, adaptor.getDst(), elementType, rewriter);
    Value dstPtr = dstMemObj.getBase();

    Value barrierPtr = nullptr;
    if (op.getBarrier()) {
      auto smemObj = LLVM::getSharedMemoryObjectFromStruct(
          loc, adaptor.getBarrier(),
          typeConverter->convertType(
              op.getBarrier().getType().getElementType()),
          rewriter);
      barrierPtr = smemObj.getBase();
    }

    // Get the source row indices for gather
    SmallVector<Value> srcRowIndices =
        unpackTensorElements(loc, adaptor.getSrcRowIndices(), rewriter,
                             op.getSrcRowIndices().getType());

    auto shapePerCTA = triton::gpu::getShapePerCTA(smemTy);

    auto srcRowIndicesType =
        cast<RankedTensorType>(op.getSrcRowIndices().getType());

    // Create the CGA layout
    triton::LinearLayout sharedLayout;
    unsigned padInterval = 0;
    unsigned padAmount = 0;
    if (paddedEnc) {
      assert(paddedEnc.getIntervals().size() == 1 &&
             paddedEnc.getPaddings().size() == 1);
      sharedLayout = paddedEnc.getLinearComponent();
      padInterval = paddedEnc.getIntervals()[0];
      padAmount = paddedEnc.getPaddings()[0];
    } else {
      sharedLayout = triton::gpu::toLinearLayout(smemTy);
    }

    auto kBlock = rewriter.getStringAttr("block");
    auto cgaLayout = sharedLayout.sublayout(
        {kBlock}, to_vector(sharedLayout.getOutDimNames()));
    auto ctaId = targetInfo.getClusterCTAId(rewriter, loc);
    int numWarps = triton::gpu::lookupNumWarps(op);

    Value multicastMask;
    if (targetInfo.supportsMultiCTALaunch()) {
      // Use the sharedLayout to compute the multicast mask because the index
      // layout only describes rows and misses information about columns.
      multicastMask = LLVM::AMD::emitCtaMulticastMask(
          rewriter, loc, ctaId, sharedLayout,
          targetInfo.getMaxMulticastMaskPopcount());
    }

    if (failed(mlir::LLVM::AMD::emitTDMGatherScatter(
            rewriter, loc, getTypeConverter(), desc, shapePerCTA, padInterval,
            padAmount, dstPtr, multicastMask, elementType, barrierPtr,
            cgaLayout, ctaId, srcRowIndices,
            /*isGather=*/true, numWarps, srcRowIndicesType)))
      return failure();

    rewriter.eraseOp(op);
    return success();
  }
};

struct StoreOpConversion : public ConvertOpToLLVMPattern<triton::StoreOp>,
                           public LoadStoreConversionBase {
  StoreOpConversion(LLVMTypeConverter &converter,
                    const AMD::TargetInfo &targetInfo,
                    ModuleAxisInfoAnalysis &axisAnalysisPass,
                    PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        LoadStoreConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::StoreOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value ptr = op.getPtr();
    Value value = op.getValue();
    Value mask = op.getMask();

    Value llPtr = adaptor.getPtr();
    Value llMask = adaptor.getMask();
    Value llValue = adaptor.getValue();

    auto loc = op->getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    MLIRContext *ctx = rewriter.getContext();

    auto valueTy = value.getType();
    Type valueElemTy =
        typeConverter->convertType(getElementTypeOrSelf(valueTy));

    // Determine the vectorization size
    unsigned vec = getVectorSize(ptr, axisAnalysisPass);
    unsigned elemsPerThread = getTotalElemsPerThread(ptr.getType());

    auto ptrElems = unpackTensorElements(loc, llPtr, rewriter, ptr.getType());
    auto valueElems =
        unpackTensorElements(loc, llValue, rewriter, op.getValue().getType());
    assert(ptrElems.size() == valueElems.size());

    SmallVector<Value> maskElems =
        getMaskElemsAndUpdateVeclen(rewriter, loc, llMask, mask, vec);

    const size_t valueElemNBits =
        std::max<int>(8, valueElemTy.getIntOrFloatBitWidth());

    auto cacheMod = op.getCache();
    const int numVecs = elemsPerThread / vec;
    auto freeVarMasks = getFreeVariableMasks(valueTy);
    Value threadPred = emitRedundantThreadPredicateNonNull(
        freeVarMasks, rewriter, loc, targetInfo);
    uint32_t regMask = freeVarMasks[str_attr("reg")];
    for (size_t vecStart = 0; vecStart < elemsPerThread; vecStart += vec) {
      if (!isCanonicalIndex(vecStart, regMask)) {
        // Don't emit store ops for redundant elements within a thread
        continue;
      }

      Value pred =
          llMask ? b.and_(threadPred, maskElems[vecStart]) : threadPred;

      auto vecTy = LLVM::getVectorType(valueElemTy, vec);

      const size_t maxWordWidth = std::max<size_t>(32, valueElemNBits);
      const size_t totalWidth = valueElemNBits * vec;
      const size_t width = std::min(totalWidth, maxWordWidth);
      const size_t nWords = std::max<size_t>(1, totalWidth / width);
      const size_t wordNElems = width / valueElemNBits;
      assert(wordNElems * nWords * numVecs == elemsPerThread);

      SmallVector<std::pair<Value, std::string>> asmArgs;
      Value ptr = ptrElems[vecStart];

      // Create the store val
      Value storeVal = packElementRangeIntoVector(
          rewriter, this->getTypeConverter(), loc, cast<VectorType>(vecTy),
          valueElems, vecStart);
      llStore(rewriter, loc, ptr, storeVal, pred, cacheMod);
    } // end vec
    rewriter.eraseOp(op);
    return success();
  }
};

struct BufferAtomicRMWOpConversion
    : public ConvertOpToLLVMPattern<triton::amdgpu::BufferAtomicRMWOp>,
      public LoadStoreConversionBase {
  BufferAtomicRMWOpConversion(LLVMTypeConverter &converter,
                              const AMD::TargetInfo &targetInfo,
                              ModuleAxisInfoAnalysis &axisAnalysisPass,
                              PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        LoadStoreConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::amdgpu::BufferAtomicRMWOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op->getLoc();
    insertAtomicOrderingBarriers(op, op.getSem(),
                                 !op->hasAttr("allocation.offset"), rewriter,
                                 targetInfo);
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    LLVM::AMD::BufferEmitter bufferEmitter(rewriter, loc, targetInfo);

    // original values
    Value ptr = op.getPtr();
    Value offset = op.getOffsets();
    Value mask = op.getMask();
    Value data = op.getValue();
    auto atomicRmwAttr = op.getAtomicRmwOp();

    Value llPtr = adaptor.getPtr();
    Value llOffset = adaptor.getOffsets();
    Value llMask = adaptor.getMask();
    Value llData = adaptor.getValue();
    Value llStride = adaptor.getStride();

    // Determine the vectorization size
    Type valueTy = data.getType();
    Type valueElemTy =
        typeConverter->convertType(getElementTypeOrSelf(valueTy));
    Type ptrType = getPointerTypeWithShape(ptr, offset);

    unsigned numElems = getTotalElemsPerThread(ptrType);
    unsigned vec = getVectorSize(ptr, offset, axisAnalysisPass);

    // v4f16 and v4bf16 variants of buffer atomics do not exist.
    // only v2f16 and v2bf16.
    if (valueElemTy.isBF16() || valueElemTy.isF16()) {
      // We clamp to the only supported vectorization width here (2).
      // In ConvertToBufferOps we check that we have a large enough vector size
      assert(vec >= 2);
      vec = 2u;
      // The max width of a buffer atomic op is 64-bits
      // Some types like F32 don't have a 2x vectorized version
    } else if (valueElemTy.isF32() || valueElemTy.isF64() ||
               valueElemTy.isInteger(32) || valueElemTy.isInteger(64)) {
      vec = 1u;
    }

    // Get the offsets and value
    SmallVector<Value> offsetElems = unpackTensorElements(
        loc, llOffset, rewriter, op.getOffsets().getType());
    SmallVector<Value> valueElems =
        unpackTensorElements(loc, llData, rewriter, data.getType());

    // Get the mask
    SmallVector<Value> maskElems =
        getMaskElemsAndUpdateVeclen(rewriter, loc, llMask, mask, vec);

    Value rsrcDesc = bufferEmitter.createResourceDescriptor(llPtr, llStride);
    SmallVector<Value> loadedVals;

    // We need to manually emit memory fences (LLVM doesn't do this for buffer
    // ops) see: https://llvm.org/docs/AMDGPUUsage.html#memory-model-gfx942
    auto memOrdering = op.getSem();
    auto memScope = op.getScope();
    if (failed(emitFence(op, rewriter, loc, memOrdering, memScope,
                         true /*preAtomic*/))) {
      return failure();
    }

    mlir::Operation *lastRMWOp;
    MLIRContext *ctx = rewriter.getContext();
    GCNBuilder waitcntBuilder;

    //    We set GLC=1, to return the old value. Atomics in GFX942 execute with
    //    either device (default) or system scope (controlled by the sc1 flag).
    //    This is distinct from the memory scope of the atomic (i.e, the memory
    //    fences which appear before/after the ops).

    // Check if the op has users, if it does we set GLC=1, otherwise GLC=0
    auto opUsers = op.getResult().getUsers();
    auto hasUsers = std::distance(opUsers.begin(), opUsers.end()) > 0;

    auto freeVarMasks = getFreeVariableMasks(valueTy);
    Value threadPred = emitRedundantThreadPredicateNonNull(
        freeVarMasks, rewriter, loc, targetInfo);
    uint32_t regMask = freeVarMasks[str_attr("reg")];
    for (size_t vecStart = 0; vecStart < numElems; vecStart += vec) {
      if (!isCanonicalIndex(vecStart, regMask)) {
        // Don't emit store ops for redundant elements within a thread
        continue;
      }

      Value pred =
          llMask ? b.and_(threadPred, maskElems[vecStart]) : threadPred;

      Type vecTy = LLVM::getVectorType(valueElemTy, vec);
      // Create the store val
      Value storeVal = packElementRangeIntoVector(
          rewriter, this->getTypeConverter(), loc, cast<VectorType>(vecTy),
          valueElems, vecStart);

      Value loadVal = bufferEmitter.emitAtomicRMW(
          atomicRmwAttr, vecTy, rsrcDesc, offsetElems[vecStart], storeVal, pred,
          hasUsers);
      // Track the last op, so we can emit a fenceop after the loop
      lastRMWOp = loadVal.getDefiningOp();

      for (size_t ii = 0; ii < vec; ++ii) {
        Value vecIdx = createIndexAttrConstant(
            rewriter, loc, getTypeConverter()->getIndexType(), ii);
        Value loaded = b.extract_element(valueElemTy, loadVal, vecIdx);
        loadedVals.push_back(loaded);
      }
    } // end vec

    // Acquire Fence post-atomic
    if (failed(emitFence(op, rewriter, lastRMWOp->getLoc(), memOrdering,
                         memScope, false /*preAtomic*/))) {
      return failure();
    }

    finalizeTensorAtomicResults(op, dyn_cast<RankedTensorType>(valueTy),
                                rewriter, loadedVals, valueElemTy, b,
                                threadPred, targetInfo, getTypeConverter());
    return success();
  }
};

struct BufferAtomicCASOpConversion
    : public ConvertOpToLLVMPattern<triton::amdgpu::BufferAtomicCASOp>,
      public LoadStoreConversionBase {
  BufferAtomicCASOpConversion(LLVMTypeConverter &converter,
                              const AMD::TargetInfo &targetInfo,
                              ModuleAxisInfoAnalysis &axisAnalysisPass,
                              PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        LoadStoreConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::amdgpu::BufferAtomicCASOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op->getLoc();
    insertAtomicOrderingBarriers(op, op.getSem(),
                                 !op->hasAttr("allocation.offset"), rewriter,
                                 targetInfo);
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    LLVM::AMD::BufferEmitter bufferEmitter(rewriter, loc, targetInfo);

    // original values
    Value ptr = op.getPtr();
    Value offset = op.getOffsets();
    Value val = op.getVal();

    Value llPtr = adaptor.getPtr();
    Value llOffset = adaptor.getOffsets();
    Value llVal = adaptor.getVal();
    Value llCmp = adaptor.getCmp();
    Value llStride = adaptor.getStride();

    // Determine the vectorization size
    Type valueTy = val.getType();
    Type valueElemTy =
        typeConverter->convertType(getElementTypeOrSelf(valueTy));
    Type ptrType = getPointerTypeWithShape(ptr, offset);

    unsigned numElems = getTotalElemsPerThread(ptrType);
    // Max supported vectorization for i32 and i64 is 1x
    // on CDNA3 and CDNA4
    // BUFFER_ATOMIC_CMPSWAP(i32) and BUFFER_ATOMIC_CMPSWAP_X2(i64)
    unsigned vec = 1u;

    // Get the offsets, val, and cmp
    SmallVector<Value> offsetElems = unpackTensorElements(
        loc, llOffset, rewriter, op.getOffsets().getType());
    SmallVector<Value> valElems =
        unpackTensorElements(loc, llVal, rewriter, op.getVal().getType());
    SmallVector<Value> cmpElems =
        unpackTensorElements(loc, llCmp, rewriter, op.getCmp().getType());

    Value rsrcDesc = bufferEmitter.createResourceDescriptor(llPtr, llStride);
    SmallVector<Value> loadedVals;

    // We need to manually emit memory fences (LLVM doesn't do this for buffer
    // ops)
    auto memOrdering = op.getSem();
    auto memScope = op.getScope();
    // Release Fence pre-atomic
    if (failed(emitFence(op, rewriter, loc, memOrdering, memScope,
                         true /*preAtomic*/))) {
      return failure();
    }

    mlir::Operation *lastCASOp;
    GCNBuilder waitcntBuilder;

    // Check if the op has users, if it does we set GLC=1, otherwise GLC=0
    auto hasUsers = !op.getResult().getUsers().empty();
    auto freeVarMasks = getFreeVariableMasks(valueTy);
    Value threadPred = emitRedundantThreadPredicateNonNull(
        freeVarMasks, rewriter, loc, targetInfo);

    for (size_t vecStart = 0; vecStart < numElems; vecStart += vec) {
      Type vecTy = LLVM::getVectorType(valueElemTy, vec);
      Value pred = threadPred;
      // Create the store val
      Value casStoreVal = packElementRangeIntoVector(
          rewriter, this->getTypeConverter(), loc, cast<VectorType>(vecTy),
          valElems, vecStart);
      // Create the cmp val
      Value casCmpVal = packElementRangeIntoVector(
          rewriter, this->getTypeConverter(), loc, cast<VectorType>(vecTy),
          cmpElems, vecStart);

      Value loadVal =
          bufferEmitter.emitAtomicCAS(vecTy, rsrcDesc, offsetElems[vecStart],
                                      casCmpVal, casStoreVal, pred, hasUsers);
      // Track the last op, so we can emit a fenceop after the loop
      lastCASOp = loadVal.getDefiningOp();

      for (size_t ii = 0; ii < vec; ++ii) {
        Value vecIdx = createIndexAttrConstant(
            rewriter, loc, getTypeConverter()->getIndexType(), ii);
        Value loaded = b.extract_element(valueElemTy, loadVal, vecIdx);
        loadedVals.push_back(loaded);
      }
    } // end vec

    // Emit post-atomic acquire fence
    if (failed(emitFence(op, rewriter, lastCASOp->getLoc(), memOrdering,
                         memScope, false /*preAtomic*/))) {
      return failure();
    }

    finalizeTensorAtomicResults(op, dyn_cast<RankedTensorType>(valueTy),
                                rewriter, loadedVals, valueElemTy, b,
                                threadPred, targetInfo, getTypeConverter());
    return success();
  }
};

struct BufferStoreOpConversion
    : public ConvertOpToLLVMPattern<triton::amdgpu::BufferStoreOp>,
      public LoadStoreConversionBase {
  BufferStoreOpConversion(LLVMTypeConverter &converter,
                          const AMD::TargetInfo &targetInfo,
                          ModuleAxisInfoAnalysis &axisAnalysisPass,
                          PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        LoadStoreConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::amdgpu::BufferStoreOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op->getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    LLVM::AMD::BufferEmitter bufferEmitter(rewriter, loc, targetInfo);

    // original values
    Value ptr = op.getPtr();
    Value offset = op.getOffsets();
    Value mask = op.getMask();
    Value data = op.getValue();
    auto cacheMod = op.getCache();

    Value llPtr = adaptor.getPtr();
    Value llOffset = adaptor.getOffsets();
    Value llMask = adaptor.getMask();
    Value llData = adaptor.getValue();
    Value llStride = adaptor.getStride();

    // Determine the vectorization size
    Type valueTy = data.getType();
    Type valueElemTy =
        typeConverter->convertType(getElementTypeOrSelf(valueTy));
    Type ptrType = getPointerTypeWithShape(ptr, offset);

    unsigned numElems = getTotalElemsPerThread(ptrType);
    unsigned vec = getVectorSize(ptr, offset, axisAnalysisPass);
    // If the op has a contiguity hint use it to increase the vector size.
    vec = std::max(vec, op.getContiguity());

    // Get the offsets and value
    SmallVector<Value> offsetElems = unpackTensorElements(
        loc, llOffset, rewriter, op.getOffsets().getType());
    SmallVector<Value> valueElems =
        unpackTensorElements(loc, llData, rewriter, data.getType());

    // Get the mask
    SmallVector<Value> maskElems =
        getMaskElemsAndUpdateVeclen(rewriter, loc, llMask, mask, vec);

    Value rsrcDesc = bufferEmitter.createResourceDescriptor(llPtr, llStride);
    MLIRContext *ctx = rewriter.getContext();
    auto freeVarMasks = getFreeVariableMasks(valueTy);
    Value threadPred = emitRedundantThreadPredicateNonNull(
        freeVarMasks, rewriter, loc, targetInfo);
    uint32_t regMask = freeVarMasks[str_attr("reg")];
    for (size_t vecStart = 0; vecStart < numElems; vecStart += vec) {
      if (!isCanonicalIndex(vecStart, regMask)) {
        // Don't emit store ops for redundant elements within a thread
        continue;
      }

      Value pred =
          llMask ? b.and_(threadPred, maskElems[vecStart]) : threadPred;

      Type vecTy = LLVM::getVectorType(valueElemTy, vec);
      // Create the store val
      Value storeVal = packElementRangeIntoVector(
          rewriter, this->getTypeConverter(), loc, cast<VectorType>(vecTy),
          valueElems, vecStart);
      bufferEmitter.emitStore(rsrcDesc, offsetElems[vecStart], storeVal, pred,
                              cacheMod);
    } // end vec

    rewriter.eraseOp(op);
    return success();
  }
};

struct AtomicCASOpConversion
    : public ConvertOpToLLVMPattern<triton::AtomicCASOp>,
      public LoadStoreConversionBase {
  AtomicCASOpConversion(LLVMTypeConverter &converter,
                        const AMD::TargetInfo &targetInfo,
                        ModuleAxisInfoAnalysis &axisAnalysisPass,
                        PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        LoadStoreConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::AtomicCASOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // extract relevant info from Module
    auto loc = op.getLoc();
    insertAtomicOrderingBarriers(op, op.getSem(),
                                 !op->hasAttr("allocation.offset"), rewriter,
                                 targetInfo);
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    MLIRContext *ctx = rewriter.getContext();

    Value llPtr = adaptor.getPtr();
    Value llCmp = adaptor.getCmp();
    Value llVal = adaptor.getVal();

    // prep data by unpacking to get data ready
    auto ptrElements =
        unpackTensorElements(loc, llPtr, rewriter, op.getPtr().getType());
    auto cmpElements =
        unpackTensorElements(loc, llCmp, rewriter, op.getCmp().getType());
    auto valElements =
        unpackTensorElements(loc, llVal, rewriter, op.getVal().getType());

    auto memOrdering = op.getSem();
    auto atomicMemOrdering = getMemoryOrdering(memOrdering);
    if (!atomicMemOrdering)
      return rewriter.notifyMatchFailure(op, "Unknown AMDGPU memory ordering");
    auto scope = getAMDGPUMemScopeStr(op.getScope());
    if (!scope)
      return rewriter.notifyMatchFailure(op, "Unknown AMDGPU memory scope");

    // deal with tensor or scalar
    auto valueTy = op.getResult().getType();
    auto tensorTy = dyn_cast<RankedTensorType>(valueTy);
    Type valueElemTy =
        tensorTy ? getTypeConverter()->convertType(tensorTy.getElementType())
                 : valueTy;
    auto valueElemNBits = valueElemTy.getIntOrFloatBitWidth();
    Type valueElemIntTy{};
    if (!valueElemTy.isSignlessInteger()) {
      valueElemIntTy = rewriter.getIntegerType(valueElemNBits);
    }
    auto elemsPerThread = getTotalElemsPerThread(op.getVal().getType());
    SmallVector<Value> resultVals(elemsPerThread);

    auto successOrdering = *atomicMemOrdering;
    auto failureOrdering = LLVM::AtomicOrdering::monotonic;
    auto scopeStr = StringRef(scope.value());

    auto freeVarMasks = getFreeVariableMasks(op.getPtr().getType());
    Value threadPred = emitRedundantThreadPredicateNonNull(
        freeVarMasks, rewriter, loc, targetInfo);
    uint32_t regMask = freeVarMasks[str_attr("reg")];

    // atomic ops
    for (size_t i = 0; i < elemsPerThread; i += 1) {
      if (tensorTy && (i & ~regMask) != i) {
        resultVals[i] = resultVals[i & ~regMask];
        continue;
      }

      Value casVal = valElements[i];
      Value casCmp = cmpElements[i];
      Value casPtr = ptrElements[i];
      if (valueElemIntTy) {
        casVal = LLVM::BitcastOp::create(rewriter, loc, valueElemIntTy, casVal);
        casCmp = LLVM::BitcastOp::create(rewriter, loc, valueElemIntTy, casCmp);
      }
      // use op
      if (tensorTy) { // for tensor
        Value undefVal = b.undef(valueElemTy);
        auto *curBlock = rewriter.getInsertionBlock();
        auto *endBlock = curBlock->splitBlock(rewriter.getInsertionPoint());
        auto *atomicBlock = rewriter.createBlock(
            curBlock->getParent(), std::next(Region::iterator(curBlock)));
        endBlock->addArgument({valueElemTy}, {loc});

        rewriter.setInsertionPointToEnd(curBlock);
        LLVM::CondBrOp::create(rewriter, loc, threadPred, atomicBlock, endBlock,
                               undefVal);

        rewriter.setInsertionPointToEnd(atomicBlock);

        auto cmpxchg = LLVM::AtomicCmpXchgOp::create(
            rewriter, loc, casPtr, casCmp, casVal, successOrdering,
            failureOrdering, scopeStr);

        // Extract the new_loaded value from the pair.
        Value ret;
        if (valueElemIntTy) {
          ret = b.extract_val(valueElemIntTy, cmpxchg, 0);
          ret = LLVM::BitcastOp::create(rewriter, loc, valueElemTy, ret);
        } else {
          ret = b.extract_val(valueElemTy, cmpxchg, 0);
        }

        LLVM::BrOp::create(rewriter, loc, ret, endBlock);
        rewriter.setInsertionPointToStart(endBlock);
        resultVals[i] = endBlock->getArgument(0);
      } else { // for scalar
        // Build blocks to bypass the atomic instruction for ~rmwMask.
        auto *curBlock = rewriter.getInsertionBlock();
        auto *endBlock = curBlock->splitBlock(rewriter.getInsertionPoint());
        auto *atomicBlock = rewriter.createBlock(
            curBlock->getParent(), std::next(Region::iterator(curBlock)));

        // Fill entry block with global memory barrier and conditional branch.
        rewriter.setInsertionPointToEnd(curBlock);
        auto tid = getThreadId(rewriter, loc);
        Value pred = b.icmp_eq(tid, b.i32_val(i));
        LLVM::CondBrOp::create(rewriter, loc, pred, atomicBlock, endBlock);

        // Build main block with atomic_cmpxchg.
        rewriter.setInsertionPointToEnd(atomicBlock);

        auto cmpxchg = LLVM::AtomicCmpXchgOp::create(
            rewriter, loc, casPtr, casCmp, casVal, successOrdering,
            failureOrdering, scopeStr);

        if (!op.getResult().use_empty()) {
          // Extract the new_loaded value from the pair.
          Value newLoaded;
          if (valueElemIntTy) {
            newLoaded = b.extract_val(valueElemIntTy, cmpxchg, 0);
            newLoaded =
                LLVM::BitcastOp::create(rewriter, loc, valueElemTy, newLoaded);
          } else {
            newLoaded = b.extract_val(valueElemTy, cmpxchg, 0);
          }
          Value atomPtr =
              getSharedMemoryBase(loc, rewriter, targetInfo, op.getOperation());
          b.store(newLoaded, atomPtr);
        }

        LLVM::BrOp::create(rewriter, loc, ValueRange(), endBlock);

        // Build the last block: synced load from shared memory, exit.
        rewriter.setInsertionPointToStart(endBlock);

        if (op.getResult().use_empty()) {
          rewriter.eraseOp(op);
          return success();
        }

        b.barrier(triton::gpu::AddrSpace::Local);
        Value atomPtr =
            getSharedMemoryBase(loc, rewriter, targetInfo, op.getOperation());
        Value ret = b.load(valueElemTy, atomPtr);
        rewriter.replaceOp(op, {ret});
        return success();
      }
    }

    finalizeTensorAtomicResults(op, tensorTy, rewriter, resultVals, valueElemTy,
                                b, threadPred, targetInfo, getTypeConverter());
    return success();
  }
};

bool supportsGlobalAtomicF16PackedAndDpp(ISAFamily isaFamily) {
  switch (isaFamily) {
  case ISAFamily::CDNA1:
  case ISAFamily::CDNA2:
  case ISAFamily::CDNA3:
  case ISAFamily::CDNA4:
    return true;
  default:
    break;
  }
  return false;
}

struct AtomicRMWOpConversion
    : public ConvertOpToLLVMPattern<triton::AtomicRMWOp>,
      public LoadStoreConversionBase {
  AtomicRMWOpConversion(LLVMTypeConverter &converter,
                        const AMD::TargetInfo &targetInfo,
                        ModuleAxisInfoAnalysis &axisAnalysisPass,
                        PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit),
        LoadStoreConversionBase(targetInfo, axisAnalysisPass) {}

  LogicalResult
  matchAndRewrite(triton::AtomicRMWOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    insertAtomicOrderingBarriers(op, op.getSem(),
                                 !op->hasAttr("allocation.offset"), rewriter,
                                 targetInfo);
    auto b = TritonLLVMOpBuilder(loc, rewriter);

    auto binOp = matchAtomicOp(op.getAtomicRmwOp());
    if (!binOp)
      return rewriter.notifyMatchFailure(op, "Unsupported RMW operation");

    auto memOrder = getMemoryOrdering(op.getSem());
    if (!memOrder)
      return rewriter.notifyMatchFailure(op, "Unsupported RMW memory order");

    auto scopeStr = getAMDGPUMemScopeStr(op.getScope());
    if (!scopeStr)
      return rewriter.notifyMatchFailure(op, "Unsupported RMW scope");

    auto emitter =
        LLVM::AMD::AtomicRMWEmitter(targetInfo, *binOp, *memOrder, *scopeStr);

    Value val = op.getVal();
    Value ptr = op.getPtr();
    Value opResult = op.getResult();
    auto atomicRmwAttr = op.getAtomicRmwOp();

    Value llPtr = adaptor.getPtr();
    Value llVal = adaptor.getVal();
    Value llMask = adaptor.getMask();

    auto valElements =
        unpackTensorElements(loc, llVal, rewriter, op.getVal().getType());
    auto ptrElements =
        unpackTensorElements(loc, llPtr, rewriter, op.getPtr().getType());
    SmallVector<Value> maskElements;
    if (llMask)
      maskElements =
          unpackTensorElements(loc, llMask, rewriter, op.getMask().getType());

    auto tensorTy = dyn_cast<RankedTensorType>(opResult.getType());
    Type valueElemTy =
        tensorTy ? getTypeConverter()->convertType(tensorTy.getElementType())
                 : opResult.getType();

    // In the case of unpaired f16 elements utilize dpp instructions to
    // accelerate atomics. Here is an algorithm of lowering
    // tt::atomicRmwOp(%ptr, %val, %mask):
    // 0. Group thread by pairs. Master thread is (tid % 2 == 0);
    // 1. All the threads send %val to (tid - 1) thread via dppUpdateOp shl, so
    //    all the masters receive value from secondary threads;
    // 2. Take into account parity in the %mask value, build control flow
    //    structures according to it;
    // 3. Generate llvm::atomicRmwOp in the threads enabled by %mask value;
    // 4. All the threads send result of generated operation to (tid + 1) thread
    //    via dppUpdateOp shl, so all secondary thread also receive their
    //    result.
    //
    // This approach enables us to use half the active threads committing atomic
    // requests to avoid generating of code providing unified access to f16
    // element and reduce contention.
    bool applyPackingF16 = false;
    auto vec = getVectorSize(ptr, axisAnalysisPass);
    if (llMask) {
      vec = std::min<unsigned>(vec, getMaskAlignment(op.getMask()));
    }

    // CDNA3/CDNA4 arch allows to accelerate its atomics with LDS reduction
    // algorithm, which is only applicable for atomics with no return. Otherwise
    // we have to deal with an additional overhead.
    bool enableIntraWaveReduce =
        llvm::is_contained({ISAFamily::CDNA3, ISAFamily::CDNA4},
                           targetInfo.getISAFamily()) &&
        tensorTy && opResult.use_empty();

    // TODO: support data types less than 32 bits
    enableIntraWaveReduce &= valueElemTy.getIntOrFloatBitWidth() >= 32;

    if (tensorTy) {
      bool isF16Ty = valueElemTy.isF16() || valueElemTy.isBF16();
      unsigned availableVecSize = isF16Ty ? 2 : 1;
      vec = std::min<unsigned>(vec, availableVecSize);
      // Force F16 packing in the case it's not coming in as packed, but the
      // ISA can support packed atomic instructions.
      applyPackingF16 =
          supportsGlobalAtomicF16PackedAndDpp(targetInfo.getISAFamily()) &&
          vec == 1 && isF16Ty && atomicRmwAttr == RMWOp::FADD &&
          !enableIntraWaveReduce;

      auto threadOrder = getThreadOrder(tensorTy);
      unsigned contigWithinLanes =
          axisAnalysisPass.getAxisInfo(ptr)->getContiguity(threadOrder.front());
      enableIntraWaveReduce &= contigWithinLanes == 1;
    }

    auto vecTy = vec_ty(valueElemTy, vec);
    auto elemsPerThread = getTotalElemsPerThread(val.getType());

    auto freeVarMasks = getFreeVariableMasks(op.getPtr().getType());
    Value threadPred = emitRedundantThreadPredicateNonNull(
        freeVarMasks, rewriter, loc, targetInfo);

    bool needLdsStaging = !tensorTy && !opResult.use_empty();
    std::optional<Value> atomicSharedMemBase =
        op->hasAttr("allocation.offset") && needLdsStaging
            ? std::optional<Value>(getSharedMemoryBase(
                  loc, rewriter, targetInfo, op.getOperation()))
            : std::nullopt;

    SmallVector<Value> resultVals(elemsPerThread);
    for (size_t i = 0; i < elemsPerThread; i += vec) {
      // TODO: in case llMask is zero we can create only one branch for all
      // elemsPerThread.
      Value rmwMask = llMask ? b.and_(threadPred, maskElements[i]) : threadPred;
      if (applyPackingF16) {
        resultVals[i] = emitter.emitPairedAtomicForEvenTID(
            rewriter, ptrElements[i], valElements[i], rmwMask);
      } else {
        Value valElement;
        if (vec == 1) {
          valElement = valElements[i];
        } else {
          Value vecVal = b.undef(vecTy);
          for (size_t ii = 0; ii < vec; ++ii)
            vecVal = b.insert_element(vecTy, vecVal, valElements[i + ii],
                                      b.i32_val(ii));
          valElement = vecVal;
        }

        // If we have a single tl.atomic_rmw that is lowered into multiple
        // llvm.atomic_rmw, and we set the ordering for each to aql_rel (the
        // default if no sem value is explicitly set in the DSL level
        // tl.atomic_add. The llvm backend will insert extra buffer invalidates
        // and L2 write backs causing a perforance degration. To avoid this we
        // set the ordering to release for the first, acquire for the last, and
        // relaxed for anything in between so that only a single set of
        // buffer_inv and buffer_wbl2 instructions are inserted by the backend
        // for any "cluster" of atomic ops.
        if ((vec > 1 || elemsPerThread > 1) &&
            op.getSem() == MemSemantic::ACQUIRE_RELEASE) {
          if (i == 0) {
            // First
            emitter.setAtomicOrdering(LLVM::AtomicOrdering::release);
          } else if (i == elemsPerThread - vec) {
            // Last
            emitter.setAtomicOrdering(LLVM::AtomicOrdering::acquire);
          } else {
            // Middle
            emitter.setAtomicOrdering(LLVM::AtomicOrdering::monotonic);
          }
        }

        Value retVal =
            emitter.emitAtomicRMW(rewriter, ptrElements[i], valElement, rmwMask,
                                  atomicSharedMemBase, enableIntraWaveReduce);

        if (tensorTy) {
          for (int ii = 0; ii < vec; ++ii) {
            resultVals[i + ii] =
                vec == 1
                    ? retVal
                    : b.extract_element(valueElemTy, retVal, b.i32_val(ii));
          }
        } else {
          if (!atomicSharedMemBase.has_value()) {
            rewriter.eraseOp(op);
            return success();
          }
          Value atomPtr = *atomicSharedMemBase;
          b.barrier(triton::gpu::AddrSpace::Local);
          Value ret = b.load(valueElemTy, atomPtr);

          rewriter.replaceOp(op, {ret});
          return success();
        }
      }
    }
    finalizeTensorAtomicResults(op, tensorTy, rewriter, resultVals, valueElemTy,
                                b, threadPred, targetInfo, getTypeConverter());
    return success();
  }
};

// Lower ttg.async_wait directly to wait_asyncmark for architectures that use
// asyncmark-based synchronization (CDNA3/CDNA4). The commit group count from
// ttg.async_wait is passed through without clamping since LLVM will compute
// the final waitcnt based on #instructions instead of #commitGroups.
struct TTGAsyncWaitOpConversion : public ConvertOpToLLVMPattern<AsyncWaitOp> {
  TTGAsyncWaitOpConversion(LLVMTypeConverter &converter,
                           const AMD::TargetInfo &targetInfo,
                           PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit), targetInfo(targetInfo) {}

  LogicalResult
  matchAndRewrite(AsyncWaitOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (!targetInfo.useAsyncMarks())
      return rewriter.notifyMatchFailure(
          op, "ttg.async_wait should have been lowered to amdg.async_wait by "
              "UpdateAsyncWaitCount for non-asyncmark targets");

    auto loc = op->getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);

    // Pass through the commit group count. LLVM will compute the correct
    // vmcnt based on asyncmark/wait_asyncmark markers.
    ROCDL::WaitAsyncmarkOp::create(rewriter, loc, op.getNum());

    // Drop the result AsyncToken
    rewriter.replaceOp(op, b.i32_val(0));
    return success();
  }

private:
  const AMD::TargetInfo &targetInfo;
};

struct AsyncWaitOpConversion
    : public ConvertOpToLLVMPattern<amdgpu::AsyncWaitOp> {
  AsyncWaitOpConversion(LLVMTypeConverter &converter,
                        const AMD::TargetInfo &targetInfo,
                        PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit), targetInfo(targetInfo) {}

  LogicalResult
  matchAndRewrite(amdgpu::AsyncWaitOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op->getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);

    if (targetInfo.getISAFamily() == ISAFamily::GFX1250) {
      // Clamp asyncCnt to 6bits(hw limit); lower means conservative
      unsigned asyncCnt = std::min(63u, op.getNumInst());
      ROCDL::WaitAsynccntOp::create(rewriter, loc, asyncCnt);
    } else {
      return rewriter.notifyMatchFailure(
          op, "async wait not supported on this target architecture");
    }

    // Drop the result AsyncToken
    rewriter.replaceOp(op, b.i32_val(0));
    return success();
  }

private:
  const AMD::TargetInfo &targetInfo;
};

struct AsyncTDMIntrinsicWaitConversion
    : public ConvertOpToLLVMPattern<triton::amdgpu::AsyncTDMIntrinsicWait> {
  AsyncTDMIntrinsicWaitConversion(LLVMTypeConverter &converter,
                                  PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit) {}

  LogicalResult
  matchAndRewrite(triton::amdgpu::AsyncTDMIntrinsicWait op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    ROCDL::WaitTensorcntOp::create(rewriter, loc, op.getCount());
    rewriter.eraseOp(op);
    return success();
  }
};

struct AsyncCommitGroupOpConversion
    : public ConvertOpToLLVMPattern<AsyncCommitGroupOp> {
  AsyncCommitGroupOpConversion(LLVMTypeConverter &converter,
                               const AMD::TargetInfo &targetInfo,
                               PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit), targetInfo(targetInfo) {}

  LogicalResult
  matchAndRewrite(AsyncCommitGroupOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op->getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);

    // Emit asyncmark for architectures that use async tracking
    if (targetInfo.useAsyncMarks())
      ROCDL::AsyncmarkOp::create(rewriter, loc);

    // Drop the result AsyncToken
    rewriter.replaceOp(op, b.i32_val(0));
    return success();
  }

private:
  const AMD::TargetInfo &targetInfo;
};

struct AsyncCopyMbarrierArriveOpConversion
    : public ConvertOpToLLVMPattern<triton::amdgpu::AsyncCopyMbarrierArriveOp> {
  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(triton::amdgpu::AsyncCopyMbarrierArriveOp op,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op->getLoc();
    TritonLLVMOpBuilder b(loc, rewriter);
    auto smemObj = LLVM::getSharedMemoryObjectFromStruct(
        loc, adaptor.getBarrier(),
        typeConverter->convertType(op.getBarrier().getType().getElementType()),
        rewriter);
    auto newOp = ROCDL::DsAtomicAsyncBarrierArriveOp::create(rewriter, loc, {},
                                                             smemObj.getBase());
    rewriter.replaceOp(op, newOp);
    return success();
  }
};

struct TDMPrefetchConversion
    : public ConvertOpToLLVMPattern<triton::amdgpu::TDMPrefetchOp> {
  TDMPrefetchConversion(LLVMTypeConverter &converter,
                        const AMD::TargetInfo &targetInfo,
                        PatternBenefit benefit)
      : ConvertOpToLLVMPattern(converter, benefit), targetInfo(targetInfo) {}
  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(triton::amdgpu::TDMPrefetchOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);

    auto tdescType = op.getDesc().getType();
    SmallVector<int64_t> blockShape = llvm::to_vector(tdescType.getShape());
    Type elementType =
        getTypeConverter()->convertType(tdescType.getElementType());
    SmallVector<Value> desc =
        mlir::LLVM::AMD::unpackTDMDescriptor(rewriter, loc, adaptor.getDesc());
    SmallVector<Value> offset = adaptor.getIndices();

    auto mod = op->getParentOfType<ModuleOp>();
    int threadsPerWarp = TritonGPUDialect::getThreadsPerWarp(mod);
    int numWarps = lookupNumWarps(op);
    int numCTAs = lookupNumCTAs(op);

    auto [laneId, warpId] = getLaneAndWarpId(rewriter, loc);
    auto ctaId = targetInfo.getClusterCTAId(rewriter, loc);

    auto offsets = mlir::LLVM::AMD::emitTDMPrefetch(
        rewriter, loc, desc, blockShape, threadsPerWarp, numWarps, numCTAs,
        offset, op.getPred(), elementType, laneId, warpId, ctaId,
        op.getSpeculative());

    // If the op has no results, just erase it
    if (op->getNumResults() == 0) {
      rewriter.eraseOp(op);
      return success();
    }

    // Return offsets
    Value resultStruct = packTensorElements(loc, getTypeConverter(), offsets,
                                            rewriter, op.getType(0));
    rewriter.replaceOp(op, {resultStruct});
    return success();
  }

private:
  const AMD::TargetInfo &targetInfo;
};
} // namespace

namespace mlir::triton::AMD {
void populateLoadStoreOpToLLVMPatterns(LLVMTypeConverter &typeConverter,
                                       const TargetInfo &targetInfo,
                                       RewritePatternSet &patterns,
                                       ModuleAxisInfoAnalysis &axisInfoAnalysis,
                                       PatternBenefit benefit) {
  patterns.add<AtomicCASOpConversion, AtomicRMWOpConversion, LoadOpConversion,
               StoreOpConversion, BufferLoadOpConversion,
               BufferLoadToLocalOpConversion, BufferStoreOpConversion,
               BufferAtomicRMWOpConversion, AsyncCopyGlobalToLocalOpConversion,
               AsyncCopyLocalToGlobalOpConversion, BufferAtomicCASOpConversion,
               AsyncTDMCopyGlobalToLocalOpConversion,
               AsyncTDMFusedCopyGlobalToLocalOpConversion,
               AsyncTDMCopyLocalToGlobalOpConversion,
               AsyncTDMScatterOpConversion, AsyncTDMGatherOpConversion>(
      typeConverter, targetInfo, axisInfoAnalysis, benefit);
  patterns.add<TTGAsyncWaitOpConversion>(typeConverter, targetInfo, benefit);
  patterns.add<AsyncWaitOpConversion>(typeConverter, targetInfo, benefit);
  patterns.add<TDMPrefetchConversion>(typeConverter, targetInfo, benefit);
  patterns.add<AsyncTDMIntrinsicWaitConversion>(typeConverter, benefit);
  patterns.add<AsyncCommitGroupOpConversion>(typeConverter, targetInfo,
                                             benefit);
  patterns.add<AsyncCopyMbarrierArriveOpConversion>(typeConverter, benefit);
}
} // namespace mlir::triton::AMD
