From c3a536719ba83164f2331f05b7f430df581cb5e6 Mon Sep 17 00:00:00 2001 From: Paul Fultz II Date: Mon, 21 Oct 2024 14:31:46 -0500 Subject: [PATCH] Add flag to dump mlir modules to mxrs (#3182) --- docs/dev/env_vars.rst | 5 +++ src/include/migraphx/program.hpp | 2 + src/include/migraphx/shape.hpp | 2 + src/program.cpp | 16 +++---- src/shape.cpp | 12 ++++++ src/targets/gpu/include/migraphx/gpu/mlir.hpp | 4 ++ src/targets/gpu/jit/mlir.cpp | 8 ++++ src/targets/gpu/mlir.cpp | 43 +++++++++++++++++++ 8 files changed, 81 insertions(+), 11 deletions(-) diff --git a/docs/dev/env_vars.rst b/docs/dev/env_vars.rst index 8c7ff31b9c2..8816af8f550 100644 --- a/docs/dev/env_vars.rst +++ b/docs/dev/env_vars.rst @@ -299,6 +299,11 @@ Enable reduction fusions in MLIR. Set to "1", "enable", "enabled", "yes", or "true" to use. Enable Split-k perf configs when tuning with MLIR. +.. envvar:: MIGRAPHX_MLIR_DUMP_TO_MXR + +Set to path where MXRs will be saved. +Dumps MLIRs module to mxr files. + CK vars ----------- diff --git a/src/include/migraphx/program.hpp b/src/include/migraphx/program.hpp index 6e9bb41f4af..87de258a44b 100644 --- a/src/include/migraphx/program.hpp +++ b/src/include/migraphx/program.hpp @@ -58,6 +58,8 @@ struct MIGRAPHX_EXPORT program { program(); + explicit program(module m); + // move constructor program(program&&) noexcept; diff --git a/src/include/migraphx/shape.hpp b/src/include/migraphx/shape.hpp index 4a5952d5eac..e15c4dece44 100644 --- a/src/include/migraphx/shape.hpp +++ b/src/include/migraphx/shape.hpp @@ -143,6 +143,8 @@ struct MIGRAPHX_EXPORT shape const std::size_t& y); }; + static std::string to_sizes_string(const std::vector& shapes); + static const std::vector& types(); static std::string name(type_t t); diff --git a/src/program.cpp b/src/program.cpp index 25cb16cc950..cac833803b3 100644 --- a/src/program.cpp +++ b/src/program.cpp @@ -77,6 +77,10 @@ struct program_impl }; program::program() : impl(std::make_unique()) { this->create_module("main"); } +program::program(module m) : impl(std::make_unique()) +{ + this->create_module("main", std::move(m)); +} program::program(program&&) noexcept = default; program::~program() noexcept = default; @@ -852,17 +856,7 @@ std::string perf_group(instruction_ref ins, bool detailed) if(detailed) { result += "<" + ins->get_shape().type_string(); - std::vector sizes; - std::transform(ins->inputs().begin(), - ins->inputs().end(), - std::back_inserter(sizes), - [&](instruction_ref input) { - std::string r = to_string_range(input->get_shape().lens(), "x"); - if(not input->get_shape().standard()) - r += ":" + to_string_range(input->get_shape().strides(), "x"); - return r; - }); - result += "(" + join_strings(sizes, ", ") + ")>"; + result += "(" + shape::to_sizes_string(to_shapes(ins->inputs())) + ")>"; } return result; } diff --git a/src/shape.cpp b/src/shape.cpp index c9899403548..fd8173f0bff 100644 --- a/src/shape.cpp +++ b/src/shape.cpp @@ -227,6 +227,18 @@ struct shape_impl std::shared_ptr copy() const { return std::make_shared(*this); } }; +std::string shape::to_sizes_string(const std::vector& shapes) +{ + std::vector sizes; + std::transform(shapes.begin(), shapes.end(), std::back_inserter(sizes), [&](const shape& s) { + std::string r = to_string_range(s.lens(), "x"); + if(not s.standard()) + r += ":" + to_string_range(s.strides(), "x"); + return r; + }); + return join_strings(sizes, ", "); +} + const std::vector& shape::types() { static const std::vector result = { diff --git a/src/targets/gpu/include/migraphx/gpu/mlir.hpp b/src/targets/gpu/include/migraphx/gpu/mlir.hpp index 7f31f43dfd4..13d8fdfa916 100644 --- a/src/targets/gpu/include/migraphx/gpu/mlir.hpp +++ b/src/targets/gpu/include/migraphx/gpu/mlir.hpp @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -67,6 +68,9 @@ MIGRAPHX_GPU_EXPORT tuning_config get_tuning_config_mlir(const context& migraphx const std::vector& inputs, bool exhaustive); +MIGRAPHX_GPU_EXPORT void +dump_mlir_to_mxr(module m, const std::vector& inputs, const fs::path& location); + } // namespace gpu } // namespace MIGRAPHX_INLINE_NS } // namespace migraphx diff --git a/src/targets/gpu/jit/mlir.cpp b/src/targets/gpu/jit/mlir.cpp index 6565ba056a1..06bb69fbd11 100644 --- a/src/targets/gpu/jit/mlir.cpp +++ b/src/targets/gpu/jit/mlir.cpp @@ -38,6 +38,8 @@ namespace migraphx { inline namespace MIGRAPHX_INLINE_NS { namespace gpu { +MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_MLIR_DUMP_TO_MXR); + static module create_pointwise_module(module_ref in_mod) { module pw_mod; @@ -209,8 +211,14 @@ struct mlir_compiler : compiler const operation&, bool exhaustive) const { + static const auto mxr_loc = string_value_of(MIGRAPHX_MLIR_DUMP_TO_MXR{}); + auto shapes = to_shapes(ins->inputs()); auto* smod = ins->module_inputs().front(); + if(not mxr_loc.empty()) + { + dump_mlir_to_mxr(*smod, ins->inputs(), mxr_loc); + } return get_tuning_config_mlir(ctx, *smod, shapes, exhaustive); } diff --git a/src/targets/gpu/mlir.cpp b/src/targets/gpu/mlir.cpp index 32c903857a8..e9d06ce0a9b 100644 --- a/src/targets/gpu/mlir.cpp +++ b/src/targets/gpu/mlir.cpp @@ -59,6 +59,8 @@ #include #include #include +#include +#include #include #include #include @@ -1063,6 +1065,23 @@ void adjust_param_shapes(module& m, const std::vector& inputs) } } +void replace_params_with_literals(module& m, const std::vector& inputs) +{ + auto names = m.get_parameter_names(); + std::sort(names.begin(), names.end()); + for(auto i : range(names.size())) + { + const auto& name = names[i]; + const auto& input = inputs[i]; + if(input->name() != "@literal") + continue; + auto param = m.get_parameter(name); + auto lit = m.add_literal(input->get_literal()); + m.replace_instruction(param, lit); + m.remove_instruction(param); + } +} + std::string dump_mlir(module m, const std::vector& inputs) { const_module_ref mr = &m; @@ -1182,6 +1201,30 @@ tuning_config get_tuning_config_mlir(const context& migraphx_ctx, return tc; } +void dump_mlir_to_mxr(module m, + const std::vector& inputs, + const fs::path& location) +{ + static std::mutex mutex; + const std::lock_guard lock(mutex); + + adjust_param_shapes(m, to_shapes(inputs)); + replace_params_with_literals(m, inputs); + std::vector sizes; + for(auto ins : iterator_for(m)) + { + if(not contains({"convolution", "dot"}, ins->name())) + continue; + sizes.insert(sizes.end(), ins->inputs().begin(), ins->inputs().end()); + } + auto name = + mlir_program::get_symbol_name(m) + "_" + shape::to_sizes_string(to_shapes(sizes)) + ".mxr"; + replace_string_inplace(name, ", ", "_"); + replace_string_inplace(name, ":", "s"); + auto f = location / name; + save(program{std::move(m)}, f.string()); +} + #else template