[PROTON-DEV] Refactor frontend interface (#5825)

Jokeren · web-flow · commit 922077218d77 · 2025-02-10T11:00:34.000-05:00
diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
@@ -631,14 +631,15 @@ void populateTritonPatterns(TritonGPUTypeConverter &typeConverter,
 }
 // Proton patterns
 // NOTE: Because Proton's inputs are scalars and not tensors this conversion
-// isn't strictly nessessary however you could envision a case where we pass in
+// isn't strictly necessary however you could envision a case where we pass in
 // tensors in for Triton object specific tracing operations in which case we
 // would need to fill in the OpConversionPattern
 void populateProtonPatterns(TritonGPUTypeConverter &typeConverter,
                             RewritePatternSet &patterns) {
   MLIRContext *context = patterns.getContext();
-  patterns.add<GenericOpPattern<triton::proton::RecordOp>>(typeConverter,
-                                                           context);
+  patterns.add<GenericOpPattern<triton::proton::RecordOp>,
+               GenericOpPattern<triton::proton::InitScopeOp>>(typeConverter,
+                                                              context);
 }
 //
 // SCF patterns
diff --git a/python/src/ir.cc b/python/src/ir.cc
@@ -1735,8 +1735,12 @@ void init_triton_ir(py::module &&m) {
            })
       // Proton Ops
       .def("create_proton_record",
-           [](TritonOpBuilder &self, bool isStart, int32_t regionId) -> void {
-             self.create<mlir::triton::proton::RecordOp>(isStart, regionId);
+           [](TritonOpBuilder &self, bool isStart, Value &scopeId) -> void {
+             self.create<mlir::triton::proton::RecordOp>(isStart, scopeId);
+           })
+      .def("create_proton_init_scope",
+           [](TritonOpBuilder &self, const std::string &name) -> Value {
+             return self.create<mlir::triton::proton::InitScopeOp>(name);
            });
 
   py::class_<PassManager>(m, "pass_manager", py::module_local())
diff --git a/python/triton/compiler/compiler.py b/python/triton/compiler/compiler.py
@@ -355,6 +355,7 @@ class CompiledKernel:
     # TODO: move out of this namespace since it's a runtime thing
     launch_enter_hook = None
     launch_exit_hook = None
+    init_handle_hook = None
 
     def __init__(self, src, metadata_group, hash):
         from collections import namedtuple
@@ -403,6 +404,8 @@ def _init_handles(self):
         # TODO: n_regs, n_spills should be metadata generated when calling `ptxas`
         self.module, self.function, self.n_regs, self.n_spills = driver.active.utils.load_binary(
             self.name, self.kernel, self.metadata.shared, device)
+        if self.init_handle_hook is not None:
+            self.init_handle_hook(self.module, self.function, self.metadata_path)
 
     def __getattribute__(self, name):
         if name == 'run':
@@ -415,11 +418,7 @@ def launch_metadata(self, grid, stream, *args):
         ret = LazyDict({"name": self.name, "function": self.function, "stream": stream})
         if not isinstance(self.src, ASTSource) or self.src.fn.launch_metadata is None:
             return ret
-        arg_dict = {}
-        arg_idx = 0
-        for i, arg_name in enumerate(self.src.fn.arg_names):
-            arg_dict[arg_name] = args[arg_idx]
-            arg_idx += 1
+        arg_dict = {name: arg for name, arg in zip(self.src.fn.arg_names, args)}
         ret.add(self.src.fn.launch_metadata, (grid, self.metadata, arg_dict))
         return ret
 
diff --git a/test/Proton/ops.mlir b/test/Proton/ops.mlir
@@ -1,12 +1,21 @@
-// RUN: triton-opt --split-input-file %s -cse -canonicalize --proton-lowering-pass | FileCheck %s
+// RUN: triton-opt --split-input-file %s | FileCheck %s
 
 module {
+  // CHECK-LABEL: proton_init_scope
+  tt.func @proton_init_scope() {
+    // CHECK: proton.init_scope "name0" : i32
+    // CHECK-NEXT: tt.return
+    %0 = proton.init_scope "name0" : i32
+    tt.return
+  }
+  // CHECK-LABEL: proton_record
   tt.func @proton_record() {
-    // CHECK: proton.record() {isStart = true, regionId = 1 : i32}
-    // CHECK-NEXT: proton.record() {isStart = false, regionId = 1 : i32}
+    // CHECK: proton.record start %0
+    // CHECK: proton.record end %0
     // CHECK-NEXT: tt.return
-    proton.record() {isStart = true, regionId = 1 : i32}
-    proton.record() {isStart = false, regionId = 1 : i32}
+    %0 = proton.init_scope "name0" : i32
+    proton.record start %0
+    proton.record end %0
     tt.return
   }
 } // end module
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp
@@ -231,6 +231,8 @@ struct ConvertTritonAMDGPUToLLVM
     mlir::triton::populatePrintOpToLLVMPattern(typeConverter, patterns,
                                                targetInfo, commonBenefit);
 
+    mlir::triton::proton::populateInitScopeOpToLLVMPattern(
+        typeConverter, patterns, commonBenefit);
     mlir::triton::proton::populateRecordOpToLLVMPattern(
         typeConverter, patterns, targetInfo, commonBenefit);
 
diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp
@@ -153,6 +153,8 @@ struct ConvertTritonGPUToLLVM
                                                     targetInfo, benefit);
     mlir::triton::populatePrintOpToLLVMPattern(typeConverter, patterns,
                                                targetInfo, benefit);
+    mlir::triton::proton::populateInitScopeOpToLLVMPattern(typeConverter,
+                                                           patterns, benefit);
     mlir::triton::proton::populateRecordOpToLLVMPattern(typeConverter, patterns,
                                                         targetInfo, benefit);
     mlir::triton::populateControlFlowOpToLLVMPattern(typeConverter, patterns,
diff --git a/third_party/proton/csrc/include/Context/Shadow.h b/third_party/proton/csrc/include/Context/Shadow.h
@@ -37,4 +37,4 @@ class ShadowContextSource : public ContextSource, public ScopeInterface {
 
 } // namespace proton
 
-#endif // PROTON_CONTEXT_CONTEXT_H_
+#endif // PROTON_CONTEXT_SHADOW_H_
diff --git a/third_party/proton/dialect/include/Dialect/Proton/IR/ProtonOps.td b/third_party/proton/dialect/include/Dialect/Proton/IR/ProtonOps.td
@@ -13,61 +13,80 @@ include "ProtonDialect.td"
 include "ProtonAttrDefs.td"
 
 class PT_Op<string mnemonic, list<Trait> traits = []> :
-    Op<Proton_Dialect, mnemonic, !listconcat(traits, [])> {
+  Op<Proton_Dialect, mnemonic, !listconcat(traits, [])> {
 }
 
 def GlobalMemory : Resource<"::mlir::triton::GlobalMemory">;
 def SharedMemory : Resource<"::mlir::triton::gpu::SharedMemory">;
 
 // Proton profiling metric.
 def MetricAttr : I32EnumAttr<
-    "Metric", "",
-    [
-        I32EnumAttrCase<"CYCLE", 0, "cycle">,
-    ]> {
-    let cppNamespace = "::mlir::triton::proton";
+  "Metric", "",
+  [
+    I32EnumAttrCase<"CYCLE", 0, "cycle">,
+  ]> {
+  let cppNamespace = "::mlir::triton::proton";
 }
 
 // Proton profiling granularity.
 def GranularityAttr : I32EnumAttr<
-    "Granularity", "",
-    [
-        I32EnumAttrCase<"WARPGROUP", 0, "warpgroup">,
-        I32EnumAttrCase<"WARP", 1, "warp">,
-    ]> {
-    let cppNamespace = "::mlir::triton::proton";
+  "Granularity", "",
+  [
+    I32EnumAttrCase<"WARPGROUP", 0, "warpgroup">,
+    I32EnumAttrCase<"WARP", 1, "warp">,
+  ]> {
+  let cppNamespace = "::mlir::triton::proton";
+}
+
+def PT_InitScopeOp : PT_Op<"init_scope", [Pure]> {
+  let summary = "Initialize a scope";
+
+  let description = [{
+    This operation initializes a scope with the given name and returns a unique id for the scope.
+
+    Example:
+
+    ```mlir
+    %scope0 = proton.init_scope "name0" : i32
+    ```
+  }];
+
+  let arguments = (
+    ins StrAttr: $scopeName
+  );
+  let results = (outs I32 : $scopeId);
+
+  let assemblyFormat = "$scopeName attr-dict `:` type($scopeId) ";
+  // hasVerifier = 1; verify that (1) each scope is used twice, and (2) the use ops are record like ops
 }
 
 def PT_RecordOp : PT_Op<"record", [
-  MemoryEffects<[MemRead<DefaultResource>]>,
-  MemoryEffects<[MemWrite<DefaultResource>]>
-  ]> {
-  let summary = "Record a GPU metric event";
+  MemoryEffects<[MemRead<DefaultResource>, MemWrite<DefaultResource>]>
+]> {
+  let summary = "Record an event";
 
   let description = [{
-    The operator records GPU event of a particular metric.
-    Essentially a marker with a region id.
+    This operation records events of a particular metric.
 
     Example:
 
     ```mlir
-    proton.record() {isStart = true, regionId = 4 : i32}
+    proton.record start %scope0
     ...
-    proton.record() {isStart = false, regionId = 4 : i32}
+    proton.record end %scope0
     ```
   }];
   let arguments = (
-    ins BoolAttr: $isStart,
-    ConfinedAttr<I32Attr, [IntNonNegative]>:$regionId
+    ins UnitAttr: $isStart,
+    I32: $scopeId
   );
 
-  let assemblyFormat = " `(` operands `)` attr-dict";
+  let assemblyFormat = "(`start` $isStart^):(`end`)? $scopeId attr-dict";
 }
 
 def PT_CircularRecordOp : PT_Op<"circular_record", [
-  MemoryEffects<[MemRead<DefaultResource>]>,
-  MemoryEffects<[MemWrite<DefaultResource>]>
-  ]> {
+  MemoryEffects<[MemRead<DefaultResource>, MemWrite<DefaultResource>]>
+]> {
   let summary = "Record a GPU metric event into a circular buffer";
 
   let description = [{
@@ -76,16 +95,19 @@ def PT_CircularRecordOp : PT_Op<"circular_record", [
     get dropped if the `data` is full.
   }];
   let arguments = (
-    ins TTG_MemDescType:$data,
+    ins UnitAttr: $isStart,
+    I32: $scopeId,
+    TTG_MemDescType:$data,
     TT_PtrLike :$indexPtr,
-    BoolAttr: $isStart,
-    ConfinedAttr<I32Attr, [IntNonNegative]>:$regionId,
     DefaultValuedAttr<MetricAttr, "Metric::CYCLE">:$metric,
     DefaultValuedAttr<GranularityAttr, "Granularity::WARPGROUP">:$granularity
   );
   let hasVerifier = 1;
 
-  let assemblyFormat = [{$data `,` $indexPtr attr-dict `:` qualified(type($data)) `,` type($indexPtr)}];
+  let assemblyFormat = [{
+    (`start` $isStart^):(`end`)? $scopeId `,` $data `,` $indexPtr attr-dict `:`
+    qualified(type($data)) `,` type($indexPtr)
+  }];
 }
 
 def PT_FinalizeOp : PT_Op<"finalize", [
diff --git a/third_party/proton/dialect/include/TritonProtonToLLVM/PatternTritonProtonOpToLLVM.h b/third_party/proton/dialect/include/TritonProtonToLLVM/PatternTritonProtonOpToLLVM.h
@@ -1,5 +1,5 @@
-#ifndef TRITON_CONVERSION_TRITONPROTON_TO_LLVM_PATTERNS_TRITON_PROTON_OP_TO_LLVM_H
-#define TRITON_CONVERSION_TRITONPROTON_TO_LLVM_PATTERNS_TRITON_PROTON_OP_TO_LLVM_H
+#ifndef TRITON_PROTON_TO_LLVM_PATTERN_TRITON_PROTON_OP_TO_LLVM_H
+#define TRITON_PROTON_TO_LLVM_PATTERN_TRITON_PROTON_OP_TO_LLVM_H
 
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 
@@ -10,7 +10,12 @@ void populateRecordOpToLLVMPattern(LLVMTypeConverter &typeConverter,
                                    RewritePatternSet &patterns,
                                    const TargetInfoBase &targetInfo,
                                    PatternBenefit benefit);
+
+void populateInitScopeOpToLLVMPattern(LLVMTypeConverter &typeConverter,
+                                      RewritePatternSet &patterns,
+                                      PatternBenefit benefit);
+
 } // namespace proton
 } // namespace mlir::triton
 
-#endif
+#endif // TRITON_PROTON_TO_LLVM_PATTERN_TRITON_PROTON_OP_TO_LLVM_H
diff --git a/third_party/proton/dialect/lib/TritonProtonToLLVM/CMakeLists.txt b/third_party/proton/dialect/lib/TritonProtonToLLVM/CMakeLists.txt
@@ -1,4 +1,5 @@
 add_triton_library(TritonProtonToLLVM
+    InitScopeOpToLLVM.cpp
     RecordOpToLLVM.cpp
     ProtonLoweringPass.cpp
 
diff --git a/third_party/proton/dialect/lib/TritonProtonToLLVM/InitScopeOpToLLVM.cpp b/third_party/proton/dialect/lib/TritonProtonToLLVM/InitScopeOpToLLVM.cpp
@@ -0,0 +1,39 @@
+#include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+#include "mlir/IR/PatternMatch.h"
+#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h"
+#include "triton/Conversion/TritonGPUToLLVM/Utility.h"
+#include "triton/Dialect/Triton/IR/Dialect.h"
+
+#include "third_party/proton/dialect/include/Dialect/Proton/IR/Dialect.h"
+#include "third_party/proton/dialect/include/TritonProtonToLLVM/PatternTritonProtonOpToLLVM.h"
+
+namespace {
+
+struct InitScopeOpConversion
+    : public ConvertOpToLLVMPattern<mlir::triton::proton::InitScopeOp> {
+  explicit InitScopeOpConversion(LLVMTypeConverter &typeConverter,
+                                 PatternBenefit benefit)
+      : mlir::ConvertOpToLLVMPattern<mlir::triton::proton::InitScopeOp>(
+            typeConverter, benefit) {}
+
+  LogicalResult
+  matchAndRewrite(mlir::triton::proton::InitScopeOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    // TODO: Use an allocate scope id pass to attach each scope id op an
+    // attribute
+    auto b = TritonLLVMOpBuilder(loc, rewriter);
+    auto i0 = b.i32_val(0);
+    rewriter.replaceOp(op, i0);
+    return success();
+  }
+};
+
+} // namespace
+
+void mlir::triton::proton::populateInitScopeOpToLLVMPattern(
+    LLVMTypeConverter &typeConverter, RewritePatternSet &patterns,
+    PatternBenefit benefit) {
+  patterns.add<InitScopeOpConversion>(typeConverter, benefit);
+}
diff --git a/third_party/proton/proton/hook.py b/third_party/proton/proton/hook.py
@@ -1,6 +1,8 @@
 from .state import enter_state, exit_state
 from .scope import enter_scope, exit_scope
 
+from triton._C.libtriton import ir
+
 COMPUTE_METADATA_SCOPE_NAME = "__proton_launch_metadata"
 
 
@@ -24,15 +26,45 @@ def exit(lazy_dict: LazyDict) -> None:
         exit_scope(triton_op=True)
 
 
-def register_triton_hook() -> None:
+class TritonInitHandleHook:
+    function_scope_ids: dict = {}
+
+    @staticmethod
+    def map_scope_ids(function, module, metadata_group):
+        if function and function not in TritonInitHandleHook.function_scope_ids:
+            ir_path = None
+            if "ttgir" in metadata_group:
+                ir_path = metadata_group["ttgir"]
+            elif "ttir" in metadata_group:
+                ir_path = metadata_group["ttir"]
+            if ir_path:
+                context = ir.context()
+                module = ir.parse_mlir_module(ir_path, context)
+                module.context = context
+
+
+def register_launch_hook() -> None:
     from triton.compiler import CompiledKernel
     if CompiledKernel.launch_enter_hook is None:
         CompiledKernel.launch_enter_hook = TritonHook.enter
         CompiledKernel.launch_exit_hook = TritonHook.exit
+    else:
+        raise RuntimeError("Triton launch hook is already registered.")
+
+
+def unregister_launch_hook() -> None:
+    from triton.compiler import CompiledKernel
+    CompiledKernel.launch_enter_hook = None
+    CompiledKernel.launch_exit_hook = None
+
+
+def register_init_handle_hook() -> None:
+    from triton.compiler import CompiledKernel
+    if CompiledKernel.init_handle_hook is not None:
+        raise RuntimeError("Triton init handle hook is already registered.")
+    CompiledKernel.init_handle_hook = TritonInitHandleHook.map_scope_ids
 
 
-def unregister_triton_hook() -> None:
+def unregister_init_handle_hook() -> None:
     from triton.compiler import CompiledKernel
-    if CompiledKernel.launch_enter_hook == TritonHook.enter:
-        CompiledKernel.launch_enter_hook = None
-        CompiledKernel.launch_exit_hook = None
+    CompiledKernel.init_handle_hook = None
diff --git a/third_party/proton/proton/language.py b/third_party/proton/proton/language.py
diff --git a/third_party/proton/proton/profile.py b/third_party/proton/proton/profile.py
diff --git a/third_party/proton/test/test_record.py b/third_party/proton/test/test_record.py

Original file line number	Diff line number	Diff line change
`@@ -37,4 +37,4 @@ class ShadowContextSource : public ContextSource, public ScopeInterface {`
`37`	`37`
`38`	`38`	`} // namespace proton`
`39`	`39`
`40`		`-#endif // PROTON_CONTEXT_CONTEXT_H_`
	`40`	`+#endif // PROTON_CONTEXT_SHADOW_H_`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`add_triton_library(TritonProtonToLLVM`
	`2`	`+ InitScopeOpToLLVM.cpp`
`2`	`3`	`RecordOpToLLVM.cpp`
`3`	`4`	`ProtonLoweringPass.cpp`
`4`	`5`