feat: new enzymexla_tt_ext dialect bindings

deps/ReactantExtra/BUILD

@@ -1447,6 +1447,24 @@ gentbl_cc_library(
     ],
 )
 
+gentbl_cc_library(
+    name = "TritonExtJLIncGen",
+    tbl_outs = [
+        (
+            [
+                "--generator=jl-op-defs",
+                "--disable-module-wrap=0",
+            ],
+            "TritonExt.jl",
+        ),
+    ],
+    tblgen = "//:mlir-jl-tblgen",
+    td_file = "@enzyme_ad//src/enzyme_ad/jax:Dialect/TritonExt/Ops.td",
+    deps = [
+        "@enzyme_ad//src/enzyme_ad/jax:TritonExtDialectTdFiles",
+    ],
+)
+
 gentbl_cc_library(
     name = "TPUJLIncGen",
     tbl_outs = [

deps/ReactantExtra/make-bindings.jl

@@ -42,6 +42,7 @@ for file in [
     "MPI.jl",
     "MemRef.jl",
     "SparseTensor.jl",
+    "TritonExt.jl",
 ]
     build_file(joinpath(src_dir, "mlir", "Dialects", file))
 end

docs/package.json

@@ -5,6 +5,7 @@
     "@types/node": "^22.13.9",
     "markdown-it": "^14.1.0",
     "markdown-it-mathjax3": "^4.3.2",
+    "patch-package": "^8.0.1",
     "vitepress": "^1.6.3",
     "vitepress-plugin-tabs": "^0.6.0"
   },

docs/src/.vitepress/config.mts

@@ -131,6 +131,7 @@ export default defineConfig({
               { text: "SparseTensor", link: "/api/dialects/sparsetensor" },
               { text: "StableHLO", link: "/api/dialects/stablehlo" },
               { text: "Triton", link: "/api/dialects/triton" },
+              { text: "TritonExt", link: "/api/dialects/tritonext" },
               { text: "TPU", link: "/api/dialects/tpu" },
               { text: "VHLO", link: "/api/dialects/vhlo" },
             ],
@@ -221,6 +222,7 @@ export default defineConfig({
               { text: "SparseTensor", link: "/api/dialects/sparsetensor" },
               { text: "StableHLO", link: "/api/dialects/stablehlo" },
               { text: "Triton", link: "/api/dialects/triton" },
+              { text: "TritonExt", link: "/api/dialects/tritonext" },
               { text: "TPU", link: "/api/dialects/tpu" },
               { text: "VHLO", link: "/api/dialects/vhlo" },
             ],

docs/src/api/dialects/tritonext.md

@@ -0,0 +1,11 @@
+```@meta
+CollapsedDocStrings = true
+```
+
+# TritonExt Dialect
+
+Provides extensions to the Triton dialect.
+
+```@autodocs
+Modules = [Reactant.MLIR.Dialects.enzymexla_tt_ext]
+```

src/mlir/Dialects/Enzyme.jl

@@ -361,12 +361,12 @@ end
 
 Extract value from batched operand at index
 """
-function extract(input::Value, index::Value; output::IR.Type, location=Location())
+function extract(input::Value; output::IR.Type, index, location=Location())
     op_ty_results = IR.Type[output,]
-    operands = Value[input, index]
+    operands = Value[input,]
     owned_regions = Region[]
     successors = Block[]
-    attributes = NamedAttribute[]
+    attributes = NamedAttribute[namedattribute("index", index),]
 
     return create_operation(
         "enzyme.extract",

src/mlir/Dialects/TritonExt.jl

@@ -0,0 +1,89 @@
+module enzymexla_tt_ext
+using ...IR
+import ...IR:
+    NamedAttribute,
+    Value,
+    Location,
+    Block,
+    Region,
+    Attribute,
+    create_operation,
+    context,
+    IndexType
+import ..Dialects: namedattribute, operandsegmentsizes
+import ...API
+
+function call(
+    gridx::Value,
+    gridy::Value,
+    gridz::Value,
+    blockx::Value,
+    blocky::Value,
+    blockz::Value,
+    clusterx::Value,
+    clustery::Value,
+    clusterz::Value,
+    inputs::Vector{Value};
+    result_0::Vector{IR.Type},
+    fn,
+    backend_config=nothing,
+    operand_layouts=nothing,
+    result_layouts=nothing,
+    arg_attrs=nothing,
+    res_attrs=nothing,
+    output_operand_aliases=nothing,
+    xla_side_effect_free=nothing,
+    location=Location(),
+)
+    op_ty_results = IR.Type[result_0...,]
+    operands = Value[
+        gridx, gridy, gridz, blockx, blocky, blockz, clusterx, clustery, clusterz, inputs...
+    ]
+    owned_regions = Region[]
+    successors = Block[]
+    attributes = NamedAttribute[namedattribute("fn", fn),]
+    !isnothing(backend_config) &&
+        push!(attributes, namedattribute("backend_config", backend_config))
+    !isnothing(operand_layouts) &&
+        push!(attributes, namedattribute("operand_layouts", operand_layouts))
+    !isnothing(result_layouts) &&
+        push!(attributes, namedattribute("result_layouts", result_layouts))
+    !isnothing(arg_attrs) && push!(attributes, namedattribute("arg_attrs", arg_attrs))
+    !isnothing(res_attrs) && push!(attributes, namedattribute("res_attrs", res_attrs))
+    !isnothing(output_operand_aliases) &&
+        push!(attributes, namedattribute("output_operand_aliases", output_operand_aliases))
+    !isnothing(xla_side_effect_free) &&
+        push!(attributes, namedattribute("xla_side_effect_free", xla_side_effect_free))
+
+    return create_operation(
+        "enzymexla_tt_ext.call",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=op_ty_results,
+        result_inference=false,
+    )
+end
+
+function module_(; sym_name, bodyRegion::Region, location=Location())
+    op_ty_results = IR.Type[]
+    operands = Value[]
+    owned_regions = Region[bodyRegion,]
+    successors = Block[]
+    attributes = NamedAttribute[namedattribute("sym_name", sym_name),]
+
+    return create_operation(
+        "enzymexla_tt_ext.module",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=op_ty_results,
+        result_inference=false,
+    )
+end
+
+end # enzymexla_tt_ext