Add a comment at the end of autotuning a Triton function

I'd like to see what config was chosen at the end of autotuning. Tracking `is the new best config` is a bit hard. PiperOrigin-RevId: 538465212
google · Jun 7, 2023 · 01ed663 · 01ed663
1 parent fa099fd
commit 01ed663
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/jaxlib/gpu/triton.cc b/jaxlib/gpu/triton.cc
@@ -296,6 +296,9 @@ class TritonAutotunedKernelCall : public TritonKernelCallBase {
     py::gil_scoped_acquire gil;
     configs_.erase(configs_.begin() + 1, configs_.end());
 
+    LOG(INFO) << "Finished autotuning function: " << name_ << " best config "
+              << configs_[0].description;
+
     // Restore aliased inputs to their original values.
     for (auto [input_idx, _, size] : input_output_aliases_) {
       CUDA_RETURN_IF_ERROR(