feat: add a pass to convert tt to ttg preserving attributes (#1604)

Author: avik-pal

src/enzyme_ad/jax/Passes/ConvertTritonToTritonGPUPreservingModuleAttributes.cpp

@@ -0,0 +1,64 @@
+#include "src/enzyme_ad/jax/Passes/Passes.h"
+
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+
+#include "triton/Conversion/TritonToTritonGPU/Passes.h"
+
+#define DEBUG_TYPE "convert-triton-to-triton-gpu-preserving-module-attributes"
+
+namespace mlir {
+namespace enzyme {
+#define GEN_PASS_DEF_CONVERTTRITONTOTRITONGPUPRESERVINGMODULEATTRIBUTESPASS
+#include "src/enzyme_ad/jax/Passes/Passes.h.inc"
+} // namespace enzyme
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::enzyme;
+
+struct ConvertTritonToTritonGPUPreservingModuleAttributesPass
+    : public mlir::enzyme::impl::
+          ConvertTritonToTritonGPUPreservingModuleAttributesPassBase<
+              ConvertTritonToTritonGPUPreservingModuleAttributesPass> {
+  using Base::Base;
+
+  void runOnOperation() override {
+    ModuleOp mod = getOperation();
+
+    int32_t numWarps = 4, threadsPerWarp = 32, numCtas = 1;
+    bool enableSourceRemat = false;
+
+    if (mod->hasAttr("enzymexla.ttg.num-ctas")) {
+      numCtas =
+          mod->getAttrOfType<IntegerAttr>("enzymexla.ttg.num-ctas").getInt();
+    }
+
+    if (mod->hasAttr("enzymexla.ttg.num-warps")) {
+      numWarps =
+          mod->getAttrOfType<IntegerAttr>("enzymexla.ttg.num-warps").getInt();
+    }
+
+    if (mod->hasAttr("enzymexla.ttg.threads-per-warp")) {
+      threadsPerWarp =
+          mod->getAttrOfType<IntegerAttr>("enzymexla.ttg.threads-per-warp")
+              .getInt();
+    }
+
+    if (mod->hasAttr("enzymexla.ttg.enable-source-remat")) {
+      enableSourceRemat = true;
+    }
+
+    OpPassManager pm;
+    pm.addPass(triton::createConvertTritonToTritonGPU(
+        {target, numWarps, threadsPerWarp, numCtas, enableSourceRemat}));
+    if (failed(runPipeline(pm, mod))) {
+      mod->emitError() << "failed to run triton passes";
+      signalPassFailure();
+      return;
+    }
+
+    return;
+  }
+};

src/enzyme_ad/jax/Passes/Passes.td

@@ -1023,4 +1023,19 @@ def SCFCPUify : Pass<"cpuify"> {
   Option<"method", "method", "std::string", /*default=*/"\"distribute\"", "Method of doing distribution">
   ];
 }
+
+def ConvertTritonToTritonGPUPreservingModuleAttributesPass : Pass<
+    "convert-triton-to-triton-gpu-preserving-module-attributes", "mlir::ModuleOp"> {
+  let summary = "Triton generally compiles a single kernel, so they can specify the number of ctas and warps. However, we want to be able to compile multiple kernels. This pass will use the attributes from the module and use that to lower to TritonGPU.";
+  let dependentDialects = [];
+  let options = [
+    Option<
+        /*C++ variable name=*/"target",
+        /*CLI argument=*/"target",
+        /*type=*/"std::string",
+        /*default=*/"\"\"",
+        /*description=*/"the GPU target, e.g., cuda:80, hip:gfx942"
+      >];
+}
+
 #endif