Add environment variable to enable mlir exhaustive tuning configs (#2103)

22d18668 · Paul Fultz II · GitHub · dfd95c55 · 22d18668
Unverified Commit 22d18668 authored Sep 11, 2023 by Paul Fultz II Committed by GitHub Sep 11, 2023
Show whitespace changes
Inline Side-by-side

Showing with 5 additions and 2 deletions

src/targets/gpu/mlir.cpp src/targets/gpu/mlir.cpp +5 -2

No files found.
--- a/src/targets/gpu/mlir.cpp
+++ b/src/targets/gpu/mlir.cpp
@@ -70,6 +70,7 @@ inline namespace MIGRAPHX_INLINE_NS {
 namespace gpu {

 MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_TRACE_MLIR);
+MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_MLIR_TUNE_EXHAUSTIVE);
 MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_MLIR_TUNING_DB);
 MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_MLIR_TUNING_CFG);

@@ -685,8 +686,10 @@ struct mlir_program
    {
        tuning_config tc;
        run_high_level_pipeline();
-        mlir_tuning_space params{
-            mlirRockTuningSpaceCreate(mmodule.get(), RocmlirTuningParamSetKindFull)};
+        auto tuning_mode = RocmlirTuningParamSetKindFull;
+        if(enabled(MIGRAPHX_MLIR_TUNE_EXHAUSTIVE{}))
+            tuning_mode = RocmlirTuningParamSetKindExhaustive;
+        mlir_tuning_space params{mlirRockTuningSpaceCreate(mmodule.get(), tuning_mode)};
        for(auto i : range(mlirRockTuningGetNumParams(params.get())))
        {
            mlir_tuning_param param{mlirRockTuningParamCreate()};