Use bounds_check_indices v2 on ROCm (#3916)

Summary:
X-link: facebookresearch/FBGEMM#1005

Pull Request resolved: #3916

This diff forces using bounds_check_indices v2 on ROCm because ROCm
has a constraint that the gridDim * blockDim has to be smaller than
2^32. The v1 kernel can be launched with gridDim * blockDim > 2^32
while the v2 kernel limits the gridDim size to 64 * # of SMs.  Thus,
its gridDim * blockDim is guaranteed to be smaller than 2^32

Reviewed By: q10, jianyuh, joebos

Differential Revision: D72334377

fbshipit-source-id: 9c955b691e4462721d500b5b643e037d71e13e0c

Author: sryap

fbgemm_gpu/codegen/utils/embedding_bounds_check_host.cpp

@@ -60,13 +60,21 @@ void bounds_check_indices_cuda(
     const int64_t info_B_num_bits,
     const int64_t info_B_mask,
     const int8_t bounds_check_version) {
+#if USE_ROCM
+  // Force using bounds_check_indices v2 on ROCm because ROCm has a constraint
+  // that the gridDim * blockDim has to be smaller than 2^32. The v1 kernel can
+  // be launched with gridDim * blockDim > 2^32 while the v2 kernel limits the
+  // gridDim size to 64 * # of SMs.  Thus, its gridDim * blockDim is guaranteed
+  // to be smaller than 2^32
+  const auto bounds_check_indices_fn = _bounds_check_indices_cuda_v2;
+#else
   TORCH_CHECK(bounds_check_version == 1 || bounds_check_version == 2);
-  const static bool use_v2 =
-      fbgemm_gpu::config::is_feature_enabled(
-          fbgemm_gpu::config::FeatureGateName::BOUNDS_CHECK_INDICES_V2) ||
-      bounds_check_version == 2;
-  const auto bounds_check_indices_fn =
-      use_v2 ? _bounds_check_indices_cuda_v2 : _bounds_check_indices_cuda_v1;
+  const static bool use_v2 = fbgemm_gpu::config::is_feature_enabled(
+      fbgemm_gpu::config::FeatureGateName::BOUNDS_CHECK_INDICES_V2);
+  const auto bounds_check_indices_fn = (use_v2 || bounds_check_version == 2)
+      ? _bounds_check_indices_cuda_v2
+      : _bounds_check_indices_cuda_v1;
+#endif
   bounds_check_indices_fn(
       rows_per_table,
       indices,

fbgemm_gpu/fbgemm_gpu/sparse_ops.py

@@ -829,6 +829,7 @@ def bounds_check_indices_abstract(
     b_t_map: Optional[torch.Tensor] = None,
     info_B_num_bits: int = -1,
     info_B_mask: int = -1,
+    bounds_check_version: int = 1,
 ) -> None:
     """
     This meta function is used to fake the bounds checking

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

@@ -716,6 +716,16 @@ def __init__(  # noqa C901
             # See:
             #   https://fb.workplace.com/groups/fbgemmusers/permalink/9438488366231860/
             cache_precision = SparseType.FP32
+            self.log("Override cache_precision=SparseType.FP32 on ROCm")
+
+            # NOTE: Use bounds_check_indices v2 on ROCm because ROCm has a
+            # constraint that the gridDim * blockDim has to be smaller than
+            # 2^32. The v1 kernel can be launched with gridDim * blockDim >
+            # 2^32 while the v2 kernel limits the gridDim size to 64 * # of
+            # SMs.  Thus, its gridDim * blockDim is guaranteed to be smaller
+            # than 2^32
+            self.bounds_check_version = 2
+            self.log("Override bounds_check_version=2 on ROCm")
         else:
             # NOTE: The changes from D65865527 are retained here until we can
             # test that the the hack also works for non-ROCm environments.