Add a host map option for a UVM tensor alloc (#3073)

Summary:
X-link: facebookresearch/FBGEMM#167

Pull Request resolved: #3073

FBGEMM-GPU provides two options to allocate a UVM buffer: (1) `malloc`
& `cudaHostRegister` and (2) `cudaMallocManaged`.  Each one of them
has a different performance implication (this is platform specific).
This diff adds an option for a user to choose a UVM buffer allocation
method in TBE.

Reviewed By: q10, spcyppt

Differential Revision: D61401239

Author: sryap

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

@@ -230,6 +230,7 @@ def apply_split_helper(
     make_dev_param: bool = False,
     dev_reshape: Optional[Tuple[int, ...]] = None,
     uvm_tensors_log: Optional[List[str]] = None,
+    uvm_host_mapped: bool = False,
 ) -> None:
     set_attr_fn(f"{prefix}_physical_placements", split.placements)
     set_attr_fn(f"{prefix}_physical_offsets", split.offsets)
@@ -313,11 +314,12 @@ def apply_split_helper(
                 f"{prefix}_uvm",
                 torch.zeros(
                     split.uvm_size,
-                    out=torch.ops.fbgemm.new_managed_tensor(
+                    out=torch.ops.fbgemm.new_unified_tensor(
                         # pyre-fixme[6]: Expected `Optional[Type[torch._dtype]]`
                         #  for 3rd param but got `Type[Type[torch._dtype]]`.
                         torch.zeros(1, device=current_device, dtype=dtype),
                         [split.uvm_size],
+                        is_host_mapped=uvm_host_mapped,
                     ),
                 ),
             )
@@ -540,6 +542,9 @@ def __init__(  # noqa C901
         multipass_prefetch_config: Optional[MultiPassPrefetchConfig] = None,
         # Global weight decay params
         global_weight_decay: Optional[GlobalWeightDecayDefinition] = None,
+        # Set to True to alloc a UVM tensor using malloc+cudaHostRegister.
+        # Set to False to use cudaMallocManaged
+        uvm_host_mapped: bool = False,
     ) -> None:
         super(SplitTableBatchedEmbeddingBagsCodegen, self).__init__()
         self.uuid = str(uuid.uuid4())
@@ -746,6 +751,8 @@ def __init__(  # noqa C901
             loc == locations[0] for loc in locations
         )
 
+        self.uvm_host_mapped = uvm_host_mapped
+
         weight_split = construct_split_state(
             embedding_specs,
             rowwise=False,
@@ -763,6 +770,7 @@ def __init__(  # noqa C901
             enforce_hbm=enforce_hbm,
             make_dev_param=optimizer == OptimType.NONE,
             dev_reshape=(-1, self.max_D) if optimizer == OptimType.NONE else None,
+            uvm_host_mapped=self.uvm_host_mapped,
         )
 
         assert optimizer not in (
@@ -936,6 +944,7 @@ def __init__(  # noqa C901
                     #  but got `Type[torch.float32]`.
                     dtype=momentum1_dtype,
                     enforce_hbm=enforce_hbm,
+                    uvm_host_mapped=self.uvm_host_mapped,
                 )
             if optimizer in (
                 OptimType.ADAM,
@@ -972,6 +981,7 @@ def __init__(  # noqa C901
                     # pyre-fixme[6]: Expected `Type[Type[torch._dtype]]` for 3rd param
                     #  but got `Type[torch.float32]`.
                     dtype=momentum2_dtype,
+                    uvm_host_mapped=self.uvm_host_mapped,
                 )
             else:
                 # NOTE: make TorchScript work!
@@ -990,6 +1000,7 @@ def __init__(  # noqa C901
                     # pyre-fixme[6]: Expected `Type[Type[torch._dtype]]` for 3rd param
                     #  but got `Type[torch.float32]`.
                     dtype=torch.float32,
+                    uvm_host_mapped=self.uvm_host_mapped,
                 )
                 self._apply_split(
                     construct_split_state(
@@ -1001,6 +1012,7 @@ def __init__(  # noqa C901
                     # pyre-fixme[6]: Expected `Type[Type[torch._dtype]]` for 3rd param
                     #  but got `Type[torch.float32]`.
                     dtype=torch.float32,
+                    uvm_host_mapped=self.uvm_host_mapped,
                 )
                 self.register_buffer(
                     "max_counter", torch.tensor([1], dtype=torch.float32)
@@ -1019,6 +1031,7 @@ def __init__(  # noqa C901
                     # pyre-fixme[6]: Expected `Type[Type[torch._dtype]]` for 3rd param
                     #  but got `Type[torch.float32]`.
                     dtype=torch.float32,
+                    uvm_host_mapped=self.uvm_host_mapped,
                 )
                 self._register_nonpersistent_buffers("row_counter")
                 self.register_buffer(
@@ -2281,6 +2294,7 @@ def _apply_split(
         enforce_hbm: bool = False,
         make_dev_param: bool = False,
         dev_reshape: Optional[Tuple[int, ...]] = None,
+        uvm_host_mapped: bool = False,
     ) -> None:
         apply_split_helper(
             self.register_buffer,
@@ -2295,6 +2309,7 @@ def _apply_split(
             make_dev_param,
             dev_reshape,
             self._uvm_tensors_log,
+            uvm_host_mapped=uvm_host_mapped,
         )
 
     def _apply_cache_state(