Modify reshapes

src/CMakeLists.txt

@@ -192,6 +192,7 @@ register_migraphx_ops(
     reduce_sum
     relu
     reshape
+    reshape_lazy
     reverse
     rnn
     rnn_last_cell_output

src/auto_contiguous.cpp

@@ -25,7 +25,6 @@
 #include <migraphx/program.hpp>
 #include <migraphx/instruction.hpp>
 #include <migraphx/make_op.hpp>
-
 #include <migraphx/iterator_for.hpp>
 
 namespace migraphx {

src/include/migraphx/instruction.hpp

@@ -81,6 +81,7 @@ struct MIGRAPHX_EXPORT instruction
 
     const std::vector<module_ref>& module_inputs() const;
 
+    /// Where this instruction is used as an input to another instruction
     const std::vector<instruction_ref>& outputs() const;
 
     friend bool operator==(const instruction& x, const instruction& y);

src/include/migraphx/op/reshape.hpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2023 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -29,7 +29,8 @@
 #include <migraphx/config.hpp>
 #include <migraphx/value.hpp>
 #include <migraphx/dyn_output.hpp>
-#include <migraphx/optional.hpp>
+
+#include <algorithm>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -45,8 +46,6 @@ struct reshape
         return pack(f(self.dims, "dims"));
     }
 
-    value attributes() const { return {{"require_std_shape", true}}; }
-
     std::string name() const { return "reshape"; }
 
     shape dyn_compute_shape(shape s0) const
@@ -97,112 +96,6 @@ struct reshape
         return {s0.type(), output_dyn_dims};
     }
 
-    template <class Iterator>
-    static auto compute_end_dim(Iterator start, Iterator last, std::size_t dim)
-    {
-        std::size_t x = 1;
-        auto it       = std::find_if(start, last, [&](auto i) {
-            x *= i;
-            return x >= dim;
-        });
-        if(x != dim)
-            return start;
-        return it;
-    }
-
-    template <class DimIterator, class StrideIterator>
-    static auto can_strides_merge(DimIterator dim_start,
-                                  DimIterator dim_last,
-                                  StrideIterator stride_start,
-                                  StrideIterator stride_last)
-    {
-        assert(std::distance(dim_start, dim_last) == std::distance(stride_start, stride_last));
-        auto cstride = *std::prev(stride_last);
-        return std::equal(std::make_reverse_iterator(dim_last),
-                          std::make_reverse_iterator(dim_start + 1),
-                          std::make_reverse_iterator(stride_last - 1),
-                          std::make_reverse_iterator(stride_start),
-                          [&](auto dim, auto stride) {
-                              cstride *= dim;
-                              return stride == cstride;
-                          });
-    }
-
-    // This will reshape the dimesions of the input shape to use the lens of
-    // `rdims`. If this can't be done without changing memory layout then it
-    // will return nullopt
-    static optional<shape> reshape_dims(const shape& input, const std::vector<std::size_t>& rdims)
-    {
-        if(input.standard())
-            return shape{input.type(), rdims};
-
-        const auto& idims    = input.lens();
-        const auto& istrides = input.strides();
-
-        std::vector<std::size_t> rstrides;
-        std::size_t i = 0;
-        std::size_t r = 0;
-        while(i < idims.size() and r < rdims.size())
-        {
-            auto idim = idims[i];
-            auto rdim = rdims[r];
-            if(rdim == idim)
-            {
-                rstrides.push_back(istrides[i]);
-            }
-            // squeeze
-            else if(rdim > idim)
-            {
-                auto start = idims.begin() + i;
-                auto it    = compute_end_dim(start, idims.end(), rdim);
-                if(it == start)
-                    return nullopt;
-                auto n = it - start;
-                assert((i + n) <= istrides.size());
-                if(not can_strides_merge(
-                       start, it + 1, istrides.begin() + i, istrides.begin() + i + n + 1))
-                    return nullopt;
-                i += n;
-                rstrides.push_back(istrides[i]);
-            }
-            // unsqueeze
-            else // if(rdim < idim)
-            {
-                auto start = rdims.begin() + i;
-                auto it    = compute_end_dim(start, rdims.end(), idim);
-                if(it == start)
-                    return nullopt;
-                auto n = it - start;
-                assert((r + n) <= rdims.size());
-                auto stride = istrides[i] * idim;
-                std::for_each(start, it + 1, [&](auto dim) {
-                    stride /= dim;
-                    rstrides.push_back(stride);
-                });
-                r += n;
-            }
-            i++;
-            r++;
-        }
-
-        // Handle trailing 1s
-        if(rstrides.size() < rdims.size() and not rstrides.empty())
-        {
-            auto stride = rstrides.back();
-            for(auto d : range(rdims.begin() + rstrides.size(), rdims.end()))
-            {
-                if(d != 1)
-                    return nullopt;
-                rstrides.push_back(stride);
-            }
-        }
-
-        if(rdims.size() != rstrides.size())
-            return nullopt;
-
-        return shape{input.type(), rdims, rstrides};
-    }
-
     shape static_compute_shape(std::vector<shape> inputs, std::size_t n_neg_dims) const
     {
         check_shapes{inputs, *this}.has(1);
@@ -232,26 +125,26 @@ struct reshape
             }
         }
 
-        auto s = reshape_dims(inputs.front(), rdims);
-        if(not s.has_value())
-            MIGRAPHX_THROW("Reshape on axis that is not packed.");
+        shape s{inputs.front().type(), rdims};
 
-        if(s->elements() != inputs.front().elements())
-            MIGRAPHX_THROW("Reshape: Wrong number of elements for reshape: reshape has " +
-                           std::to_string(s->elements()) + " elements whereas the input has " +
+        if(s.elements() != inputs.front().elements())
+            MIGRAPHX_THROW("reshape: Wrong number of elements for reshape: reshape has " +
+                           std::to_string(s.elements()) + " elements whereas the input has " +
                            std::to_string(inputs.front().elements()));
 
-        assert(s->bytes() == inputs.front().bytes());
-        return *s;
+        assert(s.bytes() == inputs.front().bytes());
+        return s;
     }
 
     shape compute_shape(std::vector<shape> inputs) const
     {
         check_shapes{inputs, *this, true}.has(1);
+
         auto n_neg_dims = std::count(dims.begin(), dims.end(), -1);
         if(n_neg_dims > 1)
-            MIGRAPHX_THROW("Reshape: Dimensions for reshape can only have one -1 dim");
-        auto s0 = inputs[0];
+            MIGRAPHX_THROW("reshape: Dimensions for reshape can only have one -1 dim");
+
+        auto s0 = inputs.front();
         if(s0.dynamic())
         {
             return dyn_compute_shape(s0);
@@ -264,10 +157,14 @@ struct reshape
 
     argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
     {
-        return args[0].reshape(dyn_out.computed_shape);
-    }
+        assert(dyn_out.computed_shape.standard());
+        argument result{dyn_out.computed_shape};
 
-    std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
+        visit_all(result, args[0])([&](auto output, auto input) {
+            std::copy(input.begin(), input.end(), output.begin());
+        });
+        return result;
+    }
 };
 
 } // namespace op