[FIX,AUTOSCHEDULER,METASCHEDULE] Handle negative extents in featuriza…

…tion (apache#12990)

Both auto_scheduler and metaschedule featurization did not lower bound
loop extents by zero, leading to negative bounds. When multiplied with
element size integer underflow occurred before converting to floating
point.

python/tvm/auto_scheduler/feature.py

@@ -260,7 +260,7 @@ def features_from_primfunc(
     cache_line_bytes: int = 64,
     max_n_bufs: Optional[int] = None,
     log_scale: bool = False,
-) -> np.ndarray:
+) -> Optional[np.ndarray]:
     """Extract performance features from a PrimFunc.
 
     Parameters
@@ -284,7 +284,7 @@ def features_from_primfunc(
 
     Returns
     -------
-    np.ndarray
+    Optional[np.ndarray]
         Output features, one row per store into a unique buffer statement in `func`.
     """
     return _ffi_api.FeaturesFromPrimFunc(
@@ -297,7 +297,7 @@ def named_features_from_primfunc(
     cache_line_bytes: int = 64,
     max_n_bufs: Optional[int] = None,
     log_scale: bool = False,
-) -> Dict[str, np.ndarray]:
+) -> Optional[Dict[str, np.ndarray]]:
     """Extract performance features and associated names from a PrimFunc.
 
     Parameters
@@ -321,7 +321,7 @@ def named_features_from_primfunc(
 
     Returns
     -------
-    Dict[str, np.ndarray]
+    Optional[Dict[str, np.ndarray]]
         Mapping from feature name to features. One element per store into a
         unique buffer statement in `func`.
     """

src/auto_scheduler/feature.cc

@@ -38,6 +38,7 @@
 #include <tvm/tir/transform.h>
 
 #include <algorithm>
+#include <cassert>
 #include <cmath>
 #include <numeric>
 #include <unordered_map>
@@ -880,6 +881,7 @@ class PerStoreFeatureExtractor : public StmtExprVisitor {
 
         ComputeRegion(acc.indices, &local_analyzer, &tmp_region);
         int64_t touched_size = ElementProduct(tmp_region);
+        touched_size = std::max<int64_t>(0, touched_size);
         buffer_regions_map[t].push_back(
             std::make_tuple(acc.acc_type, touched_size, buffer_dtypes.at(t).bytes()));
         mem_bytes += touched_size * buffer_dtypes.at(t).bytes();
@@ -917,8 +919,9 @@ class PerStoreFeatureExtractor : public StmtExprVisitor {
         lines = 1.0f;
         unique_lines = 1.0f;
       } else {
-        unique_bytes =
-            std::get<1>(for_touch_regions_[for_loop_stack_.front()][t].front()) * ele_bytes;
+        unique_bytes = static_cast<float>(
+                           std::get<1>(for_touch_regions_[for_loop_stack_.front()][t].front())) *
+                       ele_bytes;
 
         stride = 0;
         int64_t reduce_ratio = 1;

src/meta_schedule/feature_extractor/per_store_feature.cc

@@ -836,6 +836,7 @@ void Feature::SetRegion(const LoopNest& loop_nest, IntVec* for_touched_bytes,
       // while others are discarded
       int64_t numel;
       feature.access_shape = utils::RelaxAndUnion(feature.multi_indices, &numel, analyzer);
+      numel = std::max<int64_t>(0, numel);
       feature.loop_accessed_numel[i][buffer] = numel;
       touched_bytes += numel * buffer->dtype.bytes();
       (*buffer_touched_under_loop)[loop][buffer].push_back(numel);
@@ -976,7 +977,8 @@ void Feature::SubFeature::SetFeature(const LoopNest& loop_nest, int64_t cache_li
     this->lines = 1;
     this->unique_lines = 1;
   } else {
-    this->unique_bytes = this->loop_accessed_numel.front().at(buffer) * dtype_bytes;
+    this->unique_bytes =
+        static_cast<double>(this->loop_accessed_numel.front().at(buffer)) * dtype_bytes;
     this->lines = static_cast<double>(loop_nest.prod) / this->prod_non_strided_loop_extent *
                   std::min(1.0, 1.0 * this->min_stride * dtype_bytes / cache_line_bytes);
     this->lines = std::max(1.0, this->lines);

tests/python/unittest/test_auto_scheduler_feature.py

@@ -262,6 +262,17 @@ def test_dense_lowered():
     assert total_bytes_loaded > 2 * 128 * 128 * 4  # 4 bytes per float32
 
 
+@T.prim_func
+def negative_extent(A: T.Buffer[(1,), "float32"]):
+    for j in range(0, -1):
+        A[j] = A[j] + 1.0
+
+
+def test_negative_extent():
+    features = auto_scheduler.feature.named_features_from_primfunc(negative_extent)
+    assert features["B0.unique_bytes"] == 0
+
+
 if __name__ == "__main__":
     test_cpu_matmul()
     test_cpu_fusion()

tests/python/unittest/test_meta_schedule_feature_extractor_per_store_feature.py

@@ -1588,5 +1588,21 @@ def test_cpu_layout_transform():
     )
 
 
+@T.prim_func
+def negative_extent(A: T.Buffer[(1,), "float32"]):
+    for j in range(0, -1):
+        A[j] = A[j] + 1.0
+
+
+def test_negative_extent():
+    extractor = ms.feature_extractor.PerStoreFeature()
+    (features,) = extractor.extract_from(
+        _make_context(tvm.target.Target("llvm")),
+        candidates=[_make_candidate(lambda: tir.Schedule(negative_extent))],
+    )
+    named_features = dict(zip(_feature_names(), list(features.numpy()[0, :])))
+    assert named_features["B0.unique_bytes"] == 0
+
+
 if __name__ == "__main__":
     tvm.testing.main()