[TritonGPU] Allow inlining ttng ops and actually run the canonicalizer (#7018)

Author: Mogball

include/triton/Dialect/TritonGPU/Transforms/Passes.td

@@ -369,6 +369,7 @@ def TritonGPUCanonicalize: Pass<"tritongpu-canonicalize"> {
   }];
   let dependentDialects = [
     "mlir::arith::ArithDialect",
+    "mlir::cf::ControlFlowDialect",
     "mlir::scf::SCFDialect",
   ];
 }

lib/Dialect/TritonGPU/Transforms/Canonicalize.cpp

@@ -1,4 +1,5 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -32,6 +33,8 @@ void Canonicalize::runOnOperation() {
       patterns);
   ctx->getLoadedDialect<scf::SCFDialect>()->getCanonicalizationPatterns(
       patterns);
+  ctx->getLoadedDialect<cf::ControlFlowDialect>()->getCanonicalizationPatterns(
+      patterns);
   populateForOpDeadArgumentElimination(patterns);
 
   // Populate select Triton canonicalization patterns. The important patterns to
@@ -43,4 +46,6 @@ void Canonicalize::runOnOperation() {
   ExpandDimsOp::getCanonicalizationPatterns(patterns, ctx);
   ttg::WarpSpecializeOp::getCanonicalizationPatterns(patterns, ctx);
   ttng::TensorDescToTMAPtrOp::getCanonicalizationPatterns(patterns, ctx);
+
+  (void)applyPatternsGreedily(getOperation(), std::move(patterns));
 }

lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp

@@ -30,6 +30,7 @@
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/OpImplementation.h"
 #include "triton/Analysis/Utility.h"
+#include "triton/Dialect/Triton/IR/Interfaces.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h"
 #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
@@ -264,6 +265,7 @@ void TritonNvidiaGPUDialect::initialize() {
 #include "triton/Dialect/TritonNvidiaGPU/IR/Ops.cpp.inc"
       >();
   addInterfaces<TritonGPUOpAsmInterface>();
+  addInterfaces<TritonInlinerInterface>();
 }
 
 // verify TritonNvidiaGPU ops

test/TritonGPU/inline.mlir

@@ -1,4 +1,4 @@
-// RUN: triton-opt %s -allow-unregistered-dialect -inline | FileCheck %s
+// RUN: triton-opt %s -inline | FileCheck %s
 
 #smem = #ttg.shared_memory
 #shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 32}>

test/TritonNvidiaGPU/inline.mlir

@@ -0,0 +1,22 @@
+// RUN: triton-opt %s -inline | FileCheck %s
+
+#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}>
+#smem = #ttg.shared_memory
+
+module attributes {"ttg.num-warps" = 4 : i32} {
+
+// CHECK-LABEL: @inline_ttng_ops
+tt.func public @inline_ttng_ops() {
+  // CHECK-NEXT: ttg.local_alloc
+  // CHECK-NEXT: ttng.init_barrier
+  tt.call @function_with_ttng_ops() : () -> ()
+  tt.return
+}
+
+tt.func private @function_with_ttng_ops() {
+  %0 = ttg.local_alloc : () -> !ttg.memdesc<1xi64, #shared, #smem, mutable>
+  ttng.init_barrier %0, 1 : !ttg.memdesc<1xi64, #shared, #smem, mutable>
+  tt.return
+}
+
+}