Update on "NT support for narrow() on dim=0"

jbschlosser · jbschlosser · commit aa5b50362726 · 2023-09-01T13:11:35.000-04:00
Satisfies request here: #105913 (comment) [ghstack-poisoned]
diff --git a/aten/src/ATen/native/nested/NestedTensorMatmul.cpp b/aten/src/ATen/native/nested/NestedTensorMatmul.cpp
@@ -190,6 +190,27 @@ static Tensor matmul_with_bmm_nested(const Tensor& self, const Tensor& mat2) {
 
 }
 
+// nt: NT of shape (B, *, C, D)
+// other: dense tensor of shape (D, E)
+// output: NT of shape (B, *, C, E)
+static Tensor matmul_nested_with_broadcasted_dense(const Tensor& nt, const Tensor& other) {
+  // View nt buffer as 3D jagged for matmul
+  auto *nt_impl = get_nested_tensor_impl(nt);
+  auto jagged = nt_impl->get_buffer().view({-1, nt.size(2), nt.size(3)});
+  auto new_buffer = at::matmul(jagged, other);
+
+  // Wrap result into nested tensor
+  const auto E = other.size(-1);
+  const auto component_dim = nt.dim() - 1;
+  auto new_sizes = nt_impl->get_nested_sizes().clone();
+  auto new_sizes_ptr = new_sizes.data_ptr<int64_t>();
+  for (const auto i : c10::irange(nt.size(0))) {
+    new_sizes_ptr[i * component_dim + 2] = E;
+  }
+  return at::detail::make_tensor<NestedTensorImpl>(
+      new_buffer.view(-1), new_sizes);
+}
+
 // Note [nested tensor matmul]
 // This is really a generalized batched matmul dedicated to nested tensors,
 // where `self` and `mat2` have same number (>= 3) of dimensions.
@@ -200,6 +221,18 @@ static Tensor matmul_with_bmm_nested(const Tensor& self, const Tensor& mat2) {
 // for each batch dimension `self` and `mat2` must have same size.
 // TODO: Should make full matmul semantics support some day
 Tensor matmul_nested(const Tensor& self, const Tensor& mat2) {
+  // special case of NT (B, *, C, D) with broadcasted dense (D, E)
+  if (self.is_nested() &&
+          self.is_contiguous() &&
+          !mat2.is_nested() &&
+          self.dim() == 4 &&
+          mat2.dim() == 2 &&
+          get_nested_tensor_impl(self)->opt_size(2).has_value() &&
+          get_nested_tensor_impl(self)->opt_size(3).has_value() &&
+          self.size(3) == mat2.size(0)
+     ) {
+    return matmul_nested_with_broadcasted_dense(self, mat2);
+  }
   if (self.is_nested() && !mat2.is_nested()) {
     AT_ERROR("Expected both to be nested, but got a nested self and non-nested other");
   }
diff --git a/test/test_nestedtensor.py b/test/test_nestedtensor.py
@@ -1706,6 +1706,19 @@ def test_matmul(self, device, dtype):
         expect = torch.matmul(torch.nested.to_padded_tensor(nt0, 0.0), torch.nested.to_padded_tensor(nt1, 0.0))
         self.assertEqual(actual, expect)
 
+    # only supported on CUDA for now
+    @dtypes(torch.float, torch.double)
+    def test_matmul_nt_with_broadcasted_t(self, device, dtype):
+        # NT (B, *, C, D) with T (D, E) broadcasting case
+        nt = random_nt_from_dims([3, None, 4, 5], device=device, dtype=dtype)
+        t = torch.randn(5, 6, device=device, dtype=dtype)
+        output = torch.matmul(nt, t)
+
+        # should be equivalent to matmul-ing each component with the dense tensor
+        self.assertEqual(nt.size(0), output.size(0))
+        for component, out_component in zip(nt, output):
+            self.assertEqual(out_component, torch.matmul(component, t))
+
     # cannot test torch.float16 because: RuntimeError: "bmm" not implemented for 'Half'
     @dtypes(torch.float, torch.double)
     def test_matmul_noncontiguous(self, device, dtype):
@@ -2022,6 +2035,11 @@ def test_narrow(self, device, dtype):
             with self.assertRaisesRegex(RuntimeError, "only dim=0 supported for nested tensors"):
                 nt.narrow(dim=dim, start=0, length=1)
 
+        # error case: non-contiguous NT
+        _, nt_noncont = random_nt_noncontiguous_pair((2, 3, 4))
+        with self.assertRaisesRegex(RuntimeError, "only contiguous nested tensors supported"):
+            nt_noncont.narrow(dim=0, start=0, length=1)
+
     @parametrize("input_dim", [3, 4])
     def test_scaled_dot_product_attention(self, device, input_dim):
 
