Skip to content

feat(auto-tuing): add a new auto-tuning system - #879

Open
mingdaw689 wants to merge 1 commit into
InfiniTensor:masterfrom
mingdaw689:feat/auto-tuning
Open

feat(auto-tuing): add a new auto-tuning system#879
mingdaw689 wants to merge 1 commit into
InfiniTensor:masterfrom
mingdaw689:feat/auto-tuning

Conversation

@mingdaw689

@mingdaw689 mingdaw689 commented Aug 4, 2026

Copy link
Copy Markdown

Summary

  • Added an auto-tuning system that automatically selects the fastest implementation for each operator based on runtime benchmarking.
  • src/tuning_manager.h / src/tuning_manager.cc - Manages tuning cache (in-memory + persistent JSON), provides Lookup() and Record() APIs for querying and storing optimal implementations;
  • src/tuning_signature.h - Extracts tuning signatures from operator arguments (tensor shapes, dtypes, scalar parameters) to uniquely identify each configuration;
  • src/config.h - Added auto_select_ flag (default true) to Config class, allowing per-call tuning control;
  • src/operator.h - Implemented ResolveConfigOnline<Key>() that benchmarks available implementations on first call and selects the fastest one; injects into Operator::Call() before operator cache lookup;
  • scripts/generate_wrappers.py - Fixed Python binding logic: only disable auto_select_ when user explicitly passes implementation_index parameter, preserving auto-tuning for all other cases.

Motivation

In the operator library, the same operator can have different underlying implementations. For instance, vendors can provide multiple interfaces, and there are also various algorithms to choose from for handwriting. Therefore, for the same operator, how to select the backend with the best performance in different situations becomes a problem.

Type of Change

  • [√] feat — new feature / new operator / new platform
  • fix — bug fix
  • perf — performance improvement (no behavioral change)
  • refactor — code restructuring without behavior change
  • test — adding or fixing tests only
  • docs — documentation only
  • build / ci — build system or CI configuration
  • chore — tooling, formatting, or other non-code changes
  • Breaking change (requires a ! in the Conventional Commits prefix or a BREAKING CHANGE: footer)

Platforms Affected

  • CPU (WITH_CPU)
  • [√] NVIDIA (WITH_NVIDIA)
  • Iluvatar (WITH_ILUVATAR)
  • MetaX (WITH_METAX)
  • Cambricon (WITH_CAMBRICON)
  • Moore (WITH_MOORE)
  • Ascend (WITH_ASCEND)
  • PyTorch C++ bindings (WITH_TORCH)
  • Build system / CMake / CI
  • Python bindings / user-facing API

Smoke Test Result

Testing `test_pybind_default_implementation_uses_first_active_index` failed, 
because it expects the generated pybind11 code to contain `implementation_index.value_or(DefaultImplementationIndexForMul(...))`, 
however, the actual generated code uses the new auto-tuning logic:
if (implementation_index.has_value()) {
    config.set_implementation_index(*implementation_index);
  }

This is because the recently added auto-tuning feature (WITH_TUNING) has changed the behavior of the code generator: 
- New logic: Only set when the user explicitly passes in implementation_index; otherwise, keep auto_select_=true to allow the runtime to automatically select the optimal implementation 
- Old logic: Always set implementation_index and use.value_or() to provide the default value

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA
Iluvatar
MetaX
Cambricon
Moore
Ascend
Full `pytest` output (optional)
Running 82 items in this shard
...........ss..ss....ss.......................F.........ssssss..ss...... [ 87%]
....ssssss                                                               [100%]
=================================== FAILURES ===================================
__________ test_pybind_default_implementation_uses_first_active_index __________

monkeypatch = <_pytest.monkeypatch.MonkeyPatch object at 0x7f62a9079120>
tmp_path = PosixPath('/tmp/pytest-of-root/pytest-3/test_pybind_default_implementa0')

    def test_pybind_default_implementation_uses_first_active_index(monkeypatch, tmp_path):
        module = _load_generator_module()
        base_header = tmp_path / "mul.h"
        base_header.write_text(
            """
    class Mul {
     public:
      Mul(const Tensor input, const Tensor other, Tensor out);
      virtual void operator()(const Tensor input, const Tensor other, Tensor out) const = 0;
    };
    """
        )
        monkeypatch.setattr(module, "_find_base_header", lambda op_name: base_header)
    
        operator = module._Operator(
            "mul",
            constructors=[
                module._ParsedFunction(
                    [
                        module._ParsedArgument("const Tensor", "input"),
                        module._ParsedArgument("const Tensor", "other"),
                        module._ParsedArgument("Tensor", "out"),
                    ]
                )
            ],
            calls=[
                module._ParsedFunction(
                    [
                        module._ParsedArgument("const Tensor", "input"),
                        module._ParsedArgument("const Tensor", "other"),
                        module._ParsedArgument("Tensor", "out"),
                    ]
                )
            ],
        )
    
        text = module._generate_pybind11(operator)
    
        assert "std::size_t DefaultImplementationIndexForMul" in text
        assert (
            "config.set_implementation_index("
            "DefaultImplementationIndexForMul(DeviceFromPybind11Handle(input).type()))"
        ) in text
        assert "std::optional<std::size_t> implementation_index" in text
>       assert (
            "implementation_index.value_or("
            "DefaultImplementationIndexForMul(DeviceFromPybind11Handle(input).type()))"
        ) in text
E       assert 'implementation_index.value_or(DefaultImplementationIndexForMul(DeviceFromPybind11Handle(input).type()))' in '#ifndef INFINI_OPS_BINDINGS_MUL_H_\n#define INFINI_OPS_BINDINGS_MUL_H_\n\n#include <pybind11/pybind11.h>\n#include <p...), py::arg("stream") = 0, py::arg("implementation_index") = py::none());\n}\n\n}  // namespace infini::ops\n\n#endif\n'

tests/test_generate_wrappers.py:172: AssertionError
=========================== short test summary info ============================
FAILED tests/test_generate_wrappers.py::test_pybind_default_implementation_uses_first_active_index
1 failed, 61 passed, 20 skipped, 19780 deselected in 47.10s

Benchmark / Performance Impact

Notes for Reviewers

@mingdaw689
mingdaw689 requested a review from a team August 4, 2026 11:15
@mingdaw689 mingdaw689 changed the title auto-tuning system feat(auto-tuing): add a new auto-tuning system Aug 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant