Add a tuning limit for mlir (#2442)

docs/dev/env_vars.rst

@@ -260,6 +260,10 @@ Appends to tuning cfg file that could be used with rocMLIR tuning scripts.
 Set to "1", "enable", "enabled", "yes", or "true" to use.
 Do exhaustive tuning for MLIR.
 
+.. envvar:: MIGRAPHX_MLIR_TUNE_LIMIT
+
+Set to an integer greater than 1.
+Limits the number of solutions that MLIR will use for tuning.
 
 CK vars
 -----------

src/targets/gpu/mlir.cpp

@@ -73,6 +73,7 @@ namespace gpu {
 
 MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_TRACE_MLIR);
 MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_MLIR_TUNE_EXHAUSTIVE);
+MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_MLIR_TUNE_LIMIT);
 MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_MLIR_TUNING_DB);
 MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_MLIR_TUNING_CFG);
 
@@ -796,7 +797,9 @@ struct mlir_program
         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())))
+        const auto limit =
+            value_of(MIGRAPHX_MLIR_TUNE_LIMIT{}, std::numeric_limits<std::size_t>::max());
+        for(auto i : range(std::min<std::size_t>(limit, mlirRockTuningGetNumParams(params.get()))))
         {
             mlir_tuning_param param{mlirRockTuningParamCreate()};
             if(not mlirRockTuningParamGet(params.get(), i, param.get()))