feat: new triton_ext dialect

Author: avik-pal

deps/ReactantExtra/BUILD

@@ -1,4 +1,4 @@
-load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
 load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm")
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
 load("@xla//tools/toolchains/cross_compile/cc:cc_toolchain_config.bzl", "cc_toolchain_config")
@@ -1435,6 +1435,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/WORKSPACE

@@ -4,7 +4,7 @@ NSYNC_COMMIT = "82b118aa7ace3132e517e2c467f8732978cf4023"
 
 NSYNC_SHA256 = ""
 
-ENZYMEXLA_COMMIT = "52ae936cae8f7050adc26c4ed5e755200497dc86"
+ENZYMEXLA_COMMIT = "9867ac059bb2f312a1a6d559d2b41d8ba333a589"
 
 ENZYMEXLA_SHA256 = ""
 

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

src/mlir/Dialects/EnzymeXLA.jl

@@ -792,53 +792,6 @@ function stream2token(source::Value; result::IR.Type, location=Location())
     )
 end
 
-function triton_call(
-    gridx::Value,
-    gridy::Value,
-    gridz::Value,
-    shmem::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, shmem, 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.triton_call",
-        location;
-        operands,
-        owned_regions,
-        successors,
-        attributes,
-        results=op_ty_results,
-        result_inference=false,
-    )
-end
-
 function wrap(
     operand::Value;
     result=nothing::Union{Nothing,IR.Type},

src/mlir/Dialects/TritonExt.jl

@@ -0,0 +1,82 @@
+module triton_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,
+    shmem::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, shmem, 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(
+        "triton_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(
+        "triton_ext.module",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=op_ty_results,
+        result_inference=false,
+    )
+end
+
+end # triton_ext