[fbcode] Use FastCat in PT Concat implementation (#106727)

zoranzhao · pytorchmergebot · commit d392963ac4fb · 2023-08-13T22:36:51.000Z
Summary: Reimplement D48081898 and PR #106518 in fbcode first to accelerate the launching process Test Plan: All checks have been passed: https://github.com/pytorch/pytorch/actions/runs/5758987335/job/15612600466?pr=106518 (For my own learning purpose) Check out OSS PyTorch repo and test following the instructions in https://www.internalfb.com/intern/wiki/PyTorch/PyTorchDev/Workflow/PyTorch_environment_setup/ and https://www.internalfb.com/intern/wiki/PyTorch/PyTorchDev/Workflow/PyTorch_environment_setup/oss_setup_on_devserver : ``` pytest -k test_cat_out test/test_tensor_creation_ops.py -v -s ``` To submit to GitHub ``` hg amend; jf submit; ghexport ``` Differential Revision: D48082741 Pull Request resolved: #106727 Approved by: https://github.com/ezyang, https://github.com/houseroad
diff --git a/aten/src/ATen/native/TensorShape.cpp b/aten/src/ATen/native/TensorShape.cpp
@@ -549,6 +549,21 @@ std::vector<Tensor> broadcast_tensors(TensorList tensors) {
   return expand_outplace(tensors);
 }
 
+static void fastCatOutDim0(const Tensor& out, const MaterializedITensorListRef& inputs) {
+  auto outBytes = out.nbytes();
+  char* dataPtr = reinterpret_cast<char*>(out.data_ptr());
+  size_t totalBytes = 0;
+  for (const Tensor& input : inputs) {
+    TORCH_CHECK(outBytes >= totalBytes);
+    if (input.nbytes() > 0) {
+      std::memcpy(dataPtr + totalBytes, input.data_ptr(), input.nbytes());
+    }
+    totalBytes += input.nbytes();
+  }
+  TORCH_CHECK(outBytes == totalBytes);
+}
+
+
 TORCH_IMPL_FUNC(cat_out_cpu)
 (const ITensorListRef& tensors,
  int64_t dim,
@@ -564,10 +579,20 @@ TORCH_IMPL_FUNC(cat_out_cpu)
 
   auto materialized = tensors.materialize();
 
-  // fast path for single thread when both inputs and result are contiguous and not empty
   bool use_serial_kernel = result.numel() < at::internal::GRAIN_SIZE || at::get_num_threads() == 1;
   ScalarType dtype = materialized[valid].get().scalar_type();
   bool serial_dtype = at::isFloatingType(dtype);
+  // fast path for single thread when both inputs and result are contiguous and
+  // not empty, and concat dim is 0
+  if (use_serial_kernel && all_contiguous && all_same_dtype && (MemoryFormat::Contiguous == memory_format)) {
+    if (dim == 0) {
+      fastCatOutDim0(result, materialized);
+      return;
+    }
+    // TODO: Add fast cat for higher dimensions and support multi-threaded fast cat
+  }
+
+  // fast path for single thread when both inputs and result are contiguous and not empty
   if (use_serial_kernel && all_contiguous && all_same_dtype && serial_dtype) {
     cat_serial_stub(kCPU, result, materialized, dim);
     return;
diff --git a/test/test_tensor_creation_ops.py b/test/test_tensor_creation_ops.py
@@ -21,7 +21,7 @@
     onlyCPU, largeTensorTest, precisionOverride, dtypes,
     onlyCUDA, skipCPUIf, dtypesIfCUDA, skipMeta)
 from torch.testing._internal.common_dtype import (
-    all_types_and_complex_and, all_types_and, floating_and_complex_types,
+    all_types_and_complex, all_types_and_complex_and, all_types_and, floating_and_complex_types, complex_types,
     floating_types, floating_and_complex_types_and, integral_types, integral_types_and, get_all_dtypes,
     float_to_corresponding_complex_type_map
 )
@@ -514,6 +514,7 @@ def test_cat_out(self, device):
         a = torch.cat([w[:2], w[4:6]])
         b = torch.cat([w[:2], w[4:6]], out=w[6:10])
         self.assertEqual(a, b)
+        self.assertEqual(a, w[6:10])
         self.assertEqual(w[:6], y.view(-1)[:6])
 
         # Case:
@@ -542,6 +543,68 @@ def test_cat_out(self, device):
             self.assertEqual(y, expected_y)
             self.assertEqual(x, expected_x)
 
+    @dtypes(*all_types_and_complex())
+    def test_cat_out_fast_path_dim0_dim1(self, device, dtype):
+        x = torch.zeros((0), device=device, dtype=dtype)
+        if dtype in integral_types():
+            y = torch.randint(low=0, high=100, size=(4, 6), device=device, dtype=dtype)
+        else:
+            y = torch.randn((4, 6), device=device, dtype=dtype)
+        # Test concat on dimension 0
+        w = y.view(-1).clone()
+        a = torch.cat([w[:2], w[4:6]])
+        b = torch.cat([w[:2], w[4:6]], out=w[6:10])
+        # Note that there is no guarantee that slicing here will result in
+        # contiguous tensors
+        self.assertEqual(a, b)
+        self.assertEqual(a, w[6:10])
+        self.assertEqual(w[:6], y.view(-1)[:6])
+        # If inputs are contiguous tensors, then fast concat paths will be invoked
+        a_fastcat = torch.cat([w[:2].contiguous(), w[4:6].contiguous()])
+        self.assertEqual(a_fastcat, a)
+        # Test concat on dimension 1
+        w = y.clone()
+        w_slices = torch.tensor_split(w, (2, 4), dim=1)
+        # Note that the tensor in w_slices[] here may not be a contiguous
+        # tensor and we need to make sure this is not broken by fast concat
+        b = torch.cat([w_slices[0], w_slices[1]], dim=1)
+        expected_b = torch.index_select(w, 1, torch.tensor([0, 1, 2, 3], device=device))
+        self.assertEqual(b, expected_b)
+        # If inputs are contiguous tensors, then fast concat paths will be invoked
+        b_fastcat = torch.cat([w_slices[0].contiguous(), w_slices[1].contiguous()], dim=1)
+        self.assertEqual(b_fastcat, expected_b)
+        # Finally, we need to make sure backward is not broken
+        # Integral types will not have grad
+        if dtype not in integral_types():
+            a = torch.randn((4, 3), device=device, dtype=dtype, requires_grad=True)
+            b = torch.randn((2, 3), device=device, dtype=dtype, requires_grad=True)
+            c = torch.randn((5, 3), device=device, dtype=dtype, requires_grad=True)
+            d = torch.randn((5, 2), device=device, dtype=dtype, requires_grad=True)
+            expected_a_grad = torch.ones((4, 3), device=device, dtype=dtype)
+            expected_b_grad = torch.ones((2, 3), device=device, dtype=dtype)
+            expected_c_grad = torch.ones((5, 3), device=device, dtype=dtype)
+            expected_d_grad = torch.ones((5, 2), device=device, dtype=dtype)
+            # All the new tensors should be contiguous here. Let us make sure
+            # to explicitly set them contiguous to enforce fast cat
+            dim0_cat = torch.cat([a.contiguous(), b.contiguous()], dim=0)
+            if dtype in complex_types():
+                dim0_cat.sum().abs().backward()
+                self.assertEqual(a.grad.abs(), expected_a_grad.abs())
+                self.assertEqual(b.grad.abs(), expected_b_grad.abs())
+            else:
+                dim0_cat.sum().backward()
+                self.assertEqual(a.grad, expected_a_grad)
+                self.assertEqual(b.grad, expected_b_grad)
+            dim1_cat = torch.cat([c.contiguous(), d.contiguous()], dim=1)
+            if dtype in complex_types():
+                dim1_cat.sum().abs().backward()
+                self.assertEqual(c.grad.abs(), expected_c_grad.abs())
+                self.assertEqual(d.grad.abs(), expected_d_grad.abs())
+            else:
+                dim1_cat.sum().backward()
+                self.assertEqual(c.grad, expected_c_grad)
+                self.assertEqual(d.grad, expected_d_grad)
+
     def test_cat_out_channels_last(self, device):
         x = torch.randn((4, 3, 8, 8))
         y = torch.randn(x.shape)
