[tools][lh-transform] A cmd tool for applying schedules to payloads

[rolfmorel]

Convenient for development: in case you end up with a .mlir for your schedule and a .mlir for your payload, i.e. the most generic case.

[rengolin]

The point of those tools is to make sure we're using the same infra for everything. This tool doesn't use anything from the lighthouse module and replicate the functionality. We need to discuss the design of the lighthouse modules and just apply them here.

[Copilot]

Pull request overview

Adds a lightweight developer tool and accompanying lit integration to apply MLIR transform schedules to payload modules, with an example demonstrating end-to-end usage.

Changes:

• Introduce tools/lh-transform CLI to apply a transform.named_sequence from a schedule module to a payload module.
• Update lit.cfg.py to (intended) propagate PYTHONPATH to %PYTHON and expose executable lh-* tools to lit RUN lines.
• Add a schedule example .sh test and adjust the existing schedule example generator script to print payload/schedule modules separately.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tools/lh-transform	New Python CLI that parses schedule/payload MLIR and applies a named-sequence transform.
lit.cfg.py	Adds PYTHONPATH handling for %PYTHON and auto-substitutions for executable lh-* tools.
examples/schedule/transform_a_payload_according_to_a_schedule.sh	New lit shell test that generates payload/schedule MLIR files and runs lh-transform.
examples/schedule/transform_a_payload_according_to_a_schedule.py	Updates example output modes (print payload vs schedule vs full demo) to support the new .sh test.
examples/schedule/lit.local.cfg	Enables .sh tests in the schedule examples directory.

Comments suppressed due to low confidence (1)

examples/schedule/transform_a_payload_according_to_a_schedule.py:35

• The FileCheck patterns are no longer anchored to the transformed payload. Since this script prints the untransformed payload module before applying the schedule (in the no-args path), the first CHECK: lines can match the pre-transform IR and bind %[[...]] variables there, making the later transformed-only checks brittle/flaky. Consider either (a) printing only the transformed payload in the default execution path, or (b) adjusting the FileCheck patterns/RUN lines (e.g., with a dedicated check-prefix) so matching starts at a line that only appears in the transformed output.

        # NB: Do the CHECKing on the transformed output:
        # CHECK: func.func @fold_add_on_two_matmuls
        # CHECK-SAME:      (%[[MATRIX_A:.*]]: {{.*}}, %[[MATRIX_B:.*]]: {{.*}}, %[[WEIGHTS:.*]]: {{.*}})
        @func.func(matrixType, matrixType, matrixType)
        def fold_add_on_two_matmuls(matrixA, matrixB, weights):
            empty = tensor.empty(matrixType.shape, matrixType.element_type)

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

schedule_module.body.operations[0] will raise an IndexError if the schedule module parses successfully but contains no top-level operations (e.g., empty file). Consider checking for an empty body and exiting with a clear error, and/or searching the module body for a transform.named_sequence op instead of assuming it is the first op.

[rolfmorel]

The point of those tools is to make sure we're using the same infra for everything. This tool doesn't use anything from the lighthouse module and replicate the functionality. We need to discuss the design of the lighthouse modules and just apply them here.

This tool is about exposing functionality on the cmdline. For the moment you can get away with concat-ing payload and schedule files and running mlir-opt -transform-interpreter $SINGLE_FILE though that won't work when we have custom passes and ops in lighthouse. E.g. as already being introduced here: https://github.com/llvm/lighthouse/pull/63/changes#diff-2979a502b39d3d0a16a3c4c4a62870374f3aa111df32d435975238b68c3efaf8R19-R34

Once we have those custom passes and ops, this tool will depend on lighthouse modules, namely it will register the lighthouse-specific ops/passes/rewrites, e.g. those in lighthouse.dialects, before starting the interpreter (or pass manager, for that matter).

[rengolin]

Once we have those custom passes and ops, this tool will depend on lighthouse modules, namely it will register the lighthouse-specific ops/passes/rewrites, e.g. those in lighthouse.dialects, before starting the interpreter (or pass manager, for that matter).

This doesn't need to depend on other lighthouse functionality. It can be a stand-alone tool, that uses what the workload class already does. If that class is over-provisioned, then we break it down. But the point of having such tools is to exercise the lighthouse classes and to help us design them in a way that they're reusable across various tools.

[adam-smnk]

Looks like a handy experimentation tool.
How about registering it as a proper CLI command that's part of the package?

[adam-smnk]

nit: header is skipped everywhere else

[adam-smnk]

You could also add epilog with extra usage examples and documentation like what's the expected format of the files passed.

[rengolin]

It should be a similar design as of mlir-opt and mlir-runner, which just use the functionality exposed by core libraries and show how they can be used. The lh-transform and lh-opt should use the Workload class and friends to implement the core functionality (applying passes/schedules to payloads).