Address comments

Change-Id: I506bb622065e516b4e7111d461e3dac6ee821889

Author: lhutton1

src/arith/scalable_expression.cc

@@ -66,8 +66,15 @@ std::optional<int> ExtractVscaleFactor(const PrimExpr& lanes) {
   }
 }
 
+bool IsComparison(const PrimExpr& expr) {
+  return expr->IsInstance<tir::LENode>() || expr->IsInstance<tir::LTNode>() ||
+         expr->IsInstance<tir::GENode>() || expr->IsInstance<tir::GTNode>() ||
+         expr->IsInstance<tir::EQNode>() || expr->IsInstance<tir::NENode>();
+}
+
 bool CanProveVscaleExpressionFromKnownValues(arith::Analyzer* analyzer, const PrimExpr& expr,
                                              const std::vector<unsigned int>& vscale_values) {
+  ICHECK(IsComparison(expr)) << "Expected comparison but got: " << expr;
   bool can_prove_expr = true;
   for (const unsigned int vscale_value : vscale_values) {
     PrimExpr result = SubstituteVScaleWithKnownValue(expr, vscale_value);

src/arith/scalable_expression.h

@@ -65,6 +65,7 @@ std::optional<int> ExtractVscaleFactor(const PrimExpr& lanes);
            of vscale.
  * \param analyzer An analyzer instance.
  * \param expr The expression to try to prove.
+ * \param vscale_values A list of values to substitute vscale with.
  * \return Whether or not the expression can be proven with this technique.
  */
 bool CanProveVscaleExpressionFromKnownValues(arith::Analyzer* analyzer, const PrimExpr& expr,

tests/python/arith/test_arith_simplify.py

@@ -20,6 +20,7 @@
 import tvm
 import tvm.testing
 from tvm import tir
+from tvm.script import tir as T
 
 
 def test_simplify_reshape_flattened_index():
@@ -56,14 +57,20 @@ def test_simplify_symbolic_comparison():
     assert ana.can_prove((n + 31) // 32 * 32 >= i0 * 32 + i1, PS.SYMBOLIC_BOUND)
 
 
-def test_simplify_vscale_comparison_with_sve_target():
+@pytest.mark.parametrize(
+    "expression",
+    [
+        T.vscale() * 32 < T.vscale() * 64,
+        T.vscale() * 2 * (T.vscale() * 2) >= T.vscale() * 4,
+        (T.vscale() * 4 + 114) // (T.vscale() * 4) * (T.vscale() * 4) >= 115,
+        64 % T.vscale() <= T.vscale(),
+    ],
+)
+def test_simplify_vscale_comparison_with_sve_target(expression):
     ana = tvm.arith.Analyzer()
-    vs = tvm.tir.vscale()
 
     with tvm.target.Target("llvm -mtriple=aarch64-linux-gnu -mattr=+sve"):
-        assert ana.can_prove(vs * 32 < vs * 64)
-        assert ana.can_prove(vs * 2 * (vs * 2) >= vs * 4)
-        assert ana.can_prove((vs * 4 + 114) // (vs * 4) * (vs * 4) >= 115)
+        assert ana.can_prove(expression)
 
 
 def test_simplify_vscale_comparison_without_sve_target(capfd):
@@ -83,6 +90,16 @@ def test_simplify_vscale_comparison_without_sve_target(capfd):
     assert warning_msg in capture
 
 
+def test_simplify_vscale_non_comparison():
+    ana = tvm.arith.Analyzer()
+    vs = tvm.tir.vscale()
+
+    err_msg = r".*Expected comparison but got: T.vscale\(\) \* 4"
+    with pytest.raises(tvm.TVMError, match=err_msg):
+        with tvm.target.Target("llvm -mtriple=aarch64-linux-gnu -mattr=+sve"):
+            ana.can_prove(vs * 4)
+
+
 def test_regression_simplify_inf_recursion():
     ana = tvm.arith.Analyzer()
     cond = tir.Var("cond", "int32")

tests/python/tir-schedule/test_tir_schedule_split_fuse.py

@@ -697,7 +697,7 @@ def test_sve_scalable_split_assume_exact_multiple():
     If the schedule writer knows the extent of the loop to be split will always
     be a multiple of vscale, they may use `disable_predication=True` to ensure
     a predicate is not created. This can be used to ensure predication is not
-    inserted where current analysis is not powerful enough to recognise this.
+    inserted.
     """
 
     @T.prim_func
@@ -761,7 +761,7 @@ def after(a: T.handle):
     tvm.ir.assert_structural_equal(sch.mod["main"], after)
 
 
-def test_default_scalable_split(capfd):
+def test_unsupported_target_scalable_split(capfd):
     @T.prim_func
     def before(a: T.handle):
         A = T.match_buffer(a, (128,), "float32")