From 01ed663163faa7bff9890a7b9ac8096c7014022e Mon Sep 17 00:00:00 2001 From: Peter Zhizhin Date: Wed, 7 Jun 2023 06:09:53 -0700 Subject: [PATCH] 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 --- jaxlib/gpu/triton.cc | 3 +++ 1 file changed, 3 insertions(+) diff --git a/jaxlib/gpu/triton.cc b/jaxlib/gpu/triton.cc index ab9d87823b8e..57aa2bb10717 100644 --- 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(