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

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

src/include/migraphx/op/unary.hpp

+#ifndef MIGRAPHX_GUARD_OPERATORS_UNARY_HPP
+#define MIGRAPHX_GUARD_OPERATORS_UNARY_HPP
+
+#include <migraphx/op/name.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace op {
+
+template <class Derived>
+struct unary : op_name<Derived>
+{
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs}.has(1);
+        auto s = inputs.at(0);
+        if(s.packed())
+        {
+            return s;
+        }
+        else
+        {
+            return {s.type(), s.lens()};
+        }
+    }
+    argument compute(const shape& output_shape, std::vector<argument> args) const
+    {
+        argument result{output_shape};
+        visit_all(result, args[0])([&](auto output, auto input) {
+            if(input.get_shape().standard())
+            {
+                std::transform(input.begin(),
+                               input.end(),
+                               output.begin(),
+                               static_cast<const Derived&>(*this).apply());
+            }
+            else
+            {
+                shape_for_each(output.get_shape(), [&](const auto& idx) {
+                    output(idx.begin(), idx.end()) =
+                        static_cast<const Derived&>(*this).apply()(input(idx.begin(), idx.end()));
+                });
+            }
+        });
+        return result;
+    }
+};
+
+} // namespace op
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/op/unsqueeze.hpp

+#ifndef MIGRAPHX_GUARD_OPERATORS_UNSQUEEZE_HPP
+#define MIGRAPHX_GUARD_OPERATORS_UNSQUEEZE_HPP
+
+#include <array>
+#include <migraphx/operation.hpp>
+#include <migraphx/check_shapes.hpp>
+#include <migraphx/stringutils.hpp>
+#include <migraphx/streamutils.hpp>
+#include <migraphx/literal.hpp>
+#include <migraphx/shape_for_each.hpp>
+#include <migraphx/config.hpp>
+#include <cmath>
+#include <utility>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace op {
+
+struct unsqueeze
+{
+    std::vector<int64_t> axes;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.axes, "axes"));
+    }
+
+    std::string name() const { return "unsqueeze"; }
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        auto input_shape = inputs[0];
+        auto type        = input_shape.type();
+        auto old_lens    = input_shape.lens();
+
+        if(input_shape.scalar())
+            return shape{type, old_lens};
+
+        std::size_t new_size = old_lens.size() + axes.size();
+        std::vector<std::size_t> new_lens(new_size);
+        std::size_t p = 0;
+        for(std::size_t i = 0; i < new_size; i++)
+        {
+            if(std::find(axes.begin(), axes.end(), i) != axes.end())
+            {
+                new_lens[i] = 1;
+            }
+            else
+            {
+                new_lens[i] = old_lens[p++];
+            }
+        }
+        return shape{type, new_lens};
+    }
+    argument compute(shape output_shape, std::vector<argument> args) const
+    {
+        return {std::move(output_shape), std::move(args.front().data)};
+    }
+    std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
+};
+
+} // namespace op
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/operation.hpp

@@ -49,7 +49,7 @@ struct operation
     argument compute(context& ctx, const shape& output, const std::vector<argument>& input) const;
     /// An optional method to return which argument the output will alias. If
     /// there is no aliased output then -1 can be returned.
-    int output_alias(const std::vector<shape>& input) const;
+    std::ptrdiff_t output_alias(const std::vector<shape>& input) const;
     /// An optional stream operator to print the operation. When this is not
     /// implemented, it will just print the operation's name.
     friend std::ostream& operator<<(std::ostream& os, const operation& op);
@@ -69,7 +69,7 @@ auto operator<<(std::ostream& os, const T& x) -> decltype(os << x.name())
 {
     os << x.name();
     char delim = '[';
-    reflect_each(x, [&](auto& y, auto name) {
+    reflect_each(x, [&](auto&& y, auto name) {
         os << delim;
         os << name << "=";
         stream_write_value(os, y);
@@ -87,6 +87,8 @@ namespace operation_equal {
 template <class T, class U>
 auto operator==(const T& x, const U& y) -> decltype(x.name() == y.name())
 {
+    static_assert(is_reflectable<T>{} or sizeof(T) <= 1,
+                  "Missing equality operator or reflect method.");
     if(x.name() != y.name())
         return false;
     const auto& yy = any_cast<T>(y);
@@ -175,7 +177,7 @@ auto is_context_free_op(const T& x) -> decltype(is_context_free_op(
 }
 
 template <class T>
-int output_alias_op(rank<0>, const T&, const std::vector<shape>&)
+std::ptrdiff_t output_alias_op(rank<0>, const T&, const std::vector<shape>&)
 {
     return -1;
 }
@@ -188,7 +190,7 @@ auto output_alias_op(rank<1>, const T& x, const std::vector<shape>& shapes)
 }
 
 template <class T>
-int output_alias_op(const T& x, const std::vector<shape>& shapes)
+std::ptrdiff_t output_alias_op(const T& x, const std::vector<shape>& shapes)
 {
     return output_alias_op(rank<1>{}, x, shapes);
 }
@@ -239,7 +241,7 @@ auto has_finalize_op(const T&) -> decltype(has_finalize_op(rank<1>{},
  *      std::string name() const;
  *      bool is_context_free() const;
  *      bool has_finalize() const;
- *      int output_alias(const std::vector<shape>& input) const;
+ *      std::ptrdiff_t output_alias(const std::vector<shape>& input) const;
  *      void finalize(context& ctx,const shape& output,const std::vector<shape>& input) ;
  *      shape compute_shape(const std::vector<shape>& input) const;
  *      argument compute(context& ctx,const shape& output,const std::vector<argument>& input) const;
@@ -325,7 +327,7 @@ struct operation
         return (*this).private_detail_te_get_handle().has_finalize();
     }
 
-    int output_alias(const std::vector<shape>& input) const
+    std::ptrdiff_t output_alias(const std::vector<shape>& input) const
     {
         assert((*this).private_detail_te_handle_mem_var);
         return (*this).private_detail_te_get_handle().output_alias(input);
@@ -380,10 +382,10 @@ struct operation
         virtual std::shared_ptr<private_detail_te_handle_base_type> clone() const = 0;
         virtual const std::type_info& type() const                                = 0;
 
-        virtual std::string name() const                                = 0;
-        virtual bool is_context_free() const                            = 0;
-        virtual bool has_finalize() const                               = 0;
-        virtual int output_alias(const std::vector<shape>& input) const = 0;
+        virtual std::string name() const                                           = 0;
+        virtual bool is_context_free() const                                       = 0;
+        virtual bool has_finalize() const                                          = 0;
+        virtual std::ptrdiff_t output_alias(const std::vector<shape>& input) const = 0;
         virtual void
         finalize(context& ctx, const shape& output, const std::vector<shape>& input) = 0;
         virtual shape compute_shape(const std::vector<shape>& input) const           = 0;
@@ -432,7 +434,7 @@ struct operation
 
         bool has_finalize() const override { return has_finalize_op(private_detail_te_value); }
 
-        int output_alias(const std::vector<shape>& input) const override
+        std::ptrdiff_t output_alias(const std::vector<shape>& input) const override
         {
 
             return output_alias_op(private_detail_te_value, input);

src/include/migraphx/pass_manager.hpp

+#ifndef MIGRAPHX_GUARD_MIGRAPHLIB_PASS_MANAGER_HPP
+#define MIGRAPHX_GUARD_MIGRAPHLIB_PASS_MANAGER_HPP
+
+#include <list>
+#include <unordered_map>
+#include <migraphx/operation.hpp>
+#include <migraphx/literal.hpp>
+#include <migraphx/builtin.hpp>
+#include <migraphx/instruction_ref.hpp>
+#include <migraphx/target.hpp>
+#include <migraphx/tracer.hpp>
+#include <migraphx/env.hpp>
+#include <migraphx/config.hpp>
+#include <algorithm>
+#include <iostream>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+void run_passes(program& prog, const std::vector<pass>& passes, tracer trace = tracer{});
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/program.hpp

@@ -30,8 +30,16 @@ const operation& get_operation(instruction_ref ins);
 struct program
 {
     program();
+
+    // move constructor
     program(program&&) noexcept;
-    program& operator=(program&&) noexcept;
+
+    // copy constructor
+    program(const program&);
+
+    // copy assignment operator
+    program& operator=(program);
+
     ~program() noexcept;
 
     using parameter_map = std::unordered_map<std::string, argument>;
@@ -108,6 +116,7 @@ struct program
     void debug_print() const;
     void debug_print(instruction_ref ins) const;
     void debug_print(const std::vector<instruction_ref>& inss) const;
+    void print_graph(std::ostream& os) const;
 
     void dry_run(parameter_map params) const;
 
@@ -117,6 +126,9 @@ struct program
     friend bool operator==(const program& x, const program& y);
     friend bool operator!=(const program& x, const program& y) { return !(x == y); }
 
+    private:
+    void assign(const program& p);
+
     private:
     std::unique_ptr<program_impl> impl;
 };

src/include/migraphx/constant_propagate.hpp → src/include/migraphx/propagate_constant.hpp

-#ifndef MIGRAPHX_GUARD_RTGLIB_CONSTANT_PROPAGATE_HPP
-#define MIGRAPHX_GUARD_RTGLIB_CONSTANT_PROPAGATE_HPP
+#ifndef MIGRAPHX_GUARD_RTGLIB_PROPAGATE_CONSTANT_HPP
+#define MIGRAPHX_GUARD_RTGLIB_PROPAGATE_CONSTANT_HPP
 
 #include <string>
 #include <migraphx/config.hpp>
@@ -12,9 +12,9 @@ struct program;
 /**
  * Replace instructions which take all literals with a literal of the computation.
  */
-struct constant_propagate
+struct propagate_constant
 {
-    std::string name() const { return "constant_propagate"; }
+    std::string name() const { return "propagate_constant"; }
     void apply(program& p) const;
 };
 

src/include/migraphx/ranges.hpp

@@ -12,7 +12,7 @@ inline namespace MIGRAPHX_INLINE_NS {
 namespace detail {
 
 template <class String, class T>
-auto generic_find_impl(rank<2>, String&& s, const T& x) -> decltype(s.begin() + s.find(x), s.npos)
+auto generic_find_impl(rank<2>, String&& s, const T& x) -> decltype(s.npos, s.begin() + s.find(x))
 {
     auto index = s.find(x);
     if(index == s.npos)

src/include/migraphx/reflect.hpp

@@ -11,6 +11,15 @@ inline namespace MIGRAPHX_INLINE_NS {
 
 namespace detail {
 
+struct reflect_placeholder
+{
+    template <class... Ts>
+    int operator()(Ts&&...) const
+    {
+        return 0;
+    }
+};
+
 template <class T, class Selector>
 auto reflect_impl(rank<1>, T& x, Selector f) -> decltype(T::reflect(x, f))
 {
@@ -23,8 +32,53 @@ auto reflect_impl(rank<0>, T&, Selector)
     return pack();
 }
 
+template <class T>
+auto reflectable_impl(rank<1>, T&& x)
+    -> decltype(T::reflect(x, reflect_placeholder{}), std::true_type{});
+
+template <class T>
+auto reflectable_impl(rank<0>, T &&) -> decltype(std::false_type{});
+
+template <class T>
+struct remove_rvalue_reference
+{
+    using type = T;
+};
+
+template <class T>
+struct remove_rvalue_reference<T&&>
+{
+    using type = T;
+};
+
+template <class T>
+struct wrapper
+{
+    using type = typename remove_rvalue_reference<T>::type;
+    type data;
+    type get() const { return data; }
+};
+
+template <class T>
+wrapper<T> wrap(std::remove_reference_t<T>& x)
+{
+    return wrapper<T>{std::forward<T>(x)};
+}
+
+template <class... Ts>
+using auto_tuple_t = std::tuple<typename remove_rvalue_reference<Ts>::type...>;
+
+template <class... Ts>
+auto_tuple_t<Ts...> auto_tuple(Ts&&... xs)
+{
+    return auto_tuple_t<Ts...>{std::forward<Ts>(xs)...};
+}
+
 } // namespace detail
 
+template <class T>
+using is_reflectable = decltype(detail::reflectable_impl(rank<1>{}, std::declval<T>()));
+
 template <class T, class Selector>
 auto reflect(T& x, Selector f)
 {
@@ -34,17 +88,18 @@ auto reflect(T& x, Selector f)
 template <class T>
 auto reflect_tie(T& x)
 {
-    return reflect(x, [](auto&& y, auto&&...) { return std::ref(y); })(
-        [](auto&&... xs) { return std::tie(xs.get()...); });
+    return reflect(x, [](auto&& y, auto&&...) { return detail::wrap<decltype(y)>(y); })(
+        [](auto&&... xs) { return detail::auto_tuple(xs.get()...); });
 }
 
 template <class T, class F>
 void reflect_each(T& x, F f)
 {
-    return reflect(x, [](auto&& y, auto... ys) { return pack(std::ref(y), ys...); })(
-        [&](auto&&... xs) {
-            each_args([&](auto p) { p([&](auto&& y, auto... ys) { f(y.get(), ys...); }); }, xs...);
-        });
+    return reflect(x, [](auto&& y, auto... ys) {
+        return pack(detail::wrap<decltype(y)>(y), ys...);
+    })([&](auto&&... xs) {
+        each_args([&](auto p) { p([&](auto&& y, auto... ys) { f(y.get(), ys...); }); }, xs...);
+    });
 }
 
 } // namespace MIGRAPHX_INLINE_NS

src/include/migraphx/rewrite_rnn.hpp

@@ -4,7 +4,7 @@
 #include <string>
 #include <vector>
 #include <migraphx/instruction_ref.hpp>
-#include <migraphx/operators.hpp>
+#include <migraphx/operation.hpp>
 #include <migraphx/config.hpp>
 
 namespace migraphx {

src/include/migraphx/stringutils.hpp

@@ -38,8 +38,9 @@ inline std::string join_strings(Strings strings, const std::string& delim)
         return "";
 
     auto nit = std::next(it);
-    return std::accumulate(
-        nit, strings.end(), *it, [&](std::string x, std::string y) { return x + delim + y; });
+    return std::accumulate(nit, strings.end(), *it, [&](std::string x, std::string y) {
+        return std::move(x) + delim + std::move(y);
+    });
 }
 
 template <class F>

src/instruction.cpp

@@ -28,6 +28,12 @@ void instruction::replace(const shape& r)
     }
 }
 
+void instruction::replace(operation o)
+{
+    op = std::move(o);
+    recompute_shape();
+}
+
 void instruction::recompute_shape() { replace(compute_shape(op, arguments)); }
 
 void instruction::clear_arguments()
@@ -162,7 +168,24 @@ void instruction::replace_argument(instruction_ref old, instruction_ref new_ins)
     old->remove_output(*this);
 }
 
-argument instruction::eval() const
+bool instruction::can_eval() const
+{
+    if(op.name() == "@literal")
+    {
+        return true;
+    }
+    else if(is_context_free(op))
+    {
+        return std::all_of(
+            this->inputs().begin(), this->inputs().end(), [](auto arg) { return arg->can_eval(); });
+    }
+    else
+    {
+        return false;
+    }
+}
+
+argument instruction::eval(bool check_eval) const
 {
     if(op.name() == "@literal")
     {
@@ -170,14 +193,13 @@ argument instruction::eval() const
     }
     if(is_context_free(op))
     {
+        if(check_eval and not this->can_eval())
+            return {};
         std::vector<argument> args;
-        for(auto&& arg : this->inputs())
-        {
-            argument a = arg->eval();
-            if(a.empty())
-                return {};
-            args.push_back(a);
-        }
+        std::transform(this->inputs().begin(),
+                       this->inputs().end(),
+                       std::back_inserter(args),
+                       [](auto arg) { return arg->eval(false); });
         return op.compute(result, args);
     }
     return {};

src/onnx/cifar10.cpp

@@ -32,7 +32,7 @@ auto read_cifar10_images(const std::string& full_path)
             labels[i] = *pimage++;
             for(size_t j = 0; j < nbytes_per_image; j++)
             {
-                float v                        = *(pimage + j) / 255.0f;
+                float v                        = float(*(pimage + j)) / 255.0f;
                 data[i * nbytes_per_image + j] = v;
             }
         }

src/onnx/onnx.cpp

@@ -36,7 +36,6 @@ struct onnx_parser
 
     onnx_parser()
     {
-        add_generic_op("MatMul", op::dot{});
         add_generic_op("Relu", op::relu{});
         add_generic_op("Sigmoid", op::sigmoid{});
         add_generic_op("Abs", op::abs{});
@@ -64,6 +63,7 @@ struct onnx_parser
         add_variadic_op("Max", op::max{});
         add_variadic_op("Min", op::min{});
 
+        add_mem_op("Clip", &onnx_parser::parse_clip);
         add_mem_op("LRN", &onnx_parser::parse_lrn);
         add_mem_op("ImageScaler", &onnx_parser::parse_imagescaler);
         add_mem_op("LeakyRelu", &onnx_parser::parse_leaky_relu);
@@ -77,6 +77,7 @@ struct onnx_parser
         add_mem_op("Reshape", &onnx_parser::parse_reshape);
         add_mem_op("Flatten", &onnx_parser::parse_flatten);
         add_mem_op("Gemm", &onnx_parser::parse_gemm);
+        add_mem_op("MatMul", &onnx_parser::parse_matmul);
         add_mem_op("BatchNormalization", &onnx_parser::parse_batchnorm);
         add_mem_op("Softmax", &onnx_parser::parse_softmax);
         add_mem_op("LogSoftmax", &onnx_parser::parse_logsoftmax);
@@ -141,8 +142,8 @@ struct onnx_parser
                 if(broadcasted != 0)
                 {
                     uint64_t axis = parse_value(attributes.at("axis")).at<uint64_t>();
-                    auto l =
-                        prog.add_instruction(op::broadcast{axis, args[0]->get_shape()}, args[1]);
+                    auto l = prog.add_instruction(op::broadcast{axis, args[0]->get_shape().lens()},
+                                                  args[1]);
                     return prog.add_instruction(x, args[0], l);
                 }
                 return prog.add_instruction(x, args);
@@ -154,42 +155,48 @@ struct onnx_parser
         });
     }
 
+    std::vector<std::size_t> compute_broadcasted_lens(std::vector<std::size_t> s0,
+                                                      std::vector<std::size_t> s1)
+    {
+        // Example:
+        // s0 = (3,2,4,5) and s1 = (2,1,1)
+        //
+        // In this case we need to broadcast (:,1,1) portion of
+        // s1 plus broadcast the 1st dimension of s1
+        // giving output_lens = (3,2,4,5)
+        //
+        // Another example:
+        // s0 = (3,2,1,5) and s1 = (2,7,5)
+        // In this case we need to broadcast the (:,:,1:,:) axis
+        // of s0 plus the 1st dimension of s1 giving
+        // output_lens = (3,2,7,5)
+        if(s0.size() > s1.size())
+        {
+            s0.swap(s1);
+        }
+
+        std::vector<std::size_t> out_lens(s1);
+        auto offset = s1.size() - s0.size();
+        std::transform(s0.begin(),
+                       s0.end(),
+                       s1.begin() + offset,
+                       out_lens.begin() + offset,
+                       [](auto a, auto b) { return std::max(a, b); });
+
+        return out_lens;
+    }
+
     template <class T>
     instruction_ref add_broadcastable_binary_op(instruction_ref arg0, instruction_ref arg1, T x)
     {
         if(arg0->get_shape().lens() != arg1->get_shape().lens())
         {
-            // Example:
-            // s0 = (3,2,4,5) and s1 = (2,1,1)
-            //
-            // In this case we need to broadcast (:,1,1) portion of
-            // s1 plus broadcast the 1st dimension of s1
-            // giving output_lens = (3,2,4,5)
-            //
-            // Another example:
-            // s0 = (3,2,1,5) and s1 = (2,7,5)
-            // In this case we need to broadcast the (:,:,1:,:) axis
-            // of s0 plus the 1st dimension of s1 giving
-            // output_lens = (3,2,7,5)
-            //
             // Get lengths for both arguments
-            const std::vector<std::size_t>* s0 = &arg0->get_shape().lens();
-            const std::vector<std::size_t>* s1 = &arg1->get_shape().lens();
-
-            // Make sure s0 is the smaller size
-            if(s0->size() > s1->size())
-                std::swap(s0, s1);
-
-            std::vector<std::size_t> output_lens(*s1);
-            auto offset = s1->size() - s0->size();
-            std::transform(s0->begin(),
-                           s0->end(),
-                           s1->begin() + offset,
-                           output_lens.begin() + offset,
-                           [](auto a, auto b) { return std::max(a, b); });
-
-            auto l0 = prog.add_instruction(op::multibroadcast{output_lens}, arg0);
-            auto l1 = prog.add_instruction(op::multibroadcast{output_lens}, arg1);
+            auto s0       = arg0->get_shape().lens();
+            auto s1       = arg1->get_shape().lens();
+            auto out_lens = compute_broadcasted_lens(s0, s1);
+            auto l0       = prog.add_instruction(op::multibroadcast{out_lens}, arg0);
+            auto l1       = prog.add_instruction(op::multibroadcast{out_lens}, arg1);
             return prog.add_instruction(x, l0, l1);
         }
         else
@@ -201,7 +208,7 @@ struct onnx_parser
     template <class T>
     void add_generic_op(std::string name, T x)
     {
-        add_op(name, [this, x](attribute_map, std::vector<instruction_ref> args) {
+        add_op(name, [this, x](const attribute_map&, std::vector<instruction_ref> args) {
             return prog.add_instruction(x, args);
         });
     }
@@ -209,7 +216,7 @@ struct onnx_parser
     template <class T>
     void add_variadic_op(std::string name, T x)
     {
-        add_op(name, [this, x](attribute_map, std::vector<instruction_ref> args) {
+        add_op(name, [this, x](const attribute_map&, std::vector<instruction_ref> args) {
             return std::accumulate(std::next(args.begin()),
                                    args.end(),
                                    args.front(),
@@ -219,6 +226,22 @@ struct onnx_parser
         });
     }
 
+    instruction_ref parse_clip(const std::string&,
+                               const attribute_map& attributes,
+                               std::vector<instruction_ref> args)
+    {
+        op::clip op;
+        if(contains(attributes, "max"))
+        {
+            op.max_val = parse_value(attributes.at("max")).at<float>();
+        }
+        if(contains(attributes, "min"))
+        {
+            op.min_val = parse_value(attributes.at("min")).at<float>();
+        }
+        return prog.add_instruction(op, std::move(args));
+    }
+
     instruction_ref
     parse_softmax(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
     {
@@ -300,7 +323,7 @@ struct onnx_parser
         {
             uint64_t axis = 1;
             auto l1       = prog.add_instruction(op, args[0], args[1]);
-            auto l2       = prog.add_instruction(op::broadcast{axis, l1->get_shape()}, args[2]);
+            auto l2 = prog.add_instruction(op::broadcast{axis, l1->get_shape().lens()}, args[2]);
             return prog.add_instruction(op::add{}, l1, l2);
         }
         return prog.add_instruction(op, l0, args[1]);
@@ -495,25 +518,86 @@ struct onnx_parser
         auto l2 = (transb) ? prog.add_instruction(op::transpose{perm}, args[1]) : args[1];
         if(args.size() == 3)
         {
-            if(beta != 0.f)
+            if(beta != 0.f && args[2]->get_shape().elements() > 0)
             {
-                auto l3 = prog.add_instruction(op::dot{alpha}, l1, l2);
-                auto l4 = args[2];
-                if(l4->get_shape().scalar()) // ignore args[2] (no C value added to alpha*A*B)
-                    return l3;
-                if(beta != 1.f)
+                auto out_lens   = l1->get_shape().lens();
+                out_lens.back() = l2->get_shape().lens().back();
+                auto l3         = args[2];
+                auto l3_lens    = l3->get_shape().lens();
+                if(!std::equal(out_lens.begin(), out_lens.end(), l3_lens.begin(), l3_lens.end()))
                 {
-                    auto beta_val = prog.add_literal(beta);
-                    auto l5 = prog.add_instruction(op::scalar{args[2]->get_shape()}, beta_val);
-                    l4      = prog.add_instruction(op::mul{}, args[2], l5);
+                    l3 = prog.add_instruction(op::multibroadcast{out_lens}, args[2]);
                 }
-                return add_broadcastable_binary_op(l3, l4, op::add{});
+                return prog.add_instruction(op::dot{alpha, beta}, l1, l2, l3);
             }
         }
 
         return prog.add_instruction(op::dot{alpha, beta}, l1, l2);
     }
 
+    instruction_ref
+    parse_matmul(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
+    {
+        auto l0      = args[0];
+        auto l1      = args[1];
+        auto l0_lens = l0->get_shape().lens();
+        auto l1_lens = l1->get_shape().lens();
+
+        // args[0] is a vector, prepend 1 to the shape
+        bool is_a_prepended = false;
+        if(l0_lens.size() == 1)
+        {
+            is_a_prepended = true;
+            l0_lens.insert(l0_lens.begin(), 1);
+            l0 = prog.add_instruction(op::unsqueeze{{0}}, args[0]);
+        }
+
+        bool is_b_appended = false;
+        if(l1_lens.size() == 1)
+        {
+            is_b_appended = true;
+            l1_lens.push_back(1);
+            l1 = prog.add_instruction(op::unsqueeze{{1}}, args[1]);
+        }
+
+        instruction_ref bl0 = l0;
+        instruction_ref bl1 = l1;
+        if(!std::equal(l0_lens.rbegin() + 2, l0_lens.rend(), l1_lens.rbegin() + 2, l1_lens.rend()))
+        {
+            auto l0_it = l0_lens.begin() + l0_lens.size() - 2;
+            std::vector<std::size_t> l0_broadcasted_lens(l0_lens.begin(), l0_it);
+            auto l1_it = l1_lens.begin() + l1_lens.size() - 2;
+            std::vector<std::size_t> l1_broadcasted_lens(l1_lens.begin(), l1_it);
+            auto output_lens = compute_broadcasted_lens(l0_broadcasted_lens, l1_broadcasted_lens);
+            l0_broadcasted_lens = output_lens;
+            l0_broadcasted_lens.insert(l0_broadcasted_lens.end(), l0_it, l0_lens.end());
+            l1_broadcasted_lens = output_lens;
+            l1_broadcasted_lens.insert(l1_broadcasted_lens.end(), l1_it, l1_lens.end());
+            if(l0_lens != l0_broadcasted_lens)
+            {
+                bl0 = prog.add_instruction(op::multibroadcast{l0_broadcasted_lens}, l0);
+            }
+            if(l1_lens != l1_broadcasted_lens)
+            {
+                bl1 = prog.add_instruction(op::multibroadcast{l1_broadcasted_lens}, l1);
+            }
+        }
+
+        auto dot_res     = prog.add_instruction(op::dot{1.0f, 0.0f}, bl0, bl1);
+        int64_t num_axis = static_cast<int64_t>(dot_res->get_shape().lens().size());
+        if(is_a_prepended)
+        {
+            dot_res = prog.add_instruction(op::squeeze{{num_axis - 2}}, dot_res);
+            --num_axis;
+        }
+        if(is_b_appended)
+        {
+            dot_res = prog.add_instruction(op::squeeze{{num_axis - 1}}, dot_res);
+        }
+
+        return dot_res;
+    }
+
     instruction_ref
     parse_batchnorm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
     {
@@ -604,15 +688,15 @@ struct onnx_parser
             auto&& bias_floats = attributes["bias"].floats();
             bias               = std::vector<float>(bias_floats.begin(), bias_floats.end());
         }
-        auto input_shape = args.front()->get_shape();
+        auto input_lens = args.front()->get_shape().lens();
 
         auto scale_val = prog.add_literal(scale);
         auto bias_vals = prog.add_literal(
             migraphx::literal{migraphx::shape{migraphx::shape::float_type, {bias.size()}}, bias});
 
-        auto scale_tensor = prog.add_instruction(migraphx::op::scalar{input_shape}, scale_val);
+        auto scale_tensor = prog.add_instruction(migraphx::op::scalar{input_lens}, scale_val);
         auto img_scaled   = prog.add_instruction(migraphx::op::mul{}, args.front(), scale_tensor);
-        auto bias_bcast = prog.add_instruction(migraphx::op::broadcast{1, input_shape}, bias_vals);
+        auto bias_bcast   = prog.add_instruction(migraphx::op::broadcast{1, input_lens}, bias_vals);
         return prog.add_instruction(migraphx::op::add{}, img_scaled, bias_bcast);
     }
 
@@ -1294,28 +1378,26 @@ struct onnx_parser
     static literal parse_tensor(const onnx::TensorProto& t)
     {
         std::vector<std::size_t> dims(t.dims().begin(), t.dims().end());
-        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
-        if(dims.empty())
-        {
-            dims = {1};
-        }
         if(t.has_raw_data())
         {
             const std::string& s = t.raw_data();
             switch(t.data_type())
             {
             case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
-            case onnx::TensorProto::FLOAT: return literal{{shape::float_type, dims}, s.data()};
+            case onnx::TensorProto::FLOAT: return create_literal(shape::float_type, dims, s.data());
             case onnx::TensorProto::UINT8: throw std::runtime_error("");
-            case onnx::TensorProto::INT8: return literal{{shape::int32_type, dims}, s.data()};
-            case onnx::TensorProto::UINT16: return literal{{shape::int32_type, dims}, s.data()};
-            case onnx::TensorProto::INT16: return literal{{shape::int32_type, dims}, s.data()};
-            case onnx::TensorProto::INT32: return literal{{shape::int32_type, dims}, s.data()};
-            case onnx::TensorProto::INT64: return literal{{shape::int64_type, dims}, s.data()};
+            case onnx::TensorProto::INT8: return create_literal(shape::int32_type, dims, s.data());
+            case onnx::TensorProto::UINT16:
+                return create_literal(shape::int32_type, dims, s.data());
+            case onnx::TensorProto::INT16: return create_literal(shape::int32_type, dims, s.data());
+            case onnx::TensorProto::INT32: return create_literal(shape::int32_type, dims, s.data());
+            case onnx::TensorProto::INT64: return create_literal(shape::int64_type, dims, s.data());
             case onnx::TensorProto::STRING: throw std::runtime_error("");
-            case onnx::TensorProto::BOOL: return literal{{shape::int32_type, dims}, s.data()};
-            case onnx::TensorProto::FLOAT16: return literal{{shape::half_type, dims}, s.data()};
-            case onnx::TensorProto::DOUBLE: return literal{{shape::double_type, dims}, s.data()};
+            case onnx::TensorProto::BOOL: return create_literal(shape::int32_type, dims, s.data());
+            case onnx::TensorProto::FLOAT16:
+                return create_literal(shape::half_type, dims, s.data());
+            case onnx::TensorProto::DOUBLE:
+                return create_literal(shape::double_type, dims, s.data());
             case onnx::TensorProto::UINT32: throw std::runtime_error("");
             case onnx::TensorProto::UINT64: throw std::runtime_error("");
             case onnx::TensorProto::COMPLEX64: throw std::runtime_error("");
@@ -1327,21 +1409,21 @@ struct onnx_parser
         {
         case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
         case onnx::TensorProto::FLOAT:
-            return literal{{shape::float_type, dims}, t.float_data().begin(), t.float_data().end()};
+            return create_literal(shape::float_type, dims, t.float_data());
         case onnx::TensorProto::UINT8: throw std::runtime_error("");
         case onnx::TensorProto::INT8:
-            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
+            return create_literal(shape::int32_type, dims, t.int32_data());
         case onnx::TensorProto::UINT16:
-            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
+            return create_literal(shape::int32_type, dims, t.int32_data());
         case onnx::TensorProto::INT16:
-            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
+            return create_literal(shape::int32_type, dims, t.int32_data());
         case onnx::TensorProto::INT32:
-            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
+            return create_literal(shape::int32_type, dims, t.int32_data());
         case onnx::TensorProto::INT64:
-            return literal{{shape::int64_type, dims}, t.int64_data().begin(), t.int64_data().end()};
+            return create_literal(shape::int64_type, dims, t.int64_data());
         case onnx::TensorProto::STRING: throw std::runtime_error("");
         case onnx::TensorProto::BOOL:
-            return literal{{shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
+            return create_literal(shape::int32_type, dims, t.int32_data());
         case onnx::TensorProto::FLOAT16:
         {
             std::vector<uint16_t> data_uint16(t.int32_data().begin(), t.int32_data().end());
@@ -1350,11 +1432,10 @@ struct onnx_parser
                            data_uint16.end(),
                            std::back_inserter(data_half),
                            [](uint16_t raw_val) { return *reinterpret_cast<half*>(&raw_val); });
-            return literal{{shape::half_type, dims}, data_half.begin(), data_half.end()};
+            return create_literal(shape::half_type, dims, data_half);
         }
         case onnx::TensorProto::DOUBLE:
-            return literal{
-                {shape::double_type, dims}, t.double_data().begin(), t.double_data().end()};
+            return create_literal(shape::double_type, dims, t.double_data());
         case onnx::TensorProto::UINT32: throw std::runtime_error("");
         case onnx::TensorProto::UINT64: throw std::runtime_error("");
         case onnx::TensorProto::COMPLEX64: throw std::runtime_error("");
@@ -1363,6 +1444,23 @@ struct onnx_parser
         MIGRAPHX_THROW("Invalid tensor type");
     }
 
+    static literal
+    create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, const char* data)
+    {
+        // in case of scalar constants in onnx file, use dims=1 to fill initializer data
+        if(dims.empty())
+            return literal{{shape_type}, data};
+        return literal{{shape_type, dims}, data};
+    }
+
+    template <class T, MIGRAPHX_REQUIRES(not std::is_pointer<T>{})>
+    static literal create_literal(shape::type_t shape_type, const std::vector<size_t>& dims, T data)
+    {
+        if(dims.empty())
+            return literal{{shape_type}, data.begin(), data.end()};
+        return literal{{shape_type, dims}, data.begin(), data.end()};
+    }
+
     static shape parse_type(const onnx::TypeProto& t)
     {
         shape::type_t shape_type{};

src/opt/memory_coloring_impl.cpp

+#include <migraphx/op/load.hpp>
 #include "memory_coloring_impl.hpp"
 
 namespace migraphx {
@@ -62,11 +63,11 @@ bool memory_coloring_impl::allocate(interval_ptr interval)
         }
     }
 
-    long long offset = 0;
+    std::size_t offset = 0;
     while(!conflict_queue.empty())
     {
-        live_range* range     = conflict_queue.top();
-        long long iter_offset = range->offset;
+        live_range* range       = conflict_queue.top();
+        std::size_t iter_offset = range->offset;
         if(offset > iter_offset)
         {
             offset = std::max(offset, iter_offset + range->size);
@@ -96,7 +97,7 @@ void memory_coloring_impl::build()
     if(num_of_instrs == 0)
         return;
 
-    int cur_points        = num_of_instrs * 2;
+    auto cur_points       = num_of_instrs * 2;
     instruction_ref iter  = p_program->end();
     instruction_ref begin = p_program->begin();
     std::vector<instruction_ref> dead_instrs;
@@ -192,13 +193,13 @@ void memory_coloring_impl::rewrite()
                 continue;
 
             std::size_t offset = 0;
-            if(interval->get_offset() == invalid_offset)
+            if(interval->get_offset() != invalid_offset)
             {
-                assert(interval->result.bytes() == 0);
+                offset = interval->get_offset();
             }
             else
             {
-                offset = interval->get_offset();
+                assert(interval->result.bytes() == 0);
             }
 
             if(is_allocate(ins))
@@ -206,15 +207,6 @@ void memory_coloring_impl::rewrite()
                 p_program->replace_instruction(
                     ins, op::load{ins->get_shape(), offset}, scratch_param);
             }
-            else if(is_literal(ins))
-            {
-#if 0                
-                auto pre      = p_program->add_literal(ins->lit);
-                bool pre_copy = (interval->get_begin() < earliest_end_point);
-                p_program->replace_instruction(
-                    ins, write_literal{offset, pre_copy}, scratch_param, pre);
-#endif
-            }
         }
     }
     MIGRAPHX_DEBUG(dump("---After rewrite---"));

src/opt/memory_coloring_impl.hpp

@@ -3,7 +3,6 @@
 #include <migraphx/program.hpp>
 #include <migraphx/stringutils.hpp>
 #include <migraphx/instruction.hpp>
-#include <migraphx/operators.hpp>
 #include <migraphx/iterator_for.hpp>
 #include <migraphx/pass_config.hpp>
 #include <migraphx/config.hpp>
@@ -22,15 +21,15 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
-static const int invalid_offset = -1;
+static const std::size_t invalid_offset = std::numeric_limits<std::size_t>::max();
 
 struct live_range
 {
-    int begin;        // begin point in the instruction stream.
-    int end;          // end point in the instruction stream.
-    long long offset; // offset to base pointer of allocated memory trunk.
-    int vn;           // value number that identifies this live_range.
-    long long size;   // size of required memory in bytes
+    std::size_t begin;  // begin point in the instruction stream.
+    std::size_t end;    // end point in the instruction stream.
+    std::size_t offset; // offset to base pointer of allocated memory trunk.
+    std::size_t vn;     // value number that identifies this live_range.
+    std::size_t size;   // size of required memory in bytes
 #ifdef MIGRAPHX_DEBUG_OPT
     void dump();
 #endif
@@ -46,9 +45,9 @@ struct live_interval
         is_live_on_entry = false;
     }
 
-    void add_use(int use) { use_points.push_front(use); }
-    int get_begin() const { return segment.begin; }
-    int get_end() const { return segment.end; }
+    void add_use(std::size_t use) { use_points.push_front(use); }
+    std::size_t get_begin() const { return segment.begin; }
+    std::size_t get_end() const { return segment.end; }
     long long get_offset() const { return segment.offset; }
 
 #ifdef MIGRAPHX_DEBUG_OPT
@@ -56,9 +55,9 @@ struct live_interval
 #endif
 
     live_range segment;
-    int id;
-    std::list<int> use_points;
-    int def_point;
+    std::size_t id;
+    std::list<std::size_t> use_points;
+    std::size_t def_point;
     shape result;
     bool is_literal;
     bool is_live_on_entry;
@@ -112,8 +111,8 @@ struct memory_coloring_impl
     {
         if((range1.size == 0) || (range2.size == 0))
             return false;
-        long long end1 = range1.offset + range1.size - 1;
-        long long end2 = range2.offset + range2.size - 1;
+        auto end1 = range1.offset + range1.size - 1;
+        auto end2 = range2.offset + range2.size - 1;
         return ((end1 < range2.offset) || (end2 < range1.offset));
     }
     void verify();
@@ -126,8 +125,8 @@ struct memory_coloring_impl
     {
         bool operator()(const interval_ptr i1, const interval_ptr i2) const
         {
-            int len1 = i1->get_end() - i1->get_begin();
-            int len2 = i2->get_end() - i2->get_begin();
+            auto len1 = i1->get_end() - i1->get_begin();
+            auto len2 = i2->get_end() - i2->get_begin();
             if(len1 != len2)
             {
                 return (len1 < len2);
@@ -159,7 +158,7 @@ struct memory_coloring_impl
 
     int num_of_lives;
     int max_value_number;
-    long long required_bytes;
+    std::size_t required_bytes;
     // The earliest program point where an live interval ends.
     int earliest_end_point;
     // The latest program point where an live interval ends.

src/pass_manager.cpp

+#include <migraphx/program.hpp>
+#include <migraphx/pass_manager.hpp>
+#include <migraphx/stringutils.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/operators.hpp>
+#include <migraphx/target.hpp>
+#include <migraphx/env.hpp>
+#include <migraphx/ranges.hpp>
+#include <migraphx/time.hpp>
+#include <migraphx/iterator_for.hpp>
+#include <iostream>
+#include <sstream>
+#include <algorithm>
+#include <utility>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+void run_passes(program& prog, const std::vector<pass>& passes, tracer trace)
+{
+    for(auto& p : passes)
+    {
+        trace("Pass: ", p.name());
+        p.apply(prog);
+        trace(prog);
+
+#ifndef NDEBUG
+        trace("Validate ...");
+        auto invalid = prog.validate();
+        if(invalid != prog.end())
+        {
+            auto index = std::distance(prog.begin(), invalid);
+            MIGRAPHX_THROW(p.name() + " pass produces invalid program at instruction " +
+                           std::to_string(index) + ": " + invalid->name());
+        }
+        trace();
+#endif
+    }
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/program.cpp

 #include <migraphx/program.hpp>
 #include <migraphx/stringutils.hpp>
 #include <migraphx/instruction.hpp>
-#include <migraphx/operators.hpp>
+#include <migraphx/op/identity.hpp>
 #include <migraphx/target.hpp>
 #include <migraphx/env.hpp>
 #include <migraphx/ranges.hpp>
 #include <migraphx/time.hpp>
 #include <migraphx/iterator_for.hpp>
+#include <migraphx/pass_manager.hpp>
 #include <iostream>
 #include <sstream>
 #include <algorithm>
@@ -55,18 +56,23 @@ static void print_instruction(std::ostream& os,
 }
 
 template <class F>
-static void print_program(std::ostream& os, const program& p, F annonate)
+static void print_program(const program& p, F print_func)
 {
     std::unordered_map<instruction_ref, std::string> names;
     int count = 0;
 
     for(auto ins : iterator_for(p))
     {
-        std::string var_name = "@" + std::to_string(count);
+        std::string var_name;
         if(ins->name() == "@param")
         {
             var_name = any_cast<builtin::param>(ins->get_operator()).parameter;
         }
+        else
+        {
+            var_name = "@" + std::to_string(count);
+            count++;
+        }
         names.emplace(ins, var_name);
 
         // TODO: Use all_of
@@ -76,21 +82,77 @@ static void print_program(std::ostream& os, const program& p, F annonate)
             (void)arg;
         }
 
-        print_instruction(os, ins, names);
-
-        annonate(ins, names);
-
-        os << std::endl;
-
-        count++;
+        print_func(ins, names);
     }
 }
 
 program::program() : impl(std::make_unique<program_impl>()) {}
 
 program::program(program&&) noexcept = default;
-program& program::operator=(program&&) noexcept = default;
-program::~program() noexcept                    = default;
+program::~program() noexcept         = default;
+
+// copy constructor
+program::program(const program& p) { assign(p); }
+
+// copy assignment operator
+program& program::operator=(program p)
+{
+    std::swap(p.impl, this->impl);
+    return *this;
+}
+
+void program::assign(const program& p)
+{
+    // clean the current program
+    if(!impl)
+    {
+        impl = std::make_unique<program_impl>();
+    }
+    else if(!impl->instructions.empty())
+    {
+        impl->instructions.clear();
+    }
+    impl->ctx = p.impl->ctx;
+
+    std::unordered_map<instruction_ref, instruction_ref> ins_map;
+    for(auto ins : iterator_for(p))
+    {
+        instruction_ref copy_ins{};
+        if(ins->name() == "@literal")
+        {
+            auto l   = ins->get_literal();
+            copy_ins = impl->instructions.insert(impl->instructions.end(), instruction{l});
+        }
+        else if(ins->name() == "@param")
+        {
+            auto&& name = any_cast<builtin::param>(ins->get_operator()).parameter;
+            auto s      = ins->get_shape();
+            copy_ins    = impl->instructions.insert(impl->instructions.end(),
+                                                 {builtin::param{name}, std::move(s), {}});
+        }
+        else if(ins->name() == "@outline")
+        {
+            auto s = ins->get_shape();
+            copy_ins =
+                impl->instructions.insert(impl->instructions.end(), {builtin::outline{s}, s, {}});
+        }
+        else
+        {
+            // retrieve its mapped input
+            auto inputs = ins->inputs();
+            // ensure all inputs have its corresponding copy instructions
+            assert(std::all_of(
+                inputs.begin(), inputs.end(), [&](auto i) { return ins_map.count(i) > 0; }));
+            std::vector<instruction_ref> copy_inputs(inputs.size());
+            std::transform(inputs.begin(), inputs.end(), copy_inputs.begin(), [&](auto i) {
+                return ins_map[i];
+            });
+            copy_ins = add_instruction(ins->get_operator(), copy_inputs);
+        }
+
+        ins_map[ins] = copy_ins;
+    }
+}
 
 instruction_ref program::add_instruction(const operation& op, std::vector<instruction_ref> args)
 {
@@ -291,23 +353,7 @@ void program::compile(const target& t, tracer trace)
         trace = tracer{std::cout};
     trace(*this);
     trace();
-    for(auto&& p : t.get_passes(this->impl->ctx))
-    {
-        trace("Pass: ", p.name());
-        p.apply(*this);
-        trace(*this);
-#ifndef NDEBUG
-        trace("Validate ...");
-        auto invalid = this->validate();
-        if(invalid != impl->instructions.end())
-        {
-            auto index = std::distance(impl->instructions.begin(), invalid);
-            MIGRAPHX_THROW(p.name() + " pass produces invalid program at instruction " +
-                           std::to_string(index) + ": " + invalid->name());
-        }
-        trace();
-#endif
-    }
+    run_passes(*this, t.get_passes(this->impl->ctx), trace);
     auto invalid = this->validate();
     if(invalid != impl->instructions.end())
     {
@@ -475,10 +521,12 @@ void program::perf_report(std::ostream& os, std::size_t n, parameter_map params)
     double calculate_overhead_time    = total_time - total_instruction_time;
     double calculate_overhead_percent = calculate_overhead_time * 100.0 / total_time;
 
-    print_program(os, *this, [&](auto ins, auto&&) {
+    print_program(*this, [&](auto ins, const auto& names) {
+        print_instruction(std::cout, ins, names);
         double avg     = common_average(ins_vec[ins]);
         double percent = std::ceil(100.0 * avg / total_instruction_time);
         os << ": " << avg << "ms, " << percent << "%";
+        os << std::endl;
     });
 
     os << std::endl;
@@ -516,7 +564,7 @@ void program::debug_print(instruction_ref ins) const
         return;
     }
     std::stringstream ss;
-    print_program(ss, *this, [&](auto x, auto&& names) {
+    print_program(*this, [&](auto x, const auto& names) {
         if(x == ins)
         {
             print_instruction(std::cout, x, names);
@@ -531,6 +579,32 @@ void program::debug_print(const std::vector<instruction_ref>& inss) const
     std::cout << std::endl;
 }
 
+static std::string enclose_name(const std::string& name)
+{
+    return '"' + replace_string(name, "\"", "\\\"") + '"';
+}
+
+void program::print_graph(std::ostream& os) const
+{
+    os << "digraph {" << std::endl;
+    os << "\trankdir=LR;" << std::endl;
+    print_program(*this, [&](auto ins, const auto& names) {
+        os << "\t" << enclose_name(names.at(ins))
+           << "[label=" << enclose_name(to_string(ins->get_operator())) << "];";
+        os << std::endl;
+        if(!ins->inputs().empty())
+        {
+            for(auto&& arg : ins->inputs())
+            {
+                os << "\t" << enclose_name(names.at(arg)) << " -> " << enclose_name(names.at(ins));
+                os << "[label=" << enclose_name(to_string(ins->get_shape())) << "];";
+                os << std::endl;
+            }
+        }
+    });
+    os << "}" << std::endl;
+}
+
 void program::dry_run(std::unordered_map<std::string, argument> params) const
 {
     auto& ctx = this->impl->ctx;
@@ -539,14 +613,21 @@ void program::dry_run(std::unordered_map<std::string, argument> params) const
 
 void program::annotate(std::ostream& os, std::function<void(instruction_ref)> a) const
 {
-    print_program(os, *this, [&](auto ins, auto&&) { a(ins); });
+    print_program(*this, [&](auto ins, const auto& names) {
+        print_instruction(os, ins, names);
+        a(ins);
+        os << std::endl;
+    });
 }
 
 bool operator==(const program& x, const program& y) { return to_string(x) == to_string(y); }
 
 std::ostream& operator<<(std::ostream& os, const program& p)
 {
-    print_program(os, p, [](auto&&...) {});
+    print_program(p, [&](auto ins, const auto& names) {
+        print_instruction(os, ins, names);
+        os << std::endl;
+    });
     return os;
 }
 

src/propagate_constant.cpp

+#include <migraphx/propagate_constant.hpp>
+#include <migraphx/program.hpp>
+#include <migraphx/matcher.hpp>
+#include <migraphx/literal.hpp>
+#include <migraphx/functional.hpp>
+#include <unordered_set>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+bool skip_propogate(instruction_ref ins)
+{
+    if(ins->name() == "@literal")
+        return true;
+    auto&& s = ins->get_shape();
+    if(s.broadcasted() and not s.scalar())
+        return true;
+    if(s.scalar() and s.elements() != 1)
+        return true;
+    return false;
+}
+
+void propagate_constant::apply(program& p) const
+{
+    for(auto i : iterator_for(p))
+    {
+        if(i->name() != "@literal")
+            continue;
+        if(i->outputs().empty())
+            continue;
+        fix([&](auto self, auto ins) {
+            std::unordered_set<instruction_ref> children(ins->outputs().begin(),
+                                                         ins->outputs().end());
+            for(auto child : children)
+            {
+                if(skip_propogate(child))
+                {
+                    self(child);
+                    continue;
+                }
+                auto r = child->eval();
+                if(not r.empty())
+                {
+                    assert(r.get_shape() == child->get_shape());
+                    auto l = p.add_literal(r.get_shape(), r.data());
+                    self(p.replace_instruction(child, l));
+                }
+            }
+        })(i);
+    }
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx