#include "mlir/Support/LLVM.h"
#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h"
#include "triton/Conversion/TritonGPUToLLVM/Utility.h"
#include "triton/Dialect/TritonGPU/IR/Attributes.h"
#include "triton/Dialect/TritonGPU/IR/Dialect.h"
#include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h"
#include "triton/Dialect/TritonGPU/IR/Types.h"
#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
#include "triton/Tools/LayoutUtils.h"

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

Value bitOrPtrCast(Value val, Type type, TritonLLVMOpBuilder &b) {
  if (isa<LLVM::LLVMPointerType>(val.getType()) &&
      !isa<LLVM::LLVMPointerType>(type)) {
    return b.ptrtoint(type, val);
  } else {
    return b.bitcast(val, type);
  }
}

struct SplatOpConversion : public ConvertOpToLLVMPattern<triton::SplatOp> {
  using ConvertOpToLLVMPattern<triton::SplatOp>::ConvertOpToLLVMPattern;
  // Convert SplatOp or arith::ConstantOp with SplatElementsAttr to a
  // LLVM::StructType value.
  //
  // @elemType: the element type in operand.
  // @resType: the return type of the Splat-like op.
  // @constVal: a LLVM::ConstantOp or other scalar value.
  static Value convertSplatLikeOp(Type elemType, Type resType, Value constVal,
                                  const LLVMTypeConverter *typeConverter,
                                  ConversionPatternRewriter &rewriter,
                                  Location loc) {
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    auto tensorTy = cast<RankedTensorType>(resType);
    // Check the converted type for the tensor as depending on the encoding the
    // converter may pick different element types.
    auto srcType = typeConverter->convertType(tensorTy);
    if (auto structTy = dyn_cast<LLVM::LLVMStructType>(srcType))
      srcType = structTy.getBody()[0];
    // If the type sizes don't match we need to pack constants.
    if (srcType.isIntOrFloat() && constVal.getType().getIntOrFloatBitWidth() !=
                                      srcType.getIntOrFloatBitWidth()) {
      unsigned cstBitWidth = constVal.getType().getIntOrFloatBitWidth();
      unsigned srcBitWidth = srcType.getIntOrFloatBitWidth();
      assert(cstBitWidth <= srcBitWidth && srcBitWidth % cstBitWidth == 0);
      unsigned ratio = srcBitWidth / cstBitWidth;
      Type intTy = IntegerType::get(elemType.getContext(), cstBitWidth);
      VectorType vecType = VectorType::get(ratio, intTy);
      Value intCst = bitOrPtrCast(constVal, intTy, b);
      Value vec = b.undef(vecType);
      for (unsigned i = 0; i < ratio; ++i)
        vec = b.insert_element(vecType, vec, intCst, b.int_val(32, i));
      constVal = vec;
    }
    Value llSrc = bitOrPtrCast(constVal, srcType, b);
    size_t elemsPerThread = getUniqueElemsPerThread(tensorTy);
    llvm::SmallVector<Value> elems(elemsPerThread, llSrc);
    return packUniqueTensorElements(loc, typeConverter, elems, rewriter,
                                    resType);
  }
  LogicalResult matchAndRewrite(triton::SplatOp op, OpAdaptor adaptor,
                                ConversionPatternRewriter &rewriter) const {
    auto loc = op->getLoc();
    auto src = adaptor.getSrc();
    auto typeConverter = getTypeConverter();
    auto llStruct = convertSplatLikeOp(src.getType(), op.getType(), src,
                                       typeConverter, rewriter, loc);
    rewriter.replaceOp(op, {llStruct});
    return success();
  }
};

struct UnsplatOpConversion : public ConvertOpToLLVMPattern<triton::UnsplatOp> {
  using ConvertOpToLLVMPattern<triton::UnsplatOp>::ConvertOpToLLVMPattern;
  LogicalResult matchAndRewrite(triton::UnsplatOp op, OpAdaptor adaptor,
                                ConversionPatternRewriter &rewriter) const {
    auto loc = op->getLoc();
    auto scrVals = unpackUniqueTensorElements(loc, adaptor.getSrc(), rewriter);
    rewriter.replaceOp(op, scrVals[0]);
    return success();
  }
};

// This pattern helps to convert arith::ConstantOp(with SplatElementsAttr),
// the logic is the same as triton::SplatOp, so the underlying implementation
// is reused.
struct ArithConstantSplatOpConversion
    : public ConvertOpToLLVMPattern<arith::ConstantOp> {
  using ConvertOpToLLVMPattern<arith::ConstantOp>::ConvertOpToLLVMPattern;
  LogicalResult
  matchAndRewrite(arith::ConstantOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto value = op.getValue();
    auto values = dyn_cast<SplatElementsAttr>(op.getValue());
    if (!values)
      return failure();
    auto loc = op->getLoc();
    LLVM::ConstantOp arithConstantOp;
    auto elemType = values.getElementType();
    Attribute val;
    if (type::isFloat(elemType)) {
      val = values.getValues<FloatAttr>()[0];
    } else if (type::isInt(elemType)) {
      val = values.getValues<IntegerAttr>()[0];
    } else {
      llvm::errs() << "ArithConstantSplatOpConversion get unsupported type: "
                   << value.getType() << "\n";
      return failure();
    }
    // Lower FP8 constant to int8 constant since FP8 types are not supported on
    // LLVM IR.
    if (type::isFloat8(elemType))
      elemType = rewriter.getIntegerType(8);
    auto constOp = LLVM::ConstantOp::create(rewriter, loc, elemType, val);
    auto typeConverter = getTypeConverter();
    auto llStruct = SplatOpConversion::convertSplatLikeOp(
        elemType, op.getType(), constOp, typeConverter, rewriter, loc);
    rewriter.replaceOp(op, llStruct);
    return success();
  }
};

struct CatOpConversion : public ConvertOpToLLVMPattern<CatOp> {
  using OpAdaptor = typename CatOp::Adaptor;
  explicit CatOpConversion(LLVMTypeConverter &typeConverter,
                           PatternBenefit benefit = patternBenefitDefault)
      : ConvertOpToLLVMPattern<CatOp>(typeConverter, benefit) {}
  LogicalResult
  matchAndRewrite(CatOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op->getLoc();
    auto resultTy = cast<RankedTensorType>(op.getType());
    // Unpack input values.
    auto lhsVals = unpackUniqueTensorElements(loc, adaptor.getLhs(), rewriter);
    auto rhsVals = unpackUniqueTensorElements(loc, adaptor.getRhs(), rewriter);

    // concatenate (and potentially reorder) values
    SmallVector<Value> retVals;
    for (Value v : lhsVals)
      retVals.push_back(v);
    for (Value v : rhsVals)
      retVals.push_back(v);

    assert(retVals.size() == getUniqueElemsPerThread(resultTy));

    // pack and replace
    Value ret = packUniqueTensorElements(loc, getTypeConverter(), retVals,
                                         rewriter, resultTy);
    rewriter.replaceOp(op, ret);
    return success();
  }
};
struct JoinOpConversion : public ConvertOpToLLVMPattern<JoinOp> {
  using OpAdaptor = typename JoinOp::Adaptor;
  explicit JoinOpConversion(LLVMTypeConverter &typeConverter,
                            PatternBenefit benefit = patternBenefitDefault)
      : ConvertOpToLLVMPattern<JoinOp>(typeConverter, benefit) {}
  LogicalResult
  matchAndRewrite(JoinOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // We rely on the following invariants of this op (which are checked by its
    // verifier):
    //
    // - The last dimension (the one we're joining) is also the most minor
    //   dimension.
    // - The input and output encodings are the same, except the output has
    //   2 elements per thread in the last dim.
    //
    // With these invariants, join is trivial: We can count how many contiguous
    // registers belong to the same chunk then we merge the registers between
    // two different chunks.
    Location loc = op->getLoc();
    RankedTensorType dstTy = op.getType();
    auto ll = toLinearLayout(dstTy);
    int splitDim = dstTy.getRank() - 1;
    auto kReg = mlir::StringAttr::get(dstTy.getContext(), "register");
    ll = ll.removeZeroBasesAlongDim(kReg);
    const auto &bases = ll.getBases();
    const auto &regs = bases.find(kReg)->second;
    int numContiguousValues = 1;
    bool found = false;
    for (const auto &reg : regs) {
      if (reg[splitDim] == 1) {
        found = true;
        break;
      }
      numContiguousValues *= 2;
    }
    assert(found && "Join dimension is not distributed along registers.");
    SmallVector<Value> lhsVals =
        unpackUniqueTensorElements(loc, adaptor.getLhs(), rewriter);
    SmallVector<Value> rhsVals =
        unpackUniqueTensorElements(loc, adaptor.getRhs(), rewriter);
    assert(lhsVals.size() == rhsVals.size());
    SmallVector<Value> joinedVals;
    joinedVals.resize(lhsVals.size() * 2);
    for (int i = 0; i < lhsVals.size(); i += numContiguousValues) {
      for (int j = 0; j < numContiguousValues; j++) {
        joinedVals[2 * i + j] = lhsVals[i + j];
        joinedVals[2 * i + numContiguousValues + j] = rhsVals[i + j];
      }
    }
    Value ret = packUniqueTensorElements(loc, getTypeConverter(), joinedVals,
                                         rewriter, dstTy);
    rewriter.replaceOp(op, ret);
    return success();
  }
};
struct SplitOpConversion : public ConvertOpToLLVMPattern<SplitOp> {
  using OpAdaptor = typename SplitOp::Adaptor;
  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
  LogicalResult
  matchAndRewrite(SplitOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // We rely on the following invariants of this op (which are checked by its
    // verifier):
    //
    // - The layout distribute the last dimension along registers
    // - The last dimension (the one we're splitting) has sizePerThread=2,
    // threadPerWarp=1 and warpPerBlock=1.
    //
    // With these invariants, split is trivial: We can count how many contiguous
    // registers belong to the same chunk then we separate the registers between
    // two different chunks.
    auto srcTy = cast<RankedTensorType>(op.getSrc().getType());
    auto ll = toLinearLayout(srcTy);
    int splitDim = srcTy.getRank() - 1;
    auto kReg = mlir::StringAttr::get(srcTy.getContext(), "register");
    ll = ll.removeZeroBasesAlongDim(kReg);
    const auto &bases = ll.getBases();
    const auto &regs = bases.find(kReg)->second;
    int numContiguousValues = 1;
    bool found = false;
    for (const auto &reg : regs) {
      if (reg[splitDim] == 1) {
        found = true;
        break;
      }
      numContiguousValues *= 2;
    }
    assert(found && "Split dimension is not distributed along registers.");
    Location loc = op->getLoc();
    SmallVector<Value> srcVals =
        unpackUniqueTensorElements(loc, adaptor.getSrc(), rewriter);
    assert(srcVals.size() % 2 == 0);
    SmallVector<Value> outLhsVals;
    SmallVector<Value> outRhsVals;
    for (int i = 0; i < srcVals.size(); i += 2 * numContiguousValues) {
      for (int j = 0; j < numContiguousValues; j++) {
        outLhsVals.push_back(srcVals[i + j]);
        outRhsVals.push_back(srcVals[i + numContiguousValues + j]);
      }
    }
    auto resultTy = op.getResult(0).getType();
    Value retLhs = packUniqueTensorElements(loc, getTypeConverter(), outLhsVals,
                                            rewriter, resultTy);
    Value retRhs = packUniqueTensorElements(loc, getTypeConverter(), outRhsVals,
                                            rewriter, resultTy);
    rewriter.replaceOp(op, {retLhs, retRhs});
    return success();
  }
};
struct ReshapeOpConversion : public ConvertOpToLLVMPattern<ReshapeOp> {
  using OpAdaptor = typename ReshapeOp::Adaptor;
  explicit ReshapeOpConversion(LLVMTypeConverter &typeConverter,
                               PatternBenefit benefit = patternBenefitDefault)
      : ConvertOpToLLVMPattern<ReshapeOp>(typeConverter, benefit) {}
  LogicalResult
  matchAndRewrite(ReshapeOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    assert(!isExpensiveView(op.getSrc().getType(), op.getType()));
    rewriter.replaceOp(op, adaptor.getSrc());
    return success();
  }
};
struct ExpandDimsOpConversion : public ConvertOpToLLVMPattern<ExpandDimsOp> {
  using OpAdaptor = typename ExpandDimsOp::Adaptor;
  explicit ExpandDimsOpConversion(
      LLVMTypeConverter &typeConverter,
      PatternBenefit benefit = patternBenefitDefault)
      : ConvertOpToLLVMPattern<ExpandDimsOp>(typeConverter, benefit) {}
  LogicalResult
  matchAndRewrite(ExpandDimsOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    rewriter.replaceOp(op, adaptor.getSrc());
    return success();
  }
};
struct MemDescTransOpConversion
    : public ConvertOpToLLVMPattern<MemDescTransOp> {
  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
  LogicalResult
  matchAndRewrite(MemDescTransOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op->getLoc();
    auto resultTy = cast<TensorOrMemDesc>(op.getType());
    auto llvmElemTy =
        getTypeConverter()->convertType(resultTy.getElementType());
    auto srcSmemObj = getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(),
                                                      llvmElemTy, rewriter);
    auto dstSmemObj = SharedMemoryObject(
        srcSmemObj.getBases(), srcSmemObj.getBaseElemType(),
        /*offsets=*/applyPermutation(srcSmemObj.getOffsets(), op.getOrder()));
    auto retVal = getStructFromSharedMemoryObject(loc, dstSmemObj, rewriter);
    rewriter.replaceOp(op, retVal);
    return success();
  }
};

struct MemDescReshapeOpConversion
    : public ConvertOpToLLVMPattern<MemDescReshapeOp> {
  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
  LogicalResult
  matchAndRewrite(MemDescReshapeOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op->getLoc();
    auto resultTy = cast<TensorOrMemDesc>(op.getType());
    auto llvmElemTy =
        getTypeConverter()->convertType(resultTy.getElementType());
    auto srcSmemObj = getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(),
                                                      llvmElemTy, rewriter);
    SmallVector<Value> offsets = srcSmemObj.getOffsets();
    // FIXME: This should be done by composing a linear layout with its
    // reshaped counterpart.
    SmallVector<unsigned> srcShape;
    for (int64_t d : op.getSrc().getType().getShape())
      srcShape.push_back(d);
    SmallVector<unsigned> dstShape;
    for (int64_t d : op.getType().getShape())
      dstShape.push_back(d);
    Value linearOffset = LLVM::linearize(rewriter, loc, offsets, srcShape);
    SmallVector<Value> delinearizedOffset =
        LLVM::delinearize(rewriter, loc, linearOffset, dstShape);
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    auto dstSmemObj =
        SharedMemoryObject(srcSmemObj.getBases(), srcSmemObj.getBaseElemType(),
                           delinearizedOffset);
    auto retVal = getStructFromSharedMemoryObject(loc, dstSmemObj, rewriter);
    rewriter.replaceOp(op, retVal);
    return success();
  }
};

struct TransOpConversion : public ConvertOpToLLVMPattern<TransOp> {
  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
  LogicalResult
  matchAndRewrite(TransOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // By construction, TransOp::inferReturnTypes ensures that the src encoding
    // is the same as the dst encoding so that this op is a no-op.
    rewriter.replaceOp(op, adaptor.getSrc());
    return success();
  }
};

struct BroadcastOpConversion
    : public ConvertOpToLLVMPattern<triton::BroadcastOp> {
  using ConvertOpToLLVMPattern<triton::BroadcastOp>::ConvertOpToLLVMPattern;
  LogicalResult
  matchAndRewrite(triton::BroadcastOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Following the order of indices in the legacy code, a broadcast of:
    //   [s(0), s(1) ... s(k-1),    1, s(k+1), s(k+2) ... s(n-1)]
    // =>
    //   [s(0), s(1) ... s(k-1), s(k), s(k+1), s(k+2) ... s(n-1)]
    //
    // logically maps to a broadcast within a thread's scope:
    //   [cta(0)..cta(k-1),     1,cta(k+1)..cta(n-1),spt(0)..spt(k-1),
    //   1,spt(k+1)..spt(n-1)]
    // =>
    //   [cta(0)..cta(k-1),cta(k),cta(k+1)..cta(n-1),spt(0)..spt(k-1),spt(k),spt(k+1)..spt(n-1)]
    //
    // regardless of the order of the layout
    //
    Location loc = op->getLoc();
    Value src = adaptor.getSrc();
    Value result = op.getResult();
    auto srcTy = cast<RankedTensorType>(op.getSrc().getType());
    auto resultTy = cast<RankedTensorType>(result.getType());
    auto srcLayout = srcTy.getEncoding();
    auto resultLayout = resultTy.getEncoding();
    auto srcShape = srcTy.getShape();
    unsigned rank = srcTy.getRank();
    auto typeConverter = getTypeConverter();
    assert(rank == resultTy.getRank());
    auto srcOffsets = emitOffsetForLayout(srcLayout, srcTy);
    auto resultOffsets = emitOffsetForLayout(resultLayout, resultTy);
    SmallVector<Value> srcVals =
        unpackTensorElements(loc, src, rewriter, srcTy);
    std::map<SmallVector<unsigned>, Value> srcValues;
    for (size_t i = 0; i < srcOffsets.size(); i++) {
      srcValues[srcOffsets[i]] = srcVals[i];
    }
    SmallVector<Value> resultVals;
    for (size_t i = 0; i < resultOffsets.size(); i++) {
      auto offset = resultOffsets[i];
      for (size_t j = 0; j < srcShape.size(); j++)
        if (srcShape[j] == 1)
          offset[j] = 0;
      resultVals.push_back(srcValues.at(offset));
    }
    Value resultStruct =
        packTensorElements(loc, typeConverter, resultVals, rewriter, resultTy);
    rewriter.replaceOp(op, {resultStruct});
    return success();
  }
};

struct MemDescIndexOpConversion
    : public ConvertOpToLLVMPattern<triton::gpu::MemDescIndexOp> {
  using ConvertOpToLLVMPattern<
      triton::gpu::MemDescIndexOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(triton::gpu::MemDescIndexOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op->getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    auto srcTy = op.getSrc().getType();
    auto dstTy = op.getResult().getType();
    auto llvmElemTy = getTypeConverter()->convertType(srcTy.getElementType());

    // Stride is computed from dstTy (the result after dropping the leading
    // buffer dimension). The encoding's partitionDim is relative to this
    // reduced shape, so partitionDim=0 refers to the tensor's first dimension,
    // not the multi-buffer dimension that was just indexed away.
    //
    // getAllocationElems returns the correct number of fp4 elements that
    // we need to skip when we have fp4Padded=True. getShapePerCTA does not
    // account for this.
    int64_t stride =
        getAllocationElems(dstTy.getEncoding(), dstTy.getAllocShape());
    if (auto partEnc =
            dyn_cast<PartitionedSharedEncodingAttr>(dstTy.getEncoding())) {
      stride /= partEnc.getNumPartitions();
    }
    Value offset = b.mul(op.getIndex(), b.i32_val(stride));
    auto smemObj = getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(),
                                                   llvmElemTy, rewriter);
    auto prevOffsets = smemObj.getOffsets();
    SmallVector<Value> offsetVals(prevOffsets.end() - dstTy.getRank(),
                                  prevOffsets.end());

    // Apply padding based on the amount we move the base ptr
    if (auto padEnc = getPaddedEncoding(dstTy.getEncoding())) {
      auto bitwidth = dstTy.getElementTypeBitWidth();
      auto paddingShifts =
          getPaddedSharedShifts(padEnc, bitwidth, /*offsetInBytes=*/false);
      offset = applyPadding(loc, rewriter, offset, paddingShifts);
    }

    // For partitioned tensors, all partitions have the same layout structure.
    // Each partition contains portions of every "buffer" along dimension 0.
    // When we select buffer i, we advance all partition bases by the same
    // offset.
    SmallVector<Value> newBases;
    for (Value base : smemObj.getBases()) {
      auto elemPtrTy = base.getType();
      newBases.push_back(b.gep(elemPtrTy, llvmElemTy, base, offset));
    }

    // Create new smemObj with advanced base pointers
    auto newSmemObj = SharedMemoryObject(newBases, llvmElemTy, offsetVals);
    auto retVal = getStructFromSharedMemoryObject(loc, newSmemObj, rewriter);
    rewriter.replaceOp(op, retVal);
    return success();
  }
};

struct MemDescSubsliceOpConversion
    : public ConvertOpToLLVMPattern<triton::gpu::MemDescSubsliceOp> {
  using ConvertOpToLLVMPattern<
      triton::gpu::MemDescSubsliceOp>::ConvertOpToLLVMPattern;

  LogicalResult
  matchAndRewrite(triton::gpu::MemDescSubsliceOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op->getLoc();
    auto b = TritonLLVMOpBuilder(loc, rewriter);
    auto *ctx = op.getContext();
    auto srcTy = op.getSrc().getType();
    auto llvmElemTy = getTypeConverter()->convertType(srcTy.getElementType());

    auto smemObj = getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(),
                                                   llvmElemTy, rewriter);
    auto opOffsetVals = op.getOffsets();
    auto encoding = srcTy.getEncoding();
    auto layoutOffsets = dropPipeliningDim(opOffsetVals, encoding);
    SmallVector<Value> newBases = llvm::to_vector(smemObj.getBases());

    if (layoutOffsets.size() != opOffsetVals.size() &&
        opOffsetVals.front() != 0) {
      int64_t stride = getAllocationElems(
          encoding, dropPipeliningDim(srcTy.getAllocShape(), encoding));
      if (auto partEnc = dyn_cast<PartitionedSharedEncodingAttr>(encoding))
        stride /= partEnc.getNumPartitions();
      Value offset = b.i32_val(opOffsetVals.front() * stride);
      for (Value &base : newBases)
        base = b.gep(base.getType(), llvmElemTy, base, offset);
    }

    // Accumulate the logical offsets
    SmallVector<Value> offsetVals;
    for (auto [oldOffVal, opOff] :
         llvm::zip(smemObj.getOffsets(), opOffsetVals)) {
      offsetVals.push_back(b.add(oldOffVal, b.i32_val(opOff)));
    }

    // For PartitionedSharedEncoding we need to pick the right base at load
    // time. Let
    //   o     = this op's static subslice offsets (one per dim),
    //   c     = a logical base indices to the current subslice op,
    //   L     = the shared LL (inputs: offset, partition, block;
    //                          outputs: dim0, dim1, ...).
    //
    //   (1) c_src = c + o
    //   (2) verifier makes o's bits disjoint from c's bits per dim, so:
    //         c + o = c ^ o                                   (bit-disjoint)
    //   (3) L^-1 is linear, so projecting (2) onto the partition component:
    //         partition(L^-1(c_src)) = partition(L^-1(c)) ^ S
    //       where  S = partition(L^-1(o))
    //   (4) the existing lowering already computes i = partition(L^-1(c))
    //       and indexes bases[i]; from (3) the correct base is:
    //         bases[i ^ S]
    //   (5) precompute that rotation once here
    //         newBases[i] = oldBases[i ^ S]
    //
    // The offset component of (3) is already XORed in by getShmemOffset at
    // load time; only the partition component needs this fix.
    if (newBases.size() > 1) {
      LinearLayout ll = triton::gpu::toLinearLayoutIgnoringPadding(srcTy);
      auto kPartition = StringAttr::get(ctx, "partition");
      assert(ll.hasInDim(kPartition) &&
             "multiple bases require a partition input dim");
      auto dimNames = standardOutDimNames(ctx, layoutOffsets.size());
      SmallVector<std::pair<StringAttr, int32_t>> namedOffsets;
      for (auto [dim, off] : llvm::zip(dimNames, layoutOffsets))
        namedOffsets.push_back({dim, off});
      auto partitionLayout = ll.invert().sublayout(dimNames, {kPartition});
      int32_t partitionShift = partitionLayout.apply(namedOffsets)[0].second;
      SmallVector<Value> rotated(newBases.size());
      for (size_t i = 0; i < newBases.size(); ++i)
        rotated[i] = newBases[i ^ partitionShift];
      newBases = std::move(rotated);
    }

    smemObj = SharedMemoryObject(newBases, llvmElemTy, offsetVals);
    auto retVal = getStructFromSharedMemoryObject(loc, smemObj, rewriter);
    rewriter.replaceOp(op, retVal);
    return success();
  }
};

struct MemDescReinterpretOpConversion
    : public ConvertOpToLLVMPattern<MemDescReinterpretOp> {
  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;

  LogicalResult matchAndRewrite(MemDescReinterpretOp op, OpAdaptor adaptor,
                                ConversionPatternRewriter &b) const override {
    Location loc = op.getLoc();
    MemDescType srcTy = op.getSrc().getType();
    MemDescType dstTy = op.getType();
    Type srcElemTy = getTypeConverter()->convertType(srcTy.getElementType());
    Type dstElemTy = getTypeConverter()->convertType(dstTy.getElementType());

    // PartitionedSharedEncoding is not yet supported for memdesc_reinterpret
    if (isa<PartitionedSharedEncodingAttr>(srcTy.getEncoding())) {
      return b.notifyMatchFailure(
          op,
          "PartitionedSharedEncoding not yet supported in memdesc_reinterpret");
    }

    auto smemObj =
        getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(), srcElemTy, b);
    Value newBase = smemObj.getShmemAffineBase(loc, b, srcTy);
    SharedMemoryObject newObj(newBase, dstElemTy, dstTy.getRank(), loc, b);
    b.replaceOp(op, getStructFromSharedMemoryObject(loc, newObj, b));
    return success();
  }
};

} // namespace

void mlir::triton::populateViewOpToLLVMPatterns(
    LLVMTypeConverter &typeConverter, RewritePatternSet &patterns,
    PatternBenefit benefit) {
  patterns.add<ReshapeOpConversion>(typeConverter, benefit);
  patterns.add<ExpandDimsOpConversion>(typeConverter, benefit);
  patterns.add<SplatOpConversion>(typeConverter, benefit);
  patterns.add<UnsplatOpConversion>(typeConverter, benefit);
  patterns.add<ArithConstantSplatOpConversion>(typeConverter, benefit);
  patterns.add<CatOpConversion>(typeConverter, benefit);
  patterns.add<JoinOpConversion>(typeConverter, benefit);
  patterns.add<SplitOpConversion>(typeConverter, benefit);
  patterns.add<MemDescTransOpConversion, MemDescReshapeOpConversion>(
      typeConverter, benefit);
  patterns.add<TransOpConversion>(typeConverter, benefit);
  patterns.add<BroadcastOpConversion>(typeConverter, benefit);
  patterns.add<MemDescSubsliceOpConversion, MemDescIndexOpConversion>(
      typeConverter, benefit);
  patterns.add<MemDescReinterpretOpConversion>(typeConverter, benefit);
}
