#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/IR/Types.h"
#include "mlir/IR/Value.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "triton/Analysis/Utility.h"
#include "triton/Conversion/MLIRTypes.h"
#include "triton/Dialect/Triton/IR/Dialect.h"
#include "triton/Dialect/Triton/IR/OpInterfaces.h"
#include "triton/Dialect/Triton/IR/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/Transforms/DecomposeScaledBlocked.h"
#include "triton/Dialect/TritonGPU/Transforms/Passes.h"
#include "triton/Dialect/TritonGPU/Transforms/Utility.h"
#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
#include "triton/Tools/LayoutUtils.h"
#include "triton/Tools/StrUtil.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"

namespace mlir {
namespace triton {
namespace gpu {

namespace {

static bool isUnsupportedMMAv5Int8Dot(int computeCapability, DotOp op) {
  if (nvidia_gpu::TargetFeatures(computeCapability).supportsI8Tcgen05MMA())
    return false;
  auto aElemTy = op.getA().getType().getElementType();
  auto bElemTy = op.getB().getType().getElementType();
  return aElemTy.isInteger(8) && bElemTy.isInteger(8);
}

// Get the highest version supported for the hardware and the dot.
static int getMMAVersionSafe(int computeCapability, DotOp op) {
  // List supported mma version in order of preference.
  SmallVector<int> versionsSupported;
  if (computeCapability < 75) {
    versionsSupported = {1};
  } else if (computeCapability < 90) {
    versionsSupported = {2};
  } else if (computeCapability < 100) {
    versionsSupported = {3, 2};
  } else if (computeCapability < 120) {
    // Exclude consumer Blackwell (sm120)
    if (isUnsupportedMMAv5Int8Dot(computeCapability, op)) {
      versionsSupported = {2};
    } else {
      versionsSupported = {5, 2};
    }
  } else if (computeCapability < 130) {
    versionsSupported = {2};
  } else {
    assert(false && "computeCapability not supported");
  }
  for (int baseVersion : versionsSupported) {
    if (supportMMA(op, baseVersion))
      return baseVersion;
    if (baseVersion == 3) {
      auto remark = op.emitRemark()
                    << "MMA version 3 acceleration not applied due to "
                       "unsupported shapes or data types.";
      remark.attachNote() << "Target compute capability (" << computeCapability
                          << ") supports MMA v3.";
    }

    if (baseVersion == 5) {
      auto remark = op.emitRemark()
                    << "MMA version 5 acceleration not applied due to "
                       "unsupported shapes or data types.";
      remark.attachNote() << "Target compute capability (" << computeCapability
                          << ") supports MMA v5.";
    }
  }
  return 0;
}

SmallVector<unsigned> warpsPerTileV2(DotOpInterface dotOp,
                                     const ArrayRef<int64_t> shape,
                                     int numWarps) {
  auto rank = shape.size();
  // Early exit for batched matmul
  if (rank == 3)
    return getMmaV2WarpsPerCTA(shape, numWarps);

  auto filter = [&dotOp](Operation *op) {
    return op->getParentRegion() == dotOp->getParentRegion() &&
           !isa<TransOp>(op);
  };
  auto slices = mlir::getSlice(dotOp, {filter}, {filter});
  bool hasChainedDot = false;
  for (Operation *op : slices) {
    if (isa<DotOp, DotScaledOp>(op) && (op != dotOp)) {
      auto resTy = cast<RankedTensorType>(op->getResult(0).getType());
      if (resTy.getRank() != rank) {
        continue;
      }
      if (auto mmaEncoding =
              dyn_cast<NvidiaMmaEncodingAttr>(resTy.getEncoding())) {
        return to_vector(mmaEncoding.getWarpsPerCTA());
      }
      hasChainedDot = true;
    }
  }
  if (hasChainedDot) {
    if (shape[0] >= shape[1]) {
      return {(unsigned)numWarps, 1};
    } else {
      return {1, (unsigned)numWarps};
    }
  }

  return getMmaV2WarpsPerCTA(shape, numWarps);
}
SmallVector<unsigned, 2>
warpsPerTileV3(DotOpInterface dotOp, const ArrayRef<int64_t> shape,
               int numWarps, const SmallVector<unsigned, 3> &instrShape) {
  SetVector<Operation *> slices;
  mlir::getForwardSlice(dotOp.getD(), &slices);
  // Contains a chained dot. We prefer to assign warps to one axis
  // to facilitate use cases like flash attention, allowing reductions within
  // the same warp.
  if (llvm::find_if(slices, [](Operation *op) {
        return isa<mlir::triton::DotOpInterface>(op);
      }) != slices.end())
    return {(unsigned)numWarps, 1};

  // For MMAv3, the smallest indivisible unit of warp shape is (4, 1).
  SmallVector<unsigned, 2> ret = {4, 1};
  SmallVector<int64_t, 2> shapePerWarp = {16, instrShape[1]};
  do {
    if (ret[0] * ret[1] >= numWarps)
      break;
    if (shape[0] > shapePerWarp[0] * ret[0]) {
      ret[0] *= 2;
    } else {
      ret[1] *= 2;
    }
  } while (true);
  return ret;
}

// Returns a shared memory allocation that can be used by a dotMMA op for the
// given value.
static Value
getSharedMemoryMMAOperand(Value v, mlir::PatternRewriter &rewriter, int opIdx,
                          bool allowTranspose, bool isMMAv5Fp4Padded = false,
                          bool forceTranspose = false,
                          Operation *op = nullptr /*only for diagnostic*/) {
  OpBuilder::InsertionGuard g(rewriter);
  Value arg = v;
  while (auto cvtOp = arg.getDefiningOp<ConvertLayoutOp>())
    arg = cvtOp.getSrc();
  auto argType = cast<RankedTensorType>(arg.getType());
  assert(argType.getEncoding() && "unexpected tensor type");
  auto order = getOrderForMemory(argType);

  // If the MMA op doesn't support transpose pick the layout expected by the MMA
  // op.
  llvm::SmallVector<unsigned> newOrder = order;
  if (!allowTranspose) {
    if (opIdx == 1) {
      newOrder = {0, 1};
    } else {
      newOrder = {1, 0};
    }
    if (forceTranspose)
      std::swap(newOrder[0], newOrder[1]);
  }

  if (newOrder != order && op) {
    op->emitWarning("Warning: Forcing a different order [")
        << newOrder[0] << ", " << newOrder[1]
        << "] on SMEM than the register order for the operand " << opIdx
        << ". Registers will be transposed before SMEM store and the pipelined "
           "load for this operand will be disabled, so poor performance is "
           "expected. Recommendation: consider transposing the operand in "
           "global "
           "memory to remove the need to transpose the tensor in registers.";
  }

  Attribute SharedMemorySpace =
      SharedMemorySpaceAttr::get(argType.getContext());
  auto CGALayout = getCGALayout(argType.getEncoding());
  auto newLayout = NVMMASharedEncodingAttr::get(
      argType.getContext(), argType.getShape(), newOrder, CGALayout,
      argType.getElementType(), isMMAv5Fp4Padded);
  auto newType = MemDescType::get(argType.getShape(), argType.getElementType(),
                                  newLayout, SharedMemorySpace);
  rewriter.setInsertionPointAfterValue(arg);
  return LocalAllocOp::create(rewriter, arg.getLoc(), newType, arg);
}

static bool isTmemCopyCompatibleScale(RankedTensorType argType,
                                      bool usesTMAload) {
  assert(argType.getEncoding() && "unexpected tensor type");
  {
    auto shmemLayout = NVMMASharedEncodingAttr::get(
        argType.getContext(), /*swizzlingByteWidth=*/0,
        /*transposed=*/false,
        /*elementBitWidth=*/8,
        /*fp4Padded=*/false, getCGALayout(argType.getEncoding()));
    auto type = MemDescType::get(
        argType.getShape(), argType.getElementType(), shmemLayout,
        SharedMemorySpaceAttr::get(argType.getContext()));
    if (!isInnermostContiguous(type, 512)) {
      // TMEM copy expects metadata "chunks" in SMEM to be stored contiguously
      // in the innermost axes.
      return false;
    }
  }

  if (usesTMAload)
    return true;

  if (argType.getRank() != 2) {
    // TODO: Add support for higher rank when 5D coalesced load is fixed.
    return false;
  }

  auto innermostBits = argType.getDimSize(argType.getRank() - 1) * 8;
  return innermostBits % (32 * 128) == 0;
}

template <typename OpTy>
static OpTy getDefiningOpSkippingConvertLayout(Value value) {
  while (auto cvtOp = value.getDefiningOp<ConvertLayoutOp>())
    value = cvtOp.getSrc();
  return value.getDefiningOp<OpTy>();
}

static Value stripConvertLayout(Value value) {
  while (auto cvtOp = value.getDefiningOp<ConvertLayoutOp>())
    value = cvtOp.getSrc();
  return value;
}

static Value
tryCreateTmemCopyCompatibleScaleOperand(Value scale,
                                        mlir::PatternRewriter &rewriter) {
  OpBuilder::InsertionGuard g(rewriter);
  // The final reshape to conform to the logical shape requirement of
  // dot_scaled.
  auto reshape2D = getDefiningOpSkippingConvertLayout<ReshapeOp>(scale);
  if (!reshape2D)
    return Value();

  auto scaleType = dyn_cast<RankedTensorType>(reshape2D.getResult().getType());
  if (!scaleType || scaleType.getShape().size() != 2)
    return Value();

  auto scale2DShape = scaleType.getShape();
  auto blockMN = scale2DShape[0];
  auto numScales = scale2DShape[1];
  if (blockMN % 128 != 0 || numScales % 4 != 0)
    return Value();

  // Permute to expose the 32x4x4 interleaving used by scale blocks in TMEM.
  auto transOp =
      getDefiningOpSkippingConvertLayout<TransOp>(reshape2D.getSrc());
  if (!transOp || transOp.getOrder() != ArrayRef<int32_t>({0, 3, 2, 1, 4}))
    return Value();

  // A logical scale tensor has shape [block_mn, num_scales]. The tensor view
  // chain in a user kernel signals TMEMCopy compatibility by exposing the
  // underlying 128x4 scale block as
  //     --reshape
  //   [mn_copy_tiles, k_copy_tiles, 32, 4, 4]
  //     --trans {0, 3, 2, 1, 4}-->
  //   [mn_copy_tiles, 4, 32, k_copy_tiles, 4]
  //     --flatten-->
  //   [block_mn, num_scales]
  //
  // The block shape with which the scale tensor is loaded is flexible:
  // For example, [.., 2, 256], [.., 512], and [.., 32, 4, 4] are all acceptable
  // as long as a reshape exposes the same 32x4x4 blocks before the transpose.
  auto mnCopyTiles = blockMN / 128;
  auto kCopyTiles = numScales / 4;
  SmallVector<int64_t> tiledScaleShape = {mnCopyTiles, kCopyTiles, 32, 4, 4};

  Value tiledScale = stripConvertLayout(transOp.getSrc());
  auto reshapeTiled = tiledScale.getDefiningOp<ReshapeOp>();
  if (!reshapeTiled ||
      reshapeTiled.getType().getShape() != ArrayRef<int64_t>(tiledScaleShape))
    return Value();

  Value packedScale = stripConvertLayout(reshapeTiled.getSrc());
  auto packedScaleType = dyn_cast<RankedTensorType>(packedScale.getType());
  if (!packedScaleType)
    return Value();

  auto *loadOp = packedScale.getDefiningOp();
  if (!isa_and_nonnull<LoadOp, DescriptorLoadLikeOpInterface>(loadOp))
    return Value();

  bool usesTMAload = isa<DescriptorLoadLikeOpInterface>(loadOp);
  if (!isTmemCopyCompatibleScale(packedScaleType, usesTMAload))
    return Value();

  auto sharedMemory = SharedMemorySpaceAttr::get(scaleType.getContext());
  auto sinkLayout = getScaleSmemLayoutForTMEMCopy(
      scaleType.getContext(), scaleType.getShape(),
      getCGALayout(scaleType.getEncoding()));
  auto layout = SharedLinearEncodingAttr::get(scaleType.getContext(),
                                              std::move(sinkLayout),
                                              /*alignment=*/128);
  auto sinkType =
      MemDescType::get(scaleType.getShape(), scaleType.getElementType(), layout,
                       sharedMemory, /*mutableMemory=*/false);
  rewriter.setInsertionPointAfterValue(reshape2D.getResult());
  return LocalAllocOp::create(rewriter, reshape2D.getLoc(), sinkType,
                              reshape2D.getResult());
}

static Value createTmemScaleOperand(Value tensor, Attribute tmemEncoding,
                                    Attribute tensorMemorySpace, int numWarps,
                                    Location loc, PatternRewriter &rewriter) {
  auto tensorType = cast<RankedTensorType>(tensor.getType());
  auto tmemType =
      MemDescType::get(tensorType.getShape(), tensorType.getElementType(),
                       tmemEncoding, tensorMemorySpace,
                       /*mutableMemory=*/false);
  auto tmemLayout = nvidia_gpu::getDefaultLayoutForTmemLdSt(tmemType, numWarps);
  auto stagedType = tensorType.cloneWithEncoding(tmemLayout);
  Value converted = ConvertLayoutOp::create(rewriter, loc, stagedType, tensor);
  return triton::nvidia_gpu::TMEMAllocOp::create(rewriter, loc, tmemType,
                                                 /*token=*/Type(), converted);
}

SmallVector<unsigned, 3>
getWarpsPerTile(DotOpInterface dotOp, const ArrayRef<int64_t> shape,
                int version, int numWarps,
                const SmallVector<unsigned, 3> &instrShape) {
  switch (version) {
  case 2:
    return warpsPerTileV2(dotOp, shape, numWarps);
  case 3:
    return warpsPerTileV3(dotOp, shape, numWarps, instrShape);
  default:
    assert(false && "not supported version");
    return {0, 0};
  }
}

static bool bwdFilter(Operation *op) {
  return (op->hasTrait<OpTrait::Elementwise>() && isMemoryEffectFree(op)) ||
         isView(op) ||
         isa<Fp4ToFpOp, LoadOp, DescriptorLoadLikeOpInterface, BroadcastOp,
             ConvertLayoutOp>(op);
}

static bool joinPacksLastDim(JoinOp join) {
  auto type = cast<RankedTensorType>(join.getResult().getType());
  unsigned rank = type.getRank();
  if (rank < 2)
    return false;

  auto order = getOrder(type);
  // Check that the appended join axis and the previous last axis are the two
  // fastest-changing axes, so the join packs the last dimension contiguously.
  return order[0] == rank - 1 && order[1] == rank - 2;
}

static bool canWidenKWidthForJoin(int loadBitWidth,
                                  const SetVector<Operation *> &slice) {
  return llvm::any_of(slice, [loadBitWidth](Operation *op) {
    auto join = dyn_cast<JoinOp>(op);
    if (!join)
      return false;
    auto type = cast<RankedTensorType>(join.getResult().getType());
    return type.getElementTypeBitWidth() == loadBitWidth * 2 &&
           joinPacksLastDim(join);
  });
}

// Finds the bitwidth with which the value x is loaded
static int computeOrigBitWidth(Value x) {
  SetVector<Operation *> slice;
  mlir::BackwardSliceOptions opt;
  opt.omitBlockArguments = true;
  opt.filter = bwdFilter;
  (void)getBackwardSlice(x, &slice, opt);

  // TODO: This heuristic may be a bit too coarse and may need improving
  // If the chain contains a fp4 to fp16/bf16 conversion, then the original
  // bitwidth is 4.
  if (llvm::any_of(slice, [](Operation *op) { return isa<Fp4ToFpOp>(op); }))
    return 4;

  int origBitWidth = getElementTypeOrSelf(x).getIntOrFloatBitWidth();
  for (auto op : slice) {
    if (isa<LoadOp, DescriptorLoadLikeOpInterface>(op)) {
      if (auto tensorTy =
              dyn_cast<RankedTensorType>(op->getResultTypes().front())) {
        origBitWidth =
            std::min<int>(origBitWidth, tensorTy.getElementTypeBitWidth());
      }
    }
  }

  // If JoinOp occurred at least once, in backward layout propagation,
  // the kWidth will be split in half as we pass through the JoinOp.
  // Hence we divide origBitWidth by 2 here to compensate for that and
  // improve our load width.
  // This won't be optimal if there is a tree of multiple JoinOps, which
  // would require counting the max number of JoinOp's along any path.
  //
  // In the future we might want to do something like trying a large kWidth,
  // run layout backpropagation and see what's the contiguity that you
  // get at the loads that feed into it.
  //
  // This heuristic is intended for packed-K values of shape [..., K / 2]
  // joined into [..., K / 2, 2]. We also check that the bitwidth of the joined
  // value is twice the original bitwidth, meaning "one loaded storage element
  // contributes two logical K values after unpack/join".
  if (canWidenKWidthForJoin(origBitWidth, slice))
    origBitWidth /= 2;

  return origBitWidth;
}

namespace {

// Common MMA encoding creation
struct MMAEncodingResult {
  NvidiaMmaEncodingAttr mmaEnc;
  RankedTensorType newRetType;
  Value newAcc;
  int versionMajor;
  int versionMinor;
};

// Unified implementation for DotOpInterface
static MMAEncodingResult createMMAEncodingForDot(DotOpInterface dotOp,
                                                 PatternRewriter &rewriter,
                                                 int computeCapability,
                                                 int versionMajor) {
  auto oldRetType = cast<RankedTensorType>(dotOp.getD().getType());
  auto oldAType = cast<RankedTensorType>(dotOp.getA().getType());

  int numWarps = lookupNumWarps(dotOp);

  int versionMinor = computeCapability == 75 ? 1 : 0;
  // Only MMAv2 and MMAv3 rely on computing instrShape/warpsPerTile here.
  if (!(versionMajor == 2 || versionMajor == 3)) {
    return {nullptr, RankedTensorType(), Value(), versionMajor, versionMinor};
  }

  auto CGALayout = getCGALayout(oldRetType.getEncoding());
  auto retShapePerCTA = getShapePerCTA(oldRetType);
  auto instrShape = mmaVersionToInstrShape(versionMajor, retShapePerCTA,
                                           oldAType.getElementType(), numWarps);
  auto warpsPerTile = getWarpsPerTile(dotOp, retShapePerCTA, versionMajor,
                                      numWarps, instrShape);

  auto mmaEnc = NvidiaMmaEncodingAttr::get(oldRetType.getContext(),
                                           versionMajor, versionMinor,
                                           warpsPerTile, CGALayout, instrShape);
  auto newRetType = oldRetType.cloneWithEncoding(mmaEnc);

  auto oldAcc = dotOp->getOperand(2);
  auto newAcc =
      ConvertLayoutOp::create(rewriter, oldAcc.getLoc(), newRetType, oldAcc);

  return {mmaEnc, newRetType, newAcc, versionMajor, versionMinor};
}

// Common operand conversion
static Value convertDotOperandForMMA(Value v, int opIdx, int bitwidth,
                                     RankedTensorType newRetType,
                                     PatternRewriter &rewriter) {
  auto minType = bitwidth > 0 ? rewriter.getIntegerType(bitwidth) : v.getType();
  auto vType = cast<RankedTensorType>(v.getType());
  auto newVEncoding = DotOperandEncodingAttr::get(
      v.getContext(), opIdx, newRetType.getEncoding(), minType);
  auto newVType = vType.cloneWithEncoding(newVEncoding);
  return ConvertLayoutOp::create(rewriter, v.getLoc(), newVType, v);
}

} // namespace

class BlockedToMMA : public mlir::OpRewritePattern<DotOp> {
  int computeCapability;
  mutable llvm::DenseMap<Operation *, unsigned> dotOpInstNs;

public:
  BlockedToMMA(mlir::MLIRContext *context, int computeCapability, int benefit)
      : OpRewritePattern<DotOp>(context, benefit),
        computeCapability(computeCapability) {}

  mlir::LogicalResult
  matchAndRewrite(triton::DotOp dotOp,
                  mlir::PatternRewriter &rewriter) const override {
    if (computeCapability < 70)
      return failure();
    if (computeCapability < 80) {
      dotOp.emitRemark()
          << "Dot op using MMA for compute capability " << computeCapability
          << " has been deprecated. It falls back to the FMA path.";
      return failure();
    }
    // TODO: Check data-types and SM compatibility
    auto retType = dotOp.getType();
    if (!retType.getEncoding() ||
        mlir::isa<NvidiaMmaEncodingAttr>(retType.getEncoding()))
      return failure();

    Value a = dotOp.getA();
    Value b = dotOp.getB();
    auto oldAType = cast<RankedTensorType>(a.getType());
    auto oldBType = cast<RankedTensorType>(b.getType());
    auto oldRetType = cast<RankedTensorType>(dotOp.getType());

    // Enable F64 MMA only on targets with high performance F64 tensor cores.
    // Otherwise, fallback to F64 FMA for better performance.
    if ((oldAType.getElementType().isF64() ||
         oldBType.getElementType().isF64() ||
         oldRetType.getElementType().isF64()) &&
        !(computeCapability == 80 || computeCapability == 90 ||
          (computeCapability >= 100 && computeCapability < 120))) {
      return failure();
    }

    auto mmaVersion = getMMAVersionSafe(computeCapability, dotOp);
    auto mmaResult =
        createMMAEncodingForDot(dotOp, rewriter, computeCapability, mmaVersion);
    if (!(mmaResult.versionMajor >= 1 && mmaResult.versionMajor <= 3))
      return failure();

    Operation *newDot = nullptr;
    bool aFromLoad = comesFromLoadOrBlockArg(a);

    if (mmaResult.versionMajor == 3) {
      auto eltType = cast<RankedTensorType>(a.getType()).getElementType();
      bool allowTranspose = eltType.isF16() || eltType.isBF16();
      if (!aFromLoad) {
        int bitwidth = getElementTypeOrSelf(a).getIntOrFloatBitWidth();
        a = convertDotOperandForMMA(a, 0, bitwidth, mmaResult.newRetType,
                                    rewriter);
      } else {
        a = getSharedMemoryMMAOperand(a, rewriter, 0, allowTranspose,
                                      /*isMMAv5Fp4Padded=*/false,
                                      /*forceTranspose=*/false, dotOp);
      }
      b = getSharedMemoryMMAOperand(b, rewriter, 1, allowTranspose,
                                    /*isMMAv5Fp4Padded=*/false,
                                    /*forceTranspose=*/false, dotOp);

      newDot = triton::nvidia_gpu::WarpGroupDotOp::create(
          rewriter, dotOp.getLoc(), mmaResult.newRetType, a, b,
          mmaResult.newAcc, nullptr, dotOp.getInputPrecision(),
          dotOp.getMaxNumImpreciseAcc(), false);
    } else {
      int minBitwidth =
          std::min(computeOrigBitWidth(a), computeOrigBitWidth(b));
      a = convertDotOperandForMMA(a, 0, minBitwidth, mmaResult.newRetType,
                                  rewriter);
      b = convertDotOperandForMMA(b, 1, minBitwidth, mmaResult.newRetType,
                                  rewriter);
      newDot = DotOp::create(rewriter, dotOp.getLoc(), mmaResult.newRetType, a,
                             b, mmaResult.newAcc, dotOp.getInputPrecision(),
                             dotOp.getMaxNumImpreciseAcc());
    }

    rewriter.replaceOpWithNewOp<ConvertLayoutOp>(dotOp, dotOp.getType(),
                                                 newDot->getResult(0));
    return success();
  }
};

static DistributedEncodingTrait
replaceCGALayout(DistributedEncodingTrait layout,
                 const triton::gpu::CGAEncodingAttr &newCGALayout) {
  if (auto blockedLayout = mlir::dyn_cast<BlockedEncodingAttr>(layout)) {
    return BlockedEncodingAttr::get(
        layout.getContext(), blockedLayout.getSizePerThread(),
        blockedLayout.getThreadsPerWarp(), blockedLayout.getWarpsPerCTA(),
        blockedLayout.getOrder(), newCGALayout);
  } else if (auto sliceLayout = mlir::dyn_cast<SliceEncodingAttr>(layout)) {
    return SliceEncodingAttr::get(
        layout.getContext(), sliceLayout.getDim(),
        replaceCGALayout(sliceLayout.getParent(), newCGALayout));
  } else {
    llvm::report_fatal_error("not implemented");
    return layout;
  }
}

static Value splitBOperand(Value b, mlir::PatternRewriter &rewriter) {
  OpBuilder::InsertionGuard g(rewriter);
  MLIRContext *ctx = b.getContext();
  while (auto cvtOp = b.getDefiningOp<ConvertLayoutOp>())
    b = cvtOp.getSrc();
  auto loadOp = b.getDefiningOp();
  assert((isa<triton::LoadOp, triton::DescriptorLoadLikeOpInterface>(loadOp)) &&
         "expected LoadOp");
  RankedTensorType bType = cast<RankedTensorType>(b.getType());
  auto currentLayout = cast<DistributedEncodingTrait>(bType.getEncoding());
  auto kBlock = StringAttr::get(ctx, "block");
  auto dims = standardOutDimNames(ctx, 2);
  auto newCGALayout =
      CGAEncodingAttr::get(ctx, LinearLayout({{kBlock, {{0, 1}}}}, dims));
  Attribute newLayout = replaceCGALayout(currentLayout, newCGALayout);
  rewriter.setInsertionPoint(loadOp);
  for (OpOperand &operand : loadOp->getOpOperands()) {
    auto tensorType = dyn_cast<RankedTensorType>(operand.get().getType());
    if (!tensorType)
      continue;
    Value newOperand = ConvertLayoutOp::create(
        rewriter, operand.get().getLoc(),
        tensorType.cloneWithEncoding(newLayout), operand.get());
    loadOp->setOperand(operand.getOperandNumber(), newOperand);
  }
  loadOp->getResult(0).setType(bType.cloneWithEncoding(newLayout));
  Value newB = loadOp->getResult(0);
  rewriter.setInsertionPointAfter(loadOp);
  auto cvt = ConvertLayoutOp::create(rewriter, b.getLoc(), bType, newB);
  rewriter.replaceAllUsesExcept(newB, cvt.getResult(), cvt);
  return newB;
}

class BlockedToMMAv5 : public mlir::OpRewritePattern<DotOp> {
  int computeCapability;

public:
  BlockedToMMAv5(mlir::MLIRContext *context, int computeCapability, int benefit)
      : OpRewritePattern<DotOp>(context, benefit),
        computeCapability(computeCapability) {}

  mlir::LogicalResult
  matchAndRewrite(triton::DotOp dotOp,
                  mlir::PatternRewriter &rewriter) const override {
    RankedTensorType oldRetType = dotOp.getType();
    if (!oldRetType.getEncoding() ||
        mlir::isa<NvidiaMmaEncodingAttr>(oldRetType.getEncoding()))
      return failure();

    // get MMA encoding for the given number of warps
    auto retShapePerCTA = getShapePerCTA(oldRetType);
    int numWarps = lookupNumWarps(dotOp);
    auto CGALayout = getCGALayout(oldRetType.getEncoding());

    int versionMajor = getMMAVersionSafe(computeCapability, dotOp);
    if (versionMajor != 5)
      return failure();
    Location loc = dotOp.getLoc();
    // operands
    Value a = dotOp.getA();
    Value b = dotOp.getB();
    if (std::min(computeOrigBitWidth(a), computeOrigBitWidth(b)) >= 32 &&
        dotOp.getInputPrecision() != InputPrecision::TF32)
      return failure();
    auto oldAType = dotOp.getA().getType();
    // NYI: PTX 13+ requires all tcgen instructions in a kernel to have a
    // consistent CTA mode, disabling 2CTA mode for now. To re-enable,
    // change the line below to: bool useTwoCTAs = canUseTwoCTAs(dotOp);
    bool useTwoCTAs = false;
    if (useTwoCTAs) {
      b = splitBOperand(b, rewriter);
    }
    // TF32 transpose is only supported with 128 swizzle mode with 32B
    // atomicity. As we currently don't support this layout we disallow
    // transpose for TF32 inputs.
    bool allowTranspose = !dotOp.getA().getType().getElementType().isF32();
    a = getSharedMemoryMMAOperand(a, rewriter, 0, allowTranspose);
    b = getSharedMemoryMMAOperand(b, rewriter, 1, allowTranspose);
    MLIRContext *context = dotOp->getContext();
    auto instrShape = mmaVersionToInstrShape(
        versionMajor, retShapePerCTA, oldAType.getElementType(), numWarps);
    auto bitwidth = oldRetType.getElementType().getIntOrFloatBitWidth();
    unsigned colStride = 32 / bitwidth;
    Attribute accEncoding = triton::nvidia_gpu::TensorMemoryEncodingAttr::get(
        context, instrShape[0], instrShape[1], colStride, CGALayout,
        useTwoCTAs);
    Attribute tensorMemorySpace =
        triton::nvidia_gpu::TensorMemorySpaceAttr::get(context);
    MemDescType accMemDescType =
        MemDescType::get(oldRetType.getShape(), oldRetType.getElementType(),
                         accEncoding, tensorMemorySpace,
                         /*mutableMemory=*/true);
    auto newDistributedEncoding =
        nvidia_gpu::getDefaultLayoutForTmemLdSt(accMemDescType, numWarps);
    auto newAccType = oldRetType.cloneWithEncoding(newDistributedEncoding);
    Value cvtAcc =
        ConvertLayoutOp::create(rewriter, loc, newAccType, dotOp.getOperand(2));
    auto tokType = rewriter.getType<AsyncTokenType>();
    auto acc = triton::nvidia_gpu::TMEMAllocOp::create(
        rewriter, loc, accMemDescType, tokType, cvtAcc);
    auto vTrue = arith::ConstantIntOp::create(rewriter, dotOp.getLoc(), 1, 1);
    auto mma = triton::nvidia_gpu::TCGen5MMAOp::create(
        rewriter, loc, tokType, a, b, acc, acc.getToken(), /*useD=*/vTrue,
        /*pred=*/vTrue);
    mma.setTwoCtas(useTwoCTAs);

    auto ld = triton::nvidia_gpu::TMEMLoadOp::create(
        rewriter, loc, newAccType, tokType, acc, /*dep=*/mma.getToken());
    rewriter.replaceOpWithNewOp<ConvertLayoutOp>(dotOp, oldRetType, ld);
    return success();
  }
};

class ScaledBlockedToMMA : public mlir::OpRewritePattern<triton::DotScaledOp> {
  int computeCapability;

public:
  ScaledBlockedToMMA(mlir::MLIRContext *context, int computeCapability,
                     int benefit)
      : mlir::OpRewritePattern<triton::DotScaledOp>(context, benefit),
        computeCapability(computeCapability) {}

  mlir::LogicalResult
  matchAndRewrite(triton::DotScaledOp dotOp,
                  mlir::PatternRewriter &rewriter) const override {
    if (computeCapability / 10 != 12)
      return failure();

    auto numCTAs = lookupNumCTAs(rewriter);
    if (numCTAs != 1) {
      return failure();
    }
    // Skip if any scale is missing. This pattern requires both scales.
    if (!dotOp.getAScale() || !dotOp.getBScale())
      return failure();

    auto aScaleType = dotOp.getAScale().getType();
    auto bScaleType = dotOp.getBScale().getType();

    if (mlir::isa<LinearEncodingAttr>(aScaleType.getEncoding()) ||
        mlir::isa<LinearEncodingAttr>(bScaleType.getEncoding())) {
      return failure();
    }
    auto aElemType = dotOp.getAElemType();
    auto bElemType = dotOp.getBElemType();
    auto isFP8 = [&](ScaleDotElemType elemType) -> bool {
      return elemType == ScaleDotElemType::E4M3 ||
             elemType == ScaleDotElemType::E5M2;
    };
    auto isFP4 = [&](ScaleDotElemType elemType) -> bool {
      return elemType == ScaleDotElemType::E2M1;
    };

    bool isFP4xFP4 = isFP4(aElemType) && isFP4(bElemType);
    // TODO: Enable mixed-precision mxfp for sm120
    if (!((isFP8(aElemType) && isFP8(bElemType)) || isFP4xFP4)) {
      return rewriter.notifyMatchFailure(
          dotOp, "only FP8xFP8 and FP4xFP4 are supported on sm120");
    }
    if (isFP4xFP4 && (!dotOp.getLhsKPack() || !dotOp.getRhsKPack())) {
      return rewriter.notifyMatchFailure(
          dotOp, "SM120 native FP4xFP4 requires K-packed operands");
    }

    auto scaleElemType = dotOp.getAScale().getType().getElementType();
    if (scaleElemType != dotOp.getBScale().getType().getElementType()) {
      return failure();
    }

    // Common MMA encoding creation
    auto mmaResult =
        createMMAEncodingForDot(dotOp, rewriter, computeCapability, 2);

    // Operand processing
    Value a = dotOp.getA();
    Value b = dotOp.getB();
    auto oldAType = cast<RankedTensorType>(a.getType());
    auto oldBType = cast<RankedTensorType>(b.getType());

    Operation *newDot = nullptr;

    // ScaledBlockedToMMA logic
    int bitwidthA = oldAType.getElementType().getIntOrFloatBitWidth();
    int bitwidthB = oldBType.getElementType().getIntOrFloatBitWidth();
    int minBitwidth = std::min(bitwidthA, bitwidthB);

    Value newA = convertDotOperandForMMA(a, 0, minBitwidth,
                                         mmaResult.newRetType, rewriter);
    Value newB = convertDotOperandForMMA(b, 1, minBitwidth,
                                         mmaResult.newRetType, rewriter);
    const auto mmaWarps = mmaResult.mmaEnc.getWarpsPerCTA(); // [wM, wN]
    // Convert scales to Linear layout
    auto convertScale = [&](Value scale, int opIdx) -> Value {
      auto ty = cast<RankedTensorType>(scale.getType());
      SmallVector<int64_t> shape = llvm::to_vector(ty.getShape());
      MLIRContext *ctx = ty.getContext();
      auto blocked = cast<triton::gpu::BlockedEncodingAttr>(ty.getEncoding());

      auto ll = triton::gpu::getSM120DotScaledScaleLayout(
          ctx, shape, opIdx, mmaWarps, blocked.getCGALayout());
      auto newEnc = triton::gpu::LinearEncodingAttr::get(ctx, std::move(ll));
      auto newTy = RankedTensorType::get(shape, ty.getElementType(), newEnc);
      return ConvertLayoutOp::create(rewriter, scale.getLoc(), newTy, scale);
    };
    Value aScale = convertScale(dotOp.getAScale(), /*opIdx=*/0);
    Value bScale = convertScale(dotOp.getBScale(), /*opIdx=*/1);

    newDot = triton::DotScaledOp::create(
        rewriter, dotOp.getLoc(), mmaResult.newRetType, newA, newB,
        mmaResult.newAcc, aScale, bScale, dotOp.getAElemType(),
        dotOp.getBElemType(), dotOp.getFastMath(), dotOp.getLhsKPack(),
        dotOp.getRhsKPack());
    rewriter.replaceOpWithNewOp<ConvertLayoutOp>(dotOp, dotOp.getType(),
                                                 newDot->getResult(0));
    return success();
  }
};

class ScaledBlockedToMMAv5
    : public mlir::OpRewritePattern<triton::DotScaledOp> {
  int computeCapability;

public:
  ScaledBlockedToMMAv5(mlir::MLIRContext *context, int computeCapability,
                       int benefit)
      : mlir::OpRewritePattern<triton::DotScaledOp>(context, benefit),
        computeCapability(computeCapability) {}

  mlir::LogicalResult
  matchAndRewrite(triton::DotScaledOp dotOp,
                  mlir::PatternRewriter &rewriter) const override {
    RankedTensorType oldRetType = dotOp.getType();
    if (!oldRetType.getEncoding() ||
        mlir::isa<NvidiaMmaEncodingAttr>(oldRetType.getEncoding()))
      return failure();

    if (dotOp.getAScale() == nullptr || dotOp.getBScale() == nullptr) {
      return failure();
    }

    // get MMA encoding for the given number of warps
    auto retShapePerCTA = getShapePerCTA(oldRetType);
    int numWarps = lookupNumWarps(dotOp);
    auto CGALayout = getCGALayout(oldRetType.getEncoding());
    if (computeCapability < 100 || computeCapability >= 120)
      return failure();
    if (numWarps != 4 && numWarps != 8)
      return failure();
    if (retShapePerCTA[0] < 128 || retShapePerCTA[1] < 16)
      return failure();
    Location loc = dotOp.getLoc();
    // operands
    Value a = dotOp.getA();
    Value b = dotOp.getB();

    bool IsAMixedPrecFp4 = false;
    bool IsBMixedPrecFp4 = false;
    bool isAFP4 = dotOp.getAElemType() == ScaleDotElemType::E2M1;
    bool isBFP4 = dotOp.getBElemType() == ScaleDotElemType::E2M1;

    if (dotOp.getAElemType() != dotOp.getBElemType()) {
      if (isAFP4)
        IsAMixedPrecFp4 = true;
      else if (isBFP4)
        IsBMixedPrecFp4 = true;
    }

    bool isFp4MMA = isAFP4 && isBFP4;
    bool requiresFp4Padding =
        nvidia_gpu::TargetFeatures(computeCapability).requiresFp4Padding();

    // On Blackwell, if we use mixed-precision MMA we need to pad the fp4
    // operand
    // https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-packing-formats-mxf8f6f4-smem
    // On Rubin, the fp4 operand must be packed in SMEM if MMA_K = 64 is used.
    // However, MMA_K = 64 with MN-major fp4 is not supported. In this case,
    // the fp4 operand must be padded and MMA_K = 32 must be used.
    // MMA_K = 64 is also not supported when BLOCK_M = 64.
    auto blockK = dotOp.getA().getType().getShape().back() * (isAFP4 ? 2 : 1);
    auto blockM = dotOp.getA().getType().getShape()[0];
    // mxf4/mxf4nvf4 do not support MN-major operands, so fp4 x fp4 falls
    // back to mxf8f6f4 if either operand is not K-packed. The mxf8f6f4
    // shared-memory packing format requires padding for every fp4 operand,
    // even if the operand is K packed.
    bool isFp4MMAUsingMxf8f6f4 =
        isFp4MMA && (!dotOp.getLhsKPack() || !dotOp.getRhsKPack());
    bool isMMAv5Fp4PaddedLhs =
        isFp4MMAUsingMxf8f6f4 ||
        (IsAMixedPrecFp4 && (requiresFp4Padding || blockM == 64 ||
                             blockK == 32 || !dotOp.getLhsKPack()));
    bool isMMAv5Fp4PaddedRhs =
        isFp4MMAUsingMxf8f6f4 ||
        (IsBMixedPrecFp4 &&
         (requiresFp4Padding || blockK == 32 || !dotOp.getRhsKPack()));

    // For mixed-precision fp4 operands, set allowTranspose = false, to force
    // the packed axis, K, to be contiguous in SMEM
    a = getSharedMemoryMMAOperand(a, rewriter, 0,
                                  /*allowTranspose=*/!isAFP4,
                                  /*isMMAv5Fp4Padded=*/isMMAv5Fp4PaddedLhs,
                                  /*forceTranspose=*/!dotOp.getLhsKPack(),
                                  dotOp);
    b = getSharedMemoryMMAOperand(b, rewriter, 1,
                                  /*allowTranspose=*/!isBFP4,
                                  /*isMMAv5Fp4Padded=*/isMMAv5Fp4PaddedRhs,
                                  /*forceTranspose=*/!dotOp.getRhsKPack(),
                                  dotOp);

    MLIRContext *context = dotOp->getContext();
    unsigned m = 128;
    unsigned n = retShapePerCTA[1] >= 256 ? 256 : retShapePerCTA[1];

    auto bitwidth = oldRetType.getElementType().getIntOrFloatBitWidth();
    unsigned colStride = 32 / bitwidth;
    Attribute accEncoding = triton::nvidia_gpu::TensorMemoryEncodingAttr::get(
        context, m, n, colStride, CGALayout);
    Attribute tensorMemorySpace =
        triton::nvidia_gpu::TensorMemorySpaceAttr::get(context);
    MemDescType accMemDescType =
        MemDescType::get(oldRetType.getShape(), oldRetType.getElementType(),
                         accEncoding, tensorMemorySpace,
                         /*mutableMemory=*/true);
    auto newDistributedEncoding =
        nvidia_gpu::getDefaultLayoutForTmemLdSt(accMemDescType, numWarps);
    auto newAccType = oldRetType.cloneWithEncoding(newDistributedEncoding);
    Value cvtAcc =
        ConvertLayoutOp::create(rewriter, loc, newAccType, dotOp.getOperand(2));
    auto tokType = rewriter.getType<AsyncTokenType>();
    auto acc = triton::nvidia_gpu::TMEMAllocOp::create(
        rewriter, loc, accMemDescType, tokType, cvtAcc);

    RankedTensorType oldScaleAType = dotOp.getAScale().getType();
    RankedTensorType oldScaleBType = dotOp.getBScale().getType();

    auto aScaleBlockRepOrder =
        triton::nvidia_gpu::getTensorMemoryScalesBlockRepOrder(
            dotOp, /*isA=*/true, dotOp.getAElemType(), dotOp.getBElemType(),
            oldScaleAType.getElementType(), oldScaleBType.getElementType());
    auto bScaleBlockRepOrder =
        triton::nvidia_gpu::getTensorMemoryScalesBlockRepOrder(
            dotOp, /*isA=*/false, dotOp.getAElemType(), dotOp.getBElemType(),
            oldScaleAType.getElementType(), oldScaleBType.getElementType());
    Attribute scaleAEncoding =
        triton::nvidia_gpu::TensorMemoryScalesEncodingAttr::get(
            context, CGALayout, aScaleBlockRepOrder);
    Attribute scaleBEncoding =
        triton::nvidia_gpu::TensorMemoryScalesEncodingAttr::get(
            context, CGALayout, bScaleBlockRepOrder);
    Value scaleA =
        tryCreateTmemCopyCompatibleScaleOperand(dotOp.getAScale(), rewriter);
    if (!scaleA) {
      scaleA =
          createTmemScaleOperand(dotOp.getAScale(), scaleAEncoding,
                                 tensorMemorySpace, numWarps, loc, rewriter);
    }
    Value scaleB =
        tryCreateTmemCopyCompatibleScaleOperand(dotOp.getBScale(), rewriter);
    if (!scaleB) {
      scaleB =
          createTmemScaleOperand(dotOp.getBScale(), scaleBEncoding,
                                 tensorMemorySpace, numWarps, loc, rewriter);
    }

    auto vTrue = arith::ConstantIntOp::create(rewriter, dotOp.getLoc(), 1, 1);
    auto mmaOp = triton::nvidia_gpu::TCGen5MMAScaledOp::create(
        rewriter, loc, tokType, a, b, acc.getResult(), acc.getToken(), scaleA,
        scaleB, dotOp.getAElemType(), dotOp.getBElemType(),
        /*useD=*/vTrue, /*pred=*/vTrue);

    auto ld = triton::nvidia_gpu::TMEMLoadOp::create(
        rewriter, loc, newAccType, tokType, acc, mmaOp.getToken());
    rewriter.replaceOpWithNewOp<ConvertLayoutOp>(dotOp, oldRetType, ld);
    return success();
  }
};
} // namespace

static Value promoteOperand(OpBuilder &builder, Location loc, Value operand,
                            Type promotedType) {
  Type tensorPromotedType = cast<RankedTensorType>(operand.getType())
                                .cloneWith(std::nullopt, promotedType);
  Type operandElType =
      cast<RankedTensorType>(operand.getType()).getElementType();
  if (type::isFloat8(operandElType)) {
    return FpToFpOp::create(builder, loc, tensorPromotedType, operand);
  }
  return arith::ExtFOp::create(builder, loc, tensorPromotedType, operand);
}

static bool mmav2SupportsFp8Operands(int computeCapability) {
  // promote operands for sm < 89 since fp8 mma is not natively supported
  // although PTX instructions for mma v2 w/ fp8 operands exist for sm90 and
  // sm100, they are emulated as fp16 upcasts + fp16 HMMA in SASS. sm120 has
  // hardware support for fp8 operands w/ mmav2.
  return computeCapability == 89 || computeCapability / 10 == 12;
}

// promote operands of dot op if the existing combination is not natively
// supported.
static void decomposeMixedModeDotOp(ModuleOp mod, int computeCapability) {
  mod.walk([=](DotOp dotOp) -> void {
    auto D = dotOp.getD();
    OpBuilder builder(dotOp);
    Type AElType = dotOp.getA().getType().getElementType();
    Type promoteType;
    NvidiaMmaEncodingAttr mmaLayout =
        dyn_cast<NvidiaMmaEncodingAttr>(D.getType().getEncoding());
    if (mmaLayout) {
      bool isNativeFP8 = llvm::isa<Float8E5M2Type, Float8E4M3FNType>(AElType);
      // promote to f16 unless there's hardware support for fp8 operands
      if (!isNativeFP8 ||
          (isNativeFP8 && (mmav2SupportsFp8Operands(computeCapability) ||
                           mmaLayout.isHopper())))
        return;
      promoteType = builder.getF16Type();
    } else {
      // FMA case.
      Type AElType = dotOp.getA().getType().getElementType();
      Type DElType = D.getType().getElementType();
      if (AElType == DElType)
        return;
      promoteType = DElType;
    }
    Location loc = dotOp.getLoc();
    Value promotedA = promoteOperand(builder, loc, dotOp.getA(), promoteType);
    Value promotedB = promoteOperand(builder, loc, dotOp.getB(), promoteType);
    dotOp.setOperand(0, promotedA);
    dotOp.setOperand(1, promotedB);
  });
}

// Transpose scaled_dot ops that have a scale on lhs.
static void transposeDotOp(DotScaledOp dotOp) {
  OpBuilder builder(dotOp);
  Value lhs = dotOp.getA();
  std::array<int, 2> transOrder = {1, 0};
  Value lhsTransposed = TransOp::create(builder, lhs.getLoc(), lhs, transOrder);
  Value rhs = dotOp.getB();
  Value rhsTransposed = TransOp::create(builder, rhs.getLoc(), rhs, transOrder);
  Value c = dotOp.getC();
  Value cTransposed = TransOp::create(builder, c.getLoc(), c, transOrder);
  Value result = DotScaledOp::create(
      builder, dotOp.getLoc(), cTransposed.getType(), rhsTransposed,
      lhsTransposed, cTransposed, dotOp.getBScale(), dotOp.getAScale(),
      dotOp.getBElemType(), dotOp.getAElemType(), dotOp.getFastMath());
  Operation *transposedResult =
      TransOp::create(builder, result.getLoc(), result, transOrder);
  dotOp.replaceAllUsesWith(transposedResult);
  dotOp.erase();
}

static void transposeDots(ModuleOp m) {
  // TODO: extend to regular dot when it is profitable. For instance when we may
  // want to use rhs from register for mmav3.
  SmallVector<DotScaledOp> toTranspose;
  m.walk([&](DotScaledOp dotOp) -> void {
    if (dotOp.getAScale() == nullptr && dotOp.getBScale() != nullptr)
      toTranspose.push_back(dotOp);
  });
  for (DotScaledOp dotOp : toTranspose) {
    transposeDotOp(dotOp);
  }
}

#define GEN_PASS_DEF_TRITONGPUACCELERATEMATMUL
#include "triton/Dialect/TritonGPU/Transforms/Passes.h.inc"

class TritonGPUAccelerateMatmulPass
    : public impl::TritonGPUAccelerateMatmulBase<
          TritonGPUAccelerateMatmulPass> {
public:
  using impl::TritonGPUAccelerateMatmulBase<
      TritonGPUAccelerateMatmulPass>::TritonGPUAccelerateMatmulBase;

  void runOnOperation() override {
    MLIRContext *context = &getContext();
    ModuleOp m = getOperation();

    auto computeCapability = getNVIDIAComputeCapability(m);
    // We could do this generically if we manage to improve the heuristics
    // reverted in these two PRs https://github.com/triton-lang/triton/pull/5834
    // https://github.com/triton-lang/triton/pull/5837
    transposeDots(m);

    mlir::RewritePatternSet patterns(context);
    constexpr int benefitDefault = 1;
    constexpr int benefitMMAv5 = 10;
    constexpr int benefitSM120 = 10;

    patterns.add<BlockedToMMA>(context, computeCapability, benefitDefault);
    patterns.add<ScaledBlockedToMMA>(context, computeCapability, benefitSM120);
    populateDecomposeScaledBlockedPatterns(patterns, benefitDefault);
    patterns.add<BlockedToMMAv5, ScaledBlockedToMMAv5>(
        context, computeCapability, benefitMMAv5);

    if (applyPatternsGreedily(m, std::move(patterns)).failed())
      return signalPassFailure();
    // Now that we have picked the mma type, decompose dot that are not natively
    // supported.
    decomposeMixedModeDotOp(m, computeCapability);
  }
};

} // namespace gpu
} // namespace triton
} // namespace mlir
