attempt to write specialized reshape code for 1d <-> 2d (#2075)

Summary:
Pull Request resolved: #2075

This diff contains torchrec changes that, together with a separate PyTorch PR, will ensure that uniform batch test with view conversions work in the existing torchrec tests.

I am deliberately not including the changes in test_pt2_multiprocess.py that I used to test locally on my devserver. I will let the torchrec team decide if they want those changes in this diff or they prefer to do it themselves in a separate diff.

Reviewed By: IvanKobzarev

Differential Revision: D58200241

fbshipit-source-id: b68e2cebce7e75b34a36f9ecf6ff546b3a1a3d8c

torchrec/distributed/tests/test_pt2_multiprocess.py

@@ -353,7 +353,9 @@ def _test_compile_rank_fn(
             torch._dynamo.config.capture_scalar_outputs = True
             torch._dynamo.config.capture_dynamic_output_shape_ops = True
             opt_fn = torch.compile(
-                dmp, backend=torch_compile_backend, fullgraph=True, dynamic=True
+                dmp,
+                backend=torch_compile_backend,
+                fullgraph=True,
             )
             compile_out = opt_fn(kjt_for_pt2_tracing(kjt, convert_to_vb=convert_to_vb))
             torch.testing.assert_close(eager_out, compile_out)
@@ -403,6 +405,13 @@ def disable_cuda_tf32(self) -> bool:
                     _ConvertToVariableBatch.TRUE,
                     "eager",
                 ),
+                (
+                    _ModelType.EBC,
+                    ShardingType.TABLE_WISE.value,
+                    _InputType.SINGLE_BATCH,
+                    _ConvertToVariableBatch.FALSE,
+                    "eager",
+                ),
             ]
         ),
     )

torchrec/pt2/checks.py

@@ -11,6 +11,8 @@
 
 import torch
 
+from torch.fx.experimental.symbolic_shapes import guard_size_oblivious
+
 USE_TORCHDYNAMO_COMPILING_PATH: bool = False
 
 
@@ -74,3 +76,19 @@ def pt2_checks_all_is_size(list: List[int]) -> List[int]:
     for i in list:
         torch._check_is_size(i)
     return list
+
+
+def pt2_check_size_nonzero(x: torch.Tensor) -> torch.Tensor:
+    if torch.jit.is_scripting() or not is_torchdynamo_compiling():
+        return x
+
+    for i in range(x.dim()):
+        torch._check(x.size(i) > 0)
+    return x
+
+
+def pt2_guard_size_oblivious(x: bool) -> bool:
+    if torch.jit.is_scripting() or not is_torchdynamo_compiling():
+        return x
+
+    return guard_size_oblivious(x)

torchrec/sparse/jagged_tensor.py

@@ -20,8 +20,10 @@
 from torchrec.pt2.checks import (
     is_non_strict_exporting,
     is_torchdynamo_compiling,
+    pt2_check_size_nonzero,
     pt2_checks_all_is_size,
     pt2_checks_tensor_slice,
+    pt2_guard_size_oblivious,
 )
 from torchrec.streamable import Pipelineable
 
@@ -878,8 +880,10 @@ def _maybe_compute_length_per_key(
                 _length_per_key_from_stride_per_key(lengths, stride_per_key)
                 if variable_stride_per_key
                 else (
-                    torch.sum(lengths.view(-1, stride), dim=1).tolist()
-                    if lengths.numel() != 0
+                    torch.sum(
+                        pt2_check_size_nonzero(lengths.view(len(keys), stride)), dim=1
+                    ).tolist()
+                    if pt2_guard_size_oblivious(lengths.numel() != 0)
                     else [0] * len(keys)
                 )
             )