From ecf0e0d666fae34a93e94a8d4421e89b668426a5 Mon Sep 17 00:00:00 2001
From: Prakalp Srivastava <prakalp@octoml.ai>
Date: Mon, 13 Feb 2023 06:35:39 -0800
Subject: [PATCH] [Fix] Pool must return error if layout is tiled on H, W, D
 dimensions.

---
 include/tvm/topi/nn/pooling.h                 | 44 ++++++++-------
 tests/python/topi/python/test_topi_pooling.py | 55 ++++++++++++++++---
 2 files changed, 72 insertions(+), 27 deletions(-)

diff --git a/include/tvm/topi/nn/pooling.h b/include/tvm/topi/nn/pooling.h
index 3503584687fea..4b5a1914df792 100644
--- a/include/tvm/topi/nn/pooling.h
+++ b/include/tvm/topi/nn/pooling.h
@@ -214,50 +214,54 @@ inline Tensor pool_grad_impl(const Tensor& out_grad, const Tensor& x,
   }
 }
 
+/*!
+ * \brief Find index of Depth, Height or Width dimension in a layout string.
+ *
+ * \param layout The layout string
+ * \param depth_axis set as the index of depth ('D') if not nullptr.
+ * \param height_axis set as the index of height ('H') if not nullptr.
+ * \param width_axis set as the index of width ('W') if not nullptr.
+ *
+ * \return true if the layout is valid (i.e., no tiling on D, H or W dimensions, no duplicates and
+ * if the requested dimensions are found), otherwise false.
+ */
 inline bool find_depth_height_width(const std::string& layout, int* depth_axis, int* height_axis,
                                     int* width_axis) {
-  *depth_axis = -1;
-  *height_axis = -1;
-  *width_axis = -1;
+  if (depth_axis) *depth_axis = -1;
+  if (height_axis) *height_axis = -1;
+  if (width_axis) *width_axis = -1;
   int curr_idx = 0;
   for (size_t i = 0; i < layout.size(); ++i) {
     if ((layout[i] >= 'A' && layout[i] <= 'Z') || (layout[i] >= 'a' && layout[i] <= 'z')) {
       if (layout[i] == 'D') {
-        if (*depth_axis != -1) return false;
+        if (depth_axis && (*depth_axis != -1)) return false;
         *depth_axis = curr_idx;
       } else if (layout[i] == 'H') {
-        if (*height_axis != -1) return false;
+        if (height_axis && *height_axis != -1) return false;
         *height_axis = curr_idx;
       } else if (layout[i] == 'W') {
-        if (*width_axis != -1) return false;
+        if (width_axis && *width_axis != -1) return false;
         *width_axis = curr_idx;
       } else if (layout[i] == 'd' || layout[i] == 'h' || layout[i] == 'w') {
-        // do not support split on height or width, e.g., NCHW16w
+        // do not support split on height, width or depth, e.g., NCHW16w
         return false;
       }
       ++curr_idx;
     }
   }
-  if (*depth_axis == -1 || *height_axis == -1 || *width_axis == -1) return false;
+  if ((depth_axis && *depth_axis == -1) || (height_axis && *height_axis == -1) ||
+      (width_axis && *width_axis == -1))
+    return false;
   return true;
 }
 
 inline bool find_height_width(const std::string& layout, int* height_axis, int* width_axis) {
-  int dummy;
-  ICHECK_EQ(find_depth_height_width(layout, &dummy, height_axis, width_axis), false);
-  if (*height_axis != -1 && *width_axis != -1) {
-    return true;
-  }
-  return false;
+  return find_depth_height_width(layout, /*depth_axis=*/nullptr, height_axis, width_axis);
 }
 
 inline bool find_width(const std::string& layout, int* width_axis) {
-  int dummy;
-  ICHECK_EQ(find_depth_height_width(layout, &dummy, &dummy, width_axis), false);
-  if (*width_axis != -1) {
-    return true;
-  }
-  return false;
+  return find_depth_height_width(layout, /*depth_axis=*/nullptr, /*height_axis=*/nullptr,
+                                 width_axis);
 }
 
 /*!
diff --git a/tests/python/topi/python/test_topi_pooling.py b/tests/python/topi/python/test_topi_pooling.py
index b7f11de1391f9..5f8aebabc2dfa 100644
--- a/tests/python/topi/python/test_topi_pooling.py
+++ b/tests/python/topi/python/test_topi_pooling.py
@@ -17,12 +17,13 @@
 # pylint: disable=invalid-name, too-many-locals, too-many-statements, unused-argument
 """Test code for pooling"""
 import math
+import pytest
 
 import numpy as np
 import tvm
 import tvm.testing
 import tvm.topi.testing
-from tvm import te, topi
+from tvm import te, topi, TVMError
 from tvm.topi.utils import get_const_tuple
 
 _pool_schedule = {
@@ -762,10 +763,50 @@ def test_pool1d():
     verify_pool1d([1, 31, 16], [3], [3], [3], [3, 0], "max", True, layout="NWC")
 
 
+def test_pool_invalid_tiled_layout():
+
+    with pytest.raises(TVMError, match="Unsupported layout NCHWD4d"):
+        A_3d = te.placeholder([1, 16, 32, 32, 32], name="A")
+        B = topi.nn.pool3d(
+            A_3d,
+            kernel=[2, 2, 2],
+            stride=[2, 2, 2],
+            dilation=[1, 1, 1],
+            padding=[0, 0, 0, 0, 0, 0],
+            pool_type="avg",
+            ceil_mode=False,
+            count_include_pad=True,
+            layout="NCHWD4d",
+        )
+
+    with pytest.raises(TVMError, match="Unsupported layout NCHW4h4w"):
+        A_2d = te.placeholder([1, 16, 32, 32], name="A")
+        B = topi.nn.pool2d(
+            A_2d,
+            kernel=[2, 2],
+            stride=[2, 2],
+            dilation=[1, 1],
+            padding=[0, 0, 0, 0],
+            pool_type="avg",
+            ceil_mode=False,
+            count_include_pad=True,
+            layout="NCHW4h4w",
+        )
+
+    with pytest.raises(TVMError, match="Unsupported layout NCW4w"):
+        A_1d = te.placeholder([1, 16, 32], name="A")
+        B = topi.nn.pool1d(
+            A_1d,
+            kernel=[2],
+            stride=[2],
+            dilation=[1],
+            padding=[0, 0],
+            pool_type="avg",
+            ceil_mode=False,
+            count_include_pad=True,
+            layout="NCW4w",
+        )
+
+
 if __name__ == "__main__":
-    test_pool1d()
-    test_pool2d()
-    test_pool3d()
-    test_pool_grad()
-    test_global_pool()
-    test_adaptive_pool()
+    tvm.testing.main()