Log best_config for all triton kernels

Reviewed By: chenyang78

Differential Revision: D57196022

fbshipit-source-id: 16dcb60da7160df224a916c6f84692b6c84948a3

torchbenchmark/operators/gemm/operator.py

@@ -163,13 +163,20 @@ def best_config(
  self, fn_name: str, example_inputs: Any, metrics: BenchmarkOperatorMetrics
  ) -> float:
  if "triton_tutorial_matmul" in str(fn_name):
- bconfig = triton_matmul_kernel.best_config
- kwargs = deepcopy(bconfig.kwargs)
- kwargs["num_stages"] = bconfig.num_stages
- kwargs["num_warps"] = bconfig.num_warps
- dumped_str = json.dumps(kwargs)
- return dumped_str
- return ""
+ kernel = triton_matmul_kernel
+ elif "triton_ops_matmul" in str(fn_name):
+ kernel = triton.ops._matmul.kernel
+ elif "hstu_triton_matmul" in str(fn_name):
+ import hammer
+ kernel = hammer.ops.triton.triton_matmul._epilogue_mm
+ else:
+ return ""
+
+ bconfig = kernel.best_config
+ kwargs = deepcopy(bconfig.kwargs)
+ kwargs["num_stages"] = bconfig.num_stages
+ kwargs["num_warps"] = bconfig.num_warps
+ return json.dumps(kwargs)
 
  @register_metric()
  def tflops(