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

python/src/ir.cc

@@ -1814,6 +1814,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

@@ -1,6 +1,7 @@
 from triton.backends.compiler import BaseBackend, Language
 from triton._C.libtriton import ir, passes, llvm, intel
 from triton.backends.intel.driver import compile_module_from_src
+from triton.backends.intel.track import track
 from triton import knobs
 
 from dataclasses import dataclass
@@ -211,6 +212,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()
@@ -230,6 +232,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:
@@ -307,6 +310,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)
@@ -348,7 +352,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
@@ -367,6 +373,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
@@ -394,7 +401,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'