Merge branch 'refactor_dynamic_compute' of...

Merge branch 'refactor_dynamic_compute' of github.com:ROCmSoftwarePlatform/AMDMIGraphX into dyn_ref_broadcast

src/include/migraphx/op/convert.hpp

@@ -44,7 +44,7 @@ struct convert : unary<convert>
 
     shape compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, *this}.has(1);
+        check_shapes{inputs, *this, true}.has(1);
         auto input = inputs.at(0);
         if(input.dynamic())
         {

src/include/migraphx/op/elu.hpp

@@ -32,14 +32,13 @@ namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 namespace op {
 
-struct elu
+struct elu : unary<elu>
 {
-    std::string name() const { return "elu"; }
     float alpha = 1;
-    shape compute_shape(std::vector<shape> inputs) const
+
+    std::string point_op() const
     {
-        check_shapes{inputs, *this}.has(1);
-        return inputs.front();
+        return "${function:where}(${0} > 0, ${0}, ${alpha} * (migraphx::exp(${0}) - 1))";
     }
 
     template <class Self, class F>
@@ -47,6 +46,11 @@ struct elu
     {
         return pack(f(self.alpha, "alpha"));
     }
+
+    auto apply() const
+    {
+        return [&](auto x) { return x > 0 ? x : alpha * std::expm1(x); };
+    }
 };
 
 } // namespace op

src/include/migraphx/op/leaky_relu.hpp

@@ -26,12 +26,13 @@
 
 #include <migraphx/check_shapes.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/op/unary.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 namespace op {
 
-struct leaky_relu
+struct leaky_relu : unary<leaky_relu>
 {
     float alpha = 0.01;
 
@@ -41,11 +42,13 @@ struct leaky_relu
         return pack(f(self.alpha, "alpha"));
     }
 
+    std::string point_op() const { return "${function:where}(${0} > 0, ${0}, ${alpha} * ${0})"; }
+
     std::string name() const { return "leaky_relu"; }
-    shape compute_shape(std::vector<shape> inputs) const
+
+    auto apply() const
     {
-        check_shapes{inputs, *this}.has(1);
-        return inputs.front();
+        return [&](auto x) { return x > 0 ? x : x * alpha; };
     }
 };
 

src/include/migraphx/op/unary.hpp

@@ -30,6 +30,7 @@
 #include <migraphx/argument.hpp>
 #include <migraphx/stringutils.hpp>
 #include <migraphx/value.hpp>
+#include <migraphx/operation.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -62,9 +63,9 @@ struct unary : op_name<Derived>
     value attributes() const { return base_attributes(); }
     shape compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, static_cast<const Derived&>(*this)}.has(1);
+        check_shapes{inputs, static_cast<const Derived&>(*this), true}.has(1);
         auto s = inputs.at(0);
-        if(s.scalar())
+        if(s.dynamic() or s.scalar())
         {
             return s;
         }
@@ -78,9 +79,9 @@ struct unary : op_name<Derived>
         }
     }
 
-    argument compute(const shape& output_shape, std::vector<argument> args) const
+    argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
     {
-        argument result{output_shape};
+        argument result{dyn_out.computed_shape};
         result.visit([&](auto output) {
             args[0].visit([&](auto input) {
                 std::transform(input.begin(),

src/include/migraphx/operation.hpp

@@ -32,6 +32,7 @@
 #include <utility>
 #include <unordered_map>
 #include <migraphx/reflect.hpp>
+#include <migraphx/functional.hpp>
 #include <migraphx/streamutils.hpp>
 #include <migraphx/normalize_attributes.hpp>
 #include <migraphx/argument.hpp>
@@ -94,6 +95,46 @@ bool has_finalize(const operation& x);
 
 #else
 
+struct dyn_output
+{
+    // original shape from the instruction
+    shape ins_shape;
+    // shape computed at eval time using input arguments
+    shape computed_shape;
+};
+
+/**
+ * Handle dynamic and static shape at evaluation time.
+ * If converted to shape type, returns original ins_shape.
+ * If converted to dyn_output type, will compute an output shape using the input arguments.
+ */
+template <class F>
+struct compute_output_shape
+{
+    F ins_inputs;
+
+    operator dyn_output() const
+    {
+        return ins_inputs([](const auto& x, shape ins_shape, const std::vector<argument>& inputs) {
+            if(ins_shape.dynamic())
+                return dyn_output{ins_shape, compute_shape(x, to_shapes(inputs))};
+            return dyn_output{ins_shape, ins_shape};
+        });
+    }
+
+    operator shape() const
+    {
+        return ins_inputs(
+            [](const auto&, shape ins_shape, const std::vector<argument>&) { return ins_shape; });
+    }
+};
+
+template <class F>
+compute_output_shape<F> make_compute_output_shape(F f)
+{
+    return {f};
+}
+
 namespace detail {
 
 namespace operation_operators {
@@ -199,9 +240,12 @@ auto compute_op(rank<1>,
                 context& ctx,
                 const shape& output_shape,
                 const std::vector<argument>& input)
-    -> decltype(x.compute(auto_any_cast(ctx), output_shape, input))
+    -> decltype(x.compute(auto_any_cast(ctx),
+                          make_compute_output_shape(pack(x, output_shape, input)),
+                          input))
 {
-    return x.compute(auto_any_cast(ctx), output_shape, input);
+    return x.compute(
+        auto_any_cast(ctx), make_compute_output_shape(pack(x, output_shape, input)), input);
 }
 
 template <class T>
@@ -220,9 +264,9 @@ compute_op(const T& x, context& ctx, const shape& output_shape, const std::vecto
 
 template <class T>
 auto compute_op(rank<1>, const T& x, const shape& output_shape, const std::vector<argument>& input)
-    -> decltype(x.compute(output_shape, input))
+    -> decltype(x.compute(make_compute_output_shape(pack(x, output_shape, input)), input))
 {
-    return x.compute(make_compute_output_shape(x, output_shape, input), input);
+    return x.compute(make_compute_output_shape(pack(x, output_shape, input)), input);
 }
 
 template <class T>
@@ -244,9 +288,11 @@ auto compute_op(rank<1>,
                 const shape& output,
                 const std::vector<argument>& inputs,
                 const std::vector<module_ref>& module_args,
-                F f) -> decltype(x.compute(output, inputs, module_args, f))
+                F f)
+    -> decltype(
+        x.compute(make_compute_output_shape(pack(x, output, inputs)), inputs, module_args, f))
 {
-    return x.compute(output, inputs, module_args, f);
+    return x.compute(make_compute_output_shape(pack(x, output, inputs)), inputs, module_args, f);
 }
 
 template <class T, class F>
@@ -278,9 +324,17 @@ auto compute_op(rank<4>,
                 const shape& output,
                 const std::vector<argument>& inputs,
                 const std::vector<module_ref>& module_args,
-                F f) -> decltype(x.compute(auto_any_cast(ctx), output, inputs, module_args, f))
+                F f) -> decltype(x.compute(auto_any_cast(ctx),
+                                           make_compute_output_shape(pack(x, output, inputs)),
+                                           inputs,
+                                           module_args,
+                                           f))
 {
-    return x.compute(auto_any_cast(ctx), output, inputs, module_args, f);
+    return x.compute(auto_any_cast(ctx),
+                     make_compute_output_shape(pack(x, output, inputs)),
+                     inputs,
+                     module_args,
+                     f);
 }
 
 template <class T, class F>
@@ -290,9 +344,11 @@ auto compute_op(rank<3>,
                 const shape& output,
                 const std::vector<argument>& inputs,
                 const std::vector<module_ref>& module_args,
-                F f) -> decltype(x.compute(output, inputs, module_args, f))
+                F f)
+    -> decltype(
+        x.compute(make_compute_output_shape(pack(x, output, inputs)), inputs, module_args, f))
 {
-    return x.compute(output, inputs, module_args, f);
+    return x.compute(make_compute_output_shape(pack(x, output, inputs)), inputs, module_args, f);
 }
 
 template <class T, class F>
@@ -302,9 +358,10 @@ auto compute_op(rank<2>,
                 const shape& output,
                 const std::vector<argument>& inputs,
                 const std::vector<module_ref>&,
-                F) -> decltype(x.compute(output, inputs))
+                F)
+    -> decltype(x.compute(make_compute_output_shape(pack(x, output, inputs)), inputs))
 {
-    return x.compute(output, inputs);
+    return x.compute(make_compute_output_shape(pack(x, output, inputs)), inputs);
 }
 
 template <class T, class F>
@@ -314,9 +371,12 @@ auto compute_op(rank<1>,
                 const shape& output,
                 const std::vector<argument>& inputs,
                 const std::vector<module_ref>&,
-                F) -> decltype(x.compute(auto_any_cast(ctx), output, inputs))
+                F) -> decltype(x.compute(auto_any_cast(ctx),
+                                         make_compute_output_shape(pack(x, output, inputs)),
+                                         inputs))
 {
-    return x.compute(auto_any_cast(ctx), output, inputs);
+    return x.compute(
+        auto_any_cast(ctx), make_compute_output_shape(pack(x, output, inputs)), inputs);
 }
 
 template <class T, class F>
@@ -348,7 +408,8 @@ auto is_context_free_op(rank<1>,
                         const T& x,
                         const shape& output_shape,
                         const std::vector<argument>& input)
-    -> decltype(x.compute(output_shape, input), std::true_type{});
+    -> decltype(x.compute(make_compute_output_shape(pack(x, output_shape, input)), input),
+                std::true_type{});
 
 template <class T>
 auto is_context_free_op(rank<0>, const T&, const shape&, const std::vector<argument>&)
@@ -1278,59 +1339,6 @@ inline const ValueType& any_cast(const operation& x)
 
 inline bool operator!=(const operation& x, const operation& y) { return not(x == y); }
 
-// used for dynamic operators
-struct dyn_output
-{
-    // original instruction output shape
-    shape ins_shape;
-    std::function<shape()> compute_shape;
-
-    shape get_output_shape()
-    {
-        if(output_shape.element_space() == 0)
-        {
-            output_shape = compute_shape();
-        }
-        return output_shape;
-    }
-
-    private:
-    // shape computed at eval time using input arguments
-    shape output_shape;
-};
-
-/**
- * Handle dynamic and static shape at evaluation time.
- * If converted to shape type, returns original ins_shape
- * If converted to dyn_output type, will compute an output shape using the input arguments
- */
-template <class F>
-struct compute_output_shape
-{
-    F ins_inputs;
-    operator dyn_output() const
-    {
-        return unpack(
-            [](const auto& x, shape ins_shape, const std::vector<argument>& args) {
-                return dyn_output{ins_shape, [&]() { compute_shape(x, to_shapes(args)); }};
-            },
-            ins_inputs);
-    }
-
-    operator shape() const
-    {
-        return unpack(
-            [](const auto&, shape ins_shape, const std::vector<argument>&) { return ins_shape; },
-            ins_inputs);
-    }
-};
-
-template <class T>
-auto make_compute_output_shape(const T& x, shape ins_shape, const std::vector<argument>& input)
-{
-    return compute_output_shape{pack(x, ins_shape, input)};
-}
-
 inline value
 compile(operation& op, context& ctx, const shape& output_shape, const std::vector<shape>& input)
 {

src/targets/cpu/lowering.cpp

@@ -216,55 +216,6 @@ struct cpu_pad
 };
 MIGRAPHX_REGISTER_OP(cpu_pad)
 
-struct leaky_relu_op
-{
-    op::leaky_relu op;
-    std::string name() const { return "cpu::leaky_relu"; }
-    auto fcn() const
-    {
-        auto a = op.alpha;
-        return [a](auto x) { return x > 0 ? x : x * a; };
-    }
-};
-
-template <typename Op>
-struct cpu_unary2 : auto_register_op<cpu_unary2<Op>>
-{
-    cpu_unary2() = default;
-
-    template <class T>
-    cpu_unary2(T pop) : op(Op{std::move(pop)})
-    {
-    }
-
-    Op op;
-
-    template <class Self, class F>
-    static auto reflect(Self& self, F f)
-    {
-        return migraphx::reflect(self.op.op, f);
-    }
-    std::string name() const { return op.name(); }
-    shape compute_shape(const std::vector<shape>& inputs) const
-    {
-        check_shapes{inputs, *this}.has(1);
-        const auto& s = inputs.at(0);
-        return {s.type(), s.lens()};
-    }
-
-    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
-    {
-        argument result{output_shape};
-        visit_all(result, args[0])([&](auto output, auto input) {
-            assert(input.get_shape().standard());
-            std::transform(input.begin(), input.end(), output.begin(), op.fcn());
-        });
-
-        return result;
-    }
-};
-template struct cpu_unary2<leaky_relu_op>;
-
 struct cpu_rnn_var_sl_last_output
 {
     op::rnn_var_sl_last_output op;

src/targets/gpu/jit/scatternd.cpp

@@ -79,8 +79,9 @@ struct scatternd_compiler : compiler<scatternd_compiler>
     {
         assert(starts_with(op.name(), "scatternd_"));
         auto reduction = op.name().substr(10);
-        return insert(compile_op(ctx,
-                                 to_shapes({ins->inputs().begin() + 1, ins->inputs().end()}),
+        return insert(compile_op(
+            ctx,
+            to_shapes(std::vector<instruction_ref>{ins->inputs().begin() + 1, ins->inputs().end()}),
             {{"reduction", reduction}}));
     }
 

src/targets/gpu/lowering.cpp

@@ -96,9 +96,9 @@ struct miopen_apply
 
         add_extend_op("argmax");
         add_extend_op("argmin");
-        add_extend_op("elu");
+        // add_extend_op("elu");
         add_extend_op("gather");
-        add_extend_op("leaky_relu");
+        // add_extend_op("leaky_relu");
         add_extend_op("logsoftmax");
         add_extend_op("lrn");
         add_extend_op("multinomial");

src/targets/ref/lowering.cpp

@@ -32,9 +32,7 @@
 #include <migraphx/op/quant_convolution.hpp>
 #include <migraphx/op/dot.hpp>
 #include <migraphx/op/quant_dot.hpp>
-#include <migraphx/op/elu.hpp>
 #include <migraphx/op/im2col.hpp>
-#include <migraphx/op/leaky_relu.hpp>
 #include <migraphx/op/logsoftmax.hpp>
 #include <migraphx/op/loop.hpp>
 #include <migraphx/op/lrn.hpp>
@@ -234,9 +232,8 @@ struct ref_convolution : auto_register_op<ref_convolution<Op>>
         return op.normalize_compute_shape(inputs);
     }
 
-    argument compute(context&, migraphx::dyn_output dyn_output, std::vector<argument> args) const
+    argument compute(context&, shape output_shape, std::vector<argument> args) const
     {
-        shape output_shape;
         std::vector<std::size_t> padding;
         if(op.padding_mode != op::padding_mode_t::default_)
         {
@@ -252,7 +249,11 @@ struct ref_convolution : auto_register_op<ref_convolution<Op>>
         else
         {
             padding = op.padding;
-            output_shape = dyn_output.get_output_shape();
+            if(output_shape.dynamic())
+            {
+                output_shape =
+                    op.normalize_compute_shape({args.at(0).get_shape(), args.at(1).get_shape()});
+            }
         }
 
         argument result{output_shape};
@@ -507,65 +508,6 @@ struct ref_quant_gemm
 };
 MIGRAPHX_REGISTER_OP(ref_gemm)
 
-struct leaky_relu_op
-{
-    op::leaky_relu op;
-    std::string name() const { return "ref::leaky_relu"; }
-    auto fcn() const
-    {
-        auto a = op.alpha;
-        return [a](auto x) { return x > 0 ? x : x * a; };
-    }
-};
-
-struct elu_op
-{
-    op::elu op;
-    std::string name() const { return "ref::elu"; }
-    auto fcn() const
-    {
-        auto a = op.alpha;
-        return [a](auto x) { return x > 0 ? x : a * std::expm1(x); };
-    }
-};
-
-template <typename Op>
-struct ref_unary : auto_register_op<ref_unary<Op>>
-{
-    ref_unary() = default;
-
-    template <class T>
-    ref_unary(T pop) : op(Op{std::move(pop)})
-    {
-    }
-
-    Op op;
-
-    template <class Self, class F>
-    static auto reflect(Self& self, F f)
-    {
-        return migraphx::reflect(self.op.op, f);
-    }
-    std::string name() const { return op.name(); }
-    shape compute_shape(const std::vector<shape>& inputs) const
-    {
-        check_shapes{inputs, *this}.has(1);
-        const auto& s = inputs.at(0);
-        return {s.type(), s.lens()};
-    }
-
-    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
-    {
-        argument result{output_shape};
-        visit_all(result, args[0])([&](auto output, auto input) {
-            assert(input.get_shape().standard());
-            std::transform(input.begin(), input.end(), output.begin(), op.fcn());
-        });
-
-        return result;
-    }
-};
-
 template <class Op>
 struct ref_softmax : auto_register_op<ref_softmax<Op>>
 {
@@ -708,9 +650,7 @@ struct ref_apply
         apply_map["quant_dot"]   = extend_op<ref_quant_gemm, op::quant_dot>();
         apply_map["quant_convolution"] =
             extend_op<ref_convolution<op::quant_convolution>, op::quant_convolution>();
-        apply_map["elu"]        = extend_op<ref_unary<elu_op>, op::elu>();
         apply_map["im2col"]     = extend_op<ref_im2col, op::im2col>();
-        apply_map["leaky_relu"] = extend_op<ref_unary<leaky_relu_op>, op::leaky_relu>();
         apply_map["logsoftmax"] = extend_op<ref_softmax<op::logsoftmax>, op::logsoftmax>();
         apply_map["lrn"]        = extend_op<ref_lrn, op::lrn>();
         apply_map["pad"]        = extend_op<ref_pad, op::pad>();

test/onnx/gen_onnx.py

@@ -3647,6 +3647,16 @@ def neg_test():
     return ([node], [x], [y])
 
 
+@onnx_test
+def neg_dynamic_test():
+    x = helper.make_tensor_value_info('0', TensorProto.INT64, [None, 3])
+    y = helper.make_tensor_value_info('1', TensorProto.INT64, [None, 3])
+
+    node = onnx.helper.make_node('Neg', inputs=['0'], outputs=['1'])
+
+    return ([node], [x], [y])
+
+
 @onnx_test
 def nms_test():
     b = helper.make_tensor_value_info('boxes', TensorProto.FLOAT, [1, 6, 4])
@@ -5280,6 +5290,20 @@ def sinh_test():
     return ([node], [x], [y])
 
 
+@onnx_test
+def sinh_dynamic_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [None])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [None])
+
+    node = onnx.helper.make_node(
+        'Sinh',
+        inputs=['x'],
+        outputs=['y'],
+    )
+
+    return ([node], [x], [y])
+
+
 @onnx_test
 def size_float_test():
     x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [2, 3, 4])

test/onnx/onnx_test.cpp

@@ -3476,6 +3476,21 @@ TEST_CASE(neg_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(neg_dynamic_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape s{migraphx::shape::int64_type, {{1, 10, 0}, {3, 3, 0}}};
+    auto input = mm->add_parameter("0", s);
+    auto ret   = mm->add_instruction(migraphx::make_op("neg"), input);
+    mm->add_return({ret});
+
+    migraphx::onnx_options options;
+    options.default_dyn_dim_value = {1, 10, 0};
+    auto prog                     = migraphx::parse_onnx("neg_dynamic_test.onnx", options);
+    EXPECT(p == prog);
+}
+
 TEST_CASE(nms_test)
 {
     migraphx::program p;
@@ -5199,6 +5214,29 @@ TEST_CASE(sinh_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(sinh_dynamic_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape::dynamic_dimension dd{1, 10, 0};
+    std::vector<migraphx::shape::dynamic_dimension> dyn_dims;
+    dyn_dims.push_back(dd);
+    auto input = mm->add_parameter("x", migraphx::shape{migraphx::shape::float_type, dyn_dims});
+    mm->add_instruction(migraphx::make_op("sinh"), input);
+
+    migraphx::onnx_options options;
+    options.default_dyn_dim_value = dd;
+    auto prog                     = parse_onnx("sinh_dynamic_test.onnx", options);
+    auto* mm_onnx                 = prog.get_main_module();
+    auto last_ins                 = std::prev(mm_onnx->end());
+    if(last_ins->name() == "@return")
+    {
+        mm->remove_instruction(last_ins);
+    }
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(size_float_test)
 {
     migraphx::program p;

test/verify/test_elu.cpp

@@ -34,7 +34,7 @@ struct test_elu : verify_program<test_elu>
         migraphx::program p;
         auto* mm = p.get_main_module();
         auto x = mm->add_parameter("x", migraphx::shape{migraphx::shape::float_type, {4, 3, 3, 3}});
-        mm->add_instruction(migraphx::make_op("leaky_relu", {{"alpha", 1.0}}), x);
+        mm->add_instruction(migraphx::make_op("elu", {{"alpha", 0.8}}), x);
         return p;
     }
 };

test/verify/test_leaky_relu.cpp

@@ -34,7 +34,7 @@ struct test_leaky_relu : verify_program<test_leaky_relu>
         migraphx::program p;
         auto* mm = p.get_main_module();
         auto x = mm->add_parameter("x", migraphx::shape{migraphx::shape::float_type, {4, 3, 3, 3}});
-        mm->add_instruction(migraphx::make_op("leaky_relu", {{"alpha", 0.01}}), x);
+        mm->add_instruction(migraphx::make_op("leaky_relu", {{"alpha", 0.41}}), x);
         return p;
     }
 };

tools/include/operation.hpp

@@ -32,6 +32,7 @@
 #include <utility>
 #include <unordered_map>
 #include <migraphx/reflect.hpp>
+#include <migraphx/functional.hpp>
 #include <migraphx/streamutils.hpp>
 #include <migraphx/normalize_attributes.hpp>
 #include <migraphx/argument.hpp>
@@ -94,6 +95,46 @@ bool has_finalize(const operation& x);
 
 #else
 
+struct dyn_output
+{
+    // original shape from the instruction
+    shape ins_shape;
+    // shape computed at eval time using input arguments
+    shape computed_shape;
+};
+
+/**
+ * Handle dynamic and static shape at evaluation time.
+ * If converted to shape type, returns original ins_shape.
+ * If converted to dyn_output type, will compute an output shape using the input arguments.
+ */
+template <class F>
+struct compute_output_shape
+{
+    F ins_inputs;
+
+    operator dyn_output() const
+    {
+        return ins_inputs([](const auto& x, shape ins_shape, const std::vector<argument>& inputs) {
+            if(ins_shape.dynamic())
+                return dyn_output{ins_shape, compute_shape(x, to_shapes(inputs))};
+            return dyn_output{ins_shape, ins_shape};
+        });
+    }
+
+    operator shape() const
+    {
+        return ins_inputs(
+            [](const auto&, shape ins_shape, const std::vector<argument>&) { return ins_shape; });
+    }
+};
+
+template <class F>
+compute_output_shape<F> make_compute_output_shape(F f)
+{
+    return {f};
+}
+
 namespace detail {
 
 namespace operation_operators {
@@ -199,9 +240,12 @@ auto compute_op(rank<1>,
                 context& ctx,
                 const shape& output_shape,
                 const std::vector<argument>& input)
-    -> decltype(x.compute(auto_any_cast(ctx), output_shape, input))
+    -> decltype(x.compute(auto_any_cast(ctx),
+                          make_compute_output_shape(pack(x, output_shape, input)),
+                          input))
 {
-    return x.compute(auto_any_cast(ctx), output_shape, input);
+    return x.compute(
+        auto_any_cast(ctx), make_compute_output_shape(pack(x, output_shape, input)), input);
 }
 
 template <class T>
@@ -220,9 +264,9 @@ compute_op(const T& x, context& ctx, const shape& output_shape, const std::vecto
 
 template <class T>
 auto compute_op(rank<1>, const T& x, const shape& output_shape, const std::vector<argument>& input)
-    -> decltype(x.compute(output_shape, input))
+    -> decltype(x.compute(make_compute_output_shape(pack(x, output_shape, input)), input))
 {
-    return x.compute(make_compute_output_shape(x, output_shape, input), input);
+    return x.compute(make_compute_output_shape(pack(x, output_shape, input)), input);
 }
 
 template <class T>
@@ -244,9 +288,11 @@ auto compute_op(rank<1>,
                 const shape& output,
                 const std::vector<argument>& inputs,
                 const std::vector<module_ref>& module_args,
-                F f) -> decltype(x.compute(output, inputs, module_args, f))
+                F f)
+    -> decltype(
+        x.compute(make_compute_output_shape(pack(x, output, inputs)), inputs, module_args, f))
 {
-    return x.compute(output, inputs, module_args, f);
+    return x.compute(make_compute_output_shape(pack(x, output, inputs)), inputs, module_args, f);
 }
 
 template <class T, class F>
@@ -278,9 +324,17 @@ auto compute_op(rank<4>,
                 const shape& output,
                 const std::vector<argument>& inputs,
                 const std::vector<module_ref>& module_args,
-                F f) -> decltype(x.compute(auto_any_cast(ctx), output, inputs, module_args, f))
+                F f) -> decltype(x.compute(auto_any_cast(ctx),
+                                           make_compute_output_shape(pack(x, output, inputs)),
+                                           inputs,
+                                           module_args,
+                                           f))
 {
-    return x.compute(auto_any_cast(ctx), output, inputs, module_args, f);
+    return x.compute(auto_any_cast(ctx),
+                     make_compute_output_shape(pack(x, output, inputs)),
+                     inputs,
+                     module_args,
+                     f);
 }
 
 template <class T, class F>
@@ -290,9 +344,11 @@ auto compute_op(rank<3>,
                 const shape& output,
                 const std::vector<argument>& inputs,
                 const std::vector<module_ref>& module_args,
-                F f) -> decltype(x.compute(output, inputs, module_args, f))
+                F f)
+    -> decltype(
+        x.compute(make_compute_output_shape(pack(x, output, inputs)), inputs, module_args, f))
 {
-    return x.compute(output, inputs, module_args, f);
+    return x.compute(make_compute_output_shape(pack(x, output, inputs)), inputs, module_args, f);
 }
 
 template <class T, class F>
@@ -302,9 +358,10 @@ auto compute_op(rank<2>,
                 const shape& output,
                 const std::vector<argument>& inputs,
                 const std::vector<module_ref>&,
-                F) -> decltype(x.compute(output, inputs))
+                F)
+    -> decltype(x.compute(make_compute_output_shape(pack(x, output, inputs)), inputs))
 {
-    return x.compute(output, inputs);
+    return x.compute(make_compute_output_shape(pack(x, output, inputs)), inputs);
 }
 
 template <class T, class F>
@@ -314,9 +371,12 @@ auto compute_op(rank<1>,
                 const shape& output,
                 const std::vector<argument>& inputs,
                 const std::vector<module_ref>&,
-                F) -> decltype(x.compute(auto_any_cast(ctx), output, inputs))
+                F) -> decltype(x.compute(auto_any_cast(ctx),
+                                         make_compute_output_shape(pack(x, output, inputs)),
+                                         inputs))
 {
-    return x.compute(auto_any_cast(ctx), output, inputs);
+    return x.compute(
+        auto_any_cast(ctx), make_compute_output_shape(pack(x, output, inputs)), inputs);
 }
 
 template <class T, class F>
@@ -348,7 +408,8 @@ auto is_context_free_op(rank<1>,
                         const T& x,
                         const shape& output_shape,
                         const std::vector<argument>& input)
-    -> decltype(x.compute(output_shape, input), std::true_type{});
+    -> decltype(x.compute(make_compute_output_shape(pack(x, output_shape, input)), input),
+                std::true_type{});
 
 template <class T>
 auto is_context_free_op(rank<0>, const T&, const shape&, const std::vector<argument>&)
@@ -563,59 +624,6 @@ lifetime get_lifetime_op(const T&)
     return not(x == y);
 }
 
-// used for dynamic operators
-struct dyn_output
-{
-    // original instruction output shape
-    shape ins_shape;
-    std::function<shape()> compute_shape;
-
-    shape get_output_shape()
-    {
-        if(output_shape.element_space() == 0)
-        {
-            output_shape = compute_shape();
-        }
-        return output_shape;
-    }
-
-    private:
-    // shape computed at eval time using input arguments
-    shape output_shape;
-};
-
-/**
- * Handle dynamic and static shape at evaluation time.
- * If converted to shape type, returns original ins_shape
- * If converted to dyn_output type, will compute an output shape using the input arguments
- */
-template <class F>
-struct compute_output_shape
-{
-    F ins_inputs;
-    operator dyn_output() const
-    {
-        return unpack(
-            [](const auto& x, shape ins_shape, const std::vector<argument>& args) {
-                return dyn_output{ins_shape, [&]() { compute_shape(x, to_shapes(args)); }};
-            },
-            ins_inputs);
-    }
-
-    operator shape() const
-    {
-        return unpack(
-            [](const auto&, shape ins_shape, const std::vector<argument>&) { return ins_shape; },
-            ins_inputs);
-    }
-};
-
-template <class T>
-auto make_compute_output_shape(const T& x, shape ins_shape, const std::vector<argument>& input)
-{
-    return compute_output_shape{pack(x, ins_shape, input)};
-}
-
 inline value
 compile(operation& op, context& ctx, const shape& output_shape, const std::vector<shape>& input)
 {