Support torch.Tensor.view(dtype) (#2213)

Author: crcrpar

thunder/core/prims.py

@@ -279,6 +279,7 @@ class PrimIDs(Enum):
     # Memory access methods
     ITEM = auto()
     COPY_ = auto()
+    BITCAST = auto()
     #
     SINK = auto()
 
@@ -4334,6 +4335,27 @@ def copy__meta(
 copy_ = make_prim(PrimIDs.COPY_, "copy_", meta=copy__meta, tags=(OpTags.DONT_DCE,))
 
 
+def bitcast_meta(
+    src: TensorProxy,
+    dtype: dtypes.dtype,
+) -> TensorProxy:
+    shape = list(src.shape)
+    src_itemsize = src.dtype.bytes
+    dst_itemsize = dtype.bytes
+    if src_itemsize != dst_itemsize:
+        factor = dst_itemsize / src_itemsize
+        if factor > 1:
+            utils.check(
+                shape[-1] > factor and shape[-1] % factor == 0,
+                lambda: f"{src.shape[-1]=} is not divisible by {factor=}. Viewing {src.dtype=} as {dtype=}",
+            )
+        shape[-1] = int(shape[-1] / factor)
+    return TensorProxy(shape=tuple(shape), device=src.device, dtype=dtype)
+
+
+bitcast = make_prim(PrimIDs.BITCAST, "bitcast", meta=bitcast_meta)
+
+
 def sink_meta(*args, **kwargs):
     return
 

thunder/core/transforms.py

@@ -1649,6 +1649,7 @@ def zeros_like(x):
     prims.PrimIDs.FMOD: lambda x, y: (prims.fmod(x, y), (x, y)),
     prims.PrimIDs.COPY_: lambda x, y, grad_enabled: (prims.copy_(x, y, grad_enabled=grad_enabled), tuple()),
     prims.PrimIDs.CLONE: lambda x: (prims.clone(x), tuple()),
+    prims.PrimIDs.BITCAST: lambda x, dtype: (prims.bitcast(x, dtype), (x.dtype,)),
 }
 
 
@@ -1679,6 +1680,7 @@ def zeros_like(x):
     prims.PrimIDs.FMOD: lambda x, y, g: (g, -g * prims.trunc(x / y)),
     prims.PrimIDs.COPY_: lambda g: (g, None),
     prims.PrimIDs.CLONE: lambda g: g,
+    prims.PrimIDs.BITCAST: lambda x_dtype, g: (prims.bitcast(g, x_dtype), None),
 }
 
 

thunder/executors/nvfuserex_impl.py

@@ -970,6 +970,25 @@ def convert_element_type(
 
 register_supported(PrimIDs.CONVERT_ELEMENT_TYPE, convert_element_type, _convert_element_type_check)
 
+
+def _bitcast_check(src: TensorProxy, dtype: dtypes.dtype) -> bool:
+    return (
+        nvfuser_version() > LooseVersion("0.29.0")
+        and _convert_element_type_check(src, dtype)
+        and src.dtype.bytes == dtype.bytes
+    )
+
+
+# TODO: Expose bitcast in nvfuser to Python.
+# def bitcast(src: TensorProxy, dtype: dtypes.dtype, *, fd: FusionDefinition, lc_to_nv_map: dict):
+#     nva = getnv(src, fd, lc_to_nv_map)
+#     nvdtype = lcdtype_to_nvdtype(dtype)
+#
+#     return fd.ops.bitcast(nva, nvdtype)
+#
+#
+# register_supported(PrimIDs.BITCAST, bitcast, _bitcast_check)
+
 #
 # Tensor creation operations
 #

thunder/executors/torchex.py

@@ -2376,6 +2376,14 @@ def _shape_impl(t):
 _register_implementation(prims.shape, shape, checker=_always_executable)
 
 
+def _bitcast_impl(src, dtype):
+    return src.view(dtypes.to_torch_dtype(dtype))
+
+
+bitcast = ex.register_operator("bitcast", meta=prims.bitcast, fn=_bitcast_impl)
+_register_implementation(prims.bitcast, bitcast, checker=_always_executable)
+
+
 shallow_copy = ex.register_operator("shallow_copy", meta=prims.shallow_copy, fn=lambda x: x)
 _register_implementation(prims.shallow_copy, shallow_copy, checker=_always_executable)
 

thunder/tests/opinfos.py

@@ -3513,6 +3513,21 @@ def to_sample_generator(op, device, dtype, requires_grad, **kwargs):
 data_movement_ops.append(to_opinfo)
 
 
+def view_with_dtype_sample_generator(op, device, dtype, requires_grad, **kwargs):
+    make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
+
+    for dst_dtype in {torch.float32, torch.bfloat16, torch.float64} - {dtype}:
+        yield SampleInput(make((8, 8)), dtype)
+
+
+view_with_dtype_opinfo = OpInfo(
+    ltorch.view,
+    sample_input_generator=view_with_dtype_sample_generator,
+    torch_reference=torch.Tensor.view,
+)
+data_movement_ops.append(view_with_dtype_opinfo)
+
+
 def cuda_sample_generator(op, device, dtype, requires_grad, **kwargs):
     make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
 

thunder/tests/test_grad.py

@@ -408,7 +408,11 @@ def test_vjp_correctness(op, device, dtype, executor, comp):
         # sample.thunder() line below attempts to approximate those conversions
         # for non-differentiable arguments like dtypes so that the test will
         # execute properly.
-        sample = sample.thunder()  # converts torch.dtype to thunder.dtype
+        # NOTE: While `convert_element_type` is skipeed as of https://github.com/Lightning-AI/lightning-thunder/pull/2213
+        # as in https://github.com/Lightning-AI/lightning-thunder/blob/dbf6bad3/thunder/tests/opinfos.py#L3324-L3346,
+        # `torch.Tensor.view(dtype)` seems to require `torch.dtype` to be kept as is, opposite to `convert_element_type`.
+        if op.name != "view":
+            sample = sample.thunder()  # converts torch.dtype to thunder.dtype
         sample = sample.remove_singularities(op, eps)
 
         flat_op, flat_args, spec = flatten_func(op.op, sample.args, sample.kwargs)

thunder/torch/__init__.py

@@ -1495,9 +1495,12 @@ def unsqueeze(a: TensorLike, /, dim: int) -> TensorLike:
 
 
 # TODO Add type annotations
+# TODO(crcrpar): see [Return value of view creation ops]
 @torchsymbol(torch.Tensor.view, is_method=True)
 def view(a: TensorLike, /, *shape) -> TensorLike:
     shape = utils.extract_shape_from_varargs(shape)
+    if len(shape) == 1 and isinstance(shape[0], (torch.dtype, dtypes.dtype)):
+        return prims.bitcast(a, dtype=to_dtype(shape[0]))
     return reshape(a, shape)
 
 
@@ -6861,6 +6864,11 @@ def check_overlap_ops():
     _torch_to_thunder_function_map[torch.Tensor.reshape_as],
 }
 
+# TODO(crcrpar): [Return value of view creation ops]
+# Review what's more appropriate return value from the ops below.
+# Currently they return a new tensor, which obscures the nature of these ops, i.e.,
+# outputs share underlying storage with inputs. For more stable and improved in-place support
+# it'd be necessary to think about e.g. extending TensorProxy and/or DCE.
 _syms_returning_views: set[Symbol] = {
     diagonal,
     expand,