fix jagged tensor shared memory issue

Differential Revision: D39335880

fbshipit-source-id: 2f0de679212541fd98bfa48ab3b1763202a3cf4e

fbgemm_gpu/src/jagged_tensor_ops.cu

@@ -613,6 +613,7 @@ __global__ void jagged_dense_dense_elementwise_jagged_output_opt_gather_kernel_(
 bool jagged_dense_dense_elementwise_jagged_output_matches_opt(
     const int& num_jagged_dim,
     const Tensor& x_values,
+    const std::vector<Tensor>& x_offsets,
     const Tensor& y_0_reshaped,
     const Tensor& y_1_reshaped,
     const Tensor& output_values) {
@@ -641,6 +642,16 @@ bool jagged_dense_dense_elementwise_jagged_output_matches_opt(
   matches &= (y_0_reshaped.size(0) < INT_MAX);
   matches &= (y_0_reshaped.size(1) < INT_MAX);
 
+  AT_DISPATCH_INDEX_TYPES(
+      x_offsets[0].scalar_type(), "check_shared_memory", [&] {
+        auto B = y_0_reshaped.size(0);
+        // the default shared memory on V100/A100 is 48 KB from
+        // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#shared-memory-8-x
+        if ((B + 1) * sizeof(index_t) >= 48000) {
+          matches = false;
+        }
+      });
+
   return matches;
 }
 
@@ -705,7 +716,12 @@ void jagged_dense_elementwise_jagged_output_opt_(
   // Canonicalize y to 3D, collapsing jagged dimensions.
   const Tensor y_reshaped = y.view({y.size(0), -1, y.size(-1)});
   if (jagged_dense_dense_elementwise_jagged_output_matches_opt(
-          num_jagged_dim, x_values, y_reshaped, y_reshaped, output_values)) {
+          num_jagged_dim,
+          x_values,
+          x_offsets,
+          y_reshaped,
+          y_reshaped,
+          output_values)) {
     AT_DISPATCH_INDEX_TYPES(
         x_offsets[0].scalar_type(), "jagged_indices_fast_path", [=] {
           auto nnz = output_values.size(0);
@@ -871,6 +887,7 @@ void jagged_dense_dense_elementwise_jagged_output_opt_(
   if (jagged_dense_dense_elementwise_jagged_output_matches_opt(
           num_jagged_dim,
           x_values,
+          x_offsets,
           y_0_reshaped,
           y_1_reshaped,
           output_values)) {

fbgemm_gpu/test/jagged_tensor_ops_test.py

@@ -678,6 +678,34 @@ def test_dense_to_jagged_opt(
             precompute_total_L,
         )
 
+    # pyre-ignore [56]
+    @given(
+        num_jagged_dim=st.just(1),
+        outer_dense_size=st.just(8000),
+        inner_dense_size=st.just(16),
+        dtype=st.just(torch.half),
+        use_cpu=st.just(False),
+        precompute_total_L=st.booleans(),
+    )
+    @settings(verbosity=Verbosity.verbose, max_examples=1, deadline=None)
+    def test_dense_to_jagged_opt_large_batch(
+        self,
+        num_jagged_dim: int,
+        outer_dense_size: int,
+        inner_dense_size: int,
+        dtype: torch.dtype,
+        use_cpu: bool,
+        precompute_total_L: bool,
+    ) -> None:
+        self._test_dense_to_jagged(
+            num_jagged_dim,
+            outer_dense_size,
+            inner_dense_size,
+            dtype,
+            use_cpu,
+            precompute_total_L,
+        )
+
     # pyre-ignore [56]
     @given(
         num_jagged_dim=st.integers(1, 5),