From 0825fd20438192daf8112337f6a57e572bc36410 Mon Sep 17 00:00:00 2001
From: Weiming Zhao <weiming.zhao@alibaba-inc.com>
Date: Sun, 15 Aug 2021 08:32:49 +0000
Subject: [PATCH] [Bug] Fix axis attribute for reduction instrs

When the axis op is constant, it should update the axis attribute.

(cherry picked from commit 4fe94fa30cf2d181ccedf663ee755e883c14b285)
---
 lib/transforms/type_legalizer.cc                       | 10 ++++++++--
 ..._reduce_mean_negative_axes_keepdims_example_dnnl.cc |  2 +-
 ...t_reduce_mean_negative_axes_keepdims_random_dnnl.cc |  2 +-
 ...t_reduce_min_negative_axes_keepdims_example_dnnl.cc |  2 +-
 ...st_reduce_min_negative_axes_keepdims_random_dnnl.cc |  2 +-
 ...educe_mean_negative_axes_keepdims_example_popart.cc |  2 +-
 ...reduce_mean_negative_axes_keepdims_random_popart.cc |  2 +-
 ...reduce_min_negative_axes_keepdims_example_popart.cc |  2 +-
 ..._reduce_min_negative_axes_keepdims_random_popart.cc |  2 +-
 ...duce_l1_negative_axes_keep_dims_example_tensorrt.cc |  2 +-
 ...educe_l1_negative_axes_keep_dims_random_tensorrt.cc |  2 +-
 ...duce_l2_negative_axes_keep_dims_example_tensorrt.cc |  2 +-
 ...educe_l2_negative_axes_keep_dims_random_tensorrt.cc |  2 +-
 ..._sum_exp_negative_axes_keepdims_example_tensorrt.cc |  2 +-
 ...g_sum_exp_negative_axes_keepdims_random_tensorrt.cc |  2 +-
 .../test_reduce_log_sum_negative_axes_tensorrt.cc      |  2 +-
 ...uce_mean_negative_axes_keepdims_example_tensorrt.cc |  2 +-
 ...duce_mean_negative_axes_keepdims_random_tensorrt.cc |  2 +-
 ...duce_min_negative_axes_keepdims_example_tensorrt.cc |  2 +-
 ...educe_min_negative_axes_keepdims_random_tensorrt.cc |  2 +-
 ...m_square_negative_axes_keepdims_example_tensorrt.cc |  2 +-
 21 files changed, 28 insertions(+), 22 deletions(-)

diff --git a/lib/transforms/type_legalizer.cc b/lib/transforms/type_legalizer.cc
index 28ae88278..fe3826904 100644
--- a/lib/transforms/type_legalizer.cc
+++ b/lib/transforms/type_legalizer.cc
@@ -20,6 +20,7 @@
 #include <algorithm>
 #include <cmath>
 #include <limits>
+#include <type_traits>
 #include <unordered_set>
 
 #include "halo/api/halo_data.h"
@@ -475,8 +476,8 @@ static void RunOnInstruction(Conv2DTransposeInst* inst) {
   }
 }
 
-static void RunOnCommonReductionInstruction(Instruction* inst,
-                                            std::vector<int32_t> axis,
+template <typename T>
+static void RunOnCommonReductionInstruction(T* inst, std::vector<int32_t> axis,
                                             bool keep_dims) {
   const auto& input_type = inst->GetOperand(0).GetType();
   if (!input_type.IsValid()) {
@@ -528,6 +529,11 @@ static void RunOnCommonReductionInstruction(Instruction* inst,
     dt = DataType::INT32;
   }
 
+  constexpr bool is_arg_inst =
+      std::is_same<T, ArgmaxInst>() || std::is_same<T, ArgminInst>();
+  if constexpr (!is_arg_inst) { // NOLINT
+    inst->SetAxis(axis);
+  }
   inst->GetResultsTypes()[0] = halo::Type{dt, ret_shape};
 }
 
diff --git a/tests/unittests/lit_cases/test_dnnl/test_reduce_mean_negative_axes_keepdims_example_dnnl.cc b/tests/unittests/lit_cases/test_dnnl/test_reduce_mean_negative_axes_keepdims_example_dnnl.cc
index a830db822..1cfeec2c8 100644
--- a/tests/unittests/lit_cases/test_dnnl/test_reduce_mean_negative_axes_keepdims_example_dnnl.cc
+++ b/tests/unittests/lit_cases/test_dnnl/test_reduce_mean_negative_axes_keepdims_example_dnnl.cc
@@ -25,5 +25,5 @@
 // RUN: %t_dnnl.exe 0.0001 0 dnnl %data_path/test_reduce_mean_negative_axes_keepdims_example | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_mean_negative_axes_keepdims_example_dnnl.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_dnnl/test_reduce_mean_negative_axes_keepdims_random_dnnl.cc b/tests/unittests/lit_cases/test_dnnl/test_reduce_mean_negative_axes_keepdims_random_dnnl.cc
index 084efc254..2932b033b 100644
--- a/tests/unittests/lit_cases/test_dnnl/test_reduce_mean_negative_axes_keepdims_random_dnnl.cc
+++ b/tests/unittests/lit_cases/test_dnnl/test_reduce_mean_negative_axes_keepdims_random_dnnl.cc
@@ -25,5 +25,5 @@
 // RUN: %t_dnnl.exe 0.0001 0 dnnl %data_path/test_reduce_mean_negative_axes_keepdims_random | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_mean_negative_axes_keepdims_random_dnnl.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_dnnl/test_reduce_min_negative_axes_keepdims_example_dnnl.cc b/tests/unittests/lit_cases/test_dnnl/test_reduce_min_negative_axes_keepdims_example_dnnl.cc
index f937ab91c..b9684f8e6 100644
--- a/tests/unittests/lit_cases/test_dnnl/test_reduce_min_negative_axes_keepdims_example_dnnl.cc
+++ b/tests/unittests/lit_cases/test_dnnl/test_reduce_min_negative_axes_keepdims_example_dnnl.cc
@@ -25,5 +25,5 @@
 // RUN: %t_dnnl.exe 0.0001 0 dnnl %data_path/test_reduce_min_negative_axes_keepdims_example | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_min_negative_axes_keepdims_example_dnnl.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_dnnl/test_reduce_min_negative_axes_keepdims_random_dnnl.cc b/tests/unittests/lit_cases/test_dnnl/test_reduce_min_negative_axes_keepdims_random_dnnl.cc
index 64894fd19..0e5c8835e 100644
--- a/tests/unittests/lit_cases/test_dnnl/test_reduce_min_negative_axes_keepdims_random_dnnl.cc
+++ b/tests/unittests/lit_cases/test_dnnl/test_reduce_min_negative_axes_keepdims_random_dnnl.cc
@@ -25,5 +25,5 @@
 // RUN: %t_dnnl.exe 0.0001 0 dnnl %data_path/test_reduce_min_negative_axes_keepdims_random | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_min_negative_axes_keepdims_random_dnnl.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_popart/test_reduce_mean_negative_axes_keepdims_example_popart.cc b/tests/unittests/lit_cases/test_popart/test_reduce_mean_negative_axes_keepdims_example_popart.cc
index ecde897e9..30d7dfb83 100644
--- a/tests/unittests/lit_cases/test_popart/test_reduce_mean_negative_axes_keepdims_example_popart.cc
+++ b/tests/unittests/lit_cases/test_popart/test_reduce_mean_negative_axes_keepdims_example_popart.cc
@@ -25,5 +25,5 @@
 // RUN: %t_popart.exe 0.0001 0 popart %data_path/test_reduce_mean_negative_axes_keepdims_example | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_mean_negative_axes_keepdims_example_popart.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_popart/test_reduce_mean_negative_axes_keepdims_random_popart.cc b/tests/unittests/lit_cases/test_popart/test_reduce_mean_negative_axes_keepdims_random_popart.cc
index 5ff12350f..446db5378 100644
--- a/tests/unittests/lit_cases/test_popart/test_reduce_mean_negative_axes_keepdims_random_popart.cc
+++ b/tests/unittests/lit_cases/test_popart/test_reduce_mean_negative_axes_keepdims_random_popart.cc
@@ -25,5 +25,5 @@
 // RUN: %t_popart.exe 0.0001 0 popart %data_path/test_reduce_mean_negative_axes_keepdims_random | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_mean_negative_axes_keepdims_random_popart.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_popart/test_reduce_min_negative_axes_keepdims_example_popart.cc b/tests/unittests/lit_cases/test_popart/test_reduce_min_negative_axes_keepdims_example_popart.cc
index 76073ac0b..4d92cd9d1 100644
--- a/tests/unittests/lit_cases/test_popart/test_reduce_min_negative_axes_keepdims_example_popart.cc
+++ b/tests/unittests/lit_cases/test_popart/test_reduce_min_negative_axes_keepdims_example_popart.cc
@@ -25,5 +25,5 @@
 // RUN: %t_popart.exe 0.0001 0 popart %data_path/test_reduce_min_negative_axes_keepdims_example | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_min_negative_axes_keepdims_example_popart.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_popart/test_reduce_min_negative_axes_keepdims_random_popart.cc b/tests/unittests/lit_cases/test_popart/test_reduce_min_negative_axes_keepdims_random_popart.cc
index 1af6a8572..0a5e4103b 100644
--- a/tests/unittests/lit_cases/test_popart/test_reduce_min_negative_axes_keepdims_random_popart.cc
+++ b/tests/unittests/lit_cases/test_popart/test_reduce_min_negative_axes_keepdims_random_popart.cc
@@ -25,5 +25,5 @@
 // RUN: %t_popart.exe 0.0001 0 popart %data_path/test_reduce_min_negative_axes_keepdims_random | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_min_negative_axes_keepdims_random_popart.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_tensorrt/test_reduce_l1_negative_axes_keep_dims_example_tensorrt.cc b/tests/unittests/lit_cases/test_tensorrt/test_reduce_l1_negative_axes_keep_dims_example_tensorrt.cc
index 5c8badc81..994774334 100644
--- a/tests/unittests/lit_cases/test_tensorrt/test_reduce_l1_negative_axes_keep_dims_example_tensorrt.cc
+++ b/tests/unittests/lit_cases/test_tensorrt/test_reduce_l1_negative_axes_keep_dims_example_tensorrt.cc
@@ -25,5 +25,5 @@
 // RUN: %t_tensorrt.exe 0.0001 0 tensorrt %data_path/test_reduce_l1_negative_axes_keep_dims_example | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_l1_negative_axes_keep_dims_example_tensorrt.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_tensorrt/test_reduce_l1_negative_axes_keep_dims_random_tensorrt.cc b/tests/unittests/lit_cases/test_tensorrt/test_reduce_l1_negative_axes_keep_dims_random_tensorrt.cc
index 96927b5c0..d8103232e 100644
--- a/tests/unittests/lit_cases/test_tensorrt/test_reduce_l1_negative_axes_keep_dims_random_tensorrt.cc
+++ b/tests/unittests/lit_cases/test_tensorrt/test_reduce_l1_negative_axes_keep_dims_random_tensorrt.cc
@@ -25,5 +25,5 @@
 // RUN: %t_tensorrt.exe 0.0001 0 tensorrt %data_path/test_reduce_l1_negative_axes_keep_dims_random | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_l1_negative_axes_keep_dims_random_tensorrt.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_tensorrt/test_reduce_l2_negative_axes_keep_dims_example_tensorrt.cc b/tests/unittests/lit_cases/test_tensorrt/test_reduce_l2_negative_axes_keep_dims_example_tensorrt.cc
index 2e6ca5eeb..4414ad0dc 100644
--- a/tests/unittests/lit_cases/test_tensorrt/test_reduce_l2_negative_axes_keep_dims_example_tensorrt.cc
+++ b/tests/unittests/lit_cases/test_tensorrt/test_reduce_l2_negative_axes_keep_dims_example_tensorrt.cc
@@ -25,5 +25,5 @@
 // RUN: %t_tensorrt.exe 0.0001 0 tensorrt %data_path/test_reduce_l2_negative_axes_keep_dims_example | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_l2_negative_axes_keep_dims_example_tensorrt.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_tensorrt/test_reduce_l2_negative_axes_keep_dims_random_tensorrt.cc b/tests/unittests/lit_cases/test_tensorrt/test_reduce_l2_negative_axes_keep_dims_random_tensorrt.cc
index ba6e1ee45..97d24a30f 100644
--- a/tests/unittests/lit_cases/test_tensorrt/test_reduce_l2_negative_axes_keep_dims_random_tensorrt.cc
+++ b/tests/unittests/lit_cases/test_tensorrt/test_reduce_l2_negative_axes_keep_dims_random_tensorrt.cc
@@ -25,5 +25,5 @@
 // RUN: %t_tensorrt.exe 0.0001 0 tensorrt %data_path/test_reduce_l2_negative_axes_keep_dims_random | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_l2_negative_axes_keep_dims_random_tensorrt.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_tensorrt/test_reduce_log_sum_exp_negative_axes_keepdims_example_tensorrt.cc b/tests/unittests/lit_cases/test_tensorrt/test_reduce_log_sum_exp_negative_axes_keepdims_example_tensorrt.cc
index 3aa50571f..3feea04b7 100644
--- a/tests/unittests/lit_cases/test_tensorrt/test_reduce_log_sum_exp_negative_axes_keepdims_example_tensorrt.cc
+++ b/tests/unittests/lit_cases/test_tensorrt/test_reduce_log_sum_exp_negative_axes_keepdims_example_tensorrt.cc
@@ -25,5 +25,5 @@
 // RUN: %t_tensorrt.exe 0.0001 0 tensorrt %data_path/test_reduce_log_sum_exp_negative_axes_keepdims_example | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_log_sum_exp_negative_axes_keepdims_example_tensorrt.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_tensorrt/test_reduce_log_sum_exp_negative_axes_keepdims_random_tensorrt.cc b/tests/unittests/lit_cases/test_tensorrt/test_reduce_log_sum_exp_negative_axes_keepdims_random_tensorrt.cc
index ddf9cfd1a..b26df50f8 100644
--- a/tests/unittests/lit_cases/test_tensorrt/test_reduce_log_sum_exp_negative_axes_keepdims_random_tensorrt.cc
+++ b/tests/unittests/lit_cases/test_tensorrt/test_reduce_log_sum_exp_negative_axes_keepdims_random_tensorrt.cc
@@ -25,5 +25,5 @@
 // RUN: %t_tensorrt.exe 0.0001 0 tensorrt %data_path/test_reduce_log_sum_exp_negative_axes_keepdims_random | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_log_sum_exp_negative_axes_keepdims_random_tensorrt.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_tensorrt/test_reduce_log_sum_negative_axes_tensorrt.cc b/tests/unittests/lit_cases/test_tensorrt/test_reduce_log_sum_negative_axes_tensorrt.cc
index 303571973..7d5955c60 100644
--- a/tests/unittests/lit_cases/test_tensorrt/test_reduce_log_sum_negative_axes_tensorrt.cc
+++ b/tests/unittests/lit_cases/test_tensorrt/test_reduce_log_sum_negative_axes_tensorrt.cc
@@ -25,5 +25,5 @@
 // RUN: %t_tensorrt.exe 0.0001 0 tensorrt %data_path/test_reduce_log_sum_negative_axes | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_log_sum_negative_axes_tensorrt.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_tensorrt/test_reduce_mean_negative_axes_keepdims_example_tensorrt.cc b/tests/unittests/lit_cases/test_tensorrt/test_reduce_mean_negative_axes_keepdims_example_tensorrt.cc
index 30f44a47d..8e166674f 100644
--- a/tests/unittests/lit_cases/test_tensorrt/test_reduce_mean_negative_axes_keepdims_example_tensorrt.cc
+++ b/tests/unittests/lit_cases/test_tensorrt/test_reduce_mean_negative_axes_keepdims_example_tensorrt.cc
@@ -25,5 +25,5 @@
 // RUN: %t_tensorrt.exe 0.0001 0 tensorrt %data_path/test_reduce_mean_negative_axes_keepdims_example | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_mean_negative_axes_keepdims_example_tensorrt.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_tensorrt/test_reduce_mean_negative_axes_keepdims_random_tensorrt.cc b/tests/unittests/lit_cases/test_tensorrt/test_reduce_mean_negative_axes_keepdims_random_tensorrt.cc
index 578aa7f4c..2d0947aa7 100644
--- a/tests/unittests/lit_cases/test_tensorrt/test_reduce_mean_negative_axes_keepdims_random_tensorrt.cc
+++ b/tests/unittests/lit_cases/test_tensorrt/test_reduce_mean_negative_axes_keepdims_random_tensorrt.cc
@@ -25,5 +25,5 @@
 // RUN: %t_tensorrt.exe 0.0001 0 tensorrt %data_path/test_reduce_mean_negative_axes_keepdims_random | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_mean_negative_axes_keepdims_random_tensorrt.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_tensorrt/test_reduce_min_negative_axes_keepdims_example_tensorrt.cc b/tests/unittests/lit_cases/test_tensorrt/test_reduce_min_negative_axes_keepdims_example_tensorrt.cc
index ff3128188..8de8774af 100644
--- a/tests/unittests/lit_cases/test_tensorrt/test_reduce_min_negative_axes_keepdims_example_tensorrt.cc
+++ b/tests/unittests/lit_cases/test_tensorrt/test_reduce_min_negative_axes_keepdims_example_tensorrt.cc
@@ -25,5 +25,5 @@
 // RUN: %t_tensorrt.exe 0.0001 0 tensorrt %data_path/test_reduce_min_negative_axes_keepdims_example | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_min_negative_axes_keepdims_example_tensorrt.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_tensorrt/test_reduce_min_negative_axes_keepdims_random_tensorrt.cc b/tests/unittests/lit_cases/test_tensorrt/test_reduce_min_negative_axes_keepdims_random_tensorrt.cc
index 92164a154..5c12b1311 100644
--- a/tests/unittests/lit_cases/test_tensorrt/test_reduce_min_negative_axes_keepdims_random_tensorrt.cc
+++ b/tests/unittests/lit_cases/test_tensorrt/test_reduce_min_negative_axes_keepdims_random_tensorrt.cc
@@ -25,5 +25,5 @@
 // RUN: %t_tensorrt.exe 0.0001 0 tensorrt %data_path/test_reduce_min_negative_axes_keepdims_random | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_min_negative_axes_keepdims_random_tensorrt.cc.tmp.main.cc.in"
diff --git a/tests/unittests/lit_cases/test_tensorrt/test_reduce_sum_square_negative_axes_keepdims_example_tensorrt.cc b/tests/unittests/lit_cases/test_tensorrt/test_reduce_sum_square_negative_axes_keepdims_example_tensorrt.cc
index 1578da0e9..d5ff6924b 100644
--- a/tests/unittests/lit_cases/test_tensorrt/test_reduce_sum_square_negative_axes_keepdims_example_tensorrt.cc
+++ b/tests/unittests/lit_cases/test_tensorrt/test_reduce_sum_square_negative_axes_keepdims_example_tensorrt.cc
@@ -25,5 +25,5 @@
 // RUN: %t_tensorrt.exe 0.0001 0 tensorrt %data_path/test_reduce_sum_square_negative_axes_keepdims_example | FileCheck %s
 // CHECK: Result Pass
 // clang-format on
-// XFAIL: *
+
 #include "test_reduce_sum_square_negative_axes_keepdims_example_tensorrt.cc.tmp.main.cc.in"