NJT binary pointwise broadcasting support via jagged <-> padded dense conversion (#133021)

jbschlosser · pytorchmergebot · commit 888744bd36d9 · 2024-09-24T19:11:49.000Z
Related: #132695 This PR uses padded dense <-> jagged conversions to handle binary pointwise broadcasting of (NT, T) and (T, NT). This includes: * `(B, j0, D) + (1, 1, 1)` * `(B, j0, D) + (B, 1, 1)` * `(B, j0, D) + (B, 1, D)` * etc. This PR also adds (hacky) support for bool inputs to the jagged <-> padded dense conversions. The underlying CUDA kernels do not support integer / bool inputs; so the following workaround is employed: `convert input -> half, run conversion kernel, convert output -> bool`. Note that this bool support is needed specifically for the backward formula of `fmax`, and likely others. Pull Request resolved: #133021 Approved by: https://github.com/cpuhrsch
diff --git a/test/test_nestedtensor.py b/test/test_nestedtensor.py
@@ -7185,10 +7185,14 @@ def check(nt):
 
         check(nt)
 
-    @dtypes(torch.float32, torch.double, torch.half)
+    @dtypes(torch.float32, torch.double, torch.half, torch.bool)
     @parametrize("nt_dim", [2, 3, 4])
     @parametrize("requires_grad", [False, True])
     def test_to_padded_tensor(self, device, dtype, nt_dim, requires_grad):
+        if dtype is torch.bool and requires_grad:
+            # grads not supported for bool
+            return
+
         if nt_dim == 2:
             post_seq_len_shape = ()
         elif nt_dim == 3:
@@ -7198,7 +7202,9 @@ def test_to_padded_tensor(self, device, dtype, nt_dim, requires_grad):
 
         nt = torch.nested.nested_tensor(
             [
-                torch.randn(n, *post_seq_len_shape, device=device, dtype=dtype)
+                torch.randint(2, (n, *post_seq_len_shape), device=device, dtype=dtype)
+                if dtype is torch.bool
+                else torch.randn(n, *post_seq_len_shape, device=device, dtype=dtype)
                 for n in range(2, 9)
             ],
             layout=torch.jagged,
@@ -7219,7 +7225,7 @@ def test_to_padded_tensor(self, device, dtype, nt_dim, requires_grad):
         nt2 = nested_from_padded(padded, nt.offsets())
         self.assertEqual(nt, nt2)
 
-        if requires_grad:
+        if requires_grad and dtype is not torch.bool:
             # ensure gradients flow through conversions
             nt2.backward(torch.ones_like(nt2))
             self.assertEqual(nt.grad, torch.ones_like(nt))
@@ -7234,6 +7240,10 @@ def test_to_padded_tensor(self, device, dtype, nt_dim, requires_grad):
     @parametrize("nt_dim", [2, 3, 4])
     @parametrize("requires_grad", [False, True])
     def test_to_padded_tensor_compile(self, device, dtype, nt_dim, requires_grad):
+        if dtype is torch.bool and requires_grad:
+            # grads not supported for bool
+            return
+
         if nt_dim == 2:
             post_seq_len_shape = ()
         elif nt_dim == 3:
@@ -7243,7 +7253,9 @@ def test_to_padded_tensor_compile(self, device, dtype, nt_dim, requires_grad):
 
         nt = torch.nested.nested_tensor(
             [
-                torch.randn(n, *post_seq_len_shape, device=device, dtype=dtype)
+                torch.randint(2, (n, *post_seq_len_shape), device=device, dtype=dtype)
+                if dtype is torch.bool
+                else torch.randn(n, *post_seq_len_shape, device=device, dtype=dtype)
                 for n in range(2, 9)
             ],
             layout=torch.jagged,
@@ -7469,9 +7481,6 @@ def g(nt):
     # clone() on non-contiguous with holes NJTs currently use unbind(), leading to
     # data-dependent error in torch.compile
     "clone",
-    # torch._dynamo.exc.Unsupported: data dependent operator: aten._local_scalar_dense.default
-    # for inputs where min / max seqlen are not cached
-    "sum",
 }
 
 COMPARE_TENSOR_COMPONENT_EQUALITY = {
diff --git a/torch/nested/_internal/ops.py b/torch/nested/_internal/ops.py
@@ -283,8 +283,7 @@ def jagged_binary_pointwise(func, *args, **kwargs):
         lhs, rhs = (nt._values, t_squeezed) if a_is_nt else (t_squeezed, nt._values)
         return NestedTensor(func(lhs, rhs, *args[2:], **kwargs), **extracted_kwargs)
 
-    # Harder case: do manual broadcasting over unbound components
-    # when NT dim == non-NT dim
+    # Harder case: do manual broadcasting when NT dim == non-NT dim
     # ex: (B, j0, D_0, D_1) + (B, 1, D_0, D_1) -> (B, j0, D_0, D_1)
     if a.dim() == b.dim():
         # ex: (B, j0, D_0, D_1) + (1, 1, D_0, D_1) -> should
@@ -294,14 +293,39 @@ def jagged_binary_pointwise(func, *args, **kwargs):
                 mismatch_error_msg.format(func.__name__, a.shape, b.shape)
             )
 
-        # need to use offsets to broadcast across ragged dim properly
-        # NB: inefficient fallback here; Triton codegen can help this
-        # TODO: Make this work with autograd
-        outputs = []
-        for a_comp, b_comp in zip(a.unbind(), b.unbind()):
-            outputs.append(func(a_comp, b_comp, *args[2:], **kwargs))
-        new_values = torch.cat(outputs, dim=0)
-        return NestedTensor(new_values, **extracted_kwargs)
+        from .nested_tensor import _load_val_from_tensor, nested_from_padded
+
+        # handle broadcasting via padded dense -> jagged conversion
+        min_seqlen = None
+        if nt._min_seqlen_tensor is not None:
+            min_seqlen = _load_val_from_tensor(nt._min_seqlen_tensor)
+
+        max_seqlen = None
+        if nt._max_seqlen_tensor is not None:
+            max_seqlen = _load_val_from_tensor(nt._max_seqlen_tensor)
+
+        padded_max_S = max_seqlen
+        total_L = nt._values.shape[nt._ragged_idx - 1]
+        if padded_max_S is None:
+            # use upper bound on max seqlen if it's not present
+            padded_max_S = total_L
+
+        # convert dense tensor -> jagged
+        t = t.expand(
+            [x if i != nt._ragged_idx else padded_max_S for i, x in enumerate(t.shape)]
+        )
+        t_as_nt = nested_from_padded(
+            t,
+            offsets=nt._offsets,
+            ragged_idx=nt._ragged_idx,
+            sum_S=total_L,
+            min_seqlen=min_seqlen,
+            max_seqlen=max_seqlen,
+        )
+
+        # function call with two NJTs
+        lhs, rhs = (nt, t_as_nt) if a_is_nt else (t_as_nt, nt)
+        return func(lhs, rhs, *args[2:], **kwargs)
 
     # ex: (B, j0, D_0, D_1) + (A, B, 1, D_0, D_1) -> error because this breaks the invariant
     # that ragged dim is wrt left-most batch dim
@@ -1106,15 +1130,19 @@ def sum_dim_IntList(func, *args, **kwargs):
 
             # _jagged_to_padded_dense_forward requires values to be a 2D tensor
             # with the ragged dimension as the 0th dimension
+            max_seqlen = inp._values.shape[inp._ragged_idx - 1]
+            if inp._max_seqlen_tensor is not None:
+                max_seqlen = inp._max_seqlen
+
             padded = torch.ops.aten._jagged_to_padded_dense_forward(
                 values_ragged_dim_outer.reshape(values_ragged_dim_outer.shape[0], -1),
                 [inp._offsets],
-                max_lengths=[inp._max_seqlen],
+                max_lengths=[max_seqlen],
             )
 
             padded_ragged_dim_original = padded.view(
                 padded.shape[0],
-                inp._max_seqlen,
+                max_seqlen,
                 *values_ragged_dim_outer.shape[
                     1:
                 ],  # expand non-batch dimensions of padded tensor
@@ -1396,7 +1424,7 @@ def native_layer_norm_backward_default(func, *args, **kwargs):
     return (NestedTensor(d_input, **extract_kwargs(inp)), d_gamma, d_beta)
 
 
-@register_jagged_func(torch.ops.aten.select.int, "self: jt, dim: any, index: any")
+@register_jagged_func(torch.ops.aten.select.int, "self: jt_all, dim: any, index: any")
 def select_int(func, *args, **kwargs):
     _, new_kwargs = normalize_function(  # type: ignore[misc]
         func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
@@ -1412,6 +1440,11 @@ def select_int(func, *args, **kwargs):
     if new_kwargs["dim"] == 0:
         return inp.unbind()[new_kwargs["index"]]
 
+    if inp._lengths is not None:
+        raise ValueError(
+            "select(): not yet supported on dim != 0 for non-contiguous nested tensor with holes"
+        )
+
     return NestedTensor(func(inp._values, **new_kwargs), **extract_kwargs(inp))
 
 
@@ -1617,12 +1650,19 @@ def to_padded_tensor_default(func, *args, **kwargs):
     elif values.dim() == 1:
         values = values.unsqueeze(-1)
 
+    # NB: The CUDA kernel for jagged -> padded dense conversion does not support
+    # integer / bool types; work around this by casting to half.
+    is_bool = values.dtype is torch.bool
+    if is_bool and values.is_cuda:
+        values = values.to(torch.half)
     padded_out = torch.ops.aten._jagged_to_padded_dense_forward(
         values,
         [inp._offsets],
         [max_seq_len],
         new_kwargs["padding"],
     )
+    if is_bool and padded_out.is_cuda:
+        padded_out = padded_out.to(torch.bool)
 
     # shape gymnastics part 2
     if len(values_shape) > 2:
@@ -1656,9 +1696,16 @@ def _nested_from_padded_tensor_default(func, *args, **kwargs):
     elif padded.dim() < 3:
         padded = padded.unsqueeze(-1)
 
+    # NB: The CUDA kernel for padded dense -> jagged conversion does not support
+    # integer / bool types; work around this by casting to half.
+    is_bool = padded.dtype is torch.bool
+    if is_bool and padded.is_cuda:
+        padded = padded.to(torch.half)
     values = torch.ops.aten._padded_dense_to_jagged_forward(
         padded, [offsets], new_kwargs["sum_S"]
     )
+    if is_bool and values.is_cuda:
+        values = values.to(torch.bool)
 
     # shape gymnastics part 2
     if len(padded_shape) > 3:
@@ -1791,6 +1838,24 @@ def masked_select_default(func, *args, **kwargs):
     )
 
 
+@register_jagged_func(
+    torch.ops.aten._nested_select_backward.default,
+    "grad_output: t, self: jt, dim: any, index: any",
+)
+def _nested_select_backward_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    grad_output = new_kwargs.pop("grad_output")
+
+    grad_input = torch.zeros_like(inp, dtype=grad_output.dtype)
+    grad_input.select(new_kwargs["dim"], new_kwargs["index"]).copy_(grad_output)
+
+    return grad_input
+
+
 # Make the dummy available on the C++ side.
 @register_jagged_func(torch.ops.aten._nested_get_jagged_dummy.default, "self: any")
 def _nested_get_jagged_dummy(func, *args, **kwargs):
diff --git a/torch/testing/_internal/opinfo/definitions/nested.py b/torch/testing/_internal/opinfo/definitions/nested.py
@@ -76,32 +76,47 @@ def _sample_njts(device, dtype, requires_grad=False, dims=None):
 # This reference unbinds the input NJT and invokes the op on each of the components,
 # optionally wrapping the result in an NJT.
 def unbind_reference(op, sample, wrap_output_as_njt=True):
-    assert sample.input.is_nested
+    # first NJT in the arglist determines expected ragged structure
+    nt_inp = (
+        sample.input
+        if sample.input.is_nested
+        # TODO: look in kwargs too?
+        else next(a for a in sample.args if a.is_nested)
+    )
+
     out_ref_components = []
-    for i, component in enumerate(sample.input.unbind(dim=0)):
+    for i in range(nt_inp.shape[0]):
 
-        def _slice_njts(t, i=i, inp=sample.input):
+        def _slice_input(t, i=i, inp=nt_inp):
             # any NJT with the same ragged structure as the input should
-            # also be sliced to pass to the reference
+            # be sliced to pass to the reference
             if isinstance(t, torch.Tensor) and _raggedness_matches(t, inp):
                 return t[i]
+            # allow the SampleInput to tell us how to slice it for ref calculation
+            elif isinstance(t, torch.Tensor) and hasattr(t, "_batch_dim"):
+                bdim = t._batch_dim  # type: ignore[attr]
+                if t.shape[bdim] == 1:
+                    return t[0]
+                else:
+                    return t.select(bdim, i)
             else:
                 return t
 
-        args = tree_map(_slice_njts, sample.args)
-        kwargs = tree_map(_slice_njts, sample.kwargs)
+        inp = _slice_input(sample.input)
+        args = tree_map(_slice_input, sample.args)
+        kwargs = tree_map(_slice_input, sample.kwargs)
 
         from torch._prims_common import canonicalize_dims
 
         # Need to adjust dim to apply on NJT component
         if "dim" in kwargs:
-            kwargs["dim"] = canonicalize_dims(sample.input.dim(), kwargs["dim"]) - 1
+            kwargs["dim"] = canonicalize_dims(nt_inp.dim(), kwargs["dim"]) - 1
             assert kwargs["dim"] >= 0
 
         # TODO: handle this
         assert "dims" not in kwargs
 
-        out_ref_component = op.op(component, *args, **kwargs)
+        out_ref_component = op.op(inp, *args, **kwargs)
 
         # TODO: handle list / tuple / non-NJT outputs
         assert not isinstance(out_ref_component, (list, tuple))
@@ -157,10 +172,72 @@ def sample_inputs_elementwise_njt_binary(
         device=device, dtype=dtype, requires_grad=requires_grad, dims=[2, 3, 4]
     ):
         # TODO: account for non-contiguous NJTs here
-        # TODO: provide sample inputs for broadcasting cases and mixed (NT, T), (T, NT) inputs
         njt2 = torch.randn_like(njt1)
         yield SampleInput(njt1, args=(njt2,), kwargs=dict(op_kwargs))
 
+        # broadcasting case: (B, j0, ...) with (B, 1, ...)
+        t = torch.randn(
+            (njt1.shape[0], 1, *njt1.shape[2:]),
+            device=device,
+            dtype=dtype,
+            requires_grad=requires_grad,
+        )
+        # used for slicing in unbind_reference()
+        t._batch_dim = 0
+        # (NT, T)
+        yield SampleInput(njt1, args=(t,), kwargs=dict(op_kwargs))
+        # (T, NT)
+        yield SampleInput(t, args=(njt1,), kwargs=dict(op_kwargs))
+
+        # broadcasting case: (B, j0, ...) with (1, 1...)
+        t = torch.randn(
+            [1 for _ in range(njt1.dim())],
+            device=device,
+            dtype=dtype,
+            requires_grad=requires_grad,
+        )
+        # used for slicing in unbind_reference()
+        t._batch_dim = 0
+        # (NT, T)
+        yield SampleInput(njt1, args=(t,), kwargs=dict(op_kwargs))
+        # (T, NT)
+        yield SampleInput(t, args=(njt1,), kwargs=dict(op_kwargs))
+
+        # broadcasting case: (B, j0, ...) with (...)
+        t = torch.randn(
+            njt1.shape[2:], device=device, dtype=dtype, requires_grad=requires_grad
+        )
+        # (NT, T)
+        yield SampleInput(njt1, args=(t,), kwargs=dict(op_kwargs))
+        # (T, NT)
+        yield SampleInput(t, args=(njt1,), kwargs=dict(op_kwargs))
+
+        # broadcasting case: (B, j0, ...) with scalar
+        t = torch.randn((), device=device, dtype=dtype, requires_grad=requires_grad)
+        # (NT, T)
+        yield SampleInput(njt1, args=(t,), kwargs=dict(op_kwargs))
+        # (T, NT)
+        yield SampleInput(t, args=(njt1,), kwargs=dict(op_kwargs))
+
+    # mixed broadcasting case: (B, j0, 1) with (B, 1, D)
+    B = 4
+    D = 16
+    njt = random_nt_from_dims(
+        (B, None, 1),
+        device=device,
+        dtype=dtype,
+        requires_grad=requires_grad,
+        layout=torch.jagged,
+    )
+    t = torch.randn(B, 1, D, device=device, dtype=dtype, requires_grad=requires_grad)
+    # used for slicing in unbind_reference()
+    t._batch_dim = 0
+
+    # (NT, T)
+    yield SampleInput(njt, args=(t,), kwargs=dict(op_kwargs))
+    # (T, NT)
+    yield SampleInput(t, args=(njt,), kwargs=dict(op_kwargs))
+
 
 def sample_inputs_njt_reduction(
     op_info, device, dtype, requires_grad, op_kwargs=None, **kwargs
@@ -311,7 +388,7 @@ def sample_inputs_nn_functional_rms_norm(
     op_info, device, dtype, requires_grad, **kwargs
 ):
     for njt in _sample_njts(
-        device=device, dtype=dtype, requires_grad=requires_grad, dims=[2, 3, 4]
+        device=device, dtype=dtype, requires_grad=requires_grad, dims=[3, 4]
     ):
         # normalize over non-ragged dims
         for start_dim in range(2, njt.dim()):
