Merge branch 'develop' of https://github.com/ROCmSoftwarePlatform/AMDMIGraphX into mi100_opts

src/decompose.cpp

@@ -39,9 +39,7 @@ struct find_dot_add
         {
             auto alpha = p.add_literal(literal{shape{a_ins->get_shape().type()}, {dot.alpha}});
             auto alpha_broadcast = p.insert_instruction(
-                ins,
-                make_op("multibroadcast", {{"output_lens", a_ins->get_shape().lens()}}),
-                alpha);
+                ins, make_op("multibroadcast", {{"out_lens", a_ins->get_shape().lens()}}), alpha);
             a_ins = p.insert_instruction(ins, make_op("mul"), a_ins, alpha_broadcast);
         }
         auto dot_ins = p.insert_instruction(ins, make_op(ins->name(), {{"beta", 0}}), a_ins, b_ins);
@@ -51,7 +49,7 @@ struct find_dot_add
         {
             auto beta = p.add_literal(literal{shape{c_ins->get_shape().type()}, {dot.beta}});
             auto beta_broadcast = p.insert_instruction(
-                ins, make_op("multibroadcast", {{"output_lens", ins->get_shape().lens()}}), beta);
+                ins, make_op("multibroadcast", {{"out_lens", ins->get_shape().lens()}}), beta);
             c_ins = p.insert_instruction(ins, make_op("mul"), c_ins, beta_broadcast);
         }
         p.replace_instruction(ins, make_op("add"), dot_ins, c_ins);
@@ -72,9 +70,7 @@ struct find_dot_alpha
         {
             auto alpha = p.add_literal(literal{shape{a_ins->get_shape().type()}, {dot.alpha}});
             auto alpha_broadcast = p.insert_instruction(
-                ins,
-                make_op("multibroadcast", {{"output_lens", a_ins->get_shape().lens()}}),
-                alpha);
+                ins, make_op("multibroadcast", {{"out_lens", a_ins->get_shape().lens()}}), alpha);
             a_ins = p.insert_instruction(ins, make_op("mul"), a_ins, alpha_broadcast);
         }
         p.replace_instruction(ins, make_op(ins->name(), {{"beta", 0}}), a_ins, b_ins);

src/driver/main.cpp

@@ -503,6 +503,26 @@ struct op : command<op>
     }
 };
 
+struct onnx : command<onnx>
+{
+    bool show_ops = false;
+    void parse(argument_parser& ap)
+    {
+        ap(show_ops,
+           {"--list", "-l"},
+           ap.help("List all onnx operators supported by MIGraphX"),
+           ap.set_value(true));
+    }
+    void run() const
+    {
+        if(show_ops)
+        {
+            for(const auto& name : get_onnx_operators())
+                std::cout << name << std::endl;
+        }
+    }
+};
+
 struct main_command
 {
     static std::string get_command_help()

src/eliminate_data_type.cpp

@@ -3,17 +3,20 @@
 #include <migraphx/iterator_for.hpp>
 #include <migraphx/make_op.hpp>
 #include <migraphx/instruction.hpp>
+#include <migraphx/ranges.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
 void eliminate_data_type::apply(module& m) const
 {
+    static const std::vector<std::string> skip_op_names = {
+        "convert", "get_tuple_elem", "if", "loop"};
     for(auto ins : iterator_for(m))
     {
         if(ins->name()[0] == '@')
             continue;
-        if(ins->name() == "convert")
+        if(contains(skip_op_names, ins->name()))
             continue;
         auto inputs = ins->inputs();
         std::transform(inputs.begin(), inputs.end(), inputs.begin(), [&](auto i) {

src/file_buffer.cpp

@@ -6,7 +6,8 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
-std::vector<char> read_buffer(const std::string& filename)
+template <class T>
+T generic_read_file(const std::string& filename)
 {
     std::ifstream is(filename, std::ios::binary | std::ios::ate);
     std::streamsize size = is.tellg();
@@ -14,12 +15,22 @@ std::vector<char> read_buffer(const std::string& filename)
         MIGRAPHX_THROW("Invalid size for: " + filename);
     is.seekg(0, std::ios::beg);
 
-    std::vector<char> buffer(size);
-    if(!is.read(buffer.data(), size))
+    T buffer(size, 0);
+    if(!is.read(&buffer[0], size))
         MIGRAPHX_THROW("Error reading file: " + filename);
     return buffer;
 }
 
+std::vector<char> read_buffer(const std::string& filename)
+{
+    return generic_read_file<std::vector<char>>(filename);
+}
+
+std::string read_string(const std::string& filename)
+{
+    return generic_read_file<std::string>(filename);
+}
+
 void write_buffer(const std::string& filename, const char* buffer, std::size_t size)
 {
     std::ofstream os(filename);

src/generate.cpp

@@ -6,32 +6,59 @@ inline namespace MIGRAPHX_INLINE_NS {
 argument fill_argument(shape s, unsigned long value)
 {
     argument result;
-    s.visit_type([&](auto as) {
-        using type = typename decltype(as)::type;
-        auto v     = fill_tensor_data<type>(s, value);
-        result     = {s, v};
-    });
+    if(s.type() == shape::tuple_type)
+    {
+        std::vector<argument> sub_args;
+        const auto& sub_ss = s.sub_shapes();
+        std::transform(sub_ss.begin(), sub_ss.end(), std::back_inserter(sub_args), [&](auto ss) {
+            return fill_argument(ss, value);
+        });
+
+        result = argument(sub_args);
+    }
+    else
+    {
+        s.visit_type([&](auto as) {
+            using type = typename decltype(as)::type;
+            auto v     = fill_tensor_data<type>(s, value);
+            result     = {s, v};
+        });
+    }
     return result;
 }
 
 argument generate_argument(shape s, unsigned long seed)
 {
     argument result;
-    s.visit_type([&](auto as) {
-        // we use char type to store bool type internally, so bool_type
-        // needs special processing to generate data
-        if(s.type() == shape::bool_type)
-        {
-            auto v = generate_tensor_data<bool>(s, seed);
-            result = {s, v};
-        }
-        else
-        {
-            using type = typename decltype(as)::type;
-            auto v     = generate_tensor_data<type>(s, seed);
-            result     = {s, v};
-        }
-    });
+    if(s.type() == shape::tuple_type)
+    {
+        const auto& sub_ss = s.sub_shapes();
+        std::vector<argument> sub_args;
+        std::transform(sub_ss.begin(), sub_ss.end(), std::back_inserter(sub_args), [&](auto ss) {
+            return generate_argument(ss, seed);
+        });
+
+        result = argument(sub_args);
+    }
+    else
+    {
+        s.visit_type([&](auto as) {
+            // we use char type to store bool type internally, so bool_type
+            // needs special processing to generate data
+            if(s.type() == shape::bool_type)
+            {
+                auto v = generate_tensor_data<bool>(s, seed);
+                result = {s, v};
+            }
+            else
+            {
+                using type = typename decltype(as)::type;
+                auto v     = generate_tensor_data<type>(s, seed);
+                result     = {s, v};
+            }
+        });
+    }
+
     return result;
 }
 

src/include/migraphx/argument.hpp

@@ -27,29 +27,29 @@ struct argument : raw_data<argument>
 
     template <class F, MIGRAPHX_REQUIRES(std::is_pointer<decltype(std::declval<F>()())>{})>
     argument(shape s, F d)
-        : m_shape(std::move(s)),
-          m_data({[f = std::move(d)]() mutable { return reinterpret_cast<char*>(f()); }})
+        : m_shape(std::move(s))
 
     {
+        assign_buffer([f = std::move(d)]() mutable { return reinterpret_cast<char*>(f()); });
     }
     template <class T>
     argument(shape s, T* d)
-        : m_shape(std::move(s)), m_data({[d] { return reinterpret_cast<char*>(d); }})
+        : m_shape(std::move(s))
     {
+        assign_buffer([d] { return reinterpret_cast<char*>(d); });
     }
 
     template <class T>
     argument(shape s, std::shared_ptr<T> d)
-        : m_shape(std::move(s)), m_data({[d] { return reinterpret_cast<char*>(d.get()); }})
+        : m_shape(std::move(s))
     {
+        assign_buffer([d] { return reinterpret_cast<char*>(d.get()); });
     }
 
     argument(shape s, std::nullptr_t);
     
     argument(const std::vector<argument>& args);
 
-    static argument load(const shape& s, char* buffer);
-
     /// Provides a raw pointer to the data
     char* data() const;
 
@@ -68,6 +68,7 @@ struct argument : raw_data<argument>
     std::vector<argument> get_sub_objects() const;
 
     private:
+    void assign_buffer(std::function<char*()> d);
     struct data_t
     {
         std::function<char*()> get = nullptr;

src/include/migraphx/errors.hpp

@@ -49,12 +49,13 @@ inline std::string make_source_context(const std::string& file, int line, const
     return file + ":" + std::to_string(line) + ": " + fname;
 }
 
+// NOLINTNEXTLINE
+#define MIGRAPHX_MAKE_SOURCE_CTX() migraphx::make_source_context(__FILE__, __LINE__, __func__)
+
 /**
  * @brief Throw an exception with context information
  */
-#define MIGRAPHX_THROW(...)                                                                     \
-    throw migraphx::make_exception(migraphx::make_source_context(__FILE__, __LINE__, __func__), \
-                                   __VA_ARGS__)
+#define MIGRAPHX_THROW(...) throw migraphx::make_exception(MIGRAPHX_MAKE_SOURCE_CTX(), __VA_ARGS__)
 
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/include/migraphx/file_buffer.hpp

@@ -9,6 +9,7 @@ namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
 std::vector<char> read_buffer(const std::string& filename);
+std::string read_string(const std::string& filename);
 
 void write_buffer(const std::string& filename, const char* buffer, std::size_t size);
 void write_buffer(const std::string& filename, const std::vector<char>& buffer);

src/include/migraphx/matcher.hpp

@@ -263,6 +263,20 @@ matcher_result match_instruction(module& mod, instruction_ref ins, M&& m)
     return result;
 }
 
+/// Find first instance of a matching instruction in a module
+template <class M>
+match::matcher_result find_match(module& modl, M&& m)
+{
+    match::matcher_result result;
+    for(auto ins : iterator_for(modl))
+    {
+        result = match::match_instruction(modl, ins, m);
+        if(result.result != modl.end())
+            return result;
+    }
+    return result;
+}
+
 MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_TRACE_MATCHES)
 
 /// Find matches for an instruction in the module

src/include/migraphx/onnx.hpp

@@ -18,6 +18,8 @@ struct onnx_options
     bool skip_unknown_operators = false;
     /// Print program if an error occurs
     bool print_program_on_error = false;
+    /// Max iter num for the loop operator
+    int64_t max_loop_iterations = 10;
 };
 
 /// Create a program from an onnx file
@@ -29,6 +31,8 @@ program parse_onnx_buffer(const std::string& buffer, const onnx_options& options
 /// Create a program from an onnx buffer
 program parse_onnx_buffer(const void* data, std::size_t size, const onnx_options& options);
 
+std::vector<std::string> get_onnx_operators();
+
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx
 

src/include/migraphx/op/broadcast.hpp

@@ -30,7 +30,7 @@ struct broadcast
     template <class Self, class F>
     static auto reflect(Self& self, F f)
     {
-        return pack(f(self.axis, "axis"), f(self.broadcast_lens, "dims"));
+        return pack(f(self.axis, "axis"), f(self.broadcast_lens, "out_lens"));
     }
 
     std::string name() const { return "broadcast"; }

src/include/migraphx/op/capture.hpp

@@ -8,6 +8,7 @@
 #include <migraphx/literal.hpp>
 #include <migraphx/shape_for_each.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/context.hpp>
 #include <cmath>
 #include <utility>
 
@@ -29,7 +30,9 @@ struct capture
 
     shape compute_shape(std::vector<shape> inputs) const { return inputs.front(); }
 
-    argument compute(const shape&, std::vector<argument> args) const
+    // the context argument is added to prevent the op from be eliminated by
+    // constant propagation
+    argument compute(context&, const shape&, const std::vector<argument>& args) const
     {
         if(f)
         {
@@ -42,6 +45,8 @@ struct capture
 
         return args.front();
     }
+
+    std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 
 } // namespace op

src/include/migraphx/op/deconvolution.hpp

@@ -9,6 +9,8 @@
 #include <migraphx/literal.hpp>
 #include <migraphx/shape_for_each.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/dfor.hpp>
+#include <migraphx/par_dfor.hpp>
 #include <cmath>
 #include <utility>
 
@@ -70,6 +72,81 @@ struct deconvolution
         return inputs[0].with_lens(output_lens);
     }
 
+    argument compute(shape output_shape, std::vector<argument> args) const
+    {
+        argument result{output_shape};
+        auto kdims = this->kdims();
+        visit_all(result, args[0], args[1])([&](auto output, auto input, auto weights) {
+            using type = typename decltype(output)::value_type;
+
+            std::fill(output.begin(), output.end(), type{0});
+
+            auto in_lens = input.get_shape().lens();
+            auto in_n    = in_lens[0];
+            auto in_c    = in_lens[1];
+
+            auto wei   = weights.get_shape().lens();
+            auto wei_n = wei[0];
+            auto wei_c = wei[1];
+
+            auto out_lens = output_shape.lens();
+
+            std::vector<std::size_t> win_size{in_c};
+            std::copy(in_lens.begin() + 2, in_lens.end(), std::back_inserter(win_size));
+            std::copy(wei.begin() + 2, wei.end(), std::back_inserter(win_size));
+            shape win_shape{output_shape.type(), win_size};
+
+            par_dfor(in_n, wei_c)([&](int o, int k) {
+
+                shape_for_each(win_shape, [&](auto idx_win) {
+                    const int w = idx_win[0];
+
+                    auto input_dims_start = idx_win.begin() + 1;
+                    auto wei_dims_start   = idx_win.begin() + kdims + 1;
+
+                    std::vector<std::ptrdiff_t> win_start;
+                    for(std::size_t n = 0; n < kdims; ++n)
+                    {
+                        win_start.push_back(std::ptrdiff_t(*(input_dims_start + n) * stride[n]) -
+                                            std::ptrdiff_t(padding[n]));
+                    }
+
+                    const int group_id = w / (wei_n / group);
+                    const int in_ch    = group_id * wei_c + k;
+
+                    std::vector<std::ptrdiff_t> idx_out{o, in_ch};
+
+                    for(size_t n = 0; n < kdims; n++)
+                    {
+                        idx_out.push_back(win_start[n] + *(wei_dims_start + n) * dilation[n]);
+                    }
+
+                    std::vector<std::ptrdiff_t> idx_wei{w, k};
+                    std::copy(wei_dims_start, idx_win.end(), std::back_inserter(idx_wei));
+
+                    std::vector<std::ptrdiff_t> idx_in{o, w};
+                    std::copy(input_dims_start, wei_dims_start, std::back_inserter(idx_in));
+
+                    if(std::all_of(
+                           idx_out.begin() + 2, idx_out.end(), [&](auto ii) { return ii >= 0; }) and
+                       std::equal(idx_out.begin() + 2,
+                                  idx_out.end(),
+                                  out_lens.begin() + 2,
+                                  out_lens.end(),
+                                  std::less<std::ptrdiff_t>{}))
+                    {
+                        output(idx_out.begin(), idx_out.end()) +=
+                            input(idx_in.begin(), idx_in.end()) *
+                            weights(idx_wei.begin(), idx_wei.end());
+                    }
+                });
+
+            });
+
+        });
+        return result;
+    }
+
     size_t kdims() const
     {
         check_attribute_size();

src/include/migraphx/op/dequantizelinear.hpp

@@ -25,14 +25,15 @@ struct dequantizelinear
     std::string name() const { return "dequantizelinear"; }
     shape compute_shape(std::vector<shape> inputs) const
     {
-        return {shape::float_type, inputs[0].lens(), inputs[0].strides()};
+        check_shapes{inputs, *this}.same_dims();
+        return {inputs[1].type(), inputs[0].lens(), inputs[0].strides()};
     }
 
     argument compute(const shape& output_shape, std::vector<argument> args) const
     {
         auto x       = args.at(0);
         auto x_scale = args.at(1);
-        std::vector<int8_t> zeros(output_shape.elements(), 0);
+        std::vector<int8_t> zeros(output_shape.bytes(), 0);
         argument x_zero_point{{x.get_shape().type(), output_shape.lens()}, zeros.data()};
         if(args.size() == 3)
         {

src/include/migraphx/op/get_tuple_elem.hpp

@@ -45,6 +45,8 @@ struct get_tuple_elem
         assert(index < vec_args.size());
         return vec_args.at(index);
     }
+
+    std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 
 } // namespace op

src/include/migraphx/op/load.hpp

@@ -35,7 +35,7 @@ struct load
     {
         if((offset + s.bytes()) > args[0].get_shape().bytes())
             MIGRAPHX_THROW("Load access is out of bounds");
-        return argument::load(s, args[0].data() + offset);
+        return argument{s, args[0].data() + offset};
     }
     lifetime get_lifetime() const { return lifetime::borrow; }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }

src/include/migraphx/op/loop.hpp

+#ifndef MIGRAPHX_GUARD_OPERATORS_LOOP_HPP
+#define MIGRAPHX_GUARD_OPERATORS_LOOP_HPP
+
+#include <migraphx/check_shapes.hpp>
+#include <migraphx/argument.hpp>
+#include <migraphx/functional.hpp>
+#include <migraphx/config.hpp>
+#include <migraphx/module.hpp>
+#include <migraphx/run_loop.hpp>
+#include <migraphx/ranges.hpp>
+#include <cmath>
+#include <string>
+#include <utility>
+#include <set>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace op {
+
+struct loop
+{
+    int64_t max_iterations = 10;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.max_iterations, "max_iterations"));
+    }
+
+    std::string name() const { return "loop"; }
+
+    shape compute_shape(const std::vector<shape>& inputs, std::vector<module_ref> mods) const
+    {
+        check_shapes{inputs, *this}.standard();
+        if(mods.size() != 1)
+        {
+            MIGRAPHX_THROW("LOOP: operator should have one submodule.");
+        }
+
+        const auto& mod     = mods.front();
+        auto mod_out_shapes = mod->get_output_shapes();
+        auto dep_param_num  = inputs.size() - 2;
+
+        // first item of the mod output shapes is condition used in loop,
+        // which is not needed to compute output shape
+        mod_out_shapes.erase(mod_out_shapes.begin());
+        std::vector<shape> ins_out_shapes(mod_out_shapes.begin(),
+                                          mod_out_shapes.begin() + dep_param_num);
+        mod_out_shapes.erase(mod_out_shapes.begin(), mod_out_shapes.begin() + dep_param_num);
+        for(const auto& out_s : mod_out_shapes)
+        {
+            auto lens = out_s.lens();
+            lens.insert(lens.begin(), max_iterations);
+            ins_out_shapes.push_back({out_s.type(), lens});
+        }
+
+        return shape(ins_out_shapes);
+    }
+
+    struct ref_loop
+    {
+        int64_t max_iterations = 0;
+
+        template <class T>
+        void copy(context&, const argument& src, T& dst) const
+        {
+            dst = *src.cast<T>();
+        }
+
+        template <class T>
+        void copy(context&, T src, const argument& dst) const
+        {
+            *dst.cast<T>() = src;
+        }
+
+        void append(const std::vector<argument>& iter_state,
+                    const std::vector<argument>& concatenated_outputs,
+                    int iter) const
+        {
+            assert(iter_state.size() == concatenated_outputs.size());
+            for(auto i : range(iter_state.size()))
+            {
+                const auto& iter_stat = iter_state.at(i);
+                const auto& scan_out  = concatenated_outputs.at(i);
+
+                auto* in_data        = iter_stat.data();
+                auto* out_data       = scan_out.data();
+                std::size_t out_size = iter_stat.get_shape().bytes();
+                assert((iter + 1) * out_size <= scan_out.get_shape().bytes());
+                std::copy(in_data, in_data + out_size, out_data + iter * out_size);
+            }
+        }
+
+        void set_zero(context&, const std::vector<argument>& concatenated_outputs, int iter) const
+        {
+            if(iter >= max_iterations)
+                return;
+
+            for(const auto& out : concatenated_outputs)
+            {
+                auto s    = out.get_shape();
+                auto size = s.bytes() / max_iterations;
+                std::fill(out.data() + iter * size, out.data() + max_iterations * size, 0);
+            }
+        }
+
+        std::unordered_map<std::string, int> get_output_params(const module&) const { return {}; }
+    };
+
+    argument compute(context& ctx,
+                     const shape& out_shape,
+                     const std::vector<argument>& args,
+                     const std::vector<module_ref>& mods,
+                     const std::function<std::vector<argument>(
+                         module_ref&, const std::unordered_map<std::string, argument>&)>& run) const
+    {
+        // wrap up the arguments vector, so ref and gpu impl are the same
+        auto cpy_args = args;
+        bool in_cond  = args.at(1).at<bool>();
+        bool cond     = in_cond;
+        int64_t iter  = 0;
+        // insert iter and cond used in the loop
+        auto s_cond = args.at(1).get_shape();
+        auto s_iter = args.at(0).get_shape();
+        cpy_args.push_back({s_iter, &iter});
+        cpy_args.push_back({s_cond, &cond});
+        cpy_args.insert(cpy_args.end(), args.begin() + 2, args.end());
+
+        // add cond and mod outputs to the argument list
+        cpy_args.push_back(argument(s_cond));
+        cpy_args.push_back(argument(out_shape));
+
+        // run loop
+        return run_loop(ref_loop{max_iterations}, ctx, cpy_args, mods, run);
+    }
+};
+
+} // namespace op
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/op/multibroadcast.hpp

@@ -23,7 +23,7 @@ struct multibroadcast
     template <class Self, class F>
     static auto reflect(Self& self, F f)
     {
-        return pack(f(self.output_lens, "output_lens"));
+        return pack(f(self.output_lens, "out_lens"));
     }
 
     std::string name() const { return "multibroadcast"; }

src/include/migraphx/op/quantizelinear.hpp

@@ -25,6 +25,7 @@ struct quantizelinear
     std::string name() const { return "quantizelinear"; }
     shape compute_shape(std::vector<shape> inputs) const
     {
+        check_shapes{inputs, *this}.same_dims();
         if(inputs.size() == 3)
         {
             return {inputs[2].type(), inputs[0].lens(), inputs[0].strides()};
@@ -36,7 +37,7 @@ struct quantizelinear
     {
         auto x       = args.at(0);
         auto y_scale = args.at(1);
-        std::vector<int8_t> zeros(output_shape.elements(), 0);
+        std::vector<int8_t> zeros(output_shape.bytes(), 0);
         argument y_zero_point{output_shape, zeros.data()};
         if(args.size() == 3)
         {