Generalize workaround for unbacked size hints (#159)

Author: jansel

helion/_compiler/compile_environment.py

@@ -255,6 +255,11 @@ def _to_fake_tensor(self, tensor: torch.Tensor, source: Source) -> torch.Tensor:
 
     def size_hint(self, n: int | torch.SymInt) -> int:
         if isinstance(n, torch.SymInt):
+            expr = n._sympy_()
+            if any(s.name.startswith("u") for s in expr.free_symbols):
+                # If the size is a symbolic expression with unbacked symbols, then the shape environment
+                # hint will be wrong since we assign a default value to unbacked symbols.  Return a default hint.
+                return 8192
             # pyre-ignore[6]
             return int(self.shape_env.size_hint(n._sympy_()))
         assert isinstance(n, int)

helion/_compiler/type_propagation.py

@@ -987,13 +987,7 @@ def _get_hint(numel: int | torch.SymInt | AutoSize | None) -> int:
     if numel is None or isinstance(numel, AutoSize):
         # For data-dependent sizes, use arbitrary hint of 8192
         return 8192
-
-    hint = CompileEnvironment.current().size_hint(numel)
-    # If the hint is invalid (like 0), use a reasonable default
-    # This can happen when other hints cancel out in expressions
-    if hint <= 1:
-        return 8192
-    return hint
+    return CompileEnvironment.current().size_hint(numel)
 
 
 class TileIndexType(TypeInfo):