Cpu fusions using post_ops (#781)

* Add eliminate_data_type pass

* Formatting

* Auto convert quant ops

* Formatting

* Flip the order of decompose

* Compute max size differently

* Formatting

* Clamp values in convert

* Formatting

* Fix loss of precision in reduce

* Formatting

* Fix bugs in reduction

* Fix accumulator type in reference softmax implementation

* Formatting

* Update convert test

* Remove unused variables

* Remove unnecessary quant_dot check

* Formatting

* Add tests

* Formatting

* Remove unused code

* Remove duplicate ops

* Remove blaze dependency

* Use set since shape::type_t is no hashable on gcc 5

* Formatting

* Add dnnl binary op

* Formatting

* Add binary and eltwise

* Formatting

* Add softmax

* Formatting

* Remove unused operators

* Add missing files

* Formatting

* Add lrn

* Formatting

* Add deconvolution

* Formatting

* Change allocate default

* Add reorder

* Formatting

* Add reductions

* Formatting

* Sort lines

* Change literals in another loop

* Add pow operator

* Formatting

* Add pow operator

* Formatting

* Make sure shapes are packed

* Allow broadcasted inputs

* Remove unused operators

* Simplify functions

* Remove softmax

* Add sub and erf functions

* Formatting

* Fix bug

* Formatting

* Improve parallism

* Formatting

* Allow multiple batch dimensions

* Formatting

* Move literal transforms out of lowering

* Formatting

* Add gather operator

* Sort lines

* Add early exit for carry

* Formatting

* Add missing concat

* Rename macro

* Fix deep nesting

* Formatting

* Fix cppcheck issues

* Remov else

* Move attribute to typedef

* Formatting

* Disable maybe-uninitialized warning since its broken on gcc

* Add constexpr default constructor

* Formatting

* Fix compiler warnings

* Fix adjust_allocation test

* Add layernorm matcher

* Add gelu_erf matcher

* Formatting

* Add gelu_tanh matcher

* Formatting

* Remove match namespace

* Formatting

* Use matcher instead of string

* Formatting

* Add fusions

* Formatting

* Add post op field

* Formatting

* Make post_ops serializable

* Formatting

* Add eltwise fusions

* Formatting

* Fix null conversions

* Formatting

* Add fuse_ops source files

* Formatting

* Set binary post op index correctly

* Formatting

* Fix serialization bugs

* Check if used once

* Formatting

* Fix error in get_primitive_attr

* Formatting

* Add compile function

* Formatting

* Limit fusions

* Formatting

* Disable with env variable instead of using compile arg

* Formatting

* Fix implicit conversion to bool

* Declar on seperate lines

* Formatting

* Fix cppcheck issues

* Fix ICE in pack_join

* Formatting

* Use const ref

* Make enum hashable

* Formatting

* Add explicit this

* Fix merge issues

* Fix dangling ref

* Formatting

* Add test for compile

* Formatting

* Add more value tests

* Formatting
Co-authored-by: Shucai Xiao <shucai@gmail.com>
Co-authored-by: mvermeulen <5479696+mvermeulen@users.noreply.github.com>

src/value.cpp

@@ -204,7 +204,7 @@ const std::string& value::get_key() const { return key; }
 
 std::vector<value>* if_array_impl(const std::shared_ptr<value_base_impl>& x)
 {
-    if(!x)
+    if(x == nullptr)
         return nullptr;
     return x->if_array();
 }
@@ -286,10 +286,26 @@ const value* value::begin() const
 value* value::end() { return begin() + size(); }
 const value* value::end() const { return begin() + size(); }
 
-value& value::front() { return *begin(); }
-const value& value::front() const { return *begin(); }
-value& value::back() { return *std::prev(end()); }
-const value& value::back() const { return *std::prev(end()); }
+value& value::front()
+{
+    assert(this->size() > 0);
+    return *begin();
+}
+const value& value::front() const
+{
+    assert(this->size() > 0);
+    return *begin();
+}
+value& value::back()
+{
+    assert(this->size() > 0);
+    return *std::prev(end());
+}
+const value& value::back() const
+{
+    assert(this->size() > 0);
+    return *std::prev(end());
+}
 value& value::at(std::size_t i)
 {
     auto* a = if_array_impl(x);
@@ -322,8 +338,16 @@ const value& value::at(const std::string& pkey) const
         MIGRAPHX_THROW("Key not found: " + pkey);
     return *r;
 }
-value& value::operator[](std::size_t i) { return *(begin() + i); }
-const value& value::operator[](std::size_t i) const { return *(begin() + i); }
+value& value::operator[](std::size_t i)
+{
+    assert(i < this->size());
+    return *(begin() + i);
+}
+const value& value::operator[](std::size_t i) const
+{
+    assert(i < this->size());
+    return *(begin() + i);
+}
 value& value::operator[](const std::string& pkey) { return *emplace(pkey, nullptr).first; }
 
 void value::clear() { get_array_throw(x).clear(); }
@@ -409,8 +433,8 @@ template <class F>
 bool compare(const value& x, const value& y, F f)
 {
     bool result = false;
-    x.visit([&](auto&& a) {
-        y.visit([&](auto&& b) {
+    x.visit_value([&](auto&& a) {
+        y.visit_value([&](auto&& b) {
             result = compare_common_impl(rank<1>{}, f, x.get_key(), a, y.get_key(), b);
         });
     });
@@ -436,6 +460,8 @@ bool operator<=(const value& x, const value& y) { return x == y or x < y; }
 bool operator>(const value& x, const value& y) { return y < x; }
 bool operator>=(const value& x, const value& y) { return x == y or x > y; }
 
+void print_value(std::ostream& os, std::nullptr_t) { os << "null"; }
+
 template <class T>
 void print_value(std::ostream& os, const T& x)
 {
@@ -450,7 +476,6 @@ void print_value(std::ostream& os, const std::pair<T, U>& x)
     print_value(os, x.second);
 }
 
-void print_value(std::ostream& os, const std::nullptr_t&) { os << "null"; }
 void print_value(std::ostream& os, const std::vector<value>& x)
 {
     os << "{";

test/gpu/jit.cpp

@@ -122,13 +122,13 @@ TEST_CASE(code_object_hip)
     auto input_literal  = migraphx::generate_literal(input);
     auto output_literal = migraphx::transform(input_literal, [](auto x) { return x + 2; });
     auto x              = mm->add_literal(input_literal);
-    auto y              = mm->add_instruction(
-        migraphx::make_op("hip::allocate", {{"shape", migraphx::to_value(input)}}));
+    auto y              = mm->add_parameter("output", input);
     mm->add_instruction(co, x, y);
     migraphx::compile_options options;
     p.compile(migraphx::gpu::target{}, options);
 
-    auto result = migraphx::gpu::from_gpu(p.eval({}).front());
+    auto result =
+        migraphx::gpu::from_gpu(p.eval({{"output", migraphx::gpu::allocate_gpu(input)}}).front());
 
     EXPECT(result == output_literal.get_argument());
 }
@@ -149,12 +149,12 @@ TEST_CASE(compile_code_object_hip)
     auto input_literal  = migraphx::generate_literal(input);
     auto output_literal = migraphx::transform(input_literal, [](auto x) { return x + 1; });
     auto x              = mm->add_literal(input_literal);
-    auto y              = mm->add_instruction(
-        migraphx::make_op("hip::allocate", {{"shape", migraphx::to_value(input)}}));
+    auto y              = mm->add_parameter("output", input);
     mm->add_instruction(co, x, y);
     p.compile(migraphx::gpu::target{}, migraphx::compile_options{});
 
-    auto result = migraphx::gpu::from_gpu(p.eval({}).front());
+    auto result =
+        migraphx::gpu::from_gpu(p.eval({{"output", migraphx::gpu::allocate_gpu(input)}}).front());
 
     EXPECT(result == output_literal.get_argument());
 }

test/operation.cpp

@@ -46,6 +46,33 @@ struct simple_operation_no_print
     }
 };
 
+struct compilable_op
+{
+    std::string name() const { return "compilable"; }
+    migraphx::argument
+    compute(migraphx::context&, const migraphx::shape&, std::vector<migraphx::argument> args) const
+    {
+        if(args.empty())
+            return {};
+        return args.front();
+    }
+
+    migraphx::shape compute_shape(std::vector<migraphx::shape> inputs) const
+    {
+        if(inputs.empty())
+            return {};
+        return inputs.front();
+    }
+
+    int output_alias(const std::vector<migraphx::shape>&) const { return 0; }
+
+    migraphx::value
+    compile(migraphx::context&, const migraphx::shape&, const std::vector<migraphx::shape>&)
+    {
+        return {{"compiled", true}};
+    }
+};
+
 TEST_CASE(operation_copy_test)
 {
     simple_operation s{};
@@ -204,4 +231,20 @@ TEST_CASE(check_from_value2)
     EXPECT(op1 == op2);
 }
 
+TEST_CASE(compile)
+{
+    migraphx::operation op = compilable_op{};
+    migraphx::context ctx{};
+    auto v = op.compile(ctx, {}, {});
+    EXPECT(v.at("compiled").to<bool>() == true);
+}
+
+TEST_CASE(compile_non_compilable)
+{
+    migraphx::operation op = simple_operation{};
+    migraphx::context ctx{};
+    auto v = op.compile(ctx, {}, {});
+    EXPECT(v.empty());
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/value_test.cpp

@@ -486,6 +486,12 @@ TEST_CASE(value_emplace_object)
     EXPECT(v["three"].get_key() == "three");
 }
 
+TEST_CASE(value_bracket_convert_throws)
+{
+    migraphx::value v1;
+    EXPECT(test::throws([&] { v1["key"].to<std::string>(); }));
+}
+
 TEST_CASE(value_construct_object_string_value)
 {
     migraphx::value v = {{"one", "onev"}, {"two", "twov"}};
@@ -798,4 +804,35 @@ TEST_CASE(value_binary_object_conv)
     EXPECT(migraphx::equal(v["data"].get_binary(), data));
 }
 
+template <class T>
+bool is_null_type(T)
+{
+    return false;
+}
+
+bool is_null_type(std::nullptr_t) { return true; }
+
+TEST_CASE(visit_null)
+{
+    migraphx::value v;
+    EXPECT(v.is_null());
+    bool visited = false;
+    v.visit([&](auto&& x) { visited = is_null_type(x); });
+    EXPECT(visited);
+}
+
+TEST_CASE(value_or_convert)
+{
+    migraphx::value v = 1;
+    EXPECT(v.is_int64());
+    EXPECT(v.value_or(3) == 1);
+}
+
+TEST_CASE(value_or_null)
+{
+    migraphx::value v;
+    EXPECT(v.is_null());
+    EXPECT(v.value_or(3) == 3);
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/verify/run_verify.cpp

@@ -158,10 +158,12 @@ void run_verify::verify(const std::string& name, const migraphx::program& p) con
                                  detach_async([=] { return run_target(t, p, m); }, ti.parallel));
         }
 
+        assert(gold_f.valid());
         auto gold = gold_f.get();
 
         for(auto&& pp : results)
         {
+            assert(pp.second.valid());
             auto tname  = pp.first;
             auto x      = pp.second.get();
             auto cp     = x.first;

test/verify/test_add_relu_add.cpp

+
+#include "verify_program.hpp"
+#include <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_add_relu_add : verify_program<test_add_relu_add>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        auto x = mm->add_parameter("x", migraphx::shape{migraphx::shape::float_type, {1, 5, 2, 2}});
+        auto y = mm->add_parameter("y", migraphx::shape{migraphx::shape::float_type, {1, 5, 2, 2}});
+        auto z = mm->add_parameter("z", migraphx::shape{migraphx::shape::float_type, {1, 5, 2, 2}});
+        auto a = mm->add_instruction(migraphx::make_op("add"), x, y);
+        auto b = mm->add_instruction(migraphx::make_op("relu"), a);
+        auto c = mm->add_instruction(migraphx::make_op("add"), b, z);
+        mm->add_instruction(migraphx::make_op("relu"), c);
+        return p;
+    }
+};

tools/include/operation.hpp

@@ -341,6 +341,29 @@ auto has_finalize_op(const T&) -> decltype(has_finalize_op(rank<1>{},
     return {};
 }
 
+template <class T>
+auto compile_op(
+    rank<1>, T& x, context& ctx, const shape& output_shape, const std::vector<shape>& input)
+    -> decltype(x.compile(auto_any_cast(ctx), output_shape, input))
+{
+    return x.compile(auto_any_cast(ctx), output_shape, input);
+}
+
+template <class T>
+value compile_op(rank<0>, T&, context&, const shape&, const std::vector<shape>&)
+{
+    return value::object{};
+}
+
+template <class T>
+value compile_op(const T& x,
+                 context& ctx,
+                 const shape& output_shape,
+                 const std::vector<shape>& input)
+{
+    return compile_op(rank<1>{}, x, ctx, output_shape, input);
+}
+
 template <class T>
 value attributes_op(const T&)
 {
@@ -361,6 +384,12 @@ void from_value_op(T& x, const value& v)
     return migraphx::from_value(v, x);
 }
 
+template <class T>
+bool is_borrowed_op(const T&)
+{
+    return false;
+}
+
 } // namespace detail
 
 <%
@@ -374,11 +403,18 @@ void from_value_op(T& x, const value& v)
              const   = True,
              default = 'detail::need_normalization_op'),
      virtual('has_finalize', returns = 'bool', const = True, default = 'detail::has_finalize_op'),
+     virtual('is_borrowed', returns = 'bool', const = True, default = 'detail::is_borrowed_op'),
      virtual('output_alias',
              returns = 'std::ptrdiff_t',
              input   = 'const std::vector<shape>&',
              const   = True,
              default = 'detail::output_alias_op'),
+     virtual('compile',
+             returns = 'value',
+             ctx     = 'context&',
+             output  = 'const shape&',
+             input   = 'const std::vector<shape>&',
+             default = 'detail::compile_op'),
      virtual('finalize',
              ctx     = 'context&',
              output  = 'const shape&',
@@ -436,6 +472,24 @@ void from_value_op(T& x, const value& v)
     return !(x == y);
 }
 
+inline value
+compile(operation& op, context& ctx, const shape& output_shape, const std::vector<shape>& input)
+{
+    return op.compile(ctx, output_shape, input);
+}
+template <class Context>
+inline value
+compile(operation& op, Context& ctx, const shape& output_shape, const std::vector<shape>& input)
+{
+    dependent_type<context, Context> ctx2 = std::ref(ctx);
+    return compile(op, ctx2, output_shape, input);
+}
+template <class T, class Context>
+inline auto compile(T& op, Context& ctx, const shape& output_shape, const std::vector<shape>& input)
+    -> decltype(op.compile(ctx, ctx, output_shape, input))
+{
+    return op.compile(ctx, ctx, output_shape, input);
+}
 inline shape compute_shape(const operation& op, const std::vector<shape>& inputs)
 {
     return op.compute_shape(inputs);