Allowing static constraint in torch/__init__.py (#613)

Author: jjsjann123

thunder/core/prims.py

@@ -261,6 +261,8 @@ class PrimIDs(Enum):
     # Memory access methods
     ITEM = auto()
     COPY_ = auto()
+    #
+    SINK = auto()
 
 
 class OpTags(Enum):
@@ -3888,3 +3890,11 @@ def copy__meta(
 
 
 copy_ = make_prim(PrimIDs.COPY_, "copy_", meta=copy__meta, tags=(OpTags.DONT_DCE,))
+
+
+def sink_meta(*args, **kwargs):
+    return
+
+
+# TODO do we want another tag to remove this after prologue is constructed?
+sink = make_prim(PrimIDs.SINK, "sink", meta=sink_meta, tags=(OpTags.DONT_DCE,))

thunder/core/proxies.py

@@ -612,6 +612,7 @@ def known_value(self) -> bool:
 
     def make_static_constrained(self):
         baseutils.check(self.constraint != CONSTRAINT.DYNAMIC, lambda: f"dynamic NumberProxy cannot be made static")
+        baseutils.check(self.value is not None, lambda: f"static NumberProxy needs to have value")
         self.constraint = CONSTRAINT.STATIC
 
     def make_constrainable(self):

thunder/executors/pythonex.py

@@ -365,5 +365,13 @@ def _elementwise_binary_checker(a: NumberLike | TensorProxy, b: NumberLike | Ten
 ex.register_implementation(prims.sub, sub, checker=_elementwise_binary_checker)
 ex.register_implementation(prims.div, div, checker=_elementwise_binary_checker)
 
+
+def _sink(*args, **kwargs):
+    return
+
+
+sink = ex.register_operator("sink", like=prims.sink, fn=_sink)
+ex.register_implementation(prims.sink, sink, checker=_always_executable)
+
 # TODO: Restore truediv once we find it...
 # ex.register_implementation(prims.truediv, truediv, checker=_elementwise_binary_checker)

thunder/tests/test_jit_general.py

@@ -955,3 +955,18 @@ def bar(t):
         jbar = thunder.jit(bar, cache="symbolic values")
         t = torch.randn(4, device="cpu")
         jbar(t)
+
+
+def test_cache_symbolic_values_torch_device():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
+    def foo(dev, idx):
+        # NOTE dtype needs to be explicit, see issue: https://github.com/Lightning-AI/lightning-thunder/issues/621
+        return torch.ones(1, device=torch.device(dev, idx), dtype=torch.float32)
+
+    jfoo = thunder.jit(foo, cache="symbolic values")
+    expected = foo("cuda", 0)
+    actual = jfoo("cuda", 0)
+
+    assert_close(expected, actual)

thunder/torch/__init__.py

@@ -4906,6 +4906,10 @@ def torch_device(device_or_str: DeviceLike, /, index: int | None = None) -> devi
             not has_device_idx,
             lambda: f"device string must not include an index because index was passed explicitly: {device_or_str}",
         )
+        if isinstance(index, NumberProxy):
+            index.make_static_constrained()
+            prims.sink(index)
+            index = index.value
 
     return devices.Device(device_or_str, index)