Skip to content

Conversation

@jialiangqu
Copy link
Member

Implements the feature requested in PR #647 comments - allows users to specify tags directly in the decorator

Tags from the decorator are merged with auto-detected tags (from AST analysis and name-based heuristics). Added REGISTERED_BACKEND_TAGS global dict to store decorator tags, and updated the tagging system in run.py to merge all tag sources.

Apply a patch locally to test the new functionality
in tritonbench/operators/gemm/operator.py, add a dummy tag to triton backend

    @register_benchmark(tags=["dummy"])
    def triton_tutorial_matmul(self, a, b, bias) -> Callable:
        if bias is not None:
            return lambda: triton_tutorial_matmul_kernel(a, b) + bias
        else:
            return lambda: triton_tutorial_matmul_kernel(a, b)

run the analyzer to see it appears

$ python benchmarks/tagging/run.py --op gemm

{'gemm': {'pt2_triton_matmul': {'tags': ['pt2']}, 'pt2_matmul_maxautotune': {'tags': ['pt2']}, 'pt2_cutlass_matmul': {'tags': ['pt2', 'cutlass']}, 'matmul_decompose_k': {'tags': ['pt2']}, 'triton_tutorial_matmul': {'tags': ['dummy', 'triton'], 'kernels': ['matmul_kernel']}, 'matmul_partition_k': {'tags': ['triton'], 'kernels': ['_matmul_partition_k', '_reduce']}, 'triton_persistent_matmul': {'tags': []}, 'triton_tma_persistent_matmul': {'tags': []}, 'triton_tma_persistent_cached_matmul': {'tags': []}, 'triton_ops_matmul': {'tags': ['triton'], 'kernels': ['_splitk_kernel']}, 'aten_matmul': {'tags': ['aten']}, 'aten_tunableop_matmul': {'tags': ['aten']}, 'hstu_triton_matmul': {'tags': []}, 'streamk_matmul': {'tags': ['triton'], 'kernels': ['streamk_amd_gemm']}}}

we can see dummy tag appears in triton_tutorial_matmul

Implements the feature requested in PR #647 comments - allows users to specify tags directly in the decorator like @register_benchmark(tags=['triton', 'custom']).

Tags from the decorator are merged with auto-detected tags (from AST analysis and name-based heuristics). Added REGISTERED_BACKEND_TAGS global dict to store decorator tags, and updated the tagging system in run.py to merge all tag sources.

Updated documentation with examples showing how to use the new tags parameter.
@jialiangqu jialiangqu requested a review from xuzhao9 November 23, 2025 23:45
@meta-cla meta-cla bot added the cla signed label Nov 23, 2025
@jialiangqu jialiangqu marked this pull request as ready for review November 23, 2025 23:46
BASELINE_BENCHMARKS: Dict[str, str] = {}
# Store tags defined in @register_benchmark decorator
# Format: {operator_name: {backend_name: [tag1, tag2, ...]}}
REGISTERED_BACKEND_TAGS: Dict[str, Dict[str, List[str]]] = {}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we remove the use of REGISTERED_BACKEND_TAGS? Instead, we can add it as a field of class BenchmarkOperatorBackend

@xuzhao9
Copy link
Contributor

xuzhao9 commented Nov 24, 2025

Thanks for the contribution!
For maintenance, we want to minimize the use the global dictionaries, instead we want to use the BenchmarkOperatorBackend class to manage all tags of the backend. Can you help refactor the code?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants