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

Merge branch 'develop' of github.com:ROCmSoftwarePlatform/AMDMIGraphX into split_multiple_same_dyn_dim

test/py/onnx_backend_test.py

@@ -190,7 +190,6 @@ def disabled_tests_onnx_1_7_0(backend_test):
     backend_test.exclude(
         r'test_negative_log_likelihood_loss_input_shape_is_NCd1d2d3d4d5_none_no_weight_cpu'
     )
-    backend_test.exclude(r'test_qlinearconv_cpu')
     backend_test.exclude(r'test_qlinearmatmul_2D_cpu')
     backend_test.exclude(r'test_qlinearmatmul_3D_cpu')
     backend_test.exclude(r'test_range_float_type_positive_delta_expanded_cpu')
@@ -576,6 +575,8 @@ def disabled_tests_onnx_1_9_0(backend_test):
     backend_test.exclude(r'test_gru_batchwise_cpu')
     backend_test.exclude(r'test_simple_rnn_batchwise_cpu')
     # from OnnxBackendPyTorchConvertedModelTest
+    # MaxPool dialtion is partially supported on GPU by a workaround
+    # But these tests require too large allocations to work properly
     backend_test.exclude(r'test_MaxPool1d_stride_padding_dilation_cpu')
     backend_test.exclude(r'test_MaxPool2d_stride_padding_dilation_cpu')
 

test/ref/pooling.cpp

@@ -35,12 +35,13 @@ TEST_CASE(avgpool_rank3_test)
 {
     // 1D case 1, input is 3D
     migraphx::program p;
-    auto* mm   = p.get_main_module();
-    auto s     = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
-    auto op    = migraphx::op::pooling{migraphx::op::pooling_mode::average};
-    op.lengths = {2};
-    op.padding = {0};
-    op.stride  = {1};
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::average};
+    op.lengths   = {2};
+    op.padding   = {0};
+    op.stride    = {1};
+    op.dilations = {1};
 
     std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
     auto l0 = mm->add_literal(migraphx::literal{s, data});
@@ -54,6 +55,103 @@ TEST_CASE(avgpool_rank3_test)
     EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
 }
 
+TEST_CASE(avgpool_rank3_dil_test)
+{
+    // 1D case 1, input is 3D
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::average};
+    op.lengths   = {2};
+    op.padding   = {0};
+    op.stride    = {1};
+    op.dilations = {2};
+
+    std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{0.35, 0.15, 0.85, 0.3, 0.1, 0.65};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
+TEST_CASE(avgpool_rank3_dil_test2)
+{
+    // 1D case 1, input is 3D
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::average};
+    op.lengths   = {2};
+    op.padding   = {0};
+    op.stride    = {1};
+    op.dilations = {3};
+
+    std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{0.2, 0.45, 0.35};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
+TEST_CASE(avgpool_rank3_pad_test)
+{
+    // 1D case 1, input is 3D
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::average};
+    op.lengths   = {2};
+    op.padding   = {1};
+    op.stride    = {1};
+    op.dilations = {1};
+
+    std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{
+        0.3, 0.25, 0.3, 0.25, 0.1, 0.8, 0.65, 0.7, 0.5, 0.1, 0.1, 0.4, 0.4, 0.35, 0.6};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
+TEST_CASE(avgpool_rank3_pad_dil_test)
+{
+    // 1D case 1, input is 3D
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::average};
+    op.lengths   = {2};
+    op.padding   = {1};
+    op.stride    = {1};
+    op.dilations = {3};
+
+    std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{0.4, 0.2, 0.2, 0.9, 0.45, 0.5, 0.1, 0.35, 0.7};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
 TEST_CASE(avgpool_dyn_test)
 {
     // Dynamic input, no padding
@@ -65,7 +163,8 @@ TEST_CASE(avgpool_dyn_test)
                                           {{"mode", migraphx::op::pooling_mode::average},
                                            {"lengths", {2}},
                                            {"padding", {0}},
-                                           {"stride", {1}}}),
+                                           {"stride", {1}},
+                                           {"dilations", {1}}}),
                         x);
     p.compile(migraphx::make_target("ref"));
 
@@ -82,7 +181,7 @@ TEST_CASE(avgpool_dyn_test)
 
 TEST_CASE(avgpool_dyn_pad_test)
 {
-    // Dynamic input with explicit padding/
+    // Dynamic input with explicit padding
     migraphx::program p;
     auto* mm = p.get_main_module();
     auto s   = migraphx::shape{migraphx::shape::float_type, {{1, 3}, {3, 3}, {4, 4}}};
@@ -91,7 +190,8 @@ TEST_CASE(avgpool_dyn_pad_test)
                                           {{"mode", migraphx::op::pooling_mode::average},
                                            {"lengths", {2}},
                                            {"padding", {1}},
-                                           {"stride", {1}}}),
+                                           {"stride", {1}},
+                                           {"dilations", {1}}}),
                         x);
     p.compile(migraphx::make_target("ref"));
 
@@ -158,7 +258,8 @@ TEST_CASE(avgpool_dyn_auto_pad_1d_test)
                            //    padding added will be {1, 0} to make output
                            //    the same size as input
                            {"padding_mode", migraphx::op::padding_mode_t::same_lower},
-                           {"stride", {1}}}),
+                           {"stride", {1}},
+                           {"dilations", {1}}}),
         x);
     p.compile(migraphx::make_target("ref"));
 
@@ -171,8 +272,8 @@ TEST_CASE(avgpool_dyn_auto_pad_1d_test)
     result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
 
     // clang-format off
-    std::vector<float> gold{0.3, 0.25, 0.3, 0.25,
-                            0.8, 0.65, 0.7, 0.5,
+    std::vector<float> gold{0.3, 0.25, 0.3, 0.25, 
+                            0.8, 0.65, 0.7, 0.5, 
                             0.1, 0.4,  0.4, 0.35};
     // clang-format on
     EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
@@ -190,7 +291,8 @@ TEST_CASE(avgpool_dyn_pad_ceil_test)
                                            {"lengths", {2, 3}},
                                            {"padding", {1, 2}},
                                            {"ceil_mode", true},
-                                           {"stride", {1, 1}}}),
+                                           {"stride", {1, 1}},
+                                           {"dilations", {1, 1}}}),
                         x);
     p.compile(migraphx::make_target("ref"));
 
@@ -219,12 +321,13 @@ TEST_CASE(avgpool_rank3_stride2_test)
 {
     // 1D case 2, stride 2
     migraphx::program p;
-    auto* mm   = p.get_main_module();
-    auto s     = migraphx::shape{migraphx::shape::float_type, {2, 2, 4}};
-    auto op    = migraphx::op::pooling{migraphx::op::pooling_mode::average};
-    op.lengths = {2};
-    op.padding = {1};
-    op.stride  = {2};
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {2, 2, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::average};
+    op.lengths   = {2};
+    op.padding   = {1};
+    op.stride    = {2};
+    op.dilations = {1};
 
     // clang-format off
     std::vector<float> data{1.6321, -2.4186, 0.2239, -1.4232, 
@@ -252,12 +355,13 @@ TEST_CASE(avgpool_rank5_test)
 {
     // 3D, input is 5D
     migraphx::program p;
-    auto* mm   = p.get_main_module();
-    auto s     = migraphx::shape{migraphx::shape::float_type, {2, 2, 3, 3, 3}};
-    auto op    = migraphx::op::pooling{migraphx::op::pooling_mode::average};
-    op.lengths = {2, 2, 2};
-    op.padding = {0, 0, 0};
-    op.stride  = {1, 1, 1};
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {2, 2, 3, 3, 3}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::average};
+    op.lengths   = {2, 2, 2};
+    op.padding   = {0, 0, 0};
+    op.stride    = {1, 1, 1};
+    op.dilations = {1, 1, 1};
 
     std::vector<float> data{
         -0.179, -1.756, 0.651,  1.955,  1.87,   -0.604, 0.247,  0.449,  -0.137, 1.187,  1.593,
@@ -423,13 +527,14 @@ TEST_CASE(lppool_l1_norm_test)
 {
     // L1 norm test
     migraphx::program p;
-    auto* mm    = p.get_main_module();
-    auto s      = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
-    auto op     = migraphx::op::pooling{migraphx::op::pooling_mode::lpnorm};
-    op.lengths  = {2};
-    op.padding  = {0};
-    op.stride   = {1};
-    op.lp_order = 1;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::lpnorm};
+    op.lengths   = {2};
+    op.padding   = {0};
+    op.stride    = {1};
+    op.dilations = {1};
+    op.lp_order  = 1;
 
     std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
     auto l0 = mm->add_literal(migraphx::literal{s, data});
@@ -449,13 +554,14 @@ TEST_CASE(lppool_l1_norm_test)
 // {
 //     // padding too large for kernel size
 //     migraphx::program p;
-//     auto* mm    = p.get_main_module();
-//     auto s      = migraphx::shape{migraphx::shape::float_type, {1, 2, 5}};
-//     auto op     = migraphx::op::pooling{migraphx::op::pooling_mode::lpnorm};
-//     op.lengths  = {3};
-//     op.padding  = {2};
-//     op.stride   = {1};
-//     op.lp_order = 1;
+//     auto* mm     = p.get_main_module();
+//     auto s       = migraphx::shape{migraphx::shape::float_type, {1, 2, 5}};
+//     auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::lpnorm};
+//     op.lengths   = {3};
+//     op.padding   = {2};
+//     op.stride    = {1};
+//     op.dilations = {1};
+//     op.lp_order  = 1;
 
 //     std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7};
 //     auto l0 = mm->add_literal(migraphx::literal{s, data});
@@ -468,13 +574,14 @@ TEST_CASE(lppool_l2_norm_test)
 {
     // L2 norm test
     migraphx::program p;
-    auto* mm    = p.get_main_module();
-    auto s      = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
-    auto op     = migraphx::op::pooling{migraphx::op::pooling_mode::lpnorm};
-    op.lengths  = {2};
-    op.padding  = {0};
-    op.stride   = {1};
-    op.lp_order = 2;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::lpnorm};
+    op.lengths   = {2};
+    op.padding   = {0};
+    op.stride    = {1};
+    op.dilations = {1};
+    op.lp_order  = 2;
 
     std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
     auto l0 = mm->add_literal(migraphx::literal{s, data});
@@ -506,7 +613,8 @@ TEST_CASE(lppool_dyn_test)
                                           {{"mode", migraphx::op::pooling_mode::lpnorm},
                                            {"lengths", {2}},
                                            {"padding", {0}},
-                                           {"stride", {1}}}),
+                                           {"stride", {1}},
+                                           {"dilations", {1}}}),
                         x);
     p.compile(migraphx::make_target("ref"));
 
@@ -571,7 +679,8 @@ TEST_CASE(maxpool_test)
                                           {{"mode", migraphx::op::pooling_mode::max},
                                            {"padding", {0, 0}},
                                            {"stride", {2, 2}},
-                                           {"lengths", {3, 2}}}),
+                                           {"lengths", {3, 2}},
+                                           {"dilations", {1, 1}}}),
                         al);
     p.compile(migraphx::make_target("ref"));
     auto result = p.eval({}).back();
@@ -599,7 +708,8 @@ TEST_CASE(maxpool_pad_test)
                                           {{"mode", migraphx::op::pooling_mode::max},
                                            {"padding", {1, 1}},
                                            {"stride", {2, 2}},
-                                           {"lengths", {3, 2}}}),
+                                           {"lengths", {3, 2}},
+                                           {"dilations", {1, 1}}}),
                         al);
 
     //   * *  *  *                                           * *  *  *
@@ -620,12 +730,13 @@ TEST_CASE(maxpool_rank3_test0)
 {
     // 1D case 1, input is 3D
     migraphx::program p;
-    auto* mm   = p.get_main_module();
-    auto s     = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
-    auto op    = migraphx::op::pooling{migraphx::op::pooling_mode::max};
-    op.lengths = {2};
-    op.padding = {0};
-    op.stride  = {1};
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::max};
+    op.lengths   = {2};
+    op.padding   = {0};
+    op.stride    = {1};
+    op.dilations = {1};
 
     std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
     auto l0 = mm->add_literal(migraphx::literal{s, data});
@@ -643,12 +754,13 @@ TEST_CASE(maxpool_rank3_test1)
 {
     // 1D case 2, input is 3D
     migraphx::program p;
-    auto* mm   = p.get_main_module();
-    auto s     = migraphx::shape{migraphx::shape::float_type, {2, 2, 5}};
-    auto op    = migraphx::op::pooling{migraphx::op::pooling_mode::max};
-    op.lengths = {2};
-    op.padding = {0};
-    op.stride  = {2};
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {2, 2, 5}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::max};
+    op.lengths   = {2};
+    op.padding   = {0};
+    op.stride    = {2};
+    op.dilations = {1};
 
     std::vector<float> data{0.4975, -0.1226, -0.0405, -0.2861, -0.1227, -0.6186, -0.9618,
                             0.6022, -0.1912, 1.1925,  0.5493,  0.1692,  -0.8039, -1.0281,
@@ -664,6 +776,55 @@ TEST_CASE(maxpool_rank3_test1)
     EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
 }
 
+TEST_CASE(maxpool_rank3_test2)
+{
+    // 1D case 1, input is 3D
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::max};
+    op.lengths   = {2};
+    op.padding   = {0};
+    op.stride    = {1};
+    op.dilations = {2};
+
+    std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{0.4, 0.2, 0.9, 0.5, 0.1, 0.7};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
+TEST_CASE(maxpool_rank3_test4)
+{
+    // 1D case 1, input is 3D
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::max};
+    op.lengths   = {2};
+    op.padding   = {1};
+    op.stride    = {1};
+    op.dilations = {3};
+
+    std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{0.4, 0.3, 0.2, 0.9, 0.8, 0.5, 0.1, 0.6, 0.7};
+
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
 TEST_CASE(maxpool_rank3_ceil_test)
 {
     // 1D case 2, input is 3D, ceil mode
@@ -674,6 +835,7 @@ TEST_CASE(maxpool_rank3_ceil_test)
     op.lengths   = {2};
     op.padding   = {0};
     op.stride    = {2};
+    op.dilations = {1};
     op.ceil_mode = true;
 
     // clang-format off
@@ -702,12 +864,13 @@ TEST_CASE(maxpool_rank5_test)
 {
     // 3D, input is 5D
     migraphx::program p;
-    auto* mm   = p.get_main_module();
-    auto s     = migraphx::shape{migraphx::shape::float_type, {2, 2, 3, 3, 3}};
-    auto op    = migraphx::op::pooling{migraphx::op::pooling_mode::max};
-    op.lengths = {2, 2, 2};
-    op.padding = {0, 0, 0};
-    op.stride  = {2, 2, 2};
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {2, 2, 3, 3, 3}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::max};
+    op.lengths   = {2, 2, 2};
+    op.padding   = {0, 0, 0};
+    op.stride    = {2, 2, 2};
+    op.dilations = {1, 1, 1};
 
     std::vector<float> data{
         -2.8029, 0.5861,  0.7015,  0.1297,  -1.44,   -1.9472, 0.7812,  2.408,   -0.3145, 0.3405,
@@ -741,7 +904,8 @@ TEST_CASE(maxpool_dyn_test)
                                           {{"mode", migraphx::op::pooling_mode::max},
                                            {"lengths", {2}},
                                            {"padding", {0}},
-                                           {"stride", {1}}}),
+                                           {"stride", {1}},
+                                           {"dilations", {1}}}),
                         x);
     p.compile(migraphx::make_target("ref"));
 
@@ -755,3 +919,29 @@ TEST_CASE(maxpool_dyn_test)
     std::vector<float> gold{0.3, 0.4, 0.4, 0.8, 0.9, 0.9, 0.7, 0.7, 0.6};
     EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
 }
+
+TEST_CASE(maxpool_dyn_test2)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    auto s   = migraphx::shape{migraphx::shape::float_type, {{1, 4}, {3, 3}, {4, 4}}};
+    auto x   = mm->add_parameter("X", s);
+    mm->add_instruction(migraphx::make_op("pooling",
+                                          {{"mode", migraphx::op::pooling_mode::max},
+                                           {"lengths", {2}},
+                                           {"padding", {0}},
+                                           {"stride", {1}},
+                                           {"dilations", {2}}}),
+                        x);
+    p.compile(migraphx::make_target("ref"));
+
+    std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
+    migraphx::shape input_fixed_shape{migraphx::shape::float_type, {1, 3, 4}};
+    migraphx::parameter_map params;
+    params["X"] = migraphx::argument(input_fixed_shape, data.data());
+    auto result = p.eval(params).back();
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{0.4, 0.2, 0.9, 0.5, 0.1, 0.7};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}

test/rewrite_pooling_test.cpp

@@ -53,7 +53,8 @@ TEST_CASE(rewrite_pooling_test)
                                                          {{"mode", mode},
                                                           {"padding", {0, 0, 0}},
                                                           {"stride", {1, 1, 1}},
-                                                          {"lengths", {3, 4, 5}}}),
+                                                          {"lengths", {3, 4, 5}},
+                                                          {"dilations", {1, 1, 1}}}),
                                      input);
         m.add_return({ret});
         return m;
@@ -80,6 +81,483 @@ TEST_CASE(rewrite_pooling_test)
                  migraphx::make_op("reduce_max", {{"axes", {2, 3, 4}}}));
 }
 
+TEST_CASE(rewrite_pooling_dialtions_test)
+{
+    migraphx::shape s{migraphx::shape::float_type, {1, 1, 5, 5}};
+    auto pooling_program = [&](const migraphx::op::pooling_mode mode) {
+        migraphx::module m;
+        auto input = m.add_parameter("x", s);
+        auto ret   = m.add_instruction(migraphx::make_op("pooling",
+                                                         {{"mode", mode},
+                                                          {"padding", {0, 0}},
+                                                          {"stride", {1, 1}},
+                                                          {"lengths", {2, 2}},
+                                                          {"dilations", {2, 2}}}),
+                                     input);
+        m.add_return({ret});
+        return m;
+    };
+
+    auto opt_program = [&](const migraphx::op::pooling_mode mode) {
+        migraphx::module m;
+        auto input = m.add_parameter("x", s);
+        std::vector<int> indices{0, 2, 1, 3, 2, 4};
+        migraphx::shape s_indices{migraphx::shape::int32_type, {indices.size()}};
+        auto i1  = m.add_literal(migraphx::literal{s_indices, indices});
+        auto g1  = m.add_instruction(migraphx::make_op("gather", {{"axis", 2}}), input, i1);
+        auto i2  = m.add_literal(migraphx::literal{s_indices, indices});
+        auto g2  = m.add_instruction(migraphx::make_op("gather", {{"axis", 3}}), g1, i2);
+        auto ret = m.add_instruction(migraphx::make_op("pooling",
+                                                       {{"mode", mode},
+                                                        {"padding", {0, 0}},
+                                                        {"stride", {2, 2}},
+                                                        {"lengths", {2, 2}},
+                                                        {"dilations", {1, 1}}}),
+                                     g2);
+        m.add_return({ret});
+        return m;
+    };
+
+    auto test_rewrite = [&](const migraphx::op::pooling_mode mode) {
+        migraphx::module m1 = pooling_program(mode);
+        migraphx::module m2 = opt_program(mode);
+        opt_pooling(m1);
+        EXPECT(m1 == m2);
+    };
+
+    test_rewrite(migraphx::op::pooling_mode::average);
+    test_rewrite(migraphx::op::pooling_mode::max);
+}
+
+TEST_CASE(rewrite_pooling_dialtions_test2)
+{
+    migraphx::shape s{migraphx::shape::float_type, {1, 1, 5, 5, 5}};
+    auto pooling_program = [&](const migraphx::op::pooling_mode mode) {
+        migraphx::module m;
+        auto input = m.add_parameter("x", s);
+        auto ret   = m.add_instruction(migraphx::make_op("pooling",
+                                                         {{"mode", mode},
+                                                          {"padding", {0, 0, 0}},
+                                                          {"stride", {1, 1, 1}},
+                                                          {"lengths", {2, 2, 2}},
+                                                          {"dilations", {2, 2, 2}}}),
+                                     input);
+        m.add_return({ret});
+        return m;
+    };
+
+    auto opt_program = [&](const migraphx::op::pooling_mode mode) {
+        migraphx::module m;
+        auto input = m.add_parameter("x", s);
+        std::vector<int> indices{0, 2, 1, 3, 2, 4};
+        migraphx::shape s_indices{migraphx::shape::int32_type, {indices.size()}};
+        auto i1  = m.add_literal(migraphx::literal{s_indices, indices});
+        auto g1  = m.add_instruction(migraphx::make_op("gather", {{"axis", 2}}), input, i1);
+        auto i2  = m.add_literal(migraphx::literal{s_indices, indices});
+        auto g2  = m.add_instruction(migraphx::make_op("gather", {{"axis", 3}}), g1, i2);
+        auto i3  = m.add_literal(migraphx::literal{s_indices, indices});
+        auto g3  = m.add_instruction(migraphx::make_op("gather", {{"axis", 4}}), g2, i3);
+        auto ret = m.add_instruction(migraphx::make_op("pooling",
+                                                       {{"mode", mode},
+                                                        {"padding", {0, 0, 0}},
+                                                        {"stride", {2, 2, 2}},
+                                                        {"lengths", {2, 2, 2}},
+                                                        {"dilations", {1, 1, 1}}}),
+                                     g3);
+        m.add_return({ret});
+        return m;
+    };
+
+    auto test_rewrite = [&](const migraphx::op::pooling_mode mode) {
+        migraphx::module m1 = pooling_program(mode);
+        migraphx::module m2 = opt_program(mode);
+        opt_pooling(m1);
+        EXPECT(m1 == m2);
+    };
+
+    test_rewrite(migraphx::op::pooling_mode::average);
+    test_rewrite(migraphx::op::pooling_mode::max);
+}
+
+TEST_CASE(rewrite_pooling_dialtions_test3)
+{
+    migraphx::shape s{migraphx::shape::float_type, {2, 2, 5}};
+    auto pooling_program = [&]() {
+        migraphx::module m;
+
+        auto input = m.add_parameter("x", s);
+        auto ret =
+            m.add_instruction(migraphx::make_op("pooling",
+                                                {{"mode", migraphx::op::pooling_mode::average},
+                                                 {"padding", {1}},
+                                                 {"stride", {1}},
+                                                 {"lengths", {3}},
+                                                 {"dilations", {2}}}),
+                              input);
+        m.add_return({ret});
+        return m;
+    };
+
+    migraphx::module m1 = pooling_program();
+    migraphx::module m2 = m1;
+
+    opt_pooling(m1);
+    EXPECT(m1 == m2);
+}
+
+TEST_CASE(rewrite_pooling_dialtions_test4)
+{
+    migraphx::shape s{migraphx::shape::float_type, {1, 1, 5, 5}};
+    auto pooling_program = [&](const migraphx::op::pooling_mode mode) {
+        migraphx::module m;
+        auto input = m.add_parameter("x", s);
+        auto ret   = m.add_instruction(migraphx::make_op("pooling",
+                                                         {{"mode", mode},
+                                                          {"padding", {1, 0}},
+                                                          {"stride", {1, 3}},
+                                                          {"lengths", {3, 1}},
+                                                          {"dilations", {1, 2}}}),
+                                     input);
+        m.add_return({ret});
+        return m;
+    };
+
+    auto opt_program = [&](const migraphx::op::pooling_mode mode) {
+        migraphx::module m;
+        auto input = m.add_parameter("x", s);
+        std::vector<int> col_indices{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14};
+        migraphx::shape s_col_indices{migraphx::shape::int32_type, {col_indices.size()}};
+        std::vector<int> row_indices{0, 3};
+        migraphx::shape s_row_indices{migraphx::shape::int32_type, {row_indices.size()}};
+        auto p =
+            m.add_instruction(migraphx::make_op("pad",
+                                                {{"pads", {0, 0, 1, 0, 0, 0, 1, 0}},
+                                                 {"value", std::numeric_limits<float>::lowest()}}),
+                              input);
+        auto i1  = m.add_literal(migraphx::literal{s_col_indices, col_indices});
+        auto g1  = m.add_instruction(migraphx::make_op("gather", {{"axis", 2}}), p, i1);
+        auto i2  = m.add_literal(migraphx::literal{s_row_indices, row_indices});
+        auto g2  = m.add_instruction(migraphx::make_op("gather", {{"axis", 3}}), g1, i2);
+        auto ret = m.add_instruction(migraphx::make_op("pooling",
+                                                       {{"mode", mode},
+                                                        {"padding", {0, 0}},
+                                                        {"stride", {3, 1}},
+                                                        {"lengths", {3, 1}},
+                                                        {"dilations", {1, 1}}}),
+                                     g2);
+        m.add_return({ret});
+        return m;
+    };
+
+    auto test_rewrite = [&](const migraphx::op::pooling_mode mode) {
+        migraphx::module m1 = pooling_program(mode);
+        migraphx::module m2 = opt_program(mode);
+        opt_pooling(m1);
+        EXPECT(m1 == m2);
+    };
+
+    // Average won't work because of padding
+    test_rewrite(migraphx::op::pooling_mode::max);
+}
+
+TEST_CASE(rewrite_pooling_dialtions_test5)
+{
+    migraphx::shape s{migraphx::shape::float_type, {1, 1, 5, 5}};
+    auto pooling_program = [&](const migraphx::op::pooling_mode mode) {
+        migraphx::module m;
+        auto input = m.add_parameter("x", s);
+        auto ret   = m.add_instruction(migraphx::make_op("pooling",
+                                                         {{"mode", mode},
+                                                          {"padding", {0, 0}},
+                                                          {"stride", {2, 3}},
+                                                          {"lengths", {2, 1}},
+                                                          {"dilations", {1, 2}}}),
+                                     input);
+        m.add_return({ret});
+        return m;
+    };
+
+    auto opt_program = [&](const migraphx::op::pooling_mode mode) {
+        migraphx::module m;
+        auto input = m.add_parameter("x", s);
+        std::vector<int> col_indices{0, 1, 2, 3};
+        migraphx::shape s_col_indices{migraphx::shape::int32_type, {col_indices.size()}};
+        std::vector<int> row_indices{0, 3};
+        migraphx::shape s_row_indices{migraphx::shape::int32_type, {row_indices.size()}};
+        auto i1  = m.add_literal(migraphx::literal{s_col_indices, col_indices});
+        auto g1  = m.add_instruction(migraphx::make_op("gather", {{"axis", 2}}), input, i1);
+        auto i2  = m.add_literal(migraphx::literal{s_row_indices, row_indices});
+        auto g2  = m.add_instruction(migraphx::make_op("gather", {{"axis", 3}}), g1, i2);
+        auto ret = m.add_instruction(migraphx::make_op("pooling",
+                                                       {{"mode", mode},
+                                                        {"padding", {0, 0}},
+                                                        {"stride", {2, 1}},
+                                                        {"lengths", {2, 1}},
+                                                        {"dilations", {1, 1}}}),
+                                     g2);
+        m.add_return({ret});
+        return m;
+    };
+
+    auto test_rewrite = [&](const migraphx::op::pooling_mode mode) {
+        migraphx::module m1 = pooling_program(mode);
+        migraphx::module m2 = opt_program(mode);
+        opt_pooling(m1);
+        EXPECT(m1 == m2);
+    };
+
+    test_rewrite(migraphx::op::pooling_mode::average);
+    test_rewrite(migraphx::op::pooling_mode::max);
+}
+
+TEST_CASE(rewrite_avgpool_rank3_dil_test)
+{
+    // 1D case 1, input is 3D
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::average};
+    op.lengths   = {2};
+    op.padding   = {0};
+    op.stride    = {1};
+    op.dilations = {2};
+
+    std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    opt_pooling(*mm);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{0.35, 0.15, 0.85, 0.3, 0.1, 0.65};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
+TEST_CASE(rewrite_avgpool_rank3_dil_test2)
+{
+    // 1D case 1, input is 3D
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::average};
+    op.lengths   = {2};
+    op.padding   = {0};
+    op.stride    = {1};
+    op.dilations = {3};
+
+    std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    opt_pooling(*mm);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{0.2, 0.45, 0.35};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
+TEST_CASE(rewrite_avgpool_rank4_test)
+{
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 1, 5, 5}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::average};
+    op.lengths   = {2, 1};
+    op.padding   = {0, 0};
+    op.stride    = {2, 3};
+    op.dilations = {1, 2};
+
+    std::vector<float> data(25);
+    std::iota(data.begin(), data.end(), 1);
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    opt_pooling(*mm);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{3.5, 6.5, 13.5, 16.5};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
+TEST_CASE(rewrite_maxpool_rank3_test)
+{
+    // 1D case 1, input is 3D
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::max};
+    op.lengths   = {2};
+    op.padding   = {0};
+    op.stride    = {1};
+    op.dilations = {2};
+
+    std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    opt_pooling(*mm);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{0.4, 0.2, 0.9, 0.5, 0.1, 0.7};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
+TEST_CASE(rewrite_maxpool_rank3_test2)
+{
+    // 1D case 1, input is 3D
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::max};
+    op.lengths   = {2};
+    op.padding   = {1};
+    op.stride    = {1};
+    op.dilations = {3};
+
+    std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    opt_pooling(*mm);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{0.4, 0.3, 0.2, 0.9, 0.8, 0.5, 0.1, 0.6, 0.7};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
+TEST_CASE(rewrite_maxpool_rank3_test3)
+{
+    // 1D case 1, input is 3D
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 3, 4}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::max};
+    op.lengths   = {3};
+    op.padding   = {2};
+    op.stride    = {2};
+    op.dilations = {3};
+
+    std::vector<float> data{0.3, 0.2, 0.4, 0.1, 0.8, 0.5, 0.9, 0.1, 0.1, 0.7, 0.1, 0.6};
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    opt_pooling(*mm);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{0.2, 0.5, 0.7};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
+TEST_CASE(rewrite_maxpool_rank4_test)
+{
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {1, 1, 5, 5}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::max};
+    op.lengths   = {3, 1};
+    op.padding   = {1, 0};
+    op.stride    = {1, 3};
+    op.dilations = {1, 2};
+
+    std::vector<float> data(25);
+    std::iota(data.begin(), data.end(), 1);
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    opt_pooling(*mm);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{6, 9, 11, 14, 16, 19, 21, 24, 21, 24};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
+TEST_CASE(maxpool_rank5_test)
+{
+    // 3D, input is 5D
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {2, 2, 3, 3, 3}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::max};
+    op.lengths   = {2, 2, 2};
+    op.padding   = {0, 0, 0};
+    op.stride    = {1, 1, 1};
+    op.dilations = {2, 2, 2};
+
+    std::vector<float> data{
+        -2.8029, 0.5861,  0.7015,  0.1297,  -1.44,   -1.9472, 0.7812,  2.408,   -0.3145, 0.3405,
+        -0.9146, 0.0624,  1.5064,  -0.8345, 1.7977,  1.8949,  1.0073,  -0.2102, -0.042,  -0.7146,
+        0.6227,  -0.5263, -2.2598, 0.1713,  0.449,   0.5303,  -0.8622, -0.5691, 0.907,   -0.0569,
+        -1.5348, -0.4109, -0.1461, -0.5445, 0.4266,  0.2282,  1.3655,  -2.1519, 0.6068,  -0.2001,
+        -0.4702, 0.3864,  1.7083,  0.9096,  0.4286,  -1.8866, 0.7034,  0.0293,  1.4587,  0.7672,
+        -2.8614, 0.8124,  -0.053,  1.0449,  0.845,   -0.0131, 0.1139,  -0.859,  -1.2681, -0.6337,
+        -0.4644, 0.1938,  0.2889,  0.9035,  0.7118,  -0.5767, 0.4577,  -0.0549, 0.2237,  0.5756,
+        0.0677,  -0.0223, -0.329,  0.2364,  2.7666,  -0.7417, -1.3196, -0.2655, 0.1698,  -0.1777,
+        -0.9427, 2.6859,  -0.7501, 0.5175,  1.0029,  -2.6436, -0.4388, -1.2348, -0.1539, -0.6229,
+        -0.4136, 0.5085,  0.4136,  -0.6439, -1.1953, -0.406,  -0.0195, 0.1869,  -0.8664, 1.1364,
+        0.5041,  0.0647,  0.1941,  -1.0819, -0.4629, -0.5107, 0.3612,  -0.3583};
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    opt_pooling(*mm);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{0.7812, 1.0449, 2.7666, 2.6859};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
+TEST_CASE(maxpool_rank5_test2)
+{
+    // 3D, input is 5D
+    migraphx::program p;
+    auto* mm     = p.get_main_module();
+    auto s       = migraphx::shape{migraphx::shape::float_type, {2, 2, 3, 3, 3}};
+    auto op      = migraphx::op::pooling{migraphx::op::pooling_mode::max};
+    op.lengths   = {2, 2, 2};
+    op.padding   = {2, 2, 2};
+    op.stride    = {2, 2, 2};
+    op.dilations = {3, 3, 3};
+
+    std::vector<float> data{
+        -2.8029, 0.5861,  0.7015,  0.1297,  -1.44,   -1.9472, 0.7812,  2.408,   -0.3145, 0.3405,
+        -0.9146, 0.0624,  1.5064,  -0.8345, 1.7977,  1.8949,  1.0073,  -0.2102, -0.042,  -0.7146,
+        0.6227,  -0.5263, -2.2598, 0.1713,  0.449,   0.5303,  -0.8622, -0.5691, 0.907,   -0.0569,
+        -1.5348, -0.4109, -0.1461, -0.5445, 0.4266,  0.2282,  1.3655,  -2.1519, 0.6068,  -0.2001,
+        -0.4702, 0.3864,  1.7083,  0.9096,  0.4286,  -1.8866, 0.7034,  0.0293,  1.4587,  0.7672,
+        -2.8614, 0.8124,  -0.053,  1.0449,  0.845,   -0.0131, 0.1139,  -0.859,  -1.2681, -0.6337,
+        -0.4644, 0.1938,  0.2889,  0.9035,  0.7118,  -0.5767, 0.4577,  -0.0549, 0.2237,  0.5756,
+        0.0677,  -0.0223, -0.329,  0.2364,  2.7666,  -0.7417, -1.3196, -0.2655, 0.1698,  -0.1777,
+        -0.9427, 2.6859,  -0.7501, 0.5175,  1.0029,  -2.6436, -0.4388, -1.2348, -0.1539, -0.6229,
+        -0.4136, 0.5085,  0.4136,  -0.6439, -1.1953, -0.406,  -0.0195, 0.1869,  -0.8664, 1.1364,
+        0.5041,  0.0647,  0.1941,  -1.0819, -0.4629, -0.5107, 0.3612,  -0.3583};
+    auto l0 = mm->add_literal(migraphx::literal{s, data});
+    mm->add_instruction(op, l0);
+    opt_pooling(*mm);
+    p.compile(migraphx::make_target("ref"));
+    auto result = p.eval({}).back();
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{-0.8345, 1.5064,  -0.9146, 0.3405,  -1.44,   0.1297,  0.5861,  -2.8029,
+                            -0.4702, -0.2001, -2.1519, 1.3655,  -0.4109, -1.5348, 0.907,   -0.5691,
+                            -0.0549, 0.4577,  0.7118,  0.9035,  -1.2681, -0.859,  -0.0131, 0.845,
+                            -1.1953, -0.6439, 0.5085,  -0.4136, -2.6436, 1.0029,  -0.7501, 2.6859};
+    EXPECT(migraphx::verify::verify_rms_range(results_vector, gold));
+}
+
 TEST_CASE(rewrite_avepooling_na1_test)
 {
     migraphx::shape s{migraphx::shape::float_type, {2, 2, 3, 4, 5}};
@@ -92,7 +570,8 @@ TEST_CASE(rewrite_avepooling_na1_test)
                                                 {{"mode", migraphx::op::pooling_mode::average},
                                                  {"padding", {0, 1, 0}},
                                                  {"stride", {1, 1, 1}},
-                                                 {"lengths", {3, 4, 5}}}),
+                                                 {"lengths", {3, 4, 5}},
+                                                 {"dilations", {1, 1, 1}}}),
                               input);
         m.add_return({ret});
         return m;
@@ -117,7 +596,8 @@ TEST_CASE(rewrite_avepooling_na2_test)
                                                 {{"mode", migraphx::op::pooling_mode::average},
                                                  {"padding", {0, 0, 0}},
                                                  {"stride", {1, 2, 1}},
-                                                 {"lengths", {3, 4, 5}}}),
+                                                 {"lengths", {3, 4, 5}},
+                                                 {"dilations", {1, 1, 1}}}),
                               input);
         m.add_return({ret});
         return m;
@@ -141,7 +621,8 @@ TEST_CASE(rewrite_avepooling_na3_test)
                                                          {{"mode", migraphx::op::pooling_mode::max},
                                                           {"padding", {0, 0, 0}},
                                                           {"stride", {1, 1, 1}},
-                                                          {"lengths", {3, 3, 5}}}),
+                                                          {"lengths", {3, 3, 5}},
+                                                          {"dilations", {1, 1, 1}}}),
                                      input);
         m.add_return({ret});
         return m;
@@ -169,7 +650,8 @@ TEST_CASE(literal_rewrite_pooling_test)
                                                            {{"mode", mode},
                                                             {"padding", {0, 0, 0}},
                                                             {"stride", {1, 1, 1}},
-                                                            {"lengths", {3, 4, 5}}}),
+                                                            {"lengths", {3, 4, 5}},
+                                                            {"dilations", {1, 1, 1}}}),
                                        input);
         mm->add_return({ret});
         return p;

test/simplify_dyn_ops_test.cpp

@@ -155,29 +155,187 @@ TEST_CASE(after_split_dyn_broadcast_match)
     EXPECT(p0 == p1);
 }
 
-TEST_CASE(const_slice_3input)
+TEST_CASE(const_slice_2input_ends_axes)
 {
     migraphx::module m0;
     {
         migraphx::shape s{migraphx::shape::float_type, {6, 4, 4}};
-        auto input     = m0.add_parameter("data", s);
-        auto slice_ins = m0.add_instruction(
+        auto input = m0.add_parameter("data", s);
+        migraphx::shape s1{migraphx::shape::int32_type, {1}};
+        auto input_starts = m0.add_literal(migraphx::literal{s1, {0}});
+        auto slice_ins    = m0.add_instruction(
+            migraphx::make_op("slice", {{"ends", {3}}, {"axes", {0}}}), input, input_starts);
+        m0.add_return({slice_ins});
+    }
+    run_pass(m0);
+
+    migraphx::module m1;
+    {
+        migraphx::shape s{migraphx::shape::float_type, {6, 4, 4}};
+        auto input     = m1.add_parameter("data", s);
+        auto slice_ins = m1.add_instruction(
             migraphx::make_op("slice", {{"starts", {0}}, {"ends", {3}}, {"axes", {0}}}), input);
+        m1.add_return({slice_ins});
+    }
+
+    EXPECT(m0 == m1);
+}
+
+TEST_CASE(const_slice_2input_starts_axes)
+{
+    migraphx::module m0;
+    {
+        migraphx::shape s{migraphx::shape::float_type, {6, 4, 4}};
+        auto input = m0.add_parameter("data", s);
+        migraphx::shape s1{migraphx::shape::int32_type, {1}};
+        auto input_ends = m0.add_literal(migraphx::literal{s1, {3}});
+        auto slice_ins  = m0.add_instruction(
+            migraphx::make_op("slice", {{"starts", {0}}, {"axes", {0}}}), input, input_ends);
         m0.add_return({slice_ins});
     }
+    run_pass(m0);
 
     migraphx::module m1;
     {
         migraphx::shape s{migraphx::shape::float_type, {6, 4, 4}};
-        auto input = m1.add_parameter("data", s);
+        auto input     = m1.add_parameter("data", s);
+        auto slice_ins = m1.add_instruction(
+            migraphx::make_op("slice", {{"starts", {0}}, {"ends", {3}}, {"axes", {0}}}), input);
+        m1.add_return({slice_ins});
+    }
+
+    EXPECT(m0 == m1);
+}
+
+TEST_CASE(const_slice_2input_starts_ends)
+{
+    migraphx::module m0;
+    {
+        migraphx::shape s{migraphx::shape::float_type, {6, 4, 4}};
+        auto input = m0.add_parameter("data", s);
         migraphx::shape s1{migraphx::shape::int32_type, {1}};
-        auto input_starts = m1.add_literal(migraphx::literal{s1, {0}});
-        auto input_ends   = m1.add_literal(migraphx::literal{s1, {3}});
-        auto slice_ins    = m1.add_instruction(
+        auto input_axes = m0.add_literal(migraphx::literal{s1, {0}});
+        auto slice_ins  = m0.add_instruction(
+            migraphx::make_op("slice", {{"starts", {0}}, {"ends", {3}}}), input, input_axes);
+        m0.add_return({slice_ins});
+    }
+    run_pass(m0);
+
+    migraphx::module m1;
+    {
+        migraphx::shape s{migraphx::shape::float_type, {6, 4, 4}};
+        auto input     = m1.add_parameter("data", s);
+        auto slice_ins = m1.add_instruction(
+            migraphx::make_op("slice", {{"starts", {0}}, {"ends", {3}}, {"axes", {0}}}), input);
+        m1.add_return({slice_ins});
+    }
+
+    EXPECT(m0 == m1);
+}
+
+TEST_CASE(const_slice_3input_axes_only)
+{
+    migraphx::module m0;
+    {
+        migraphx::shape s{migraphx::shape::float_type, {6, 4, 4}};
+        auto input = m0.add_parameter("data", s);
+        migraphx::shape s1{migraphx::shape::int32_type, {1}};
+        auto input_starts = m0.add_literal(migraphx::literal{s1, {0}});
+        auto input_ends   = m0.add_literal(migraphx::literal{s1, {3}});
+        auto slice_ins    = m0.add_instruction(
             migraphx::make_op("slice", {{"axes", {0}}}), input, input_starts, input_ends);
+        m0.add_return({slice_ins});
+    }
+    run_pass(m0);
+
+    migraphx::module m1;
+    {
+        migraphx::shape s{migraphx::shape::float_type, {6, 4, 4}};
+        auto input     = m1.add_parameter("data", s);
+        auto slice_ins = m1.add_instruction(
+            migraphx::make_op("slice", {{"starts", {0}}, {"ends", {3}}, {"axes", {0}}}), input);
+        m1.add_return({slice_ins});
+    }
+
+    EXPECT(m0 == m1);
+}
+
+TEST_CASE(const_slice_3input_ends_only)
+{
+    migraphx::module m0;
+    {
+        migraphx::shape s{migraphx::shape::float_type, {6, 4, 4}};
+        auto input = m0.add_parameter("data", s);
+        migraphx::shape s1{migraphx::shape::int32_type, {1}};
+        auto input_starts = m0.add_literal(migraphx::literal{s1, {0}});
+        auto input_axes   = m0.add_literal(migraphx::literal{s1, {0}});
+        auto slice_ins    = m0.add_instruction(
+            migraphx::make_op("slice", {{"ends", {3}}}), input, input_starts, input_axes);
+        m0.add_return({slice_ins});
+    }
+    run_pass(m0);
+
+    migraphx::module m1;
+    {
+        migraphx::shape s{migraphx::shape::float_type, {6, 4, 4}};
+        auto input     = m1.add_parameter("data", s);
+        auto slice_ins = m1.add_instruction(
+            migraphx::make_op("slice", {{"starts", {0}}, {"ends", {3}}, {"axes", {0}}}), input);
+        m1.add_return({slice_ins});
+    }
+
+    EXPECT(m0 == m1);
+}
+
+TEST_CASE(const_slice_3inputs_starts_only)
+{
+    migraphx::module m0;
+    {
+        migraphx::shape s{migraphx::shape::float_type, {6, 4, 4}};
+        auto input = m0.add_parameter("data", s);
+        migraphx::shape s1{migraphx::shape::int32_type, {1}};
+        auto input_ends = m0.add_literal(migraphx::literal{s1, {3}});
+        auto input_axes = m0.add_literal(migraphx::literal{s1, {0}});
+        auto slice_ins  = m0.add_instruction(
+            migraphx::make_op("slice", {{"starts", {0}}}), input, input_ends, input_axes);
+        m0.add_return({slice_ins});
+    }
+    run_pass(m0);
+
+    migraphx::module m1;
+    {
+        migraphx::shape s{migraphx::shape::float_type, {6, 4, 4}};
+        auto input     = m1.add_parameter("data", s);
+        auto slice_ins = m1.add_instruction(
+            migraphx::make_op("slice", {{"starts", {0}}, {"ends", {3}}, {"axes", {0}}}), input);
+        m1.add_return({slice_ins});
+    }
+
+    EXPECT(m0 == m1);
+}
+
+TEST_CASE(const_slice_2input_ends_axes_dyn)
+{
+    migraphx::module m0;
+    {
+        migraphx::shape s{migraphx::shape::float_type, {{6, 6}, {2, 4, {2, 4}}, {2, 4, {2, 4}}}};
+        auto input = m0.add_parameter("data", s);
+        migraphx::shape s1{migraphx::shape::int32_type, {1}};
+        auto input_starts = m0.add_literal(migraphx::literal{s1, {0}});
+        auto slice_ins    = m0.add_instruction(
+            migraphx::make_op("slice", {{"ends", {3}}, {"axes", {0}}}), input, input_starts);
+        m0.add_return({slice_ins});
+    }
+    run_pass(m0);
+
+    migraphx::module m1;
+    {
+        migraphx::shape s{migraphx::shape::float_type, {{6, 6}, {2, 4, {2, 4}}, {2, 4, {2, 4}}}};
+        auto input     = m1.add_parameter("data", s);
+        auto slice_ins = m1.add_instruction(
+            migraphx::make_op("slice", {{"starts", {0}}, {"ends", {3}}, {"axes", {0}}}), input);
         m1.add_return({slice_ins});
     }
-    run_pass(m1);
 
     EXPECT(m0 == m1);
 }
@@ -386,4 +544,31 @@ TEST_CASE(static_dimensions_of_to_constant_alloc_reshape)
     EXPECT(m0 == m1);
 }
 
+TEST_CASE(const_alloc_fill)
+{
+    migraphx::module m0;
+    {
+        migraphx::shape val_shape{migraphx::shape::int64_type, {1}, {0}};
+        std::vector<int64_t> lit_data = {3};
+        auto value_lit                = m0.add_literal(migraphx::literal{val_shape, lit_data});
+        migraphx::shape lit_s{migraphx::shape::int64_type, {3}};
+        auto output_dim_lit = m0.add_literal(migraphx::literal{lit_s, {3, 4, 4}});
+        auto alloc_ins      = m0.add_instruction(
+            migraphx::make_op("allocate", {{"buf_type", migraphx::shape::int64_type}}),
+            output_dim_lit);
+        auto ret = m0.add_instruction(migraphx::make_op("fill"), value_lit, alloc_ins);
+        m0.add_return({ret});
+    }
+    run_pass(m0);
+
+    migraphx::module m1;
+    {
+        migraphx::shape lit_shape{migraphx::shape::int64_type, {3, 4, 4}};
+        std::vector<int64_t> lit_data(3 * 4 * 4, 3);
+        auto ret = m1.add_literal(migraphx::literal{lit_shape, lit_data});
+        m1.add_return({ret});
+    }
+    EXPECT(m0 == m1);
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/simplify_qdq_test.cpp

@@ -788,6 +788,7 @@ TEST_CASE(conv_pooling_dot)
                                                   {"padding", {0, 0, 0, 0}},
                                                   {"stride", {1, 1}},
                                                   {"lengths", {7, 7}},
+                                                  {"dilations", {1, 1}},
                                                   {"ceil_mode", 0}}),
                                a1);
         auto fl  = m1.add_instruction(migraphx::make_op("flatten", {{"axis", 1}}), ap);
@@ -835,6 +836,7 @@ TEST_CASE(conv_pooling_dot)
                                                   {"padding", {0, 0, 0, 0}},
                                                   {"stride", {1, 1}},
                                                   {"lengths", {7, 7}},
+                                                  {"dilations", {1, 1}},
                                                   {"ceil_mode", 0}}),
                                a1);
         auto fl         = m2.add_instruction(migraphx::make_op("flatten", {{"axis", 1}}), ap);
@@ -896,6 +898,7 @@ TEST_CASE(mobilenet_snippet)
                                                   {"padding", {0, 0, 0, 0}},
                                                   {"stride", {1, 1}},
                                                   {"lengths", {7, 7}},
+                                                  {"dilations", {1, 1}},
                                                   {"ceil_mode", 0}}),
                                d6);
         auto q3  = add_quantize_op(mm, "quantizelinear", ap, scale, zero);

test/verify/test_avg_pooling_1d.cpp

@@ -35,7 +35,7 @@ struct test_avg_pooling_1d : verify_program<test_avg_pooling_1d>
         auto* mm = p.get_main_module();
         auto input =
             mm->add_parameter("x", migraphx::shape{migraphx::shape::float_type, {1, 3, 5}});
-        auto op = migraphx::op::pooling{migraphx::op::pooling_mode::average, {0}, {1}, {3}};
+        auto op = migraphx::op::pooling{migraphx::op::pooling_mode::average, {0}, {1}, {3}, {1}};
         mm->add_instruction(op, input);
         return p;
     }

test/verify/test_avg_pooling_3d.cpp

@@ -36,7 +36,7 @@ struct test_avg_pooling_3d : verify_program<test_avg_pooling_3d>
         auto input =
             mm->add_parameter("x", migraphx::shape{migraphx::shape::float_type, {1, 3, 5, 5, 5}});
         auto op = migraphx::op::pooling{
-            migraphx::op::pooling_mode::average, {1, 1, 1}, {3, 3, 3}, {3, 3, 3}};
+            migraphx::op::pooling_mode::average, {1, 1, 1}, {3, 3, 3}, {3, 3, 3}, {1, 1, 1}};
         mm->add_instruction(op, input);
         return p;
     }

test/verify/test_avg_pooling_3d_opt.cpp

@@ -36,7 +36,7 @@ struct test_avg_pooling_3d_opt : verify_program<test_avg_pooling_3d_opt>
         auto input =
             mm->add_parameter("x", migraphx::shape{migraphx::shape::float_type, {4, 2, 3, 3, 3}});
         auto op = migraphx::op::pooling{
-            migraphx::op::pooling_mode::average, {0, 0, 0}, {1, 1, 1}, {3, 3, 3}};
+            migraphx::op::pooling_mode::average, {0, 0, 0}, {1, 1, 1}, {3, 3, 3}, {1, 1, 1}};
         mm->add_instruction(op, input);
         return p;
     }

test/verify/test_avg_pooling_ceil_3d.cpp

@@ -37,7 +37,7 @@ struct test_avg_pooling_ceil_3d : verify_program<test_avg_pooling_ceil_3d>
         auto input =
             mm->add_parameter("x", migraphx::shape{migraphx::shape::float_type, {1, 3, 5, 5, 5}});
         auto op = migraphx::op::pooling{
-            migraphx::op::pooling_mode::average, {1, 1, 1}, {3, 3, 3}, {3, 3, 3}, true};
+            migraphx::op::pooling_mode::average, {1, 1, 1}, {3, 3, 3}, {3, 3, 3}, {1, 1, 1}, true};
         mm->add_instruction(op, input);
         return p;
     }

test/verify/test_avg_pooling_pad.cpp

@@ -36,7 +36,7 @@ struct test_avg_pooling_pad : verify_program<test_avg_pooling_pad>
         auto* mm = p.get_main_module();
         auto input =
             mm->add_parameter("x", migraphx::shape{migraphx::shape::float_type, {1, 3, 7}});
-        auto op = migraphx::op::pooling{migraphx::op::pooling_mode::average, {2}, {1}, {3}};
+        auto op = migraphx::op::pooling{migraphx::op::pooling_mode::average, {2}, {1}, {3}, {1}};
         mm->add_instruction(op, input);
         return p;
     }

test/verify/test_concat_pooling.cpp

@@ -47,7 +47,8 @@ struct test_concat_pooling : verify_program<test_concat_pooling>
                                                   {{"mode", migraphx::op::pooling_mode::average},
                                                    {"padding", {0, 0}},
                                                    {"stride", {1, 1}},
-                                                   {"lengths", {8, 8}}}),
+                                                   {"lengths", {8, 8}},
+                                                   {"dilations", {1, 1}}}),
                                 concat_t);
         mm->add_instruction(migraphx::make_op("relu"), pooling);
         return p;

test/verify/test_conv_bn_relu_pooling.cpp

@@ -76,7 +76,8 @@ struct test_conv_bn_relu_pooling : verify_program<test_conv_bn_relu_pooling>
                                               {{"mode", migraphx::op::pooling_mode::average},
                                                {"padding", {1, 1}},
                                                {"stride", {2, 2}},
-                                               {"lengths", {3, 3}}}),
+                                               {"lengths", {3, 3}},
+                                               {"dilations", {1, 1}}}),
                             relu);
         return p;
     }

test/verify/test_conv_bn_relu_pooling2.cpp

@@ -92,7 +92,8 @@ struct test_conv_bn_relu_pooling2 : verify_program<test_conv_bn_relu_pooling2>
                                               {{"mode", migraphx::op::pooling_mode::average},
                                                {"padding", {1, 1}},
                                                {"stride", {2, 2}},
-                                               {"lengths", {3, 3}}}),
+                                               {"lengths", {3, 3}},
+                                               {"dilations", {1, 1}}}),
                             relu);
         return p;
     }

test/verify/test_max_pooling_ceil_3d.cpp

@@ -36,7 +36,7 @@ struct test_max_pooling_ceil_3d : verify_program<test_max_pooling_ceil_3d>
         auto input =
             mm->add_parameter("x", migraphx::shape{migraphx::shape::float_type, {1, 3, 5, 5, 5}});
         auto op = migraphx::op::pooling{
-            migraphx::op::pooling_mode::max, {1, 1, 1}, {3, 3, 3}, {3, 3, 3}, true};
+            migraphx::op::pooling_mode::max, {1, 1, 1}, {3, 3, 3}, {3, 3, 3}, {1, 1, 1}, true};
         mm->add_instruction(op, input);
         return p;
     }

tools/install_prereqs.sh

@@ -53,7 +53,8 @@ else
     python3-pip \
     python3-venv \
     rocblas-dev \
-    rocm-cmake
+    rocm-cmake \
+    libtbb-dev
 fi