[Proton][Dialect] Middle-end Proton operator definitions (#5754)

This PR is a follow-up of #5677 for proton compiler mid-end support,
focusing on op definition. Specifically, we 1. added the tablegen
definitions of the proton mid-end operators, 2. removed attributes of
the front-end `RecordOp` (make it a true marker), 3. cleaned up the
dialect's macro.

---------

Co-authored-by: Yuanwei Fang <[email protected]>

Author: 2 people

third_party/proton/dialect/include/Dialect/Proton/IR/Dialect.h

@@ -1,12 +1,14 @@
-#ifndef TRITON_DIALECT_PROTON_IR_DIALECT_H_
-#define TRITON_DIALECT_PROTON_IR_DIALECT_H_
+#ifndef DIALECT_PROTON_IR_DIALECT_H_
+#define DIALECT_PROTON_IR_DIALECT_H_
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/PatternMatch.h"
 #include "proton/dialect/include/Dialect/Proton/IR/Dialect.h.inc"
 #include "proton/dialect/include/Dialect/Proton/IR/OpsEnums.h.inc"
+#include "triton/Dialect/Triton/IR/Dialect.h"
+#include "triton/Dialect/TritonGPU/IR/Dialect.h"
 
 #define GET_ATTRDEF_CLASSES
 #include "proton/dialect/include/Dialect/Proton/IR/ProtonAttrDefs.h.inc"
@@ -20,4 +22,4 @@ namespace proton {} // namespace proton
 } // namespace triton
 } // namespace mlir
 
-#endif // TRITON_DIALECT_PROTON_IR_DIALECT_H_
+#endif // DIALECT_PROTON_IR_DIALECT_H_

third_party/proton/dialect/include/Dialect/Proton/IR/ProtonOps.td

@@ -4,17 +4,21 @@
 include "mlir/IR/OpBase.td"
 include "mlir/IR/EnumAttr.td"
 include "triton/Dialect/Triton/IR/TritonTypes.td"
+include "triton/Dialect/TritonGPU/IR/TritonGPUTypes.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "triton/Dialect/Triton/IR/TritonInterfaces.td"
 include "ProtonDialect.td"
 include "ProtonAttrDefs.td"
 
-class TT_Proton_Op<string mnemonic, list<Trait> traits = []> :
+class PT_Op<string mnemonic, list<Trait> 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", "",
@@ -34,32 +38,89 @@ def GranularityAttr : I32EnumAttr<
     let cppNamespace = "::mlir::triton::proton";
 }
 
-def TT_RecordOp : TT_Proton_Op<"record", [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
-  let summary = "Record a GPU hardware event";
+def PT_RecordOp : PT_Op<"record", [
+  MemoryEffects<[MemRead<DefaultResource>]>,
+  MemoryEffects<[MemWrite<DefaultResource>]>
+  ]> {
+  let summary = "Record a GPU metric event";
 
   let description = [{
-    The operator records GPU events from performance counters.
-    Currently only cycle counter is supported.
+    The operator records GPU event of a particular metric.
+    Essentially a marker with a region id.
 
     Example:
 
     ```mlir
     proton.record() {isStart = true, regionId = 4 : i32}
     ...
     proton.record() {isStart = false, regionId = 4 : i32}
-    ...
-    proton.record() {isStart = true, regionId = 1 : i32, granularity = 1 : i32}
-    ...
-    proton.record() {isStart = false, regionId = 1 : i32, granularity = 1 : i32}
     ```
   }];
   let arguments = (
     ins BoolAttr: $isStart,
+    ConfinedAttr<I32Attr, [IntNonNegative]>:$regionId
+  );
+
+  let assemblyFormat = " `(` operands `)` attr-dict";
+}
+
+def PT_CircularRecordOp : PT_Op<"circular_record", [
+  MemoryEffects<[MemRead<DefaultResource>]>,
+  MemoryEffects<[MemWrite<DefaultResource>]>
+  ]> {
+  let summary = "Record a GPU metric event into a circular buffer";
+
+  let description = [{
+    Records a metric event into a circular buffer backed by the internal memory `data`.
+    The circular buffer indexing `indexPtr` is automatically maintained. Older events
+    get dropped if the `data` is full.
+  }];
+  let arguments = (
+    ins TTG_MemDescType:$data,
+    TT_PtrLike :$indexPtr,
+    BoolAttr: $isStart,
     ConfinedAttr<I32Attr, [IntNonNegative]>:$regionId,
     DefaultValuedAttr<MetricAttr, "Metric::CYCLE">:$metric,
     DefaultValuedAttr<GranularityAttr, "Granularity::WARPGROUP">:$granularity
   );
-  let assemblyFormat = " `(` operands `)` attr-dict";
+  let hasVerifier = 1;
+
+  let assemblyFormat = [{$data `,` $indexPtr attr-dict `:` qualified(type($data)) `,` type($indexPtr)}];
+}
+
+def PT_FinalizeOp : PT_Op<"finalize", [
+  MemoryEffects<[MemRead<SharedMemory>]>,
+  MemoryEffects<[MemRead<GlobalMemory>]>,
+  MemoryEffects<[MemWrite<GlobalMemory>]>
+  ]> {
+  let summary = "Finalize the intra kernel profiler";
+
+  let description = [{
+    Finalize the intra kernel profiler, writing back the metadata and profile to the global memory.
+  }];
+  let arguments = (
+    ins TTG_MemDescType:$data,
+    TT_PtrLike :$indexPtr,
+    TT_PtrLike :$ptr,
+    I32Attr :$size
+  );
+
+  let assemblyFormat = [{$data `,` $indexPtr `,` $ptr attr-dict `:` qualified(type($data)) `,` type($indexPtr) `,` type($ptr)}];
+}
+
+def PT_InitOp : PT_Op<"init", [
+  MemoryEffects<[MemAlloc<GlobalMemory>]>
+  ]> {
+    let summary = "Initialize the intra kernel profiler";
+
+    let description = [{
+        Stack allocation and initialization for the intra kernel profiler.
+        `indexPtr` stores the number of entries proton recorded (zero initialized).
+        We expect `indexPtr` to be register promoted during the LLVM lowering phase.
+    }];
+    let arguments = (ins);
+    let results = (outs TT_PtrLike :$indexPtr);
+    let assemblyFormat = "attr-dict `:` type($indexPtr)";
 }
 
 #endif // PROTON_OPS

third_party/proton/dialect/lib/Dialect/Proton/IR/Ops.cpp

@@ -18,14 +18,12 @@ namespace mlir {
 namespace triton {
 namespace proton {
 
-// -- RecordOp --
-void RecordOp::getEffects(
-    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
-        &effects) {
-  effects.emplace_back(MemoryEffects::Write::get(),
-                       SideEffects::DefaultResource::get());
-  effects.emplace_back(MemoryEffects::Read::get(),
-                       SideEffects::DefaultResource::get());
+// -- CircularRecordOp --
+LogicalResult CircularRecordOp::verify() {
+  // TODO(fywkevin): checks the following:
+  // 1. circular buffer size power of 2.
+  // 2. function's noinline is false.
+  return success();
 }
 
 } // namespace proton