From d4cb4ea5d7c3be30b426b42fcd958763e255a748 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Tue, 20 Jun 2023 12:19:24 -0500 Subject: [PATCH] [Bugfix][Arith] Avoid flaky test with TryCombineSplitFromSameSource This commit resolves a flaky test failure that was introduced in https://github.com/apache/tvm/pull/15081. The unit test `tests/python/unittest/test_meta_schedule_schedule_rule_mlt_tc.py::test_padded_matmul_relu` failed approximately 30% of the time. The error was due to changes in `IterMapRewriter::NormalizeToIterWithOffset`. When attempting a simplfication with `TryCombineSplitFromSameSource`, if the returned `Optional` is defined, but the `args.size() < 1` check fails, then the `expr` argument is still overwritten, and presumably can cause a later `TryFuseIters(expr)` to fail. This commit replaces `expr = opt.value();` with `auto combined = opt.value();`, preserving the original argument. --- src/arith/iter_affine_map.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/arith/iter_affine_map.cc b/src/arith/iter_affine_map.cc index 377f8bb7c9b1..9a603d0a5c93 100644 --- a/src/arith/iter_affine_map.cc +++ b/src/arith/iter_affine_map.cc @@ -723,8 +723,10 @@ class IterMapRewriter : public ExprMutator { // We are normalizing a regular iter if (expr->args.size() < 1) return expr; if (auto opt = TryCombineSplitFromSameSource(expr)) { - expr = opt.value(); - if (expr->args.size() < 1) return expr; + auto combined = opt.value(); + if (combined->args.size() < 1) { + return combined; + } } Optional opt = TryFuseIters(expr, check_level_); if (opt.defined()) {