Add more helper methods for TBE benchmarking (pytorch#828)

Summary:
Pull Request resolved: facebookresearch/FBGEMM#828

X-link: pytorch#3747

- Integrate TBEDataConfig into TBE device benchmark CLI

Reviewed By: sryap

Differential Revision: D70296132

fbshipit-source-id: 7ebc3971296d6bba246070137d2c67383bd69ae3

Author: q10

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_common.py

@@ -34,6 +34,20 @@ class EmbeddingLocation(enum.IntEnum):
     MTIA = 4
 
 
+def str_to_embedding_location(key: str) -> EmbeddingLocation:
+    lookup = {
+        "device": EmbeddingLocation.DEVICE,
+        "managed": EmbeddingLocation.MANAGED,
+        "managed_caching": EmbeddingLocation.MANAGED_CACHING,
+        "host": EmbeddingLocation.HOST,
+        "mtia": EmbeddingLocation.MTIA,
+    }
+    if key in lookup:
+        return lookup[key]
+    else:
+        raise ValueError(f"Cannot parse value into EmbeddingLocation: {key}")
+
+
 class CacheAlgorithm(enum.Enum):
     LRU = 0
     LFU = 1
@@ -57,6 +71,21 @@ class PoolingMode(enum.IntEnum):
     MEAN = 1
     NONE = 2
 
+    def do_pooling(self) -> bool:
+        return self is not PoolingMode.NONE
+
+
+def str_to_pooling_mode(key: str) -> PoolingMode:
+    lookup = {
+        "sum": PoolingMode.SUM,
+        "mean": PoolingMode.MEAN,
+        "none": PoolingMode.NONE,
+    }
+    if key in lookup:
+        return lookup[key]
+    else:
+        raise ValueError(f"Cannot parse value into PoolingMode: {key}")
+
 
 class BoundsCheckMode(enum.IntEnum):
     # Raise an exception (CPU) or device-side assert (CUDA)