A tracking utility for gathering the compile and/or runtime time, size, profiling and other statistics

.github/workflows/triton-benchmarks.yml

@@ -68,6 +68,7 @@ env:
   VERIFY: ${{ (github.event_name == 'pull_request' || github.event_name == 'schedule' || inputs.verify) && '1' || '0' }}
   TAG: ${{ inputs.tag || (github.event_name == 'pull_request' && format('pr-{0}', github.event.number)) || (github.event_name == 'schedule' && 'ci') || 'test' }}
   N_RUNS: ${{ inputs.n_runs || '1' }}
+  TRITON_TRACK_DUMP: "$PWD/reports/track"
 
 jobs:
   build:

python/src/ir.cc

@@ -1816,6 +1816,44 @@ void init_triton_ir(py::module &&m) {
              self.printAsTextualPipeline(os);
              return str;
            })
+      .def("enable_timing",
+           [](PassManager &self, py::function cb) {
+             struct CallBackStrategy : OutputStrategy {
+               py::function cb;
+
+               CallBackStrategy(py::function cb)
+                   : OutputStrategy(llvm::errs()), cb(cb) {}
+
+               void printHeader(const TimeRecord &total) override {}
+
+               void printFooter() override {}
+
+               void printTime(const TimeRecord &time,
+                              const TimeRecord &total) override {}
+
+               void printListEntry(StringRef name, const TimeRecord &time,
+                                   const TimeRecord &total,
+                                   bool lastEntry = false) override {
+                 cb(std::string(name), time.wall, 0);
+               }
+
+               void printTreeEntry(unsigned indent, StringRef name,
+                                   const TimeRecord &time,
+                                   const TimeRecord &total) override {
+                 cb(std::string(name), time.wall, 1);
+               }
+
+               void printTreeEntryEnd(unsigned indent,
+                                      bool lastEntry = false) override {
+                 cb(std::string(""), 0., 2);
+               }
+             };
+
+             auto tm = std::make_unique<mlir::DefaultTimingManager>();
+             tm->setOutput(std::make_unique<CallBackStrategy>(cb));
+             tm->setEnabled(true);
+             self.enableTiming(std::move(tm));
+           })
       .def(
           "run",
           [](PassManager &self, ModuleOp &mod) {

third_party/intel/backend/compiler.py

@@ -13,6 +13,7 @@
 import os
 import subprocess
 from pathlib import Path
+from .track import track
 
 
 @dataclass
@@ -207,6 +208,7 @@ def get_split_barrier_scope(opt):
         return split_barriers_scope
 
     @staticmethod
+    @track
     def make_ttir(mod, metadata, opt):
         pm = ir.pass_manager(mod.context)
         pm.enable_debug()
@@ -226,6 +228,7 @@ def make_ttir(mod, metadata, opt):
         return mod
 
     @staticmethod
+    @track
     def make_ttgir(mod, metadata, opt, properties):
         cluster_info = intel.ClusterInfo()
         if opt.cluster_dims is not None:
@@ -303,6 +306,7 @@ def gluon_to_ttgir(self, src, metadata, options):
         return mod
 
     @staticmethod
+    @track
     def make_llir(src, metadata, options):
         mod = src
         # TritonGPU -> LLVM-IR (MLIR)
@@ -340,7 +344,9 @@ def make_llir(src, metadata, options):
             paths = [path for (name, path) in options.extern_libs]
             llvm.link_extern_libs(llvm_mod, paths)
 
-        intel.optimize_module(llvm_mod, llvm.OPTIMIZE_O3)
+        with track("optimize_module") as tr:
+            intel.optimize_module(llvm_mod, llvm.OPTIMIZE_O3, tr.callback("passes"))
+
         intel.post_process_llir(llvm_mod)
 
         # Get some metadata
@@ -357,6 +363,7 @@ def make_llir(src, metadata, options):
         return ret
 
     @staticmethod
+    @track
     def make_spv(src, metadata, options, device_arch):
         spirv, name = intel.translate_to_spirv(src)
         metadata["name"] = name
@@ -384,7 +391,7 @@ def make_spv(src, metadata, options, device_arch):
         metadata["generate_native_code"] = options.generate_native_code
 
         if options.generate_native_code:
-            with tempfile.TemporaryDirectory() as temp_dir:
+            with track("generate_native_code"), tempfile.TemporaryDirectory() as temp_dir:
                 with tempfile.NamedTemporaryFile(mode='wb', suffix='.spv', dir=temp_dir, delete=False) as fsrc:
                     fsrc.write(spirv)
                 fbin = fsrc.name + '.o'