refine the implementation

src/auto_contiguous.cpp

@@ -46,6 +46,24 @@ void auto_contiguous::apply(module& p) const
             p.replace_instruction(ins, c);
         }
     }
+
+    // if ops used as output param are alias 0, add a contiguous for the output
+    // so return outputs with standard shape
+    if (last->name() == "@return")
+    {
+        auto inputs = last->inputs();
+        for (auto ins : inputs)
+        {
+            if (ins->name() == "contiguous") continue;
+
+            auto ins_alias = ins->get_operator().output_alias({});
+            if (ins_alias == 0)
+            {
+                auto cont_ins = p.insert_instruction(last, make_op("contiguous"), ins);
+                p.replace_instruction(ins, cont_ins);
+            }
+        }
+    }
 }
 
 } // namespace MIGRAPHX_INLINE_NS

src/eliminate_contiguous.cpp

@@ -39,12 +39,6 @@ static bool try_compute_shape(instruction_ref ins,
             return false;
         }
 
-        if(std::any_of(
-               outputs.begin(), outputs.end(), [](auto o) { return o->name() == "@return"; }))
-        {
-            return false;
-        }
-
         for(auto output : outputs)
         {
             auto args = output->inputs();

src/simplify_reshapes.cpp

@@ -120,6 +120,16 @@ struct find_nop_reshapes
     void apply(module& p, const match::matcher_result& mr) const
     {
         auto ins = mr.result;
+        if (ins->name() == "contiguous")
+        {
+            auto& outputs = ins->outputs();
+            if(std::any_of(outputs.begin(), outputs.end(), [&](auto o) {
+                return o->name() == "@return";
+            }))
+            {
+                return;
+            }
+        }
         p.replace_instruction(ins, ins->inputs().front());
     }
 };