[Benchmark] Add embedding benchmark (#248)

Author: yf225

benchmark/README.md

@@ -2,7 +2,7 @@
 
 Performance comparison between Helion, torch.compile, Triton, and PyTorch eager is done by leveraging [TritonBench](https://github.com/pytorch-labs/tritonbench).
 
-Currently supported kernels for performance comparison are in `benchmark/`.
+Currently supported kernels for performance comparison are listed in `KERNEL_MAPPINGS` in `benchmark/run.py`.
 
 To run the benchmark:
 

benchmark/run.py

@@ -1,6 +1,6 @@
 """Performance comparison between Helion, torch.compile, Triton, and PyTorch eager by leveraging TritonBench.
 
-Currently supported kernels are in `benchmark/`.
+Currently supported kernels are listed in `KERNEL_MAPPINGS` in `benchmark/run.py`.
 
 Usage:
 $ python benchmark/run.py [tritonbench args...] --kernel <kernel_name>
@@ -23,6 +23,7 @@
 KERNEL_MAPPINGS: dict[str, tuple[str, str]] = {
     # <tritonbench_op_name>: (<helion_kernel_module_path>, <helion_kernel_function_name>)
     "vector_add": ("examples.add", "add"),
+    "embedding": ("examples.embedding", "embedding_tritonbench"),
 }
 
 

examples/embedding.py

@@ -24,6 +24,13 @@ def embedding(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
     return out.view(*x.size(), embedding_dim)
 
 
+def embedding_tritonbench(
+    V: int, D: int, inp: torch.Tensor, shared_weight: torch.Tensor
+) -> torch.Tensor:
+    """Wrapper for tritonbench that matches its interface."""
+    return embedding(inp, shared_weight)
+
+
 def main() -> None:
     num_embeddings, embedding_dim = 16, 64
     x = torch.randint(0, num_embeddings, [256, 32], device="cuda", dtype=torch.int32)