#include "ir.h"
#include <nanobind/nanobind.h>
#include <nanobind/stl/function.h>
#include <nanobind/stl/map.h>
#include <nanobind/stl/optional.h>
#include <nanobind/stl/pair.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/tuple.h>
#include <nanobind/stl/vector.h>

#include <optional>
#include <stdexcept>

#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/DialectRegistry.h"
#include "mlir/IR/Types.h"
#include "third_party/amd/include/Dialect/TritonAMDGPU/IR/Dialect.h"
#include "third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.h"
#include "third_party/amd/lib/TritonAMDGPUTransforms/Utility.h"
#include "triton/Analysis/Utility.h"
#include "triton/Dialect/Gluon/IR/Dialect.h"
#include "triton/Dialect/Triton/IR/Dialect.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/TritonGPU/Transforms/Utility.h"
#include "triton/Dialect/TritonInstrument/IR/Dialect.h"
#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
#include "triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h"
#include "triton/Tools/GenericSwizzling.h"
#include "triton/Tools/LayoutUtils.h"
#include "triton/Tools/LinearLayout.h"
#include "llvm/ADT/StringSwitch.h"
#include "llvm/Support/MathExtras.h"

using namespace mlir;
namespace py = nanobind;
namespace tt = triton;
namespace ttg = triton::gpu;
namespace ttng = triton::nvidia_gpu;
namespace gluon = mlir::triton::gluon;
namespace ttag = mlir::triton::amdgpu;

namespace {

ttg::CGAEncodingAttr
buildCgaLayoutAttr(MLIRContext *ctx,
                   const std::vector<std::vector<int32_t>> &layout,
                   unsigned rank) {
  auto kBlock = StringAttr::get(ctx, "block");
  tt::LinearLayout::BasesT bases;
  bases[kBlock] = layout;
  auto outDims = tt::standardOutDimNames(ctx, rank);
  tt::LinearLayout ll(std::move(bases), outDims);
  return ttg::CGAEncodingAttr::get(ctx, std::move(ll));
}

std::vector<std::vector<int32_t>>
getCgaLayoutBases(ttg::CGAEncodingAttr layout) {
  std::vector<std::vector<int32_t>> result;
  auto ctx = layout.getContext();
  auto block = StringAttr::get(ctx, "block");
  const auto &basesMap = layout.getLinearLayout().getBases();
  auto it = basesMap.find(block);
  assert(it != basesMap.end());
  return it->second;
}

// Helper to check if an MLIR type or attribute has a verifier method.
template <typename AttrOrType>
constexpr auto hasVerifier(AttrOrType t) -> decltype(t.verifyInvariants, true) {
  return true;
}
constexpr auto hasVerifier(...) { return false; }

// Print a diagnostic without its location. The frontend will attach the AST
// location to the error message.
void printDiagStr(llvm::raw_ostream &os, const Diagnostic &diag) {
  for (const DiagnosticArgument &arg : diag.getArguments())
    arg.print(os);
  os << "\n";
  for (const Diagnostic &note : diag.getNotes())
    printDiagStr(os, note);
}

struct GluonOpBuilder : public TritonOpBuilder {
  GluonOpBuilder(MLIRContext *context, std::string arch = "")
      : TritonOpBuilder(context), arch(arch) {}

  bool isRubin() const { return arch == "sm107"; }

  // Construct an attribute or type while calling its verifier. Error messages
  // are intercepted and sent back to Python via a C++ exception.
  template <typename AttrOrType, typename... ArgTs>
  std::enable_if_t<hasVerifier(AttrOrType()), AttrOrType>
  getChecked(ArgTs &&...args) {
    // Set up a scoped handler to intercept errors.
    std::string msg;
    llvm::raw_string_ostream os(msg);
    ScopedDiagnosticHandler handler(
        getContext(), [&](Diagnostic &diag) { printDiagStr(os, diag); });

    auto result =
        AttrOrType::getChecked([&] { return mlir::emitError(getLastLoc()); },
                               std::forward<ArgTs>(args)...);
    if (!result)
      throw std::runtime_error(os.str());
    return result;
  }

  // A variant of the above due to issues with C++ overload resolution and how
  // MLIR sets up the default `getChecked` implementation.
  template <typename AttrOrType, typename... ArgTs>
  std::enable_if_t<hasVerifier(AttrOrType()), AttrOrType>
  getChecked(MLIRContext *ctx, ArgTs &&...args) {
    // Set up a scoped handler to intercept errors.
    std::string msg;
    llvm::raw_string_ostream os(msg);
    ScopedDiagnosticHandler handler(
        getContext(), [&](Diagnostic &diag) { printDiagStr(os, diag); });

    if (failed(AttrOrType::verifyInvariants(
            [&] { return mlir::emitError(getLastLoc()); }, args...)))
      throw std::runtime_error(os.str());

    return AttrOrType::get(ctx, std::forward<ArgTs>(args)...);
  }

  // Fallback method for types or attributes that do not have a verifier.
  template <typename AttrOrType, typename... ArgTs>
  std::enable_if_t<!hasVerifier(AttrOrType()), AttrOrType>
  getChecked(ArgTs &&...args) {
    return AttrOrType::get(std::forward<ArgTs>(args)...);
  }

private:
  std::string arch;
};

struct GluonLayouts {
  py::handle AutoLayout;
  py::handle CoalescedLayout;
  py::handle BlockedLayout;
  py::handle SliceLayout;
  py::handle DistributedLinearLayout;
  py::handle DotOperandLayout;
  py::handle NVMMADistributedLayout;
  py::handle RubinTensorMemoryScalesLayout;
  py::handle TensorMemoryScalesLayout;
  py::handle TensorMemoryLayout;
  py::handle NVMMASharedLayout;
  py::handle SwizzledSharedLayout;
  py::handle SharedLinearLayout;
  py::handle AMDMFMALayout;
  py::handle AMDWMMALayout;
  py::handle PaddedSharedLayout;
  py::handle PartitionedSharedLayout;

  GluonLayouts() {
    auto layouts =
        py::module_::import_("triton.experimental.gluon.language._layouts");
    auto amdLayouts =
        py::module_::import_("triton.experimental.gluon.language.amd._layouts");
    auto blackwellLayouts = py::module_::import_(
        "triton.experimental.gluon.language.nvidia.blackwell");
    auto rubinLayouts =
        py::module_::import_("triton.experimental.gluon.language.nvidia.rubin");
    AutoLayout = py::object(layouts.attr("AutoLayout")).release();
    CoalescedLayout = py::object(layouts.attr("CoalescedLayout")).release();
    BlockedLayout = py::object(layouts.attr("BlockedLayout")).release();
    SliceLayout = py::object(layouts.attr("SliceLayout")).release();
    DistributedLinearLayout =
        py::object(layouts.attr("DistributedLinearLayout")).release();
    DotOperandLayout = py::object(layouts.attr("DotOperandLayout")).release();
    NVMMADistributedLayout =
        py::object(layouts.attr("NVMMADistributedLayout")).release();
    TensorMemoryScalesLayout =
        py::object(blackwellLayouts.attr("TensorMemoryScalesLayout")).release();
    RubinTensorMemoryScalesLayout =
        py::object(rubinLayouts.attr("TensorMemoryScalesLayout")).release();
    TensorMemoryLayout =
        py::object(blackwellLayouts.attr("TensorMemoryLayout")).release();
    NVMMASharedLayout = py::object(layouts.attr("NVMMASharedLayout")).release();
    SwizzledSharedLayout =
        py::object(layouts.attr("SwizzledSharedLayout")).release();
    SharedLinearLayout =
        py::object(layouts.attr("SharedLinearLayout")).release();
    AMDMFMALayout = py::object(amdLayouts.attr("AMDMFMALayout")).release();
    AMDWMMALayout = py::object(amdLayouts.attr("AMDWMMALayout")).release();
    PaddedSharedLayout =
        py::object(layouts.attr("PaddedSharedLayout")).release();
    auto cdna5Layouts = py::module_::import_(
        "triton.experimental.gluon.language.amd.cdna5._layouts");
    PartitionedSharedLayout =
        py::object(cdna5Layouts.attr("PartitionedSharedLayout")).release();

    auto core = py::module_::import_("triton.language.core");
  }
};

bool isConvertLayoutTrivial(RankedTensorType dstTy, Value value) {
  auto srcTy = cast<RankedTensorType>(value.getType());
  if (srcTy.getEncoding() == dstTy.getEncoding())
    return true;
  // Fail safe on unresolved layouts.
  if (isa<gluon::AutoEncodingAttr>(srcTy.getEncoding()))
    return false;
  if (isa<gluon::AutoEncodingAttr>(dstTy.getEncoding()))
    return false;

  // Check concrete layouts.
  triton::LinearLayout cvt = minimalCvtLayout(srcTy, dstTy);
  auto dims = llvm::to_vector(cvt.getInDimNames());
  return dims.empty() || (dims.size() == 1 && dims.front() == "register");
}

template <typename R>
std::vector<llvm::ValueTypeFromRangeType<R>> toStdVector(R &&range) {
  return {range.begin(), range.end()};
}

py::object layoutToGluon(Attribute layout, bool isRubin = false) {
  static GluonLayouts layouts;
  if (auto blocked = dyn_cast<ttg::BlockedEncodingAttr>(layout)) {
    auto cgaBases = getCgaLayoutBases(blocked.getCGALayout());
    return layouts.BlockedLayout(toStdVector(blocked.getSizePerThread()),
                                 toStdVector(blocked.getThreadsPerWarp()),
                                 toStdVector(blocked.getWarpsPerCTA()),
                                 toStdVector(blocked.getOrder()), cgaBases);
  } else if (auto sliced = dyn_cast<ttg::SliceEncodingAttr>(layout)) {
    return layouts.SliceLayout(sliced.getDim(),
                               layoutToGluon(sliced.getParent(), isRubin));
  } else if (auto linearEnc = dyn_cast<ttg::LinearEncodingTrait>(layout)) {
    const auto &ll = linearEnc.getLinearLayout();
    auto ctx = layout.getContext();
    auto kReg = mlir::StringAttr::get(ctx, "register");
    auto kLane = mlir::StringAttr::get(ctx, "lane");
    auto kWarp = mlir::StringAttr::get(ctx, "warp");
    auto kBlock = mlir::StringAttr::get(ctx, "block");
    return layouts.DistributedLinearLayout(
        ll.getBases().lookup(kReg), ll.getBases().lookup(kLane),
        ll.getBases().lookup(kWarp), ll.getBases().lookup(kBlock),
        toStdVector(ll.getOutDimSizes()));
  } else if (auto dotOp = dyn_cast<ttg::DotOperandEncodingAttr>(layout)) {
    return layouts.DotOperandLayout(dotOp.getOpIdx(),
                                    layoutToGluon(dotOp.getParent(), isRubin),
                                    dotOp.getKWidth());
  } else if (auto mma = dyn_cast<ttg::NvidiaMmaEncodingAttr>(layout)) {
    auto cgaBases = getCgaLayoutBases(mma.getCGALayout());
    return layouts.NVMMADistributedLayout(
        std::vector<unsigned>{mma.getVersionMajor(), mma.getVersionMinor()},
        toStdVector(mma.getWarpsPerCTA()), toStdVector(mma.getInstrShape()),
        cgaBases);
  } else if (auto nvmma = dyn_cast<ttg::NVMMASharedEncodingAttr>(layout)) {
    auto cgaLayout = nvmma.getCGALayout();
    auto cgaBases = getCgaLayoutBases(cgaLayout);
    return layouts.NVMMASharedLayout(nvmma.getSwizzlingByteWidth(),
                                     nvmma.getElementBitWidth(),
                                     cgaLayout.getRank(), nvmma.getTransposed(),
                                     nvmma.getFp4Padded(), cgaBases);
  } else if (auto swizzled =
                 dyn_cast<ttg::SwizzledSharedEncodingAttr>(layout)) {
    auto cgaBases = getCgaLayoutBases(swizzled.getCGALayout());
    return layouts.SwizzledSharedLayout(
        swizzled.getVec(), swizzled.getPerPhase(), swizzled.getMaxPhase(),
        toStdVector(swizzled.getOrder()), cgaBases);
  } else if (auto sharedLl = dyn_cast<ttg::SharedLinearEncodingAttr>(layout)) {
    const auto &ll = sharedLl.getLinearLayout();
    auto ctx = layout.getContext();
    auto kOffset = mlir::StringAttr::get(ctx, "offset");
    auto kBlock = mlir::StringAttr::get(ctx, "block");
    return layouts.SharedLinearLayout(
        toStdVector(ll.getBases().lookup(kOffset)),
        toStdVector(ll.getBases().lookup(kBlock)), sharedLl.getAlignment());
  } else if (auto autoEnc = dyn_cast<gluon::AutoEncodingAttr>(layout)) {
    return layouts.AutoLayout();
  } else if (auto autoEnc = dyn_cast<gluon::CoalescedEncodingAttr>(layout)) {
    return layouts.CoalescedLayout();
  } else if (auto amdMfma = dyn_cast<ttg::AMDMfmaEncodingAttr>(layout)) {
    auto cgaBases = getCgaLayoutBases(amdMfma.getCGALayout());
    return layouts.AMDMFMALayout(
        amdMfma.getVersion(), toStdVector(amdMfma.getInstrShape()),
        amdMfma.getIsTransposed(), toStdVector(amdMfma.getWarpsPerCTA()),
        amdMfma.getElementBitWidth(), toStdVector(amdMfma.getTilesPerWarp()),
        cgaBases);
  } else if (auto amdWmma = dyn_cast<ttg::AMDWmmaEncodingAttr>(layout)) {
    auto cgaBases = getCgaLayoutBases(amdWmma.getCGALayout());
    const auto &ctaLayout = amdWmma.getCtaLayout();
    auto ctx = layout.getContext();
    auto kReg = mlir::StringAttr::get(ctx, "register");
    auto kWarp = mlir::StringAttr::get(ctx, "warp");
    return layouts.AMDWMMALayout(
        amdWmma.getVersion(), amdWmma.getIsTransposed(),
        ctaLayout.getBases().lookup(kWarp), ctaLayout.getBases().lookup(kReg),
        toStdVector(amdWmma.getInstrShape()), cgaBases, amdWmma.getRank());
  } else if (auto paddedShared =
                 dyn_cast<ttg::PaddedSharedEncodingAttr>(layout)) {
    auto *ctx = paddedShared.getContext();
    std::vector<std::pair<unsigned, unsigned>> intervalPaddingPairs;
    for (auto [interval, padding] :
         llvm::zip(paddedShared.getIntervals(), paddedShared.getPaddings())) {
      intervalPaddingPairs.push_back({interval, padding});
    }
    auto kOffset = mlir::StringAttr::get(ctx, "offset");
    auto kBlock = mlir::StringAttr::get(ctx, "block");
    const auto &ll = paddedShared.getLinearComponent();
    auto outDims = llvm::to_vector(ll.getOutDimNames());

    auto ofstBases = ll.getBases().lookup(kOffset);
    auto ofstLL = triton::LinearLayout({{kOffset, ofstBases}}, outDims);
    auto blkLL = divideLeft(ll, ofstLL);
    assert(blkLL.has_value());
    auto blkBases = blkLL->getBases().lookup(kBlock);
    auto shape = toStdVector(ll.getOutDimSizes());
    return layouts.PaddedSharedLayout(intervalPaddingPairs, ofstBases, blkBases,
                                      shape);
  } else if (auto partitioned =
                 dyn_cast<ttg::PartitionedSharedEncodingAttr>(layout)) {
    py::object partitionLayout =
        layoutToGluon(partitioned.getPartitionLayout(), isRubin);
    return layouts.PartitionedSharedLayout(
        partitioned.getNumPartitions(), partitioned.getNumGroups(),
        partitioned.getPartitionDim(), partitionLayout);
  } else if (auto tmemScales =
                 dyn_cast<ttng::TensorMemoryScalesEncodingAttr>(layout)) {
    auto cgaLayout = getCgaLayoutBases(tmemScales.getCGALayout());
    if (isRubin) {
      const char *blockRepOrder =
          tmemScales.getBlockRepOrder() ==
                  ttng::TensorMemoryScalesBlockRepOrder::K_THEN_MN
              ? "kThenMn"
              : "mnThenK";
      return layouts.RubinTensorMemoryScalesLayout(cgaLayout, blockRepOrder);
    }
    return layouts.TensorMemoryScalesLayout(cgaLayout);
  } else if (auto tmem = dyn_cast<ttng::TensorMemoryEncodingAttr>(layout)) {
    return layouts.TensorMemoryLayout(
        std::vector<unsigned>{tmem.getBlockM(), tmem.getBlockN()},
        tmem.getColStride(), getCgaLayoutBases(tmem.getCGALayout()),
        tmem.getTwoCTAs(), tmem.getFp4Padded());
  }

  throw py::value_error("Unhandled encoding encountered");
}

template <typename CondT> void check(CondT &&cond, const char *msg) {
  if (!std::forward<CondT>(cond))
    throw py::value_error(msg);
}

} // namespace

void init_gluon_ir(py::module_ &m) {
  using ret = py::rv_policy;

  py::enum_<ttng::TMEMLoadReduceModifier>(m, "TMEM_LOAD_REDUCE_MODIFIER")
      .value("MIN", ttng::TMEMLoadReduceModifier::MIN)
      .value("MAX", ttng::TMEMLoadReduceModifier::MAX)
      .export_values();

  py::class_<GluonOpBuilder, TritonOpBuilder>(m, "GluonOpBuilder",
                                              py::dynamic_attr())
      .def(py::init<MLIRContext *, std::string>(), py::arg("context"),
           py::arg("arch") = "")
      .def("get_op_builder", &GluonOpBuilder::getBuilder, ret::reference)
      .def("get_distributed_ty",
           [](GluonOpBuilder &self, Type &elementType,
              std::vector<int64_t> &shape, Attribute layout) -> Type {
             return self.getChecked<RankedTensorType>(shape, elementType,
                                                      layout);
           })
      .def("get_shared_mem_desc_ty",
           [](GluonOpBuilder &self, Type &elementType,
              std::vector<int64_t> &shape, Attribute layout,
              std::vector<int64_t> &allocShape) -> Type {
             auto ctx = self.getContext();
             return self.getChecked<ttg::MemDescType>(
                 shape, elementType, layout,
                 ttg::SharedMemorySpaceAttr::get(ctx),
                 /*mutableMemory=*/true,
                 /*allocShape=*/allocShape);
           })
      .def("get_tensor_mem_desc_ty",
           [](GluonOpBuilder &self, Type &elementType,
              std::vector<int64_t> &shape, Attribute layout,
              std::vector<int64_t> &allocShape) -> Type {
             auto ctx = self.getContext();
             return self.getChecked<ttg::MemDescType>(
                 shape, elementType, layout,
                 ttng::TensorMemorySpaceAttr::get(ctx),
                 /*mutableMemory=*/true,
                 /*allocShape=*/allocShape);
           })
      .def("get_blocked_layout",
           [](GluonOpBuilder &self, std::vector<unsigned> &sizePerThread,
              std::vector<unsigned> &threadsPerWarp,
              std::vector<unsigned> &warpsPerCta, std::vector<unsigned> &order,
              std::vector<std::vector<int32_t>> &cgaBases) -> Attribute {
             auto ctx = self.getContext();
             unsigned rank = order.size();
             auto cgaLayout = buildCgaLayoutAttr(ctx, cgaBases, rank);
             return self.getChecked<ttg::BlockedEncodingAttr>(
                 ctx, sizePerThread, threadsPerWarp, warpsPerCta, order,
                 cgaLayout);
           })
      .def("get_slice_layout",
           [](GluonOpBuilder &self, unsigned dim,
              Attribute parent) -> Attribute {
             auto ctx = self.getContext();
             auto dist = cast<ttg::DistributedEncodingTrait>(parent);
             return self.getChecked<ttg::SliceEncodingAttr>(ctx, dim, dist);
           })
      .def("get_distributed_linear_layout",
           [](GluonOpBuilder &self, std::vector<std::vector<int>> regBases,
              std::vector<std::vector<int>> laneBases,
              std::vector<std::vector<int>> warpBases,
              std::vector<std::vector<int>> blockBases,
              std::vector<int64_t> shape) -> Attribute {
             auto ctx = self.getContext();
             auto kReg = mlir::StringAttr::get(ctx, "register");
             auto kLane = mlir::StringAttr::get(ctx, "lane");
             auto kWarp = mlir::StringAttr::get(ctx, "warp");
             auto kBlock = mlir::StringAttr::get(ctx, "block");
             auto outDims = tt::standardOutDimPairs(ctx, shape);
             auto ll = tt::LinearLayout({{kReg, regBases},
                                         {kLane, laneBases},
                                         {kWarp, warpBases},
                                         {kBlock, blockBases}},
                                        outDims,
                                        /*requiresSurjective=*/true);
             if (ttg::isPermutationMatrixLayout(ll))
               return ttg::LinearEncodingAttr::get(ctx, std::move(ll));
             return ttg::GenericLinearEncodingAttr::get(ctx, std::move(ll));
           })
      .def("to_linear_layout",
           [](GluonOpBuilder &self, Attribute layout,
              std::vector<int64_t> &shape) -> py::object {
             auto ctx = self.getContext();
             auto linearLayout = ttg::toLinearLayout(shape, layout);

             if (isa<ttg::DistributedEncodingTrait>(layout)) {
               auto attr =
                   ttg::LinearEncodingAttr::get(ctx, std::move(linearLayout));
               return layoutToGluon(attr, self.isRubin());
             }
             if (isa<ttg::SharedEncodingTrait>(layout)) {
               auto alignment =
                   cast<ttg::SharedEncodingTrait>(layout).getAlignment();
               auto attr = ttg::SharedLinearEncodingAttr::get(
                   ctx, std::move(linearLayout), alignment);
               return layoutToGluon(attr, self.isRubin());
             }

             // TensorMemory encodings: keep the LinearLayout but wrap as
             // print-only Python object carrying row/col bases -> dim0/dim1.
             auto inNamesRange = linearLayout.getInDimNames();
             auto inNames = llvm::to_vector(inNamesRange);
             bool isTmemLayout =
                 ((inNames.size() == 2 || inNames.size() == 3) &&
                  inNames[0].str() == "row" && inNames[1].str() == "col" &&
                  (inNames.size() == 2 || inNames[2].str() == "block"));
             if (!isTmemLayout)
               throw std::invalid_argument(
                   "Unsupported layout in to_linear_layout");

             // Build Py _TensorMemoryLinearLayout(row_bases, col_bases, shape,
             // repr)
             py::object tmemCls =
                 py::module_::import_(
                     "triton.experimental.gluon.language.nvidia.blackwell")
                     .attr("_TensorMemoryLinearLayout");
             auto bases = linearLayout.getBases();
             auto rowBases = bases[mlir::StringAttr::get(ctx, "row")];
             auto colBases = bases[mlir::StringAttr::get(ctx, "col")];
             auto outDims = linearLayout.getOutDims();
             std::vector<int> shapeVec;
             for (auto &od : outDims)
               shapeVec.push_back(od.second);

             py::object pyObj = tmemCls(py::cast(rowBases), py::cast(colBases),
                                        py::cast(shapeVec));
             return pyObj;
           })
      .def("get_dot_operand_layout",
           [](GluonOpBuilder &self, unsigned opIdx, Attribute parent,
              unsigned kWidth) -> Attribute {
             return self.getChecked<ttg::DotOperandEncodingAttr>(
                 self.getContext(), opIdx, parent, kWidth);
           })
      .def("get_mma_layout",
           [](GluonOpBuilder &self, std::vector<unsigned> &version,
              std::vector<unsigned> &warpsPerCta,
              std::vector<std::vector<int32_t>> &cgaBases,
              std::vector<unsigned> &instrShape) -> Attribute {
             auto ctx = self.getContext();
             unsigned rank = warpsPerCta.size();
             auto cgaLayout = buildCgaLayoutAttr(ctx, cgaBases, rank);
             return self.getChecked<ttg::NvidiaMmaEncodingAttr>(
                 ctx, version[0], version[1], warpsPerCta, cgaLayout,
                 instrShape);
           })
      .def("get_amd_mfma_layout",
           [](GluonOpBuilder &self, unsigned version,
              std::vector<unsigned> &warpsPerCta,
              std::vector<unsigned> &instrShape, bool transposed,
              std::vector<std::vector<int32_t>> &cgaBases,
              std::vector<unsigned> &tilesPerWarp,
              unsigned elementBitWidth) -> Attribute {
             auto ctx = self.getContext();
             unsigned rank = warpsPerCta.size();
             auto cgaLayout = buildCgaLayoutAttr(ctx, cgaBases, rank);
             return ttg::AMDMfmaEncodingAttr::get(
                 ctx, version, warpsPerCta, instrShape, transposed, cgaLayout,
                 tilesPerWarp, elementBitWidth);
           })
      .def("get_amd_wmma_layout",
           [](GluonOpBuilder &self, unsigned version, bool transposed,
              std::vector<std::vector<int32_t>> &warpBases,
              std::vector<std::vector<int32_t>> &regBases,
              std::vector<std::vector<int32_t>> &cgaBases,
              std::vector<unsigned> &instrShape, unsigned rank) -> Attribute {
             auto ctx = self.getContext();
             auto kReg = mlir::StringAttr::get(ctx, "register");
             auto kWarp = mlir::StringAttr::get(ctx, "warp");
             auto ctaLayout =
                 tt::LinearLayout({{kReg, regBases}, {kWarp, warpBases}},
                                  tt::standardOutDimNames(ctx, rank));
             auto cgaLayout = buildCgaLayoutAttr(ctx, cgaBases, rank);
             return ttg::AMDWmmaEncodingAttr::get(
                 ctx, version, ctaLayout, transposed, cgaLayout, instrShape);
           })
      .def("get_padded_shared_layout",
           [](GluonOpBuilder &self, std::vector<unsigned> &intervals,
              std::vector<unsigned> &paddings,
              std::vector<std::vector<int>> &offsetBases,
              std::vector<std::vector<int>> &blockBases,
              std::vector<int64_t> &shape) -> Attribute {
             auto ctx = self.getContext();
             auto rank = shape.size();
             auto kOffset = mlir::StringAttr::get(ctx, "offset");
             auto kBlock = mlir::StringAttr::get(ctx, "block");
             auto outDims = tt::standardOutDimNames(ctx, rank);
             auto ll = tt::LinearLayout({{kOffset, offsetBases}}, outDims) *
                       tt::LinearLayout({{kBlock, blockBases}}, outDims);
             return ttg::PaddedSharedEncodingAttr::get(ctx, intervals, paddings,
                                                       std::move(ll));
           })
      .def("get_shared_linear_layout",
           [](GluonOpBuilder &self, std::vector<std::vector<int>> &offsetBases,
              std::vector<std::vector<int>> &blockBases,
              unsigned alignment) -> Attribute {
             auto ctx = self.getContext();
             auto kOffset = mlir::StringAttr::get(ctx, "offset");
             auto kBlock = mlir::StringAttr::get(ctx, "block");
             auto outDims = tt::standardOutDimNames(ctx, offsetBases[0].size());
             auto ll = tt::LinearLayout(
                 {{kOffset, offsetBases}, {kBlock, blockBases}}, outDims);
             return self.getChecked<ttg::SharedLinearEncodingAttr>(
                 ctx, std::move(ll), alignment);
           })
      .def("get_nvmma_shared_layout",
           [](GluonOpBuilder &self, unsigned swizzleByteWidth,
              unsigned elementBitwidth, bool transposed, bool fp4Padded,
              std::vector<std::vector<int32_t>> &cgaBases,
              unsigned rank) -> Attribute {
             auto ctx = self.getContext();
             auto cgaLayout = buildCgaLayoutAttr(ctx, cgaBases, rank);
             return self.getChecked<ttg::NVMMASharedEncodingAttr>(
                 ctx, swizzleByteWidth, transposed, elementBitwidth, fp4Padded,
                 cgaLayout);
           })
      .def("get_auto_layout",
           [](GluonOpBuilder &self) -> Attribute {
             return self.getChecked<gluon::AutoEncodingAttr>(self.getContext());
           })
      .def("get_coalesced_layout",
           [](GluonOpBuilder &self) -> Attribute {
             return self.getChecked<gluon::CoalescedEncodingAttr>(
                 self.getContext());
           })
      .def("get_swizzled_shared_layout",
           [](GluonOpBuilder &self, int vec, int perPhase, int maxPhase,
              std::vector<unsigned> &order,
              std::vector<std::vector<int32_t>> &cgaBases) -> Attribute {
             auto ctx = self.getContext();
             unsigned rank = order.size();
             auto cgaLayout = buildCgaLayoutAttr(ctx, cgaBases, rank);
             return self.getChecked<ttg::SwizzledSharedEncodingAttr>(
                 ctx, vec, perPhase, maxPhase, order, cgaLayout);
           })
      .def("get_partitioned_shared_layout",
           [](GluonOpBuilder &self, unsigned numPartitions, unsigned numGroups,
              unsigned partitionDim, Attribute partitionLayout) -> Attribute {
             auto ctx = self.getContext();
             auto sharedLayout =
                 cast<ttg::SharedEncodingTrait>(partitionLayout);
             return self.getChecked<ttg::PartitionedSharedEncodingAttr>(
                 ctx, numPartitions, numGroups, partitionDim, sharedLayout);
           })
      .def("get_tensor_memory_layout",
           [](GluonOpBuilder &self, std::vector<unsigned> &block,
              unsigned colStride, std::vector<std::vector<int32_t>> &cgaBases,
              bool twoCTAs, bool fp4Padded) -> Attribute {
             auto ctx = self.getContext();
             check(block.size() == 2, "expected a 2D block");
             auto cgaLayout = buildCgaLayoutAttr(ctx, cgaBases, /*rank=*/2);
             return self.getChecked<ttng::TensorMemoryEncodingAttr>(
                 ctx, block[0], block[1], colStride, cgaLayout, twoCTAs,
                 fp4Padded);
           })
      .def(
          "get_tensor_memory_scales_layout",
          [](GluonOpBuilder &self, std::vector<std::vector<int32_t>> &cgaBases,
             const std::string &blockRepOrder) -> Attribute {
            auto ctx = self.getContext();
            auto cgaLayout = buildCgaLayoutAttr(ctx, cgaBases, /*rank=*/2);
            auto repOrder =
                blockRepOrder == "mnThenK"
                    ? ttng::TensorMemoryScalesBlockRepOrder::MN_THEN_K
                : blockRepOrder == "kThenMn"
                    ? ttng::TensorMemoryScalesBlockRepOrder::K_THEN_MN
                    : throw py::value_error(
                          "block_rep_order must be either 'mnThenK' or "
                          "'kThenMn'");
            return self.getChecked<ttng::TensorMemoryScalesEncodingAttr>(
                ctx, cgaLayout, repOrder);
          },
          py::arg("cga_bases"), py::arg("block_rep_order") = "mnThenK")
      .def("get_shape_from_tensor",
           [](GluonOpBuilder &self, Value tensor) -> std::vector<int64_t> {
             auto ty = dyn_cast<RankedTensorType>(tensor.getType());
             return ty.getShape();
           })
      .def("get_gluon_layout_from_tensor",
           [](GluonOpBuilder &self, Value tensor) -> py::object {
             auto ty = dyn_cast<RankedTensorType>(tensor.getType());
             check(ty.getEncoding(), "expected a tensor with an encoding");
             return layoutToGluon(ty.getEncoding(), self.isRubin());
           })
      .def("get_gluon_layout_from_memdesc",
           [](GluonOpBuilder &self, Value memdesc) -> py::object {
             auto ty = dyn_cast<ttg::MemDescType>(memdesc.getType());
             check(ty.getEncoding(), "expected a memdesc with an encoding");
             return layoutToGluon(ty.getEncoding(), self.isRubin());
           })
      .def("get_tensor_descriptor_layout_type",
           [](GluonOpBuilder &self, Type blockType, bool isSigned,
              Attribute layout) -> Type {
             auto blockTy = cast<RankedTensorType>(blockType);
             return triton::TensorDescType::get(blockTy.getShape(),
                                                blockTy.getElementType(),
                                                layout, isSigned);
           })
      .def("get_tensor_descriptor_im2col_layout_type",
           [](GluonOpBuilder &self, Type blockType, bool isSigned,
              Attribute layout) -> Type {
             auto blockTy = cast<RankedTensorType>(blockType);
             return triton::nvidia_gpu::TensorDescIm2ColType::get(
                 blockTy.getShape(), blockTy.getElementType(), layout,
                 isSigned);
           })
      .def("is_convert_layout_trivial",
           [](GluonOpBuilder &self, Type resultTy, Value value) -> bool {
             auto dstTy = cast<RankedTensorType>(resultTy);
             return isConvertLayoutTrivial(dstTy, value);
           })
      .def(
          "create_histogram",
          [](GluonOpBuilder &self, Value operand, int numBins,
             std::optional<Value> mask, Attribute layout) -> Value {
            auto *ctx = self.getContext();
            auto resultTy =
                RankedTensorType::get({static_cast<int64_t>(numBins)},
                                      IntegerType::get(ctx, 32), layout);
            if (!mask) {
              return self.create<triton::HistogramOp>(resultTy, operand);
            } else {
              return self.create<triton::HistogramOp>(resultTy, operand, *mask);
            }
          },
          py::arg("operand"), py::arg("numBins"), py::arg("mask").none(),
          py::arg("layout"))
      .def("create_cat",
           [](GluonOpBuilder &self, Value &lhs, Value &rhs,
              Type retType) -> Value {
             return self.create<triton::CatOp>(retType, lhs, rhs);
           })
      .def("create_experimental_fpsan_embed",
           [](GluonOpBuilder &self, Value &src, Type &dstType) -> Value {
             return self.create<triton::instrument::ExperimentalFPSanEmbedOp>(
                 dstType, src);
           })
      .def("create_experimental_fpsan_unembed",
           [](GluonOpBuilder &self, Value &src, Type &dstType) -> Value {
             return self.create<triton::instrument::ExperimentalFPSanUnembedOp>(
                 dstType, src);
           })
      .def("create_packed_arith",
           [](GluonOpBuilder &self, Type resultType,
              const std::string &operation,
              std::vector<Value> operands) -> Value {
             auto kind = ttng::symbolizePackedArithOpKind(operation);
             check(kind.has_value(), "unknown packed arithmetic operation");
             auto resultTensorType = dyn_cast<RankedTensorType>(resultType);
             if (!resultTensorType) {
               auto operandType =
                   cast<RankedTensorType>(operands.front().getType());
               auto inferred = ttg::inferFp4ToFpResultType(
                   operandType, resultType, operandType.getRank() - 1,
                   self.getLastLoc());
               check(succeeded(inferred), "cannot infer packed FP4 layout");
               resultTensorType = *inferred;
             }
             return self.create<ttng::PackedArithOp>(resultTensorType, *kind,
                                                     operands);
           })
      .def("create_fp4_to_fp",
           [](GluonOpBuilder &self, Value src, Type elemType,
              int axis) -> Value {
             return self.create<ttg::Fp4ToFpOp>(
                 cast<TypedValue<RankedTensorType>>(src), elemType, axis);
           })
      .def("create_async_copy_global_to_local",
           [](GluonOpBuilder &self, Value smem, Value pointer, Value mask,
              Value other, tt::CacheModifier cacheModifier,
              tt::EvictionPolicy evictionPolicy, bool isVolatile) {
             self.create<ttg::AsyncCopyGlobalToLocalOp>(
                 pointer, smem, mask, other, cacheModifier, evictionPolicy,
                 isVolatile);
           })
      .def("create_async_copy_local_to_global",
           [](GluonOpBuilder &self, Value smem, Value pointer, Value mask,
              tt::CacheModifier cacheModifier,
              tt::EvictionPolicy evictionPolicy) {
             self.create<ttag::AsyncCopyLocalToGlobalOp>(
                 smem, pointer, mask, cacheModifier, evictionPolicy);
           })
      .def("create_async_copy_mbarrier_arrive",
           [](GluonOpBuilder &self, Value mbarrier, bool incrementCount) {
             self.create<ttng::AsyncCopyMbarrierArriveOp>(mbarrier,
                                                          !incrementCount);
           })
      .def("create_async_commit_group",
           [](GluonOpBuilder &self) {
             ValueRange tokens;
             self.create<ttg::AsyncCommitGroupOp>(tokens);
           })
      .def("create_async_wait_group",
           [](GluonOpBuilder &self, int num) {
             ValueRange tokens;
             self.create<ttg::AsyncWaitOp>(tokens, num);
           })
      .def("create_convert_layout",
           [](GluonOpBuilder &self, Type resultTy, Value value) -> Value {
             return self.create<ttg::ConvertLayoutOp>(resultTy, value);
           })
      .def("create_local_alloc",
           [](GluonOpBuilder &self, Type resultTy) -> Value {
             return self.create<ttg::LocalAllocOp>(resultTy);
           })
      .def("create_local_alloc",
           [](GluonOpBuilder &self, Type resultTy, Value value) -> Value {
             return self.create<ttg::LocalAllocOp>(resultTy, value);
           })
      .def("create_local_store",
           [](GluonOpBuilder &self, Value memDesc, Value value) {
             self.create<ttg::LocalStoreOp>(value, memDesc);
           })
      .def(
          "create_async_shared_store",
          [](GluonOpBuilder &self, Value memDesc, Value value, Value mbarrier) {
            self.create<ttng::AsyncSharedStoreOp>(value, memDesc, mbarrier);
          })
      .def("create_local_load",
           [](GluonOpBuilder &self, Type resultTy, Value memDesc) -> Value {
             return self.create<ttg::LocalLoadOp>(resultTy, memDesc);
           })
      .def("create_local_gather",
           [](GluonOpBuilder &self, Type resultTy, Value memDesc, Value indices,
              int32_t axis) -> Value {
             auto ctx = self.getContext();
             auto i32Ty = IntegerType::get(ctx, 32);
             auto axisAttr = IntegerAttr::get(i32Ty, axis);
             return self.create<ttg::LocalGatherOp>(resultTy, memDesc, indices,
                                                    axisAttr);
           })
      .def("create_local_scatter",
           [](GluonOpBuilder &self, Value memDesc, Value values, Value indices,
              int32_t axis) {
             auto ctx = self.getContext();
             auto i32Ty = IntegerType::get(ctx, 32);
             auto axisAttr = IntegerAttr::get(i32Ty, axis);
             self.create<ttg::LocalScatterOp>(memDesc, values, indices,
                                              axisAttr);
           })
      .def(
          "create_local_atomic_scatter_rmw",
          [](GluonOpBuilder &self, tt::RMWOp rmwOp, Value memDesc, Value values,
             Value indices, std::optional<Value> mask, int32_t axis) -> Value {
            return self.create<ttg::LocalAtomicScatterRMWOp>(
                rmwOp, memDesc, values, indices, mask.value_or(Value()), axis);
          },
          py::arg("rmwOp"), py::arg("memDesc"), py::arg("values"),
          py::arg("indices"), py::arg("mask").none(), py::arg("axis"))
      .def("get_shared_bank_conflicts",
           [](GluonOpBuilder &self, Attribute regLayoutAttr,
              Attribute sharedLayoutAttr, std::vector<int64_t> &shape,
              int bitwidth) -> int {
             auto regLayout = ttg::toLinearLayout(shape, regLayoutAttr);
             auto smemLayout = ttg::toLinearLayout(shape, sharedLayoutAttr);
             auto mod = self.getBuilder()
                            .getInsertionBlock()
                            ->getParentOp()
                            ->getParentOfType<ModuleOp>();
             auto arch = getAMDArch(mod);
             if (!arch.has_value())
               return ttg::bankConflictsMemDesc(regLayout, smemLayout,
                                                bitwidth);
             tt::AMD::TargetInfo targetInfo(arch->str());
             int numBanks = targetInfo.getSharedMemoryBanks();
             auto vecBitwidth = std::max<int32_t>(
                 32, smemLayout.getInDimSize(StringAttr::get(
                         smemLayout.getInDimNames().begin()->getContext(),
                         "vector")) *
                         bitwidth);
             auto [dstTile, srcTile] =
                 targetInfo.getSharedLdStTiles(vecBitwidth);
             (void)srcTile;
             assert(srcTile.laneAddr.empty() &&
                    "srcTile.laneAddr should be empty");
             return ttg::bankConflictsMemDesc(regLayout, smemLayout, bitwidth,
                                              numBanks, dstTile);
           })
      .def(
          "create_local_dealloc",
          [](GluonOpBuilder &self, Value memDesc) -> Operation * {
            return self.create<ttg::LocalDeallocOp>(memDesc);
          },
          ret::reference)

      .def("create_memdesc_index",
           [](GluonOpBuilder &self, Type resultType, Value src,
              Value index) -> Value {
             return self.create<ttg::MemDescIndexOp>(resultType, src, index);
           })
      .def("create_memdesc_subslice",
           [](GluonOpBuilder &self, Type resultType, Value src,
              std::vector<int32_t> &offsets) -> Value {
             return self.create<ttg::MemDescSubsliceOp>(resultType, src,
                                                        offsets);
           })
      .def("create_memdesc_trans",
           [](GluonOpBuilder &self, Value src,
              std::vector<int> &order) -> Value {
             return self.create<ttg::MemDescTransOp>(src, order);
           })
      .def("create_memdesc_reshape",
           [](GluonOpBuilder &self, Value src,
              std::vector<int64_t> &shape) -> Value {
             return self.create<ttg::MemDescReshapeOp>(src, shape);
           })
      .def("create_memdesc_reinterpret",
           [](GluonOpBuilder &self, Type resultType, Value src) -> Value {
             return self.create<ttg::MemDescReinterpretOp>(resultType, src);
           })
      .def("create_set_auto_layout",
           [](GluonOpBuilder &self, Attribute layout, Value value) -> Value {
             return self.create<gluon::SetAutoLayoutOp>(layout, value);
           })
      .def("create_split",
           [](GluonOpBuilder &self, Value &a) -> py::tuple {
             auto argTy = cast<RankedTensorType>(a.getType());
             auto enc = argTy.getEncoding();
             triton::SplitOp op;
             if (isa<gluon::AutoEncodingAttr, gluon::CoalescedEncodingAttr>(
                     enc)) {
               op = self.create<triton::SplitOp>(a);
             } else {
               auto ctx = argTy.getContext();
               auto sliceEnc = ttg::SliceEncodingAttr::get(
                   ctx, argTy.getRank() - 1,
                   cast<ttg::DistributedEncodingTrait>(enc));
               auto resTy =
                   RankedTensorType::get(ArrayRef(argTy.getShape()).drop_back(),
                                         argTy.getElementType(), sliceEnc);
               op = self.create<triton::SplitOp>(TypeRange{resTy, resTy}, a);
             }
             return py::make_tuple(op->getResult(0), op->getResult(1));
           })
      .def("create_warpgroup_mma",
           [](GluonOpBuilder &self, Value a, Value b, Value acc, Value useAcc,
              triton::InputPrecision precision = triton::InputPrecision::IEEE,
              int maxNumImpreciseAcc = 0, bool isAsync = false) -> Value {
             return self.create<ttng::WarpGroupDotOp>(
                 a, b, acc, useAcc, precision, maxNumImpreciseAcc, isAsync);
           })
      .def("create_warpgroup_mma_wait",
           [](GluonOpBuilder &self, std::vector<Value> &deps, int pendings) {
             std::vector<Value> results;
             auto wait = self.create<ttng::WarpGroupDotWaitOp>(deps, pendings);
             llvm::append_range(results, wait.getResults());
             return results;
           })
      .def(
          "create_tmem_alloc",
          [](GluonOpBuilder &self, Type resultTy,
             std::optional<Value> value) -> Value {
            return self.create<ttng::TMEMAllocOp>(resultTy,
                                                  value.value_or(Value{}));
          },
          py::arg("resultTy"), py::arg("value").none())
      .def("create_tmem_store",
           [](GluonOpBuilder &self, Value memDesc, Value value, Value pred) {
             self.create<ttng::TMEMStoreOp>(memDesc, value, pred);
           })
      .def(
          "create_tmem_load",
          [](GluonOpBuilder &self, Type resultTy, Value memDesc,
             std::optional<ttng::TMEMLoadReduceModifier> redOp, bool useAbs,
             tt::PropagateNan propagateNan) -> py::object {
            ttng::TMEMLoadReduceModifierAttr redOpAttr = nullptr;
            BoolAttr absAttr = nullptr;
            BoolAttr nanAttr = nullptr;

            if (redOp) {
              redOpAttr = ttng::TMEMLoadReduceModifierAttr::get(
                  self.getContext(), redOp.value());
              if (useAbs)
                absAttr = self.getBuilder().getBoolAttr(true);
              if (propagateNan != tt::PropagateNan::NONE)
                nanAttr = self.getBuilder().getBoolAttr(true);
            }

            auto op = self.create<ttng::TMEMLoadOp>(
                resultTy, /*token=*/Type(), memDesc, /*dep=*/Value(), redOpAttr,
                absAttr, nanAttr);

            if (redOp) {
              Value result = op.getResult();
              Value red = op.getRed();
              auto redTy = cast<RankedTensorType>(red.getType());
              py::object redLayout =
                  layoutToGluon(redTy.getEncoding(), self.isRubin());
              return py::make_tuple(result, red, redLayout);
            }
            Value result = op.getResult();
            return py::cast(result);
          },
          py::arg("resultTy"), py::arg("memDesc"),
          (py::arg("redOp").none() = py::none()), py::arg("useAbs") = false,
          py::arg("propagateNan") = tt::PropagateNan::NONE)
      .def("create_tmem_copy",
           [](GluonOpBuilder &self, Value src, Value dst) {
             self.create<ttng::TMEMCopyOp>(src, dst);
           })
      .def("create_tmem_subslice",
           [](GluonOpBuilder &self, Type resultTy, Value memDesc, int offset,
              int dim) -> Value {
             return self.create<ttng::TMEMSubSliceOp>(resultTy, memDesc, offset,
                                                      dim);
           })
      .def("create_mbarrier_init",
           [](GluonOpBuilder &self, Value memDesc, int count) {
             self.create<ttng::InitBarrierOp>(memDesc, count);
           })
      .def("create_mbarrier_inval",
           [](GluonOpBuilder &self, Value memDesc) {
             self.create<ttng::InvalBarrierOp>(memDesc);
           })
      .def(
          "create_mbarrier_expect",
          [](GluonOpBuilder &self, Value memDesc, int bytes, Value pred,
             std::optional<int> fromCTA) {
            IntegerAttr fromCTAAttr =
                fromCTA ? self.getBuilder().getI32IntegerAttr(*fromCTA)
                        : IntegerAttr();
            self.create<ttng::BarrierExpectOp>(memDesc, bytes, pred,
                                               fromCTAAttr);
          },
          py::arg("memDesc"), py::arg("bytes"), py::arg("pred"),
          (py::arg("from_cta").none() = py::none()))
      .def("create_mbarrier_wait",
           [](GluonOpBuilder &self, Value memDesc, Value phase, Value pred,
              std::vector<Value> &deps) {
             self.create<ttng::WaitBarrierOp>(memDesc, phase, pred, deps);
           })
      .def(
          "create_mbarrier_arrive",
          [](GluonOpBuilder &self, Value memDesc, uint32_t count, Value pred,
             std::optional<int> fromCTA, uint32_t multicastCTA) {
            IntegerAttr fromCTAAttr =
                fromCTA ? self.getBuilder().getI32IntegerAttr(*fromCTA)
                        : IntegerAttr();
            self.create<ttng::ArriveBarrierOp>(memDesc, count, pred,
                                               fromCTAAttr, multicastCTA);
          },
          py::arg("memDesc"), py::arg("count"), py::arg("pred"),
          (py::arg("from_cta").none() = py::none()),
          py::arg("multicast_cta") = 0)
      .def(
          "create_cluster_barrier",
          [](GluonOpBuilder &self,
             bool relaxed) { self.create<ttng::ClusterBarrierOp>(relaxed); },
          py::arg("relaxed") = false)
      // CLC (Cluster Launch Control) ops - SM100+
      .def("create_clc_try_cancel",
           [](GluonOpBuilder &self, Value result, Value mbarrier) {
             self.create<ttng::CLCTryCancelOp>(result, mbarrier);
           })
      .def("create_clc_load_result",
           [](GluonOpBuilder &self, Value result) -> Value {
             return self.create<ttng::CLCLoadResultOp>(result);
           })
      .def("create_clc_is_canceled",
           [](GluonOpBuilder &self, Value clcResult) -> Value {
             return self.create<ttng::CLCIsCanceledOp>(clcResult);
           })
      .def("create_clc_get_program_id",
           [](GluonOpBuilder &self, Value clcResult, int dim) -> Value {
             return self.create<ttng::CLCGetProgramIdOp>(clcResult, dim);
           })
      .def("create_tcgen05_mma",
           [](GluonOpBuilder &self, Value a, Value b, Value acc, Value useAcc,
              Value pred, std::vector<Value> &mbarriers,
              std::vector<Value> &mbarrier_preds, bool two_ctas,
              bool multicast) {
             Value accDep;
             auto tokType = self.getBuilder().getType<ttg::AsyncTokenType>();
             self.create<ttng::TCGen5MMAOp>(tokType, a, b, acc, accDep, useAcc,
                                            pred, two_ctas, multicast,
                                            mbarriers, mbarrier_preds);
           })
      .def("create_tcgen05_mma_scaled",
           [](GluonOpBuilder &self, Value a, Value b, Value acc, Value aScale,
              Value bScale, tt::ScaleDotElemType aType,
              tt::ScaleDotElemType bType, Value useAcc, Value pred,
              std::vector<Value> &mbarriers, std::vector<Value> &mbarrier_preds,
              bool two_ctas, bool multicast) {
             Value accDep;
             auto tokType = self.getBuilder().getType<ttg::AsyncTokenType>();
             self.create<ttng::TCGen5MMAScaledOp>(
                 tokType, a, b, acc, accDep, aScale, bScale, aType, bType,
                 useAcc, pred, mbarriers, mbarrier_preds, two_ctas,
                 /*isAsync=*/false, multicast);
           })
      .def("create_tcgen05_commit",
           [](GluonOpBuilder &self, Value &barrier, Value &pred,
              std::vector<Value> &descs) {
             self.create<ttng::TCGen5CommitOp>(barrier, pred, descs);
           })

      .def(
          "create_async_tma_copy_global_to_local",
          [](GluonOpBuilder &self, Value descPtr, std::vector<Value> &coord,
             Value barrier, Value result, Value pred, bool multicast,
             std::optional<std::vector<Value>> offsets) {
            multicast &=
                ttng::hasCGABroadcast(cast<ttg::MemDescType>(result.getType()));
            ValueRange offsetsRange =
                offsets.has_value() ? ValueRange(*offsets) : ValueRange{};
            self.create<ttng::AsyncTMACopyGlobalToLocalOp>(
                descPtr, coord, offsetsRange, barrier, result, pred, multicast);
          },
          py::arg("descPtr"), py::arg("coord"), py::arg("barrier"),
          py::arg("result"), py::arg("pred"), py::arg("multicast"),
          py::arg("offsets").none())
      .def("create_async_tma_copy_local_to_global",
           [](GluonOpBuilder &self, Value descPtr, std::vector<Value> &coord,
              Value src) {
             self.create<ttng::AsyncTMACopyLocalToGlobalOp>(descPtr, coord,
                                                            src);
           })
      .def("create_async_tma_reduce",
           [](GluonOpBuilder &self, triton::DescriptorReduceKind kind,
              Value descPtr, std::vector<Value> &coord, Value src) {
             self.create<ttng::AsyncTMAReduceOp>(kind, descPtr, coord, src);
           })
      .def("create_async_tma_store_wait",
           [](GluonOpBuilder &self, int pendings, bool readOnly) {
             self.create<ttng::TMAStoreWaitOp>(pendings, readOnly);
           })
      .def(
          "create_async_tma_gather",
          [](GluonOpBuilder &self, Value descPtr, Value xOffsets, Value yOffset,
             Value barrier, Value result, Value pred, bool multicast) {
            multicast &=
                ttng::hasCGABroadcast(cast<ttg::MemDescType>(result.getType()));
            self.create<ttng::AsyncTMAGatherOp>(
                descPtr, xOffsets, yOffset, barrier, result, pred, multicast);
          })
      .def("create_async_tma_scatter",
           [](GluonOpBuilder &self, Value descPtr, Value xOffsets,
              Value yOffset, Value src) {
             self.create<ttng::AsyncTMAScatterOp>(descPtr, xOffsets, yOffset,
                                                  src);
           })
      .def("create_fence_async_shared",
           [](GluonOpBuilder &self, bool bCluster) -> OpState {
             return self.create<ttng::FenceAsyncSharedOp>(bCluster);
           })

      .def("create_broadcast",
           [](TritonOpBuilder &self, Value &arg, Type retTy) -> Value {
             return self.create<tt::BroadcastOp>(retTy, arg);
           })
      .def(
          "create_warp_return",
          [](GluonOpBuilder &self)
              -> Operation * { return self.create<ttg::WarpReturnOp>(); },
          ret::reference)
      .def(
          "create_warp_yield",
          [](GluonOpBuilder &self, std::vector<Value> &values) -> Operation * {
            return self.create<ttg::WarpYieldOp>(values);
          },
          ret::reference)
      .def(
          "create_warp_specialize_partitions",
          [](GluonOpBuilder &self, std::vector<Value> &explicitCaptures,
             int numPartitions) -> Operation * {
            return self.create<ttg::WarpSpecializePartitionsOp>(
                explicitCaptures, numPartitions);
          },
          ret::reference)
      .def("create_warp_specialize",
           [](GluonOpBuilder &self, std::vector<Type> &resultTypes,
              std::vector<int> &partitionNumWarps) {
             return self.create<ttg::WarpSpecializeOp>(resultTypes,
                                                       partitionNumWarps);
           })
      .def("create_buffer_load",
           [](GluonOpBuilder &self, Type resultType, Value ptr, Value offsets,
              Value mask, Value other, tt::CacheModifier cache) -> Value {
             return self.create<ttag::BufferLoadOp>(resultType, ptr, offsets,
                                                    Value() /*stride*/, cache,
                                                    mask, other);
           })
      .def("create_buffer_store",
           [](GluonOpBuilder &self, Value storedValue, Value ptr, Value offsets,
              Value mask, tt::CacheModifier cache) {
             self.create<ttag::BufferStoreOp>(storedValue, ptr, offsets,
                                              Value() /*stride*/, cache, mask);
           })
      .def("create_buffer_atomic_rmw",
           [](GluonOpBuilder &self, tt::RMWOp op, Value ptr, Value offsets,
              Value value, tt::MemSemantic sem, tt::MemSyncScope scope,
              Value mask) -> Value {
             return self.create<ttag::BufferAtomicRMWOp>(
                 value.getType(), op, ptr, offsets, value, Value() /*stride*/,
                 sem, scope, mask);
           })
      .def("create_buffer_load_to_local",
           [](GluonOpBuilder &self, Value dest, Value ptr, Value offsets,
              Value mask, Value other, Value stride,
              tt::CacheModifier cacheModifier) {
             self.create<ttag::BufferLoadToLocalOp>(
                 dest, ptr, offsets, mask, other, stride, cacheModifier);
           })
      .def("create_local_load_packed_transposed",
           [](GluonOpBuilder &self, Type resultType, Value memDesc) -> Value {
             return self.create<ttag::LocalLoadPackedTransposedOp>(resultType,
                                                                   memDesc);
           })
      .def("create_scaled_upcast_fp4",
           [](GluonOpBuilder &self, Value input, Value scale, Type elemType,
              int axis) -> Value {
             return self.create<ttag::ScaledUpcastFp4Op>(input, scale, elemType,
                                                         axis);
           })
      .def("create_scaled_upcast_fp8",
           [](GluonOpBuilder &self, Type resultType, Value input,
              Value scale) -> Value {
             return self.create<ttag::ScaledUpcastFp8Op>(resultType, input,
                                                         scale);
           })
      .def("create_extract_slice",
           [](GluonOpBuilder &self, Type resultType, Value source,
              std::vector<int64_t> &offsets) -> Value {
             auto offsetsAttr = self.getBuilder().getDenseI64ArrayAttr(offsets);
             return self.create<ttag::ExtractSliceOp>(resultType, source,
                                                      offsetsAttr);
           })
      .def("create_make_tensor_descriptor",
           [](TritonOpBuilder &self, Type resultTy, Value &base,
              std::vector<Value> &shape, std::vector<Value> &strides,
              tt::PaddingOption paddingOption) -> Value {
             return self.create<tt::MakeTensorDescOp>(resultTy, base, shape,
                                                      strides, paddingOption);
           })
      .def(
          "create_async_tdm_copy_global_to_local",
          [](GluonOpBuilder &self, Value descPtr, Value result, Value barrier,
             tt::CacheModifier cacheModifier,
             std::optional<uint32_t> warpUsedHint) {
            IntegerAttr hintAttr;
            if (warpUsedHint.has_value())
              hintAttr = self.getBuilder().getI32IntegerAttr(
                  static_cast<int32_t>(*warpUsedHint));
            self.create<ttag::AsyncTDMCopyGlobalToLocalOp>(
                descPtr, result, barrier, cacheModifier, hintAttr);
          },
          py::arg("descPtr"), py::arg("result"), py::arg("barrier"),
          py::arg("cacheModifier"),
          (py::arg("warpUsedHint").none() = py::none()))
      .def(
          "create_async_tdm_fused_copy_global_to_local",
          [](GluonOpBuilder &self, std::vector<Value> &descs,
             std::vector<Value> &dests, std::vector<int32_t> &warpUsedHints,
             tt::CacheModifier cacheModifier) {
            auto tokType = self.getBuilder().getType<ttg::AsyncTokenType>();
            auto hintAttr =
                self.getBuilder().getDenseI32ArrayAttr(warpUsedHints);
            self.create<ttag::AsyncTDMFusedCopyGlobalToLocalOp>(
                tokType, descs, dests, hintAttr, cacheModifier);
          },
          py::arg("descs"), py::arg("dests"), py::arg("warpUsedHints"),
          py::arg("cacheModifier") = tt::CacheModifier::NONE)
      .def("create_async_tdm_copy_local_to_global",
           [](GluonOpBuilder &self, Value descPtr, Value src, Value barrier,
              tt::CacheModifier cacheModifier) {
             self.create<ttag::AsyncTDMCopyLocalToGlobalOp>(
                 descPtr, src, barrier, cacheModifier);
           })
      .def("create_update_tensor_descriptor",
           [](GluonOpBuilder &self, Value descPtr,
              std::vector<Value> &addOffsets, std::vector<Value> &setBounds,
              Value pred, bool clampBounds) -> Value {
             Value res = self.create<ttag::UpdateTensorDescriptorOp>(
                 descPtr.getType(), descPtr, ValueRange(addOffsets),
                 ValueRange(setBounds), pred);
             if (clampBounds)
               res.getDefiningOp()->setAttr("clamp_bounds",
                                            self.getBuilder().getUnitAttr());
             return res;
           })
      .def("create_async_tdm_scatter",
           [](GluonOpBuilder &self, Value descPtr, Value dstRowIndices,
              Value src, Value barrier) {
             self.create<ttag::AsyncTDMScatterOp>(descPtr, dstRowIndices, src,
                                                  barrier);
           })
      .def("create_async_tdm_gather",
           [](GluonOpBuilder &self, Value descPtr, Value srcRowIndices,
              Value dst, Value barrier) {
             self.create<ttag::AsyncTDMGatherOp>(descPtr, srcRowIndices, dst,
                                                 barrier);
           })
      .def("create_tdm_prefetch",
           [](GluonOpBuilder &self, Value descPtr, std::vector<Value> &indices,
              Value pred, bool speculative, bool returnOffsets) -> Value {
             auto op = self.create<ttag::TDMPrefetchOp>(
                 descPtr, indices, pred, speculative,
                 returnOffsets ? UnitAttr::get(self.getContext()) : nullptr);
             return returnOffsets ? op->getResult(0) : nullptr;
           })
      .def("create_async_tdm_wait",
           [](GluonOpBuilder &self, int num) {
             ValueRange tokens;
             self.create<ttag::AsyncTDMWait>(tokens, num);
           })
      .def("create_async_copy_lds_barrier_arrive",
           [](GluonOpBuilder &self, Value mbarrier) {
             self.create<ttag::AsyncCopyMbarrierArriveOp>(mbarrier);
           })
      .def("create_lds_barrier_init",
           [](GluonOpBuilder &self, Value memDesc, int count) {
             self.create<ttag::InitBarrierOp>(memDesc, count);
           })
      .def("create_lds_barrier_wait",
           [](GluonOpBuilder &self, Value memDesc, Value phase) {
             self.create<ttag::WaitBarrierOp>(memDesc, phase);
           })
      .def("create_lds_barrier_arrive",
           [](GluonOpBuilder &self, Value memDesc, int count) -> Value {
             return self.create<ttag::ArriveBarrierOp>(memDesc, count);
           })
      .def("create_amd_cluster_arrive",
           [](GluonOpBuilder &self) {
             self.create<ttag::ClusterBarrierArriveOp>();
           })
      .def("create_amd_cluster_wait",
           [](GluonOpBuilder &self) {
             self.create<ttag::ClusterBarrierWaitOp>();
           })
      .def("create_warp_pipeline_border",
           [](GluonOpBuilder &self, const std::string &marker, int priority) {
             auto border =
                 self.create<ROCDL::SchedBarrier>(ROCDL::SchedGroupMask::none);
             auto ctx = self.getContext();
             border->setAttr("triton.warp_pipeline.border",
                             StringAttr::get(ctx, marker));
             if (priority > -1) {
               auto i32Ty = IntegerType::get(ctx, 32);
               border->setAttr("triton.warp_pipeline.priority",
                               IntegerAttr::get(i32Ty, priority));
             }
           });

  m.def("compute_tmem_reg_layout",
        [](py::object elementTyObj, std::vector<int64_t> shape,
           std::vector<int64_t> allocShape, py::object layoutObj,
           unsigned numWarps, const std::string &atomName) -> py::object {
          DialectRegistry registry;
          registry.insert<triton::TritonDialect, ttg::TritonGPUDialect,
                          ttng::TritonNvidiaGPUDialect, gluon::GluonDialect>();
          MLIRContext context(MLIRContext::Threading::DISABLED);
          context.appendDialectRegistry(registry);
          context.loadAllAvailableDialects();

          GluonOpBuilder builder(&context);
          auto builderObj = py::cast(&builder, py::rv_policy::reference);

          auto elementType =
              py::cast<Type>(elementTyObj.attr("to_ir")(builderObj));
          auto layoutAttr =
              py::cast<Attribute>(layoutObj.attr("_to_ir")(builderObj));
          auto ctx = builder.getContext();
          auto memDescTy = builder.getChecked<ttg::MemDescType>(
              shape, elementType, layoutAttr,
              ttng::TensorMemorySpaceAttr::get(ctx),
              /*mutableMemory=*/true, allocShape);

          auto maybeAtom =
              llvm::StringSwitch<std::optional<ttng::TMemAccessAtom>>(atomName)
                  .Case("32x32b", ttng::TMemAccessAtom::I32x32b)
                  .Case("16x64b", ttng::TMemAccessAtom::I16x64b)
                  .Case("16x128b", ttng::TMemAccessAtom::I16x128b)
                  .Case("16x256b", ttng::TMemAccessAtom::I16x256b)
                  .Case("16x32bx2", ttng::TMemAccessAtom::I16x32bx2)
                  .Default(std::nullopt);
          if (!maybeAtom)
            throw std::invalid_argument("unknown TMEM access atom: " +
                                        atomName);
          auto atom = *maybeAtom;
          if (atom == ttng::TMemAccessAtom::I16x32bx2)
            throw std::invalid_argument(
                "Atom 16x32bx2 is inferred implicitly and cannot be requested "
                "explicitly");
          if (numWarps < 4 || !llvm::isPowerOf2_32(numWarps))
            throw std::invalid_argument(
                "numWarps must be a power of two and >= 4");

          auto layout =
              ttng::getDistributedLayoutForTmemLdSt(memDescTy, atom, numWarps);
          if (!layout)
            return py::none();

          auto attr = ttg::LinearEncodingAttr::get(ctx, std::move(*layout));
          return layoutToGluon(attr);
        });

  m.def(
      "make_cga_layout",
      [](std::vector<unsigned> ctasPerCga, std::vector<unsigned> ctaSplitNum,
         std::vector<unsigned> ctaOrder) -> std::vector<std::vector<int32_t>> {
        DialectRegistry registry;
        registry.insert<triton::TritonDialect, ttg::TritonGPUDialect>();
        MLIRContext ctx(MLIRContext::Threading::DISABLED);
        ctx.appendDialectRegistry(registry);
        ctx.loadAllAvailableDialects();
        auto attr = ttg::CGAEncodingAttr::fromSplitParams(
            &ctx, ctasPerCga, ctaSplitNum, ctaOrder);
        return getCgaLayoutBases(attr);
      });

  m.def("get_amd_mfma_scale_layout",
        [](unsigned opIdx, std::vector<int64_t> &shape, unsigned mfmaMDim,
           std::vector<unsigned> &tilesPerWarp,
           std::vector<unsigned> &warpsPerCTA) -> py::object {
          DialectRegistry registry;
          registry.insert<triton::TritonDialect, ttg::TritonGPUDialect,
                          ttng::TritonNvidiaGPUDialect, gluon::GluonDialect>();
          MLIRContext ctx(MLIRContext::Threading::DISABLED);
          ctx.appendDialectRegistry(registry);
          ctx.loadAllAvailableDialects();

          auto ll = ttg::chooseScaledMfmaScaleLayout(
              &ctx, opIdx, shape, mfmaMDim, tilesPerWarp, warpsPerCTA);
          auto attr = ttg::LinearEncodingAttr::get(&ctx, std::move(ll));
          return layoutToGluon(attr);
        });

  m.def("compute_amd_efficient_padded_shared_layout",
        [](unsigned opIdx, unsigned kWidth, unsigned mfmaVersion,
           std::vector<unsigned> &mfmaWarpsPerCTA,
           std::vector<unsigned> &mfmaInstrShape, bool mfmaTransposed,
           std::vector<unsigned> &mfmaTilesPerWarp,
           unsigned mfmaElementBitWidth,
           std::vector<std::vector<int32_t>> &mfmaCgaBases,
           std::vector<int64_t> &shape, unsigned elemBitWidth,
           bool isKContig) -> py::object {
          DialectRegistry registry;
          registry.insert<triton::TritonDialect, ttg::TritonGPUDialect,
                          ttng::TritonNvidiaGPUDialect, gluon::GluonDialect>();
          MLIRContext ctx(MLIRContext::Threading::DISABLED);
          ctx.appendDialectRegistry(registry);
          ctx.loadAllAvailableDialects();

          if (elemBitWidth != 4 && elemBitWidth != 8 && elemBitWidth != 16)
            return py::none();
          auto elemTy = mlir::IntegerType::get(&ctx, elemBitWidth);

          auto mfmaCga =
              buildCgaLayoutAttr(&ctx, mfmaCgaBases, mfmaWarpsPerCTA.size());
          auto mfmaEnc = ttg::AMDMfmaEncodingAttr::get(
              &ctx, mfmaVersion, mfmaWarpsPerCTA, mfmaInstrShape,
              mfmaTransposed, mfmaCga, mfmaTilesPerWarp, mfmaElementBitWidth);
          if (!mfmaEnc)
            return py::none();
          auto dotOpEnc =
              ttg::DotOperandEncodingAttr::get(&ctx, opIdx, mfmaEnc, kWidth);
          if (!dotOpEnc)
            return py::none();

          auto srcTy = RankedTensorType::get(shape, elemTy, dotOpEnc);

          // sharedOrder[0] is the fast dim; for K-contig that's the K axis.
          unsigned kDimIndex = opIdx == 0 ? 1u : 0u;
          SmallVector<unsigned, 2> sharedOrder =
              isKContig ? SmallVector<unsigned, 2>{kDimIndex, 1u - kDimIndex}
                        : SmallVector<unsigned, 2>{1u - kDimIndex, kDimIndex};

          auto attr = composePaddedLayoutForAsyncCopyCDNA4(
              dotOpEnc, cast<ttg::TensorOrMemDesc>(Type(srcTy)), sharedOrder,
              /*useAsyncCopy=*/true, /*warpSize=*/64);
          if (!attr)
            return py::none();
          return layoutToGluon(attr);
        });

  m.def("get_amd_wmma_scale_layout",
        [](unsigned opIdx, std::vector<int64_t> &shape, unsigned wmmaMDim,
           unsigned wmmaNDim, bool isTransposed, unsigned scaleFactor,
           std::vector<std::vector<int32_t>> &regBases,
           std::vector<std::vector<int32_t>> &warpBases,
           std::vector<std::vector<int32_t>> &cgaBases) -> py::object {
          DialectRegistry registry;
          registry.insert<triton::TritonDialect, ttg::TritonGPUDialect,
                          ttng::TritonNvidiaGPUDialect, gluon::GluonDialect>();
          MLIRContext ctx(MLIRContext::Threading::DISABLED);
          ctx.appendDialectRegistry(registry);
          ctx.loadAllAvailableDialects();

          auto rank = shape.size();
          auto kReg = mlir::StringAttr::get(&ctx, "register");
          auto kWarp = mlir::StringAttr::get(&ctx, "warp");
          auto ctaLayout =
              tt::LinearLayout({{kReg, regBases}, {kWarp, warpBases}},
                               tt::standardOutDimNames(&ctx, rank));
          auto cgaLayout = buildCgaLayoutAttr(&ctx, cgaBases, rank);
          auto ll = ttg::chooseScaledWmmaScaleLayout(
              &ctx, opIdx, shape, wmmaMDim, wmmaNDim, isTransposed, scaleFactor,
              ctaLayout, cgaLayout);
          // A swizzled (partition-aware) WMMA produces a non-permutation scale
          // layout, which only GenericLinearEncodingAttr can represent.
          Attribute attr =
              ttg::isPermutationMatrixLayout(ll)
                  ? Attribute(ttg::LinearEncodingAttr::get(&ctx, ll))
                  : Attribute(ttg::GenericLinearEncodingAttr::get(&ctx, ll));
          return layoutToGluon(attr);
        });

  m.def("get_layout_view",
        [](py::object layout, std::vector<int64_t> shape,
           bool useHwView) -> std::string {
          DialectRegistry registry;
          registry.insert<triton::TritonDialect, ttg::TritonGPUDialect,
                          ttng::TritonNvidiaGPUDialect, gluon::GluonDialect>();
          MLIRContext ctx(MLIRContext::Threading::DISABLED);
          ctx.appendDialectRegistry(registry);
          ctx.loadAllAvailableDialects();

          GluonOpBuilder builder(&ctx);
          auto builderObj = py::cast(&builder, py::rv_policy::reference);
          Attribute attr =
              py::cast<Attribute>(layout.attr("_to_ir")(builderObj));

          if (isa<gluon::AutoEncodingAttr>(attr))
            throw py::value_error("AutoLayout cannot be visualized");
          if (isa<gluon::CoalescedEncodingAttr>(attr))
            throw py::value_error("CoalescedLayout cannot be visualized");
          if (isa<ttg::PaddedSharedEncodingAttr>(attr))
            throw py::value_error("PaddedSharedLayout cannot be visualized: "
                                  "toLinearLayout not implemented");

          auto ll = ttg::toLinearLayout(shape, attr);
          if (isa<ttg::DistributedEncodingTrait>(attr)) {
            return ttg::getDistributedLayoutStr(ll, useHwView);
          } else {
            return ttg::getSharedLayoutStr(ll, useHwView);
          }
        });

  py::class_<ttg::WarpSpecializeOp, OpState>(m, "WarpSpecializeOp")
      .def("get_default_region", &ttg::WarpSpecializeOp::getDefaultRegion,
           ret::reference)
      .def("get_partition_op_holder",
           &ttg::WarpSpecializeOp::getPartitionOpHolder, ret::reference)
      .def("set_requested_registers", [](ttg::WarpSpecializeOp &self,
                                         std::vector<int> &requestedRegisters) {
        self.setRequestedRegisters(requestedRegisters);
      });
}
