GridSample Linear Sampler Refactor (#4067)

This PR refactors the linear sampler used in the GridSample operator to use less literals. Also adds support and tests for when channels > 1.

Note: This PR does not change the nearest and bicubic samplers. Although they are currently not used in any models we support, they should likely be refactored as well sometime in the future.

Author: eddieliao

src/onnx/parse_gridsample.cpp

@@ -165,26 +165,26 @@ struct grid_sampler
     static instruction_ref concat_on_first_dim(const onnx_parser::node_info& info,
                                                std::vector<instruction_ref> instructions)
     {
-        return std::accumulate(
-            std::next(instructions.begin()),
-            instructions.end(),
-            instructions.front(),
-            [&info](auto& ret, auto& ins) {
-                return info.add_instruction(make_op("concat", {{"axis", 0}}), ret, ins);
-            });
+        return std::accumulate(std::next(instructions.begin()),
+                               instructions.end(),
+                               instructions.front(),
+                               [&info](auto& ret, auto& ins) {
+                                   return info.add_instruction(
+                                       make_op("concat", {{"axis", 0}}), ret, ins);
+                               });
     }
 
     static instruction_ref concat_on_dim(const onnx_parser::node_info& info,
                                          std::array<instruction_ref, 4> instructions,
                                          int64_t dim)
     {
-        return std::accumulate(
-            std::next(instructions.begin()),
-            instructions.end(),
-            instructions.front(),
-            [&info, &dim](auto& ret, auto& ins) {
-                return info.add_instruction(make_op("concat", {{"axis", dim}}), ret, ins);
-            });
+        return std::accumulate(std::next(instructions.begin()),
+                               instructions.end(),
+                               instructions.front(),
+                               [&info, &dim](auto& ret, auto& ins) {
+                                   return info.add_instruction(
+                                       make_op("concat", {{"axis", dim}}), ret, ins);
+                               });
     }
 
     bool has_border_padding() const { return m_padding == "border"; }
@@ -296,26 +296,31 @@ struct linear_sampler : grid_sampler
 
     instruction_ref sample(const onnx_parser::node_info& info)
     {
-        std::vector<instruction_ref> weight_indices;
-        std::vector<instruction_ref> xy_indices;
-        std::vector<instruction_ref> nc_values;
-
-        const static auto nhw_shape = migraphx::shape{migraphx::shape::int64_type, {1, 3}};
-        dfor(m_batch, m_out_height, m_out_width)([&](auto n, auto h, auto w) {
-            auto nhw = info.add_literal(migraphx::literal{nhw_shape, {n, h, w}});
-            weight_indices.push_back(nhw);
-            for(size_t c = 0; c < m_channel; c++)
-            {
-                xy_indices.push_back(nhw);
-                nc_values.push_back(info.add_literal(migraphx::literal{m_nc_shape, {n, c}}));
-            }
+        std::vector<float> xy_indices_data;
+        std::vector<float> weight_indices_data;
+        std::vector<float> nc_values_data;
+        dfor(m_batch, m_out_height, m_out_width, m_channel)([&](auto n, auto h, auto w, auto c) {
+            xy_indices_data.push_back(n);
+            xy_indices_data.push_back(h);
+            xy_indices_data.push_back(w);
+            weight_indices_data.push_back(n);
+            weight_indices_data.push_back(h);
+            weight_indices_data.push_back(w);
+            nc_values_data.push_back(n);
+            nc_values_data.push_back(c);
         });
-
-        auto xy_indices_t = concat_on_first_dim(info, xy_indices);
-        auto y0_samples   = info.add_instruction(make_op("gathernd"), m_floor_y, xy_indices_t);
-        auto x0_samples   = info.add_instruction(make_op("gathernd"), m_floor_x, xy_indices_t);
-        auto y1_samples   = info.add_instruction(make_op("gathernd"), m_ceil_y, xy_indices_t);
-        auto x1_samples   = info.add_instruction(make_op("gathernd"), m_ceil_x, xy_indices_t);
+        size_t num_indices  = m_batch * m_out_height * m_out_width * m_channel;
+        auto xy_indices_t   = info.add_literal(migraphx::literal{
+            migraphx::shape{migraphx::shape::float_type, {num_indices, 3}}, xy_indices_data});
+        auto weight_index_t = info.add_literal(migraphx::literal{
+            migraphx::shape{migraphx::shape::float_type, {num_indices, 3}}, weight_indices_data});
+        auto nc             = info.add_literal(migraphx::literal{
+            migraphx::shape{migraphx::shape::float_type, {num_indices, 2}}, nc_values_data});
+
+        auto y0_samples = info.add_instruction(make_op("gathernd"), m_floor_y, xy_indices_t);
+        auto x0_samples = info.add_instruction(make_op("gathernd"), m_floor_x, xy_indices_t);
+        auto y1_samples = info.add_instruction(make_op("gathernd"), m_ceil_y, xy_indices_t);
+        auto x1_samples = info.add_instruction(make_op("gathernd"), m_ceil_x, xy_indices_t);
 
         auto validate_samples = [&](auto& samples, auto& max) {
             auto clip       = info.add_common_op("clip", samples, m_zero_l, max);
@@ -338,8 +343,6 @@ struct linear_sampler : grid_sampler
         x1_samples = info.add_instruction(
             make_op("reshape", {{"dims", {x1_samples->get_shape().elements(), 1}}}), x1_samples);
 
-        auto nc = concat_on_first_dim(info, nc_values);
-
         auto make_corner_indices = [&](auto& x, auto& y) {
             auto hw = info.add_instruction(make_op("concat", {{"axis", 1}}), y, x);
             return info.add_instruction(make_op("concat", {{"axis", 1}}), nc, hw);
@@ -356,10 +359,6 @@ struct linear_sampler : grid_sampler
             info.add_common_op("logical_and", x1_validation, y1_validation)};
 
         std::array<instruction_ref, 4> corner_samples;
-        auto weight_index_t = concat_on_first_dim(info, weight_indices);
-        weight_index_t      = info.add_instruction(
-            make_op("reshape", {{"dims", {weight_indices.size(), 3}}}), weight_index_t);
-
         std::transform(corner_indices.begin(),
                        corner_indices.end(),
                        corner_validations.begin(),
@@ -528,13 +527,13 @@ struct bicubic_sampler : grid_sampler
                 make_op("reshape", {{"dims", {corner_weight->get_shape().elements(), 1}}}),
                 corner_weight);
         });
-        auto weights_t = std::accumulate(
-            std::next(corner_weights.begin()),
-            corner_weights.end(),
-            corner_weights.front(),
-            [&info](auto& acc, auto& ins) {
-                return info.add_instruction(make_op("concat", {{"axis", 1}}), acc, ins);
-            });
+        auto weights_t = std::accumulate(std::next(corner_weights.begin()),
+                                         corner_weights.end(),
+                                         corner_weights.front(),
+                                         [&info](auto& acc, auto& ins) {
+                                             return info.add_instruction(
+                                                 make_op("concat", {{"axis", 1}}), acc, ins);
+                                         });
         return info.add_instruction(make_op("reshape", {{"dims", out_lens}}), weights_t);
     }
 
@@ -664,7 +663,7 @@ struct parse_gridsample : op_parser<parse_gridsample>
                           std::vector<instruction_ref> args) const
     {
         bool align_corners = false;
-        // Note: defult mode can be linear or bilinear depending on the onnx version
+        // Note: default mode can be linear or bilinear depending on the onnx version
         std::string mode         = "linear";
         std::string padding_mode = "zeros";
 

test/onnx/gen_onnx.py

@@ -4462,6 +4462,41 @@ def gridsample_test():
 
     return ([node], [x, grid], [y])
 
+@onnx_test()
+def gridsample_channel_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [1, 3, 4, 4])
+    grid = helper.make_tensor_value_info('grid', TensorProto.FLOAT,
+                                         [1, 6, 6, 2])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 3, 6, 6])
+
+    node = onnx.helper.make_node(
+        "GridSample",
+        inputs=["x", "grid"],
+        outputs=["y"],
+        mode="bilinear",
+        padding_mode="border",
+        align_corners=1,
+    )
+
+    return ([node], [x, grid], [y])
+
+@onnx_test()
+def gridsample_512x512_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [1, 3, 512, 512])
+    grid = helper.make_tensor_value_info('grid', TensorProto.FLOAT,
+                                         [1, 512, 512, 2])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 3, 512, 512])
+
+    node = onnx.helper.make_node(
+        "GridSample",
+        inputs=["x", "grid"],
+        outputs=["y"],
+        mode="bilinear",
+        padding_mode="border",
+        align_corners=1,
+    )
+
+    return ([node], [x, grid], [y])
 
 @onnx_test()
 def gridsample_half_test():

test/onnx/gridsample_512x512_test.onnx

@@ -0,0 +1,26 @@
+	gridsample_512x512_test:�

+^
+
x
+grid
y"
+GridSample*
+align_corners
�
*
+mode"bilinear�
*
+padding_mode"border�
gridsample_512x512_testZ
+
x
+

+

+
+�
+�Z 
+grid
+

+

+�
+�
+b
+
y
+

+

+
+�
+�B

test/onnx/gridsample_channel_test.onnx

@@ -0,0 +1,26 @@
+	gridsample_channel_test:�

+^
+
x
+grid
y"
+GridSample*
+align_corners
�
*
+mode"bilinear�
*
+padding_mode"border�
gridsample_channel_testZ
+
x
+

+

+
+
+Z
+grid
+

+

+
+
+b
+
y
+

+

+
+
+B