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

Summary:
X-link: facebookresearch/FBGEMM#828


- Integrate TBEDataConfig into TBE device benchmark CLI

Reviewed By: sryap

Differential Revision: D70296132

Author: q10

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_common.py

@@ -33,6 +33,30 @@ class EmbeddingLocation(enum.IntEnum):
     HOST = 3
     MTIA = 4
 
+    @classmethod
+    # pyre-ignore [3]
+    def from_str(cls, key: str):
+        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 {cls}: {key}")
+
+    def __str__(self) -> str:
+        return {
+            EmbeddingLocation.DEVICE: "device",
+            EmbeddingLocation.MANAGED: "managed",
+            EmbeddingLocation.MANAGED_CACHING: "managed_caching",
+            EmbeddingLocation.HOST: "host",
+            EmbeddingLocation.MTIA: "mtia",
+        }[self]
+
 
 class CacheAlgorithm(enum.Enum):
     LRU = 0
@@ -57,6 +81,29 @@ class PoolingMode(enum.IntEnum):
     MEAN = 1
     NONE = 2
 
+    @classmethod
+    # pyre-ignore [3]
+    def from_str(cls, key: str):
+        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 {cls}: {key}")
+
+    def __str__(self) -> str:
+        return {
+            PoolingMode.SUM: "sum",
+            PoolingMode.MEAN: "mean",
+            PoolingMode.NONE: "none",
+        }[self]
+
+    def do_pooling(self) -> bool:
+        return self is not PoolingMode.NONE
+
 
 class BoundsCheckMode(enum.IntEnum):
     # Raise an exception (CPU) or device-side assert (CUDA)