up

Author: yf225

helion/_compiler/compile_environment.py

@@ -112,6 +112,9 @@ def __init__(
             collections.Counter()
         )
         self.specialized_vars: set[sympy.Symbol] = set()
+        # Maps stride symbol -> (param_name, dim) for stride specialization
+        self.stride_symbols: dict[sympy.Symbol, tuple[str, int]] = {}
+        self.specialized_strides: set[tuple[str, int]] = set()
         self.loop_dependency_checker = LoopDependencyChecker()
         self._symint_cache: dict[object, torch.SymInt] = {}
         self.device_load_count = (
@@ -469,6 +472,13 @@ def _to_fake_tensor(self, tensor: torch.Tensor, source: Source) -> torch.Tensor:
                     self.debug_shape_renames[s._sympy_()] = sympy.Symbol(
                         f"{source.local_name}_size{i}", integer=True
                     )
+        # Record stride symbols for specialization tracking
+        if isinstance(source, LocalSource):
+            for i in range(result.ndim):
+                st = result.stride(i)
+                if isinstance(st, torch.SymInt):
+                    for sym in st._sympy_().free_symbols:
+                        self.stride_symbols[sym] = (source.local_name, i)
         return result
 
     def size_hint(self, n: int | torch.SymInt) -> int:

helion/_compiler/device_function.py

@@ -602,10 +602,15 @@ def tensor_size(self, fake_value: torch.Tensor, dim: int) -> Argument:
         return self._tensor_property(TensorSizeArg, fake_value, dim, "size")
 
     def tensor_stride(self, fake_value: torch.Tensor, dim: int) -> Argument:
-        if (
-            isinstance(v := fake_value.stride(dim), int)
-            and CompileEnvironment.current().settings.static_shapes
-        ):
+        from torch._dynamo.source import LocalSource
+
+        v = fake_value.stride(dim)
+        env = CompileEnvironment.current()
+        # Check if this specific stride was specialized
+        source = env.input_sources.get(fake_value)
+        if isinstance(source, LocalSource) and (source.local_name, dim) in env.specialized_strides:
+            return StaticShape(int(v))
+        if isinstance(v, int) and env.settings.static_shapes:
             return StaticShape(v)
         return self._tensor_property(TensorStrideArg, fake_value, dim, "stride")
 

helion/_compiler/type_propagation.py

@@ -645,7 +645,7 @@ def propagate_call(
         attr = self.attr()
         if attr in {"dim", "ndimension"} and not (args or kwargs):
             return TypeInfo.from_example(self.tensor.fake_value.ndim, origin)
-        if attr in {"shape", "size"} and not kwargs:
+        if attr in {"shape", "size", "stride"} and not kwargs:
             fn = getattr(self.tensor.fake_value, attr)
             try:
                 return TypeInfo.from_example(

helion/exc.py

@@ -186,7 +186,7 @@ class SpecializeOnDevice(BaseError):
 
 
 class SpecializeArgType(BaseError):
-    message = "hl.specialize() must be called on a size from an input tensor, got: {}"
+    message = "hl.specialize() must be called on a size or stride from an input tensor, got: {}"
 
 
 class StackTensorcOnHost(BaseError):

helion/language/constexpr.py

@@ -87,7 +87,22 @@ def _(value: TypeInfo, *, origin: Origin) -> TypeInfo:
     env = CompileEnvironment.current()
 
     def handle_symint(symint: torch.SymInt) -> int:
-        env.specialized_vars.update(symint._sympy_().free_symbols)
+        from torch._dynamo.source import TensorProperty
+        from torch._dynamo.source import TensorPropertySource
+
+        syms = symint._sympy_().free_symbols
+        env.specialized_vars.update(syms)
+        # Track which strides were specialized (only stride-only symbols)
+        for sym in syms:
+            if sym in env.stride_symbols:
+                # Check if this symbol is also a size via var_to_sources
+                sources = env.shape_env.var_to_sources.get(sym, [])
+                is_size = any(
+                    isinstance(s, TensorPropertySource) and s.prop == TensorProperty.SIZE
+                    for s in sources
+                )
+                if not is_size:
+                    env.specialized_strides.add(env.stride_symbols[sym])
         return symint.__int__()
 
     specialized = _convert_specializable(proxy, on_symint=handle_symint)

helion/runtime/kernel.py

@@ -624,12 +624,14 @@ def _specialize_extra(self) -> list[Callable[[Sequence[object]], Hashable]]:
 
         def make_extractor(v: Source) -> Callable[[Sequence[object]], Hashable]:
             if isinstance(v, TensorPropertySource):
-                assert v.prop == TensorProperty.SIZE
                 index = v.idx
                 assert index is not None
                 inner = make_extractor(v.base)
-
-                return lambda args: cast("torch.Tensor", inner(args)).size(index)
+                if v.prop == TensorProperty.SIZE:
+                    return lambda args: cast("torch.Tensor", inner(args)).size(index)
+                if v.prop == TensorProperty.STRIDE:
+                    return lambda args: cast("torch.Tensor", inner(args)).stride(index)
+                raise exc.SpecializeArgType(v)
             if isinstance(v, LocalSource):
                 index = arg_name_to_index[v.local_name]
                 return operator.itemgetter(index)

test/test_examples.expected

@@ -1674,18 +1674,18 @@ def _helion_fused_linear_jsd_kernel(student_logits, teacher_logits, loss, temper
     # src[fused_linear_jsd.py:N]: teacher_div = torch.nn.functional.kl_div(
     # src[fused_linear_jsd.py:N]:     torch.log(m), teacher_prob, reduction="none", log_target=True
     # src[fused_linear_jsd.py:N]: ).sum(dim=-1)
-    v_17 = teacher_prob_1 - v_16
-    v_18 = libdevice.exp(teacher_prob_1)
-    v_19 = v_18 * v_17
+    v_17 = libdevice.exp(teacher_prob_1)
+    v_18 = teacher_prob_1 - v_16
+    v_19 = v_17 * v_18
     teacher_div = tl.cast(tl.sum(v_19, 1), tl.float32)
     # src[fused_linear_jsd.py:N]: torch.log(m), student_prob, reduction="none", log_target=True
     v_20 = tl_math.log(v_15)
     # src[fused_linear_jsd.py:N]: student_div = torch.nn.functional.kl_div(
     # src[fused_linear_jsd.py:N]:     torch.log(m), student_prob, reduction="none", log_target=True
     # src[fused_linear_jsd.py:N]: ).sum(dim=-1)
-    v_21 = student_prob_1 - v_20
-    v_22 = libdevice.exp(student_prob_1)
-    v_23 = v_22 * v_21
+    v_21 = libdevice.exp(student_prob_1)
+    v_22 = student_prob_1 - v_20
+    v_23 = v_21 * v_22
     student_div = tl.cast(tl.sum(v_23, 1), tl.float32)
     # src[fused_linear_jsd.py:N]: batch_loss = student_div + beta * (teacher_div - student_div)
     v_24 = teacher_div - student_div