[Transform][Python] Introduce transform.ffi.callback

[rolfmorel]

This transform op allows for invoking an external handler routine which can, for example, implement an analysis or transform in whichever way - and programming language - is convenient.

This PR provides full facilities for registering callbacks from Python and has test cases to show how MLIR's Python bindings can be used to easily implement at least some transforms.

[Copilot]

Pull Request Overview

This PR introduces the new transform.tune.callback functionality that enables invoking external Python callbacks from MLIR transforms. Key changes include:

• Updating test configuration (lit.cfg.py) to support Python tests.
• Adding Python modules and test cases (transform_callback.py, common.py, and tune.py) for callback registration and execution.
• Integrating a C++ callback handler (TPPDialects.cpp) and updating dialect transform ops to support the new callback op.

Reviewed Changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
test/lit.cfg.py	Updated testing configuration to include Python files and PYTHONPATH updates.
test/Python/transform_callback.py	Added test cases for transform callback behavior.
python/mlir/tpp/sched/common.py	Introduced decorators for callback registration with clear runtime mapping.
python/mlir/dialects/transform/tune.py	Added a new callback helper wrapping TuneCallbackOp creation.
python/TPPDialects.cpp	Registered the Python callback handler via C++ bindings.
lib/TPP/Dialect/Tune/TransformOps/TuneTransformOps.cpp	Implemented TuneCallbackOp operations and error handling.
include/TPP/Dialect/Tune/TuneTransformOps.td	Updated tablegen definitions with documentation for TuneCallbackOp.
include/TPP/Dialect/Tune/TuneTransformOps.h	Added header declarations for the new callback handler type.

Comments suppressed due to low confidence (1)

python/mlir/tpp/sched/common.py:34

• [nitpick] Consider adding a docstring to the 'callback' decorator to clearly document its purpose, expected input types, and behavior. This would help maintainability and improve the understandability of the callback registration process.

def callback(function: Callable):

[rolfmorel]

Failed Tests (4):
  TPP_OPT :: BF16/matmul-vnni.mlir
  TPP_OPT :: Conversion/VectorToXsmm/vector-to-transpose.mlir
  TPP_OPT :: Passes/DefaultPipeline/default-tpp-passes.mlir
  TPP_OPT :: Python/transform_callback.py

Hmm - all fail with segfaults. These tests all succeed for me on the compile server. 😕

Will need to think about it. Might be to do with enabling some Python environment vars in the lit.cfg.py as otherwise tpp-sched (which is implicated in the segfaults) is not touched at all by this PR.

[rolfmorel]

Will need to think about it. Might be to do with enabling some Python environment vars in the lit.cfg.py as otherwise tpp-sched (which is implicated in the segfaults) is not touched at all by this PR.

Issue was reference counting of the callback object. This is now fixed.

IMO this PR is ready to go (e.g. to be used for nano-tile selection algorithm prototyping).

For reference purposes, another instance of passing callbacks around: https://reviews.llvm.org/D116568

[adam-smnk]

Looks neat, I'll have to take it for a spin 😎

A general question about the whole Python-MLIR interaction.
I take that the callback signature is reified into IR. How/Where does the runtime find the actual callback definition to call?

[rolfmorel]

I take that the callback signature is reified into IR.

Yes - automatically converted from the handles' types on the invocation of a Python callback. There's a test case with an example: https://github.com/libxsmm/tpp-mlir/pull/1064/files#diff-7aa62724b21b998da9cf032da6e6b77bbdb664258a5dca850f9509a5459646f6R95

How/Where does the runtime find the actual callback definition to call?

A global func pointer gets set when a Python callback handler is provided. When the op executes it just reaches for the global and calls the handler.

The main thing that a proper solution should have is that this global gets moved to hang of the MlirContext. Though in general I think of how the FFI mechanism works as an implementation detail (e.g. the user could not use the Py bindings and use something else to set the global or just swap out the entire mechanism - or even just extremely late bind some C++ code into a transform op).