forked from pytorch/benchmark
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add --op-collection option (pytorch#2503)
Summary: This PR add `--op-collection` to tritonbench. It can run multiple ops in defined operator collections. The default collection includes all ops not included in other collections. Operator collections are defined in `torchbenchmark/operators_collection/`. For each collection, you should define a `get_operators` function to return operators included in this collection. Pull Request resolved: pytorch#2503 Reviewed By: xuzhao9 Differential Revision: D64359380 Pulled By: FindHao fbshipit-source-id: c66dd254a3c8b70c112d9b7774482813e0236789
- Loading branch information
1 parent
9e670cd
commit 58f3b1f
Showing
5 changed files
with
123 additions
and
3 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,71 @@ | ||
import importlib | ||
import pathlib | ||
from typing import List | ||
|
||
OP_COLLECTION_PATH = "operators_collection" | ||
|
||
|
||
def list_operator_collections() -> List[str]: | ||
""" | ||
List the available operator collections. | ||
This function retrieves the list of available operator collections by scanning the directories | ||
in the current path that contain an "__init__.py" file. | ||
Returns: | ||
List[str]: A list of names of the available operator collections. | ||
""" | ||
p = pathlib.Path(__file__).parent | ||
# only load the directories that contain a "__init__.py" file | ||
collection_paths = sorted( | ||
str(child.absolute()) | ||
for child in p.iterdir() | ||
if child.is_dir() and child.joinpath("__init__.py").exists() | ||
) | ||
filtered_collections = [pathlib.Path(path).name for path in collection_paths] | ||
return filtered_collections | ||
|
||
|
||
def list_operators_by_collection(op_collection: str = "default") -> List[str]: | ||
""" | ||
List the operators from the specified operator collections. | ||
This function retrieves the list of operators from the specified operator collections. | ||
If the collection name is "all", it retrieves operators from all available collections. | ||
If the collection name is not specified, it defaults to the "default" collection. | ||
Args: | ||
op_collection (str): Names of the operator collections to list operators from. | ||
It can be a single collection name or a comma-separated list of names. | ||
Special value "all" retrieves operators from all collections. | ||
Returns: | ||
List[str]: A list of operator names from the specified collection(s). | ||
Raises: | ||
ModuleNotFoundError: If the specified collection module is not found. | ||
AttributeError: If the specified collection module does not have a 'get_operators' function. | ||
""" | ||
|
||
def _list_all_operators(collection_name: str): | ||
try: | ||
module_name = f".{collection_name}" | ||
module = importlib.import_module(module_name, package=__name__) | ||
if hasattr(module, "get_operators"): | ||
return module.get_operators() | ||
else: | ||
raise AttributeError( | ||
f"Module '{module_name}' does not have a 'get_operators' function" | ||
) | ||
except ModuleNotFoundError: | ||
raise ModuleNotFoundError(f"Module '{module_name}' not found") | ||
|
||
if op_collection == "all": | ||
collection_names = list_operator_collections() | ||
else: | ||
collection_names = op_collection.split(",") | ||
|
||
all_operators = [] | ||
for collection_name in collection_names: | ||
all_operators.extend(_list_all_operators(collection_name)) | ||
return all_operators |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,5 @@ | ||
from torchbenchmark.operators import list_operators | ||
|
||
|
||
def get_operators(): | ||
return list_operators() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,28 @@ | ||
from torchbenchmark.operators_collection.all import get_operators as get_all_operators | ||
from torchbenchmark.operators_collection.liger import ( | ||
get_operators as get_liger_operators, | ||
) | ||
|
||
|
||
def get_operators(): | ||
""" | ||
Retrieve the list of operators for the default collection. | ||
This function retrieves the list of operators for the default collection by | ||
comparing the operators from the 'all' collection and the 'liger' collection. | ||
It returns a list of operators that are present in the 'all' collection but | ||
not in the 'liger' collection. | ||
In the future, if we add more operator collections, we will need to update | ||
this function to exclude desired operators in other collections. | ||
other_collections = list_operator_collections() | ||
to_remove = set(other_collections).union(liger_operators) | ||
return [item for item in all_operators if item not in to_remove] | ||
Returns: | ||
List[str]: A list of operator names for the default collection. | ||
""" | ||
all_operators = get_all_operators() | ||
liger_operators = get_liger_operators() | ||
return [item for item in all_operators if item not in liger_operators] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,5 @@ | ||
liger_operators = ["FusedLinearCrossEntropy"] | ||
|
||
|
||
def get_operators(): | ||
return liger_operators |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters