diff --git a/src/api/api.cpp b/src/api/api.cpp index c8239dbfead..3d8a77cd68b 100644 --- a/src/api/api.cpp +++ b/src/api/api.cpp @@ -174,6 +174,11 @@ static void set_exhaustive_tune_flag(compile_options& options, bool value) options.exhaustive_tune = value; } +static void set_mlss_use_specific_ops(compile_options& options, const char* value) +{ + options.mlss_use_specific_ops = value == nullptr ? "" : value; +} + static void set_file_format(file_options& options, const char* format) { options.format = format; } static void set_default_dim_value(onnx_options& options, size_t value) @@ -2248,6 +2253,19 @@ migraphx_compile_options_set_exhaustive_tune_flag(migraphx_compile_options_t com return api_error_result; } +extern "C" migraphx_status +migraphx_compile_options_set_mlss_use_specific_ops(migraphx_compile_options_t compile_options, + const char* value) +{ + auto api_error_result = migraphx::try_([&] { + if(compile_options == nullptr) + MIGRAPHX_THROW(migraphx_status_bad_param, + "Bad parameter compile_options: Null pointer"); + migraphx::set_mlss_use_specific_ops((compile_options->object), (value)); + }); + return api_error_result; +} + extern "C" migraphx_status migraphx_parse_onnx(migraphx_program_t* out, const char* name, migraphx_onnx_options_t options) { diff --git a/src/api/include/migraphx/migraphx.h b/src/api/include/migraphx/migraphx.h index 7a85274ac24..6519f9a6083 100644 --- a/src/api/include/migraphx/migraphx.h +++ b/src/api/include/migraphx/migraphx.h @@ -600,6 +600,9 @@ migraphx_compile_options_set_fast_math(migraphx_compile_options_t compile_option MIGRAPHX_C_EXPORT migraphx_status migraphx_compile_options_set_exhaustive_tune_flag( migraphx_compile_options_t compile_options, bool value); +MIGRAPHX_C_EXPORT migraphx_status migraphx_compile_options_set_mlss_use_specific_ops( + migraphx_compile_options_t compile_options, const char* value); + MIGRAPHX_C_EXPORT migraphx_status migraphx_parse_onnx(migraphx_program_t* out, const char* name, migraphx_onnx_options_t options); diff --git a/src/api/include/migraphx/migraphx.hpp b/src/api/include/migraphx/migraphx.hpp index 82512c6c430..27ee964573a 100644 --- a/src/api/include/migraphx/migraphx.hpp +++ b/src/api/include/migraphx/migraphx.hpp @@ -1204,6 +1204,13 @@ struct compile_options : MIGRAPHX_HANDLE_BASE(compile_options) { call(&migraphx_compile_options_set_exhaustive_tune_flag, this->get_handle_ptr(), value); } + + /// Force specific ops (comma-separated, e.g. "conv") onto AMDMLSS. Empty lets + /// MIGraphX decide. Mirrors the MIGRAPHX_MLSS_USE_SPECIFIC_OPS env var. + void set_mlss_use_specific_ops(const char* value) + { + call(&migraphx_compile_options_set_mlss_use_specific_ops, this->get_handle_ptr(), value); + } }; /// A program represents the all computation graphs to be compiled and executed diff --git a/src/api/migraphx.py b/src/api/migraphx.py index 7e90471cb50..eb1d64eb2ac 100644 --- a/src/api/migraphx.py +++ b/src/api/migraphx.py @@ -436,6 +436,9 @@ def compile_options(h): h.method('set_exhaustive_tune_flag', api.params(value='bool'), invoke='migraphx::set_exhaustive_tune_flag($@)') + h.method('set_mlss_use_specific_ops', + api.params(value='const char*'), + invoke='migraphx::set_mlss_use_specific_ops($@)') api.add_function('migraphx_parse_onnx', diff --git a/src/include/migraphx/compile_options.hpp b/src/include/migraphx/compile_options.hpp index 64eb3018cf3..93ae194a300 100644 --- a/src/include/migraphx/compile_options.hpp +++ b/src/include/migraphx/compile_options.hpp @@ -24,6 +24,7 @@ #ifndef MIGRAPHX_GUARD_RTGLIB_COMPILE_OPTIONS_HPP #define MIGRAPHX_GUARD_RTGLIB_COMPILE_OPTIONS_HPP +#include #include #include @@ -41,6 +42,10 @@ struct compile_options bool fast_math = true; bool exhaustive_tune = false; + // Comma-separated list of ops to force onto AMDMLSS (e.g. "conv"). Empty lets + // MIGraphX decide. Mirrors the MIGRAPHX_MLSS_USE_SPECIFIC_OPS env var. + std::string mlss_use_specific_ops = ""; + tracer trace{}; }; diff --git a/src/targets/gpu/fuse_mlss.cpp b/src/targets/gpu/fuse_mlss.cpp index 3190b60b89c..34f5c7e4ac8 100644 --- a/src/targets/gpu/fuse_mlss.cpp +++ b/src/targets/gpu/fuse_mlss.cpp @@ -54,14 +54,19 @@ MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_MLSS_USE_SPECIFIC_OPS); #ifdef MIGRAPHX_USE_AMDMLSS -static bool mlss_specific_op(std::string_view op_name) +static bool op_in_list(const std::string& list, std::string_view op_name) { - static const auto options = - split_string(string_value_of(MIGRAPHX_MLSS_USE_SPECIFIC_OPS{}, ""), ','); + const auto options = split_string(list, ','); return std::any_of( options.begin(), options.end(), [&](const auto& opt) { return opt == op_name; }); } +static bool mlss_specific_op(std::string_view op_name) +{ + static const std::string env = string_value_of(MIGRAPHX_MLSS_USE_SPECIFIC_OPS{}, ""); + return op_in_list(env, op_name); +} + // --------------------------------------------------------------------------- // Helper: create an mlss_conv_op intermediate node and replace the matched // instruction(s). The JIT compiler (jit/mlss_conv.cpp) will later convert @@ -382,7 +387,7 @@ struct find_mlss_conv_bias_leaky_relu void fuse_mlss::apply(module_pass_manager& mpm) const { #ifdef MIGRAPHX_USE_AMDMLSS - if(enable_conv or mlss_specific_op("conv")) + if(op_in_list(use_specific_ops, "conv") or mlss_specific_op("conv")) { // Match most-specific patterns first to avoid partial consumption. match::find_matches(mpm, find_mlss_conv_bias_relu{ctx}); diff --git a/src/targets/gpu/include/migraphx/gpu/fuse_mlss.hpp b/src/targets/gpu/include/migraphx/gpu/fuse_mlss.hpp index 3cb0e20f7cd..4cc6fd11e52 100644 --- a/src/targets/gpu/include/migraphx/gpu/fuse_mlss.hpp +++ b/src/targets/gpu/include/migraphx/gpu/fuse_mlss.hpp @@ -37,9 +37,9 @@ namespace gpu { struct MIGRAPHX_GPU_EXPORT fuse_mlss { context* ctx = nullptr; - // Force-enable conv fusion regardless of env-var configuration; used by tests - // so they don't have to mutate the process environment. - bool enable_conv = false; + // Comma-separated list of ops to force onto AMDMLSS (e.g. "conv"), supplied via + // compile_options. Takes effect in addition to MIGRAPHX_MLSS_USE_SPECIFIC_OPS. + std::string use_specific_ops = ""; std::string name() const { return "gpu::fuse_mlss"; } void apply(module_pass_manager& mpm) const; }; diff --git a/src/targets/gpu/target.cpp b/src/targets/gpu/target.cpp index ee34d406cf2..c8abd06b58d 100644 --- a/src/targets/gpu/target.cpp +++ b/src/targets/gpu/target.cpp @@ -184,7 +184,7 @@ struct pipeline_factory .flash_decoding_enabled = mlir_flash_decoding_enabled()}), dead_code_elimination{}, optimize_module{}, - fuse_mlss{get_context()}, + fuse_mlss{.ctx = get_context(), .use_specific_ops = options.mlss_use_specific_ops}, fuse_pointwise_reduce{}, dead_code_elimination{}, #ifndef _WIN32 diff --git a/test/gpu/fuse_mlss.cpp b/test/gpu/fuse_mlss.cpp index be9f968d3c3..bf826ddcf56 100644 --- a/test/gpu/fuse_mlss.cpp +++ b/test/gpu/fuse_mlss.cpp @@ -43,7 +43,7 @@ static migraphx::gpu::context& get_context() static void run_pass(migraphx::program& p) { migraphx::run_passes(p, - {migraphx::gpu::fuse_mlss{&get_context(), /*enable_conv=*/true}, + {migraphx::gpu::fuse_mlss{.ctx = &get_context(), .use_specific_ops = "conv"}, migraphx::dead_code_elimination{}}); } diff --git a/tools/api/api.cpp b/tools/api/api.cpp index 5caf415ed63..971a8e3f086 100644 --- a/tools/api/api.cpp +++ b/tools/api/api.cpp @@ -174,6 +174,11 @@ static void set_exhaustive_tune_flag(compile_options& options, bool value) options.exhaustive_tune = value; } +static void set_mlss_use_specific_ops(compile_options& options, const char* value) +{ + options.mlss_use_specific_ops = value == nullptr ? "" : value; +} + static void set_file_format(file_options& options, const char* format) { options.format = format; } static void set_default_dim_value(onnx_options& options, size_t value)