[TE][CreatePrimFunc] Fix loop carried dependency case with nested blo…

…ck levels (#17474) Nested level primfunc block generation for axis dependencies Co-authored-by: wrongtest <[email protected]>
apache · Nov 14, 2024 · 370ec6a · 370ec6a
1 parent 59a2256
commit 370ec6a
Show file tree

Hide file tree

Showing 3 changed files with 437 additions and 124 deletions.
diff --git a/python/tvm/script/ir_builder/tir/ir.py b/python/tvm/script/ir_builder/tir/ir.py
@@ -522,6 +522,11 @@ def _as_range(dom: Union[ir.Range, List[PrimExpr]]) -> ir.Range:
     if isinstance(dom, ir.Range):
         return dom
     if isinstance(dom, (list, tuple)):
+        from tvm.arith import Analyzer  # pylint: disable=import-outside-toplevel
+
+        extent = Analyzer().simplify(dom[1] - dom[0])
+        if isinstance(extent, tir.IntImm):
+            return ir.Range.from_min_extent(dom[0], extent)
         return ir.Range(dom[0], dom[1])
     if hasattr(dom, "dtype"):
         return ir.Range(IntImm(dom.dtype, 0), dom)