Loop operator (#853)

Add Loop operator for opset version 13.
Notes: 1) Default max iteration number is 10 if no max iteration number is provided
2) To change the max iter number, a user can set the max_loop_iterations in the onnx_option struct when parsing a model.
3) The returned shape of the scan output is from the max_loop_iterations even the actual loop num is less than that. This issue also applies to other operators like NonZero and NonMaxSuppression. A issue #948 is created to track this and to be resolved later.
Co-authored-by: Paul <pfultz2@yahoo.com>
Co-authored-by: mvermeulen <5479696+mvermeulen@users.noreply.github.com>

src/targets/gpu/device/fill.cpp

+#include <migraphx/gpu/device/fill.hpp>
+#include <migraphx/gpu/device/nary.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+namespace device {
+
+void fill(hipStream_t stream, const argument& result, unsigned long val)
+{
+    nary(stream, result)([=]() __device__ { return val; });
+}
+
+} // namespace device
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/targets/gpu/hip.cpp

@@ -169,12 +169,26 @@ void gpu_copy(context& ctx, const argument& src, const argument& dst)
 
 void copy_to_gpu(context& ctx, const argument& src, const argument& dst)
 {
-    gpu_copy(ctx, register_on_gpu(src), dst);
+    if(src.get_shape() == dst.get_shape() and dst.get_shape().packed())
+    {
+        hip_async_copy(ctx, src, dst, hipMemcpyHostToDevice);
+    }
+    else
+    {
+        gpu_copy(ctx, register_on_gpu(src), dst);
+    }
 }
 
 void copy_from_gpu(context& ctx, const argument& src, const argument& dst)
 {
-    gpu_copy(ctx, src, register_on_gpu(dst));
+    if(src.get_shape() == dst.get_shape() and dst.get_shape().packed())
+    {
+        hip_async_copy(ctx, src, dst, hipMemcpyDeviceToHost);
+    }
+    else
+    {
+        gpu_copy(ctx, src, register_on_gpu(dst));
+    }
 }
 
 argument get_preallocation(context& ctx, const std::string& id)

src/targets/gpu/include/migraphx/gpu/device/fill.hpp

+#ifndef MIGRAPHX_GUARD_RTGLIB_DEVICE_FILL_HPP
+#define MIGRAPHX_GUARD_RTGLIB_DEVICE_FILL_HPP
+
+#include <migraphx/argument.hpp>
+#include <migraphx/config.hpp>
+#include <hip/hip_runtime_api.h>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+namespace device {
+
+void fill(hipStream_t stream, const argument& result, unsigned long val);
+
+} // namespace device
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/targets/gpu/include/migraphx/gpu/loop.hpp

+#ifndef MIGRAPHX_GUARD_RTGLIB_LOOP_HPP
+#define MIGRAPHX_GUARD_RTGLIB_LOOP_HPP
+
+#include <migraphx/argument.hpp>
+#include <migraphx/reflect.hpp>
+#include <migraphx/op/loop.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+struct context;
+
+struct hip_loop
+{
+    op::loop op;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return migraphx::reflect(self.op, f);
+    }
+
+    std::string name() const { return "gpu::loop"; }
+    shape compute_shape(std::vector<shape> inputs, std::vector<module_ref> mods) const;
+    argument
+    compute(context& ctx,
+            const shape& output_shape,
+            const std::vector<argument>& args,
+            const std::vector<module_ref>& mods,
+            const std::function<std::vector<argument>(
+                module_ref&, const std::unordered_map<std::string, argument>&)>& run) const;
+    std::ptrdiff_t output_alias(const std::vector<shape>& shapes) const
+    {
+        return shapes.size() - 1;
+    }
+};
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/targets/gpu/loop.cpp

+#include <migraphx/run_loop.hpp>
+#include <migraphx/gpu/loop.hpp>
+#include <migraphx/gpu/context.hpp>
+#include <migraphx/gpu/device/fill.hpp>
+#include <unordered_map>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+shape hip_loop::compute_shape(std::vector<shape> inputs, std::vector<module_ref> mods) const
+{
+    auto input_num = (inputs.size() - 2) / 2;
+    inputs.erase(inputs.begin() + input_num, inputs.end());
+    return op.compute_shape(inputs, std::move(mods));
+}
+
+struct gpu_loop
+{
+    int64_t max_iterations = 0;
+
+    template <class T>
+    void copy(context& ctx, const argument& src, T& dst) const
+    {
+        argument arg_dst{src.get_shape(), &dst};
+        copy_from_gpu(ctx, src, arg_dst);
+    }
+
+    template <class T>
+    void copy(context& ctx, T src, const argument& dst) const
+    {
+        argument arg_src{dst.get_shape(), &src};
+        copy_to_gpu(ctx, arg_src, dst);
+    }
+
+    void append(const std::vector<argument>&, const std::vector<argument>&, int) const {}
+
+    void set_zero(context& ctx, const std::vector<argument>& concatenated_outputs, int iter) const
+    {
+        if(iter >= max_iterations)
+            return;
+
+        auto elem_num = max_iterations - iter;
+        for(const auto& out : concatenated_outputs)
+        {
+            auto s    = out.get_shape();
+            auto size = s.bytes() / max_iterations;
+            auto lens = s.lens();
+            lens[0]   = elem_num;
+            shape ss{s.type(), lens};
+            assert(ss.bytes() + iter * size <= out.get_shape().bytes());
+            device::fill(ctx.get_stream().get(), argument(ss, out.data() + iter * size), 0);
+        }
+    }
+
+    std::unordered_map<std::string, int> get_output_params(const module& m) const
+    {
+        auto get_output_index = [](const std::string& name) {
+            std::string out_prefix = "#output_";
+            auto loc               = name.find(out_prefix);
+            if(loc != std::string::npos)
+            {
+                int index = std::stoi(name.substr(loc + out_prefix.size()));
+                return index;
+            }
+
+            return -1;
+        };
+
+        const auto& param_names = m.get_parameter_names();
+        std::unordered_map<std::string, int> result;
+        for(const auto& name : param_names)
+        {
+            auto index = get_output_index(name);
+            if(index == -1)
+                continue;
+            result[name] = index;
+        }
+
+        return result;
+    }
+};
+
+argument
+hip_loop::compute(context& ctx,
+                  const shape&,
+                  const std::vector<argument>& args,
+                  const std::vector<module_ref>& mods,
+                  const std::function<std::vector<argument>(
+                      module_ref&, const std::unordered_map<std::string, argument>&)>& run) const
+{
+    return run_loop(gpu_loop{op.max_iterations}, ctx, args, mods, run);
+}
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/targets/gpu/lowering.cpp

@@ -188,6 +188,7 @@ struct miopen_apply
         add_batch_norm_inference_op();
         add_neg_op();
         add_if_op();
+        add_loop_op();
     }
 
     void copy_params()
@@ -469,9 +470,48 @@ struct miopen_apply
             return mod->replace_instruction(ins, ins->get_operator(), inputs, mod_args);
         });
     }
+
+    // replace the loop operator with gpu_loop operator
+    void add_loop_op()
+    {
+        apply_map.emplace("loop", [=](instruction_ref ins) {
+            std::vector<instruction_ref> inputs = ins->inputs();
+            // copy max_iter from gpu to cpu
+            auto cpu_max_iter =
+                mod->insert_instruction(ins, make_op("hip::copy_from_gpu"), inputs.at(0));
+            auto cpu_cond =
+                mod->insert_instruction(ins, make_op("hip::copy_from_gpu"), inputs.at(1));
+            auto synced_max_iter =
+                mod->insert_instruction(ins, make_op("hip::sync_stream"), cpu_max_iter, cpu_cond);
+            inputs.at(0)     = synced_max_iter;
+            inputs.at(1)     = cpu_cond;
+            auto copy_inputs = inputs;
+            std::transform(
+                copy_inputs.begin(), copy_inputs.end(), std::back_inserter(inputs), [&](auto in) {
+                    return mod->insert_instruction(
+                        ins, make_op("hip::allocate", {{"shape", to_value(in->get_shape())}}));
+                });
+
+            auto mod_args = ins->module_inputs();
+            auto output   = insert_allocation(ins, ins->get_shape());
+
+            const auto* sub_mod = mod_args.front();
+            auto cond_out       = mod->insert_instruction(
+                ins,
+                make_op("hip::allocate",
+                        {{"shape", to_value(sub_mod->get_output_shapes().front())}}));
+            // add cond and mod outputs to the argument list
+            inputs.push_back(cond_out);
+            inputs.push_back(output);
+
+            return mod->replace_instruction(
+                ins, make_op("gpu::loop", ins->get_operator().to_value()), inputs, mod_args);
+        });
+    }
 };
 
 void lowering::apply(module& m) const { miopen_apply{&m, this}.apply(); }
+
 } // namespace gpu
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/targets/ref/lowering.cpp

@@ -10,10 +10,10 @@
 #include <migraphx/op/dot.hpp>
 #include <migraphx/op/quant_dot.hpp>
 #include <migraphx/op/elu.hpp>
-#include <migraphx/op/if_op.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>
 #include <migraphx/op/pad.hpp>
 #include <migraphx/op/pooling.hpp>

test/api/test_cpu.cpp

@@ -163,4 +163,16 @@ TEST_CASE(get_main_module)
     p.print();
 }
 
+TEST_CASE(set_loop_default_iter_num)
+{
+    migraphx::onnx_options option;
+    option.set_default_loop_iterations(15);
+    auto p                             = migraphx::parse_onnx("loop_default_test.onnx", option);
+    auto out_shapes                    = p.get_output_shapes();
+    std::vector<std::size_t> out_lens0 = {1};
+    EXPECT(out_shapes[0].lengths() == out_lens0);
+    std::vector<std::size_t> out_lens1 = {15, 1};
+    EXPECT(out_shapes[1].lengths() == out_lens1);
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/api/test_gpu.cpp

@@ -74,4 +74,62 @@ TEST_CASE(if_pl_test)
     }
 }
 
+TEST_CASE(loop_test)
+{
+    auto run_prog = [&](int64_t max_iter_num) {
+        migraphx::onnx_options parse_options;
+        parse_options.set_default_loop_iterations(max_iter_num);
+        auto p             = migraphx::parse_onnx("loop_default_test.onnx", parse_options);
+        auto shapes_before = p.get_output_shapes();
+        migraphx_compile_options options;
+        options.offload_copy = true;
+        p.compile(migraphx::target("gpu"), options);
+        auto shapes_after = p.get_output_shapes();
+        CHECK(shapes_before.size() == 2);
+        CHECK(bool{shapes_before.front() == shapes_after.front()});
+
+        migraphx::program_parameters pp;
+        auto param_shapes     = p.get_parameter_shapes();
+        auto aas              = param_shapes["a"];
+        std::vector<float> xd = {1.0f};
+        pp.add("a", migraphx::argument(aas, xd.data()));
+        auto bbs              = param_shapes["b"];
+        std::vector<float> yd = {2.0};
+        pp.add("b", migraphx::argument(bbs, yd.data()));
+
+        auto outputs = p.eval(pp);
+        auto output  = outputs[0];
+        auto lens    = output.get_shape().lengths();
+        auto elem_num =
+            std::accumulate(lens.begin(), lens.end(), 1, std::multiplies<std::size_t>());
+        float* data_ptr = reinterpret_cast<float*>(output.data());
+        std::vector<std::vector<float>> ret;
+        ret.push_back({data_ptr, data_ptr + elem_num});
+
+        output   = outputs[1];
+        lens     = output.get_shape().lengths();
+        elem_num = std::accumulate(lens.begin(), lens.end(), 1, std::multiplies<std::size_t>());
+        data_ptr = reinterpret_cast<float*>(output.data());
+        ret.push_back({data_ptr, data_ptr + elem_num});
+
+        return ret;
+    };
+
+    {
+        auto result_vector       = run_prog(10);
+        std::vector<float> gold0 = {2.0f};
+        EXPECT(result_vector.at(0) == gold0);
+        std::vector<float> gold1 = {-2, 4, 0, 0, 0, 0, 0, 0, 0, 0};
+        EXPECT(result_vector.at(1) == gold1);
+    }
+
+    {
+        auto result_vector       = run_prog(15);
+        std::vector<float> gold0 = {2.0f};
+        EXPECT(result_vector.at(0) == gold0);
+        std::vector<float> gold1 = {-2, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
+        EXPECT(result_vector.at(1) == gold1);
+    }
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/onnx/gen_onnx.py

@@ -2346,6 +2346,84 @@ def logsoftmax_nonstd_input_test():
     return ([node0, node1], [x], [y])
 
 
+@onnx_test
+def loop_default_test():
+    body = helper.make_graph([
+        helper.make_node("Add", ["a", "b_in"], ["my_local"]),
+        helper.make_node("Sub", ["a", "b_in"], ["a_sub_b_in"]),
+        helper.make_node("Greater", ["my_local", "a_sub_b_in"],
+                         ["keep_going"]),
+        helper.make_node("Add", ["a_sub_b_in", "a_sub_b_in"],
+                         ["user_defined_vals"]),
+    ], "body", [
+        helper.make_tensor_value_info('iteration_num', TensorProto.INT64, []),
+        helper.make_tensor_value_info('keep_going_inp', TensorProto.BOOL, []),
+        helper.make_tensor_value_info('b_in', TensorProto.FLOAT, [])
+    ], [
+        helper.make_tensor_value_info('keep_going', TensorProto.BOOL, []),
+        helper.make_tensor_value_info('a_sub_b_in', TensorProto.FLOAT, []),
+        helper.make_tensor_value_info('my_local', TensorProto.FLOAT, []),
+        helper.make_tensor_value_info('user_defined_vals', TensorProto.FLOAT,
+                                      []),
+    ])
+
+    node = helper.make_node(
+        "Loop",
+        inputs=["", "", "b"],
+        outputs=["b_loop", "my_local_loop", "user_defined_vals_loop"],
+        body=body)
+
+    a = helper.make_tensor_value_info('a', TensorProto.FLOAT, [])
+    b = helper.make_tensor_value_info('b', TensorProto.FLOAT, [])
+
+    b_loop = helper.make_tensor_value_info('b_loop', TensorProto.FLOAT, [])
+    uout = helper.make_tensor_value_info('user_defined_vals_loop',
+                                         TensorProto.FLOAT, [2, 1])
+
+    return ([node], [a, b], [b_loop, uout])
+
+
+@onnx_test
+def loop_test():
+    body = helper.make_graph([
+        helper.make_node("Add", ["a", "b_in"], ["my_local"]),
+        helper.make_node("Sub", ["a", "b_in"], ["a_sub_b_in"]),
+        helper.make_node("Greater", ["my_local", "a_sub_b_in"],
+                         ["keep_going"]),
+        helper.make_node("Add", ["a_sub_b_in", "a_sub_b_in"],
+                         ["user_defined_vals"]),
+    ], "body", [
+        helper.make_tensor_value_info('iteration_num', TensorProto.INT64, [1]),
+        helper.make_tensor_value_info('keep_going_inp', TensorProto.BOOL, [1]),
+        helper.make_tensor_value_info('b_in', TensorProto.FLOAT, [1])
+    ], [
+        helper.make_tensor_value_info('keep_going', TensorProto.BOOL, [1]),
+        helper.make_tensor_value_info('a_sub_b_in', TensorProto.FLOAT, [1]),
+        helper.make_tensor_value_info('my_local', TensorProto.FLOAT, [1]),
+        helper.make_tensor_value_info('user_defined_vals', TensorProto.FLOAT,
+                                      [1]),
+    ])
+
+    node = helper.make_node(
+        "Loop",
+        inputs=["max_trip_count", "keep_going_cond", "b"],
+        outputs=["b_loop", "my_local_loop", "user_defined_vals_loop"],
+        body=body)
+
+    a = helper.make_tensor_value_info('a', TensorProto.FLOAT, [1])
+    b = helper.make_tensor_value_info('b', TensorProto.FLOAT, [1])
+    cond = helper.make_tensor_value_info('keep_going_cond', TensorProto.BOOL,
+                                         [1])
+    iter = helper.make_tensor_value_info('max_trip_count', TensorProto.INT64,
+                                         [1])
+
+    b_loop = helper.make_tensor_value_info('b_loop', TensorProto.FLOAT, [1])
+    uout = helper.make_tensor_value_info('user_defined_vals_loop',
+                                         TensorProto.FLOAT, [2, 1])
+
+    return ([node], [iter, cond, a, b], [b_loop, uout])
+
+
 @onnx_test
 def lrn_test():
     x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 28, 24, 24])

test/onnx/loop_test.onnx

+	loop_test:
+
+max_trip_count
+keep_going_cond
+
bb_loop
my_local_loopuser_defined_vals_loop"Loop*
+body2
+
+
a
+b_inmy_local"Add
+
+
a
+b_in
+a_sub_b_in"Sub
++
+my_local
+
+a_sub_b_in
+keep_going"Greater
+0
+
+a_sub_b_in
+
+a_sub_b_inuser_defined_vals"AddbodyZ
+iteration_num
+
+
+
Z
+keep_going_inp
+
+	
+
Z
+b_in
+
+

+
b
+
+keep_going
+
+	
+
b
+
+a_sub_b_in
+
+

+
b
+my_local
+
+

+
b
+user_defined_vals
+
+

+

	loop_testZ
+max_trip_count
+
+
+
Z
+keep_going_cond
+
+	
+
Z
+
a
+
+

+
Z
+
b
+
+

+
b
+b_loop
+
+

+
b(
+user_defined_vals_loop
+

+
+
B
\ No newline at end of file

test/onnx/onnx_test.cpp

@@ -2010,6 +2010,82 @@ TEST_CASE(logsoftmax_nonstd_input_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(loop_default_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+
+    migraphx::shape su{migraphx::shape::float_type};
+    auto a = mm->add_parameter("a", su);
+    auto b = mm->add_parameter("b", su);
+    migraphx::shape si{migraphx::shape::int64_type};
+    auto max_iter = mm->add_literal(migraphx::literal(si, {10}));
+    migraphx::shape sc{migraphx::shape::bool_type};
+    auto icond = mm->add_literal(migraphx::literal(sc, {1}));
+    mm->add_instruction(migraphx::make_op("undefined"));
+
+    auto* body = p.create_module("Loop_3_loop");
+    body->add_parameter("iteration_num", {migraphx::shape::int64_type});
+    body->add_parameter("keep_going_inp", {migraphx::shape::bool_type});
+    auto var = body->add_parameter("b_in", su);
+
+    auto ad = body->add_instruction(migraphx::make_op("add"), a, var);
+    auto sb = body->add_instruction(migraphx::make_op("sub"), a, var);
+    auto gt = body->add_instruction(migraphx::make_op("greater"), ad, sb);
+    auto cv = body->add_instruction(
+        migraphx::make_op("convert", {{"target_type", migraphx::shape::bool_type}}), gt);
+    auto ad1 = body->add_instruction(migraphx::make_op("add"), sb, sb);
+    body->add_return({cv, sb, ad, ad1});
+
+    auto lp = mm->add_instruction(
+        migraphx::make_op("loop", {{"max_iterations", 10}}), {max_iter, icond, b}, {body});
+    auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), lp);
+    mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), lp);
+    auto r2 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 2}}), lp);
+    mm->add_return({r0, r2});
+
+    auto prog = migraphx::parse_onnx("loop_default_test.onnx");
+
+    EXPECT(p == prog);
+}
+
+TEST_CASE(loop_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape si{migraphx::shape::int64_type, {1}};
+    auto max_iter = mm->add_parameter("max_trip_count", si);
+    migraphx::shape sc{migraphx::shape::bool_type, {1}};
+    auto icond = mm->add_parameter("keep_going_cond", sc);
+    migraphx::shape su{migraphx::shape::float_type, {1}};
+    auto a = mm->add_parameter("a", su);
+    auto b = mm->add_parameter("b", su);
+
+    auto* body = p.create_module("Loop_4_loop");
+    body->add_parameter("iteration_num", si);
+    body->add_parameter("keep_going_inp", sc);
+    auto var = body->add_parameter("b_in", su);
+
+    auto ad = body->add_instruction(migraphx::make_op("add"), a, var);
+    auto sb = body->add_instruction(migraphx::make_op("sub"), a, var);
+    auto gt = body->add_instruction(migraphx::make_op("greater"), ad, sb);
+    auto cv = body->add_instruction(
+        migraphx::make_op("convert", {{"target_type", migraphx::shape::bool_type}}), gt);
+    auto ad1 = body->add_instruction(migraphx::make_op("add"), sb, sb);
+    body->add_return({cv, sb, ad, ad1});
+
+    auto lp = mm->add_instruction(
+        migraphx::make_op("loop", {{"max_iterations", 10}}), {max_iter, icond, b}, {body});
+    auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), lp);
+    mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), lp);
+    auto r2 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 2}}), lp);
+    mm->add_return({r0, r2});
+
+    auto prog = migraphx::parse_onnx("loop_test.onnx");
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(lrn_test)
 {
     migraphx::program p;

test/ref_loop_test.cpp

+#include <iostream>
+#include <vector>
+#include <cmath>
+#include <migraphx/literal.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/quantization.hpp>
+#include <migraphx/ref/target.hpp>
+#include <migraphx/verify.hpp>
+#include <migraphx/make_op.hpp>
+
+#include "test.hpp"
+
+static auto run_prog(int64_t iter_num, bool cond, int64_t ini_val)
+{
+    migraphx::shape si{migraphx::shape::int64_type};
+    migraphx::shape s{migraphx::shape::int64_type, {1}};
+    migraphx::shape sc{migraphx::shape::bool_type};
+
+    auto create_program = [&]() {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+
+        auto in_iter = mm->add_parameter("iter_num", si);
+        auto in_cond = mm->add_parameter("ccond", sc);
+        auto in_val  = mm->add_parameter("val", s);
+
+        auto* body = p.create_module("loop_module");
+        auto iter  = body->add_parameter("#loop_module_in_0", si);
+        body->add_parameter("#loop_module_in_1", sc);
+        auto in_v               = body->add_parameter("#loop_module_in_2", s);
+        std::vector<int64_t> vd = {3};
+        auto l                  = body->add_literal(migraphx::literal(si, vd));
+        auto ad                 = body->add_instruction(migraphx::make_op("add"), iter, l);
+        auto val                = body->add_instruction(migraphx::make_op("add"), in_v, ad);
+        auto eq                 = body->add_instruction(migraphx::make_op("equal"), iter, l);
+        auto beq                = body->add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::bool_type}}), eq);
+        auto neq = body->add_instruction(migraphx::make_op("not"), beq);
+        body->add_return({neq, val, val});
+
+        auto rl = mm->add_instruction(migraphx::make_op("loop", {{"max_iterations", 10}}),
+                                      {in_iter, in_cond, in_val},
+                                      {body});
+        auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), rl);
+        auto r1 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), rl);
+        mm->add_return({r0, r1});
+
+        return p;
+    };
+
+    auto p = create_program();
+    p.compile(migraphx::ref::target{});
+    migraphx::parameter_map pp;
+    pp["iter_num"] = migraphx::argument(si, &iter_num);
+    pp["ccond"]    = migraphx::argument(sc, &cond);
+    pp["val"]      = migraphx::argument(s, &ini_val);
+    auto rets      = p.eval(pp);
+
+    std::vector<std::vector<int64_t>> res;
+    for(auto& arg : rets)
+    {
+        std::vector<int64_t> vec;
+        arg.visit([&](auto v) { vec.assign(v.begin(), v.end()); });
+        res.push_back(vec);
+    }
+
+    return res;
+}
+
+TEST_CASE(loop_test1)
+{
+    auto ress                      = run_prog(10, true, 1);
+    std::vector<int64_t> gold_last = {19};
+    EXPECT(ress.front() == gold_last);
+    std::vector<int64_t> gold_concat = {4, 8, 13, 19, 0, 0, 0, 0, 0, 0};
+    EXPECT(ress.back() == gold_concat);
+}
+
+TEST_CASE(loop_test2)
+{
+    auto ress                      = run_prog(4, true, 1);
+    std::vector<int64_t> gold_last = {19};
+    EXPECT(ress.front() == gold_last);
+    std::vector<int64_t> gold_concat = {4, 8, 13, 19, 0, 0, 0, 0, 0, 0};
+    EXPECT(ress.back() == gold_concat);
+}
+
+TEST_CASE(loop_test3)
+{
+    auto ress                      = run_prog(3, true, 1);
+    std::vector<int64_t> gold_last = {13};
+    EXPECT(ress.front() == gold_last);
+    std::vector<int64_t> gold_concat = {4, 8, 13, 0, 0, 0, 0, 0, 0, 0};
+    EXPECT(ress.back() == gold_concat);
+}
+
+TEST_CASE(loop_test4)
+{
+    auto ress                      = run_prog(5, true, 2);
+    std::vector<int64_t> gold_last = {20};
+    EXPECT(ress.front() == gold_last);
+    std::vector<int64_t> gold_concat = {5, 9, 14, 20, 0, 0, 0, 0, 0, 0};
+    EXPECT(ress.back() == gold_concat);
+}
+
+int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/run_loop_test.cpp

+#include <iostream>
+#include <vector>
+#include <cmath>
+#include <migraphx/literal.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/quantization.hpp>
+#include <migraphx/ref/target.hpp>
+#include <migraphx/shape.hpp>
+#include <migraphx/verify.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/run_loop.hpp>
+#include <migraphx/check_shapes.hpp>
+#include <migraphx/functional.hpp>
+#include <migraphx/op/loop.hpp>
+#include <basic_ops.hpp>
+#include "test.hpp"
+
+struct copy_op
+{
+    std::string name() const { return "copy"; }
+
+    migraphx::shape compute_shape(std::vector<migraphx::shape> inputs) const
+    {
+        return inputs.front();
+    }
+
+    migraphx::argument
+    compute(migraphx::context&, const migraphx::shape&, std::vector<migraphx::argument> args) const
+    {
+        visit_all(args[0], args[1])([&](auto input, auto output) {
+            std::copy(input.begin(), input.end(), output.begin());
+        });
+
+        return args[1];
+    }
+
+    int output_alias(const std::vector<migraphx::shape>&) const { return 0; }
+};
+
+struct test_loop_op
+{
+    int64_t max_iterations = 10;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return migraphx::pack(f(self.max_iterations, "max_iterations"));
+    }
+
+    std::string name() const { return "test_loop_op"; }
+
+    migraphx::shape compute_shape(const std::vector<migraphx::shape>& inputs,
+                                  std::vector<migraphx::module_ref> mods) const
+    {
+        migraphx::check_shapes{inputs, *this}.standard();
+        if(mods.size() != 1)
+        {
+            MIGRAPHX_THROW("LOOP: operator should have one submodule.");
+        }
+
+        const auto& mod     = mods.front();
+        auto mod_out_shapes = mod->get_output_shapes();
+        auto dep_param_num  = inputs.size() - 2;
+
+        // first item of the mod output shapes is condition used in loop,
+        // which is not needed to compute output shape
+        mod_out_shapes.erase(mod_out_shapes.begin());
+        std::vector<migraphx::shape> ins_out_shapes(mod_out_shapes.begin(),
+                                                    mod_out_shapes.begin() + dep_param_num);
+        mod_out_shapes.erase(mod_out_shapes.begin(), mod_out_shapes.begin() + dep_param_num);
+        for(const auto& out_s : mod_out_shapes)
+        {
+            auto lens = out_s.lens();
+            lens.insert(lens.begin(), max_iterations);
+            ins_out_shapes.push_back({out_s.type(), lens});
+        }
+
+        return migraphx::shape(ins_out_shapes);
+    }
+
+    struct test_loop : public migraphx::op::loop::ref_loop
+    {
+        test_loop(int64_t iter_num) { max_iterations = iter_num; }
+
+        std::unordered_map<std::string, int> get_output_params(const migraphx::module& m) const
+        {
+            auto get_output_index = [](const std::string& name) {
+                std::string out_prefix = "#output_";
+                auto loc               = name.find(out_prefix);
+                if(loc != std::string::npos)
+                {
+                    int index = std::stoi(name.substr(loc + out_prefix.size()));
+                    return index;
+                }
+
+                return -1;
+            };
+
+            const auto& param_names = m.get_parameter_names();
+            std::unordered_map<std::string, int> result;
+            for(const auto& name : param_names)
+            {
+                auto index = get_output_index(name);
+                if(index == -1)
+                    continue;
+                result[name] = index;
+            }
+
+            return result;
+        }
+    };
+
+    migraphx::argument
+    compute(migraphx::context& ctx,
+            const migraphx::shape& out_shape,
+            const std::vector<migraphx::argument>& args,
+            const std::vector<migraphx::module_ref>& mods,
+            const std::function<std::vector<migraphx::argument>(
+                migraphx::module_ref&, const std::unordered_map<std::string, migraphx::argument>&)>&
+                run) const
+    {
+        // wrap up the arguments vector, so ref and gpu impl are the same
+        auto cpy_args = args;
+        bool in_cond  = args.at(1).at<bool>();
+        bool cond     = in_cond;
+        int64_t iter  = 0;
+        // insert iter and cond used in the loop
+        auto s_cond = args.at(1).get_shape();
+        auto s_iter = args.at(0).get_shape();
+        cpy_args.push_back({s_iter, &iter});
+        cpy_args.push_back({s_cond, &cond});
+        cpy_args.insert(cpy_args.end(), args.begin() + 2, args.end());
+        // add cond and mod outputs to the argument list
+        cpy_args.push_back(migraphx::argument(s_cond));
+        cpy_args.push_back(migraphx::argument(out_shape));
+        // run loop
+        return run_loop(test_loop{max_iterations}, ctx, cpy_args, mods, run);
+    }
+};
+
+static auto create_program(int64_t max_loop_iterations = 10)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape si{migraphx::shape::int64_type};
+    migraphx::shape s{migraphx::shape::int64_type, {1}};
+    migraphx::shape sc{migraphx::shape::bool_type};
+
+    auto in_iter = mm->add_parameter("iter_num", si);
+    auto in_cond = mm->add_parameter("ccond", sc);
+    auto in_val  = mm->add_parameter("val", s);
+
+    auto* body = p.create_module("loop_module");
+    auto iter  = body->add_parameter("#loop_module_in_0", si);
+    body->add_parameter("#loop_module_in_1", sc);
+    auto in_v               = body->add_parameter("#loop_module_in_2", s);
+    std::vector<int64_t> vd = {3};
+    auto l                  = body->add_literal(migraphx::literal(si, vd));
+    auto ad                 = body->add_instruction(migraphx::make_op("add"), iter, l);
+    auto val                = body->add_instruction(migraphx::make_op("add"), in_v, ad);
+    auto eq                 = body->add_instruction(migraphx::make_op("equal"), iter, l);
+    auto beq                = body->add_instruction(
+        migraphx::make_op("convert", {{"target_type", migraphx::shape::bool_type}}), eq);
+    auto neq                     = body->add_instruction(migraphx::make_op("not"), beq);
+    std::string out_param_prefix = "loop_module:#output_";
+    auto out0  = body->add_parameter(out_param_prefix + std::to_string(0), neq->get_shape());
+    auto r_neq = body->add_instruction(copy_op{}, neq, out0);
+    auto out2  = body->add_parameter(out_param_prefix + std::to_string(2), val->get_shape());
+    auto r_val = body->add_instruction(copy_op{}, val, out2);
+    body->add_return({r_neq, r_val, r_val});
+
+    auto rl =
+        mm->add_instruction(test_loop_op{max_loop_iterations}, {in_iter, in_cond, in_val}, {body});
+    auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), rl);
+    auto r1 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), rl);
+    mm->add_return({r0, r1});
+
+    return p;
+};
+
+static auto run_prog(migraphx::program p, int64_t iter_num, bool cond, int64_t ini_val)
+{
+    migraphx::shape si{migraphx::shape::int64_type};
+    migraphx::shape s{migraphx::shape::int64_type, {1}};
+    migraphx::shape sc{migraphx::shape::bool_type};
+
+    p.compile(migraphx::ref::target{});
+    migraphx::parameter_map pp;
+    pp["iter_num"] = migraphx::argument(si, &iter_num);
+    pp["ccond"]    = migraphx::argument(sc, &cond);
+    pp["val"]      = migraphx::argument(s, &ini_val);
+    auto rets      = p.eval(pp);
+
+    std::vector<std::vector<int64_t>> res;
+    for(auto& arg : rets)
+    {
+        std::vector<int64_t> vec;
+        arg.visit([&](auto v) { vec.assign(v.begin(), v.end()); });
+        res.push_back(vec);
+    }
+
+    return res;
+}
+
+TEST_CASE(loop_test1)
+{
+    auto p                         = create_program();
+    auto ress                      = run_prog(p, 10, true, 1);
+    std::vector<int64_t> gold_last = {19};
+    EXPECT(ress.front() == gold_last);
+    std::vector<int64_t> gold_concat = {4, 8, 13, 19, 0, 0, 0, 0, 0, 0};
+    EXPECT(ress.back() == gold_concat);
+}
+
+TEST_CASE(loop_test2)
+{
+    auto p                         = create_program(12);
+    auto ress                      = run_prog(p, 4, true, 1);
+    std::vector<int64_t> gold_last = {19};
+    EXPECT(ress.front() == gold_last);
+    std::vector<int64_t> gold_concat = {4, 8, 13, 19, 0, 0, 0, 0, 0, 0, 0, 0};
+    EXPECT(ress.back() == gold_concat);
+}
+
+TEST_CASE(loop_test3)
+{
+    auto p                         = create_program(3);
+    auto ress                      = run_prog(p, 3, true, 1);
+    std::vector<int64_t> gold_last = {13};
+    EXPECT(ress.front() == gold_last);
+    std::vector<int64_t> gold_concat = {4, 8, 13};
+    EXPECT(ress.back() == gold_concat);
+}
+
+TEST_CASE(loop_test4)
+{
+    auto p                         = create_program(20);
+    auto ress                      = run_prog(p, 5, true, 2);
+    std::vector<int64_t> gold_last = {20};
+    EXPECT(ress.front() == gold_last);
+    std::vector<int64_t> gold_concat = {5, 9, 14, 20, 0, 0, 0, 0, 0, 0,
+                                        0, 0, 0,  0,  0, 0, 0, 0, 0, 0};
+    EXPECT(ress.back() == gold_concat);
+}
+
+int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/verify/test_loop.cpp

+
+#include "verify_program.hpp"
+#include <migraphx/literal.hpp>
+#include <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_loop : verify_program<test_loop>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        migraphx::shape si{migraphx::shape::int64_type};
+        migraphx::shape s{migraphx::shape::int64_type, {1}};
+        migraphx::shape sc{migraphx::shape::bool_type};
+        int64_t iter_num = 10;
+        auto in_iter     = mm->add_literal(migraphx::literal(si, {iter_num}));
+        auto in_cond     = mm->add_parameter("ccond", sc);
+        int64_t value    = 5;
+        auto in_val      = mm->add_literal(migraphx::literal(s, {value}));
+
+        auto* body = p.create_module("loop_module");
+        auto iter  = body->add_parameter("iter_num", si);
+        body->add_parameter("cond", sc);
+        auto in_v               = body->add_parameter("input", s);
+        std::vector<int64_t> vd = {3};
+        auto l                  = body->add_literal(migraphx::literal(si, vd));
+        auto ad                 = body->add_instruction(migraphx::make_op("add"), iter, l);
+        auto val                = body->add_instruction(migraphx::make_op("add"), in_v, ad);
+        auto eq                 = body->add_instruction(migraphx::make_op("equal"), iter, l);
+        auto beq                = body->add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::bool_type}}), eq);
+        auto neq = body->add_instruction(migraphx::make_op("not"), beq);
+        body->add_return({neq, val, val});
+
+        auto rl = mm->add_instruction(
+            migraphx::make_op("loop", {{"max_iterations", 8}}), {in_iter, in_cond, in_val}, {body});
+        auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), rl);
+        auto r1 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), rl);
+        mm->add_return({r0, r1});
+
+        return p;
+    }
+};

tools/api/api.cpp

@@ -93,6 +93,11 @@ void set_default_dim_value(onnx_options& options, size_t value)
     options.default_dim_value = value;
 }
 
+void set_default_loop_iterations(onnx_options& options, int64_t value)
+{
+    options.max_loop_iterations = value;
+}
+
 void set_nhwc(tf_options& options, bool is_nhwc) { options.is_nhwc = is_nhwc; }
 
 void set_default_dim_value(tf_options& options, size_t value) { options.batch_size = value; }

tools/include/operation.hpp

@@ -256,6 +256,18 @@ argument compute_op(const T& x,
     return compute_op(rank<1>{}, x, output, inputs, module_args, f);
 }
 
+template <class T, class F>
+auto compute_op(rank<4>,
+                const T& x,
+                context& ctx,
+                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))
+{
+    return x.compute(auto_any_cast(ctx), output, inputs, module_args, f);
+}
+
 template <class T, class F>
 auto compute_op(rank<3>,
                 const T& x,
@@ -313,7 +325,7 @@ argument compute_op(const T& x,
                     const std::vector<module_ref>& module_args,
                     F f)
 {
-    return compute_op(rank<3>{}, x, ctx, output, inputs, module_args, f);
+    return compute_op(rank<4>{}, x, ctx, output, inputs, module_args, f);
 }
 
 template <class T>