From 2de74a3ceed7005c945672bac2d72708ca27bc73 Mon Sep 17 00:00:00 2001 From: Rolf Morel Date: Tue, 11 Nov 2025 04:48:38 -0800 Subject: [PATCH 1/8] [transform] Simple Python and cmd interface for applying schedules Simple wrapper around upstream's transform_interpreter API. Invokable directly from Python and via the commandline on separate payload and schedule files. --- ...and_schedule_and_apply_latter_to_former.py | 93 +++++++++++++++++++ python/lighthouse/transform/__init__.py | 1 + python/lighthouse/transform/__main__.py | 24 +++++ python/lighthouse/transform/main.py | 14 +++ 4 files changed, 132 insertions(+) create mode 100644 python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py create mode 100644 python/lighthouse/transform/__init__.py create mode 100644 python/lighthouse/transform/__main__.py create mode 100644 python/lighthouse/transform/main.py diff --git a/python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py b/python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py new file mode 100644 index 0000000..b56fbc5 --- /dev/null +++ b/python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py @@ -0,0 +1,93 @@ +import tempfile +import subprocess + +from mlir import ir +from mlir.dialects import arith, func, linalg, tensor, transform +from mlir.dialects.transform import structured + +from lighthouse import transform as lh_transform + + +def example_payload() -> ir.Module: + payload = ir.Module.create() + with ir.InsertionPoint(payload.body): + matrixType = ir.RankedTensorType.get([16, 16], ir.F32Type.get()) + + @func.func(matrixType, matrixType) + def fold_add_on_two_matmuls(matrixA, matrixB): + splat_float = ir.FloatAttr.get(ir.F32Type.get(), 1.111111) + splat_attr = ir.DenseElementsAttr.get_splat(matrixType, splat_float) + weights = arith.constant(matrixType, splat_attr) + c0 = arith.constant(ir.F32Type.get(), 0.0) + empty = tensor.empty(matrixType.shape, matrixType.element_type) + zero_init = linalg.fill(c0, outs=[empty]) + A_x_weights = linalg.matmul(matrixA, weights, outs=[zero_init]) + empty2 = tensor.empty(matrixType.shape, matrixType.element_type) + zero_init2 = linalg.fill(c0, outs=[empty2]) + B_x_weights = linalg.matmul(matrixB, weights, outs=[zero_init2]) + added = linalg.add(A_x_weights, B_x_weights, outs=[empty]) + return added + + return payload + + +def example_schedule() -> ir.Module: + schedule = ir.Module.create() + schedule.operation.attributes["transform.with_named_sequence"] = ir.UnitAttr.get() + with ir.InsertionPoint(schedule.body): + named_seq = transform.NamedSequenceOp( # TODO: fix snake_case wrapper upstream + sym_name="__transform_main", + input_types=[transform.any_op_t()], + result_types=[], + arg_attrs=[{"transform.readonly": ir.UnitAttr.get()}], + ) + + with ir.InsertionPoint(named_seq.body): + func = structured.MatchOp.match_op_names( + named_seq.bodyTarget, ["func.func"] + ) # TODO: fix syntax upstream + with ir.InsertionPoint( + transform.apply_patterns(func).patterns.blocks.append() + ): # TODO: fix snake_case wrapper upstream + ir.Operation.create( + "transform.apply_patterns.linalg.fold_add_into_dest" + ) # TODO: expose dedicated builder upstream + transform.yield_([]) + return schedule + + +with ir.Context(), ir.Location.unknown(): + payload_module = example_payload() + print("NOTE: example payload module:") + print(payload_module) + schedule_module = example_schedule() + print("NOTE: example schedule module:") + print(schedule_module) + + print("NOTE: output of applying schedule to payload directly within Python process:") + schedule = schedule_module.body.operations[0] + lh_transform.apply(schedule, payload_module) + print(payload_module) + + with tempfile.NamedTemporaryFile( + "w", prefix="payload_" + ) as payload_file, tempfile.NamedTemporaryFile( + "w", prefix="schedule_" + ) as schedule_file: + print(payload_module, file=payload_file, flush=True) + print("NOTE: Have dumped payload to temp file:", payload_file.name) + print(schedule_module, file=schedule_file, flush=True) + print("NOTE: Have dumped schedule to temp file:", schedule_file.name) + + cmdline = [ + "python", + "-m", + "lighthouse.transform", + schedule_file.name, + payload_file.name, + ] + print("NOTE: output of applying schedule to payload from commandline:", *cmdline) + subprocess.run(cmdline) + print( + f"NOTE: cleaning-up temp files: {payload_file.name}, {schedule_file.name}" + ) diff --git a/python/lighthouse/transform/__init__.py b/python/lighthouse/transform/__init__.py new file mode 100644 index 0000000..01cd434 --- /dev/null +++ b/python/lighthouse/transform/__init__.py @@ -0,0 +1 @@ +from .main import apply diff --git a/python/lighthouse/transform/__main__.py b/python/lighthouse/transform/__main__.py new file mode 100644 index 0000000..0dd855b --- /dev/null +++ b/python/lighthouse/transform/__main__.py @@ -0,0 +1,24 @@ +import argparse +import sys + +from mlir import ir + +from .. import transform as lh_transform + + +if __name__ == "__main__": + ArgParser = argparse.ArgumentParser() + ArgParser.add_argument("schedule") + ArgParser.add_argument("payload") + args = ArgParser.parse_args(sys.argv[1:]) + + with ir.Context(), ir.Location.unknown(): + with open(args.schedule) as f: + schedule_module = ir.Module.parse(f.read()) + with open(args.payload) as f: + payload_module = ir.Module.parse(f.read()) + + schedule = schedule_module.body.operations[0] + lh_transform.apply(schedule, payload_module) + + print(payload_module) diff --git a/python/lighthouse/transform/main.py b/python/lighthouse/transform/main.py new file mode 100644 index 0000000..c5a0c97 --- /dev/null +++ b/python/lighthouse/transform/main.py @@ -0,0 +1,14 @@ +from mlir import ir +from mlir.dialects.transform import interpreter as transform_interpreter + + +def apply(schedule: ir.Operation | ir.OpView, payload: ir.Module) -> None: + assert schedule.parent and "transform.with_named_sequence" in schedule.parent.attributes + assert "transform.with_named_sequence" in schedule.parent.attributes + assert isinstance(schedule.parent.attributes["transform.with_named_sequence"], ir.UnitAttr) + + transform_interpreter.apply_named_sequence( + payload_root=payload, + transform_root=schedule, + transform_module=schedule.parent, + ) From 3a51e0494fc1ef3318778c55bb84d05149b39b70 Mon Sep 17 00:00:00 2001 From: Rolf Morel Date: Sat, 15 Nov 2025 10:56:28 -0800 Subject: [PATCH 2/8] Improvements and address feedback Now requires eb9d56c (or later) of llvm-project --- ...and_schedule_and_apply_latter_to_former.py | 31 +++++++++---------- .../{transform => schedule}/__main__.py | 12 +++---- python/lighthouse/transform/__init__.py | 1 - python/lighthouse/transform/main.py | 14 --------- 4 files changed, 21 insertions(+), 37 deletions(-) rename python/lighthouse/{transform => schedule}/__main__.py (53%) delete mode 100644 python/lighthouse/transform/__init__.py delete mode 100644 python/lighthouse/transform/main.py diff --git a/python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py b/python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py index b56fbc5..e2949d3 100644 --- a/python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py +++ b/python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py @@ -5,8 +5,6 @@ from mlir.dialects import arith, func, linalg, tensor, transform from mlir.dialects.transform import structured -from lighthouse import transform as lh_transform - def example_payload() -> ir.Module: payload = ir.Module.create() @@ -35,7 +33,7 @@ def example_schedule() -> ir.Module: schedule = ir.Module.create() schedule.operation.attributes["transform.with_named_sequence"] = ir.UnitAttr.get() with ir.InsertionPoint(schedule.body): - named_seq = transform.NamedSequenceOp( # TODO: fix snake_case wrapper upstream + named_seq = transform.named_sequence( sym_name="__transform_main", input_types=[transform.any_op_t()], result_types=[], @@ -46,9 +44,7 @@ def example_schedule() -> ir.Module: func = structured.MatchOp.match_op_names( named_seq.bodyTarget, ["func.func"] ) # TODO: fix syntax upstream - with ir.InsertionPoint( - transform.apply_patterns(func).patterns.blocks.append() - ): # TODO: fix snake_case wrapper upstream + with ir.InsertionPoint(transform.apply_patterns(func).patterns): ir.Operation.create( "transform.apply_patterns.linalg.fold_add_into_dest" ) # TODO: expose dedicated builder upstream @@ -64,16 +60,17 @@ def example_schedule() -> ir.Module: print("NOTE: example schedule module:") print(schedule_module) - print("NOTE: output of applying schedule to payload directly within Python process:") + print( + "NOTE: output of applying schedule to payload directly within Python process:" + ) schedule = schedule_module.body.operations[0] - lh_transform.apply(schedule, payload_module) + schedule.apply(payload_module) print(payload_module) - with tempfile.NamedTemporaryFile( - "w", prefix="payload_" - ) as payload_file, tempfile.NamedTemporaryFile( - "w", prefix="schedule_" - ) as schedule_file: + with ( + tempfile.NamedTemporaryFile("w", prefix="payload_") as payload_file, + tempfile.NamedTemporaryFile("w", prefix="schedule_") as schedule_file, + ): print(payload_module, file=payload_file, flush=True) print("NOTE: Have dumped payload to temp file:", payload_file.name) print(schedule_module, file=schedule_file, flush=True) @@ -82,12 +79,14 @@ def example_schedule() -> ir.Module: cmdline = [ "python", "-m", - "lighthouse.transform", + "lighthouse.schedule", schedule_file.name, payload_file.name, ] - print("NOTE: output of applying schedule to payload from commandline:", *cmdline) - subprocess.run(cmdline) + print( + "NOTE: output of applying schedule to payload from commandline:", *cmdline + ) + print(subprocess.run(cmdline, capture_output=True).stdout.decode()) print( f"NOTE: cleaning-up temp files: {payload_file.name}, {schedule_file.name}" ) diff --git a/python/lighthouse/transform/__main__.py b/python/lighthouse/schedule/__main__.py similarity index 53% rename from python/lighthouse/transform/__main__.py rename to python/lighthouse/schedule/__main__.py index 0dd855b..96177dd 100644 --- a/python/lighthouse/transform/__main__.py +++ b/python/lighthouse/schedule/__main__.py @@ -2,14 +2,13 @@ import sys from mlir import ir - -from .. import transform as lh_transform +from mlir.dialects import transform if __name__ == "__main__": - ArgParser = argparse.ArgumentParser() - ArgParser.add_argument("schedule") - ArgParser.add_argument("payload") + ArgParser = argparse.ArgumentParser(prog="lighthouse.transform") + ArgParser.add_argument("schedule", help="file which contains the MLIR schedule module") + ArgParser.add_argument("payload", help="file which contains the MLIR payload module") args = ArgParser.parse_args(sys.argv[1:]) with ir.Context(), ir.Location.unknown(): @@ -19,6 +18,7 @@ payload_module = ir.Module.parse(f.read()) schedule = schedule_module.body.operations[0] - lh_transform.apply(schedule, payload_module) + assert(isinstance(schedule, transform.NamedSequenceOp)) + schedule.apply(payload_module) print(payload_module) diff --git a/python/lighthouse/transform/__init__.py b/python/lighthouse/transform/__init__.py deleted file mode 100644 index 01cd434..0000000 --- a/python/lighthouse/transform/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .main import apply diff --git a/python/lighthouse/transform/main.py b/python/lighthouse/transform/main.py deleted file mode 100644 index c5a0c97..0000000 --- a/python/lighthouse/transform/main.py +++ /dev/null @@ -1,14 +0,0 @@ -from mlir import ir -from mlir.dialects.transform import interpreter as transform_interpreter - - -def apply(schedule: ir.Operation | ir.OpView, payload: ir.Module) -> None: - assert schedule.parent and "transform.with_named_sequence" in schedule.parent.attributes - assert "transform.with_named_sequence" in schedule.parent.attributes - assert isinstance(schedule.parent.attributes["transform.with_named_sequence"], ir.UnitAttr) - - transform_interpreter.apply_named_sequence( - payload_root=payload, - transform_root=schedule, - transform_module=schedule.parent, - ) From 19920ba6bedba9ba38a1aba4e55c0bdbb75cf7e0 Mon Sep 17 00:00:00 2001 From: Rolf Morel Date: Tue, 18 Nov 2025 11:30:43 -0800 Subject: [PATCH 3/8] Save --- ...and_schedule_and_apply_latter_to_former.py | 63 ++++++++++++------- 1 file changed, 39 insertions(+), 24 deletions(-) diff --git a/python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py b/python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py index e2949d3..2fd029b 100644 --- a/python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py +++ b/python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py @@ -1,72 +1,87 @@ +# RUN: %PYTHON %s | FileCheck %s --check-prefixes=CHECK + import tempfile import subprocess -from mlir import ir +from mlir.ir import Context, Location, InsertionPoint, Operation, Module +from mlir.ir import RankedTensorType, F32Type, FloatAttr, DenseElementsAttr, UnitAttr from mlir.dialects import arith, func, linalg, tensor, transform from mlir.dialects.transform import structured -def example_payload() -> ir.Module: - payload = ir.Module.create() - with ir.InsertionPoint(payload.body): - matrixType = ir.RankedTensorType.get([16, 16], ir.F32Type.get()) +def example_payload() -> Module: + print("NOTE: example payload module:") + payload = Module.create() + with InsertionPoint(payload.body): + matrixType = RankedTensorType.get([16, 16], F32Type.get()) + # NB: Do the CHECKing on the transformed output + # CHECK-LABEL: output of applying schedule to payload + # CHECK: func.func @fold_add_on_two_matmuls + # CHECK-SAME: (%[[MATRIX_A:.*]]: {{.*}}, %[[MATRIX_B:.*]]: {{.*}}) @func.func(matrixType, matrixType) def fold_add_on_two_matmuls(matrixA, matrixB): - splat_float = ir.FloatAttr.get(ir.F32Type.get(), 1.111111) - splat_attr = ir.DenseElementsAttr.get_splat(matrixType, splat_float) + splat_float = FloatAttr.get(F32Type.get(), 1.111111) + splat_attr = DenseElementsAttr.get_splat(matrixType, splat_float) + # CHECK: %[[WEIGHTS:.*]] = arith.constant splat weights = arith.constant(matrixType, splat_attr) - c0 = arith.constant(ir.F32Type.get(), 0.0) + c0 = arith.constant(F32Type.get(), 0.0) empty = tensor.empty(matrixType.shape, matrixType.element_type) + # CHECK: %[[ZERO_INIT:.*]] = linalg.fill zero_init = linalg.fill(c0, outs=[empty]) + # CHECK: %[[A_X_WEIGHTS]] = linalg.matmul ins(%[[MATRIX_A]], %[[WEIGHTS]]: {{.*}}) outs(%[[ZERO_INIT]] A_x_weights = linalg.matmul(matrixA, weights, outs=[zero_init]) empty2 = tensor.empty(matrixType.shape, matrixType.element_type) zero_init2 = linalg.fill(c0, outs=[empty2]) + # CHECK: %[[RES]] = linalg.matmul ins(%[[MATRIX_B]], %[[WEIGHTS]]: {{.*}}) outs(%[[A_X_WEIGHTS]] B_x_weights = linalg.matmul(matrixB, weights, outs=[zero_init2]) + # CHECK-NOT: linalg.add added = linalg.add(A_x_weights, B_x_weights, outs=[empty]) + # CHECK: func.yield %[[RES]] return added + print(payload) return payload -def example_schedule() -> ir.Module: - schedule = ir.Module.create() - schedule.operation.attributes["transform.with_named_sequence"] = ir.UnitAttr.get() - with ir.InsertionPoint(schedule.body): +def example_schedule() -> Module: + print("NOTE: example schedule module:") + schedule_module = Module.create() + schedule_module.operation.attributes["transform.with_named_sequence"] = UnitAttr.get() + with InsertionPoint(schedule_module.body): + # CHECK-LABEL: transform.named_sequence @__transform_main named_seq = transform.named_sequence( sym_name="__transform_main", input_types=[transform.any_op_t()], - result_types=[], - arg_attrs=[{"transform.readonly": ir.UnitAttr.get()}], + arg_attrs=[{"transform.readonly": UnitAttr.get()}], ) - with ir.InsertionPoint(named_seq.body): + with InsertionPoint(named_seq.body): func = structured.MatchOp.match_op_names( named_seq.bodyTarget, ["func.func"] ) # TODO: fix syntax upstream - with ir.InsertionPoint(transform.apply_patterns(func).patterns): - ir.Operation.create( + with InsertionPoint(transform.apply_patterns(func).patterns): + Operation.create( "transform.apply_patterns.linalg.fold_add_into_dest" ) # TODO: expose dedicated builder upstream transform.yield_([]) - return schedule + + print(schedule_module) + return schedule_module -with ir.Context(), ir.Location.unknown(): +with Context(), Location.unknown(): payload_module = example_payload() - print("NOTE: example payload module:") - print(payload_module) schedule_module = example_schedule() - print("NOTE: example schedule module:") - print(schedule_module) + schedule = schedule_module.body.operations[0] print( "NOTE: output of applying schedule to payload directly within Python process:" ) - schedule = schedule_module.body.operations[0] schedule.apply(payload_module) print(payload_module) + # Demonstrate apply a schedule from file to a payload from file with ( tempfile.NamedTemporaryFile("w", prefix="payload_") as payload_file, tempfile.NamedTemporaryFile("w", prefix="schedule_") as schedule_file, From 08ee4a176170eb9e9a935ce6125447bd21142700 Mon Sep 17 00:00:00 2001 From: Rolf Morel Date: Tue, 18 Nov 2025 14:21:23 -0800 Subject: [PATCH 4/8] Address feedback --- ...form_a_payload_according_to_a_schedule.py} | 47 ++++++++++++------- python/lighthouse/schedule/__main__.py | 18 +++++-- 2 files changed, 45 insertions(+), 20 deletions(-) rename python/examples/{generate_payload_and_schedule_and_apply_latter_to_former.py => schedule/transform_a_payload_according_to_a_schedule.py} (70%) diff --git a/python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py b/python/examples/schedule/transform_a_payload_according_to_a_schedule.py similarity index 70% rename from python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py rename to python/examples/schedule/transform_a_payload_according_to_a_schedule.py index 2fd029b..eee0e36 100644 --- a/python/examples/generate_payload_and_schedule_and_apply_latter_to_former.py +++ b/python/examples/schedule/transform_a_payload_according_to_a_schedule.py @@ -1,4 +1,9 @@ -# RUN: %PYTHON %s | FileCheck %s --check-prefixes=CHECK +# RUN: %PYTHON %s | FileCheck %s + +# A basic example of generating a payload, a schedule, and applying the latter +# to the former. Shows how to do it from Python and from the cmd given the +# payload and schedule are .mlir files. Run this file to see the concrete +# schedule IR, pre-transform payload IR and transformed payload IR. import tempfile import subprocess @@ -10,34 +15,39 @@ def example_payload() -> Module: + """Example payload where the results of two matmuls are summed together. + + Can be re-written so that the second matmul accumulates top of the the result of the first. + """ + print("NOTE: example payload module:") payload = Module.create() with InsertionPoint(payload.body): matrixType = RankedTensorType.get([16, 16], F32Type.get()) - # NB: Do the CHECKing on the transformed output - # CHECK-LABEL: output of applying schedule to payload + # NB: Do the CHECKing on the transformed output: + # CHECK-LABEL: result of applying schedule to payload # CHECK: func.func @fold_add_on_two_matmuls # CHECK-SAME: (%[[MATRIX_A:.*]]: {{.*}}, %[[MATRIX_B:.*]]: {{.*}}) @func.func(matrixType, matrixType) def fold_add_on_two_matmuls(matrixA, matrixB): splat_float = FloatAttr.get(F32Type.get(), 1.111111) splat_attr = DenseElementsAttr.get_splat(matrixType, splat_float) - # CHECK: %[[WEIGHTS:.*]] = arith.constant splat + # CHECK: %[[WEIGHTS:.*]] = arith.constant dense<1.11 weights = arith.constant(matrixType, splat_attr) c0 = arith.constant(F32Type.get(), 0.0) empty = tensor.empty(matrixType.shape, matrixType.element_type) # CHECK: %[[ZERO_INIT:.*]] = linalg.fill zero_init = linalg.fill(c0, outs=[empty]) - # CHECK: %[[A_X_WEIGHTS]] = linalg.matmul ins(%[[MATRIX_A]], %[[WEIGHTS]]: {{.*}}) outs(%[[ZERO_INIT]] + # CHECK: %[[A_X_WEIGHTS:.*]] = linalg.matmul ins(%[[MATRIX_A]], %[[WEIGHTS]]{{.*}}) outs(%[[ZERO_INIT]] A_x_weights = linalg.matmul(matrixA, weights, outs=[zero_init]) empty2 = tensor.empty(matrixType.shape, matrixType.element_type) zero_init2 = linalg.fill(c0, outs=[empty2]) - # CHECK: %[[RES]] = linalg.matmul ins(%[[MATRIX_B]], %[[WEIGHTS]]: {{.*}}) outs(%[[A_X_WEIGHTS]] + # CHECK: %[[RES:.*]] = linalg.matmul ins(%[[MATRIX_B]], %[[WEIGHTS]]{{.*}}) outs(%[[A_X_WEIGHTS]] B_x_weights = linalg.matmul(matrixB, weights, outs=[zero_init2]) # CHECK-NOT: linalg.add added = linalg.add(A_x_weights, B_x_weights, outs=[empty]) - # CHECK: func.yield %[[RES]] + # CHECK: return %[[RES]] return added print(payload) @@ -45,14 +55,17 @@ def fold_add_on_two_matmuls(matrixA, matrixB): def example_schedule() -> Module: + """Most basic schedule which doesn't just wrap a pass -- wraps a single rewrite pattern.""" print("NOTE: example schedule module:") schedule_module = Module.create() - schedule_module.operation.attributes["transform.with_named_sequence"] = UnitAttr.get() + schedule_module.operation.attributes["transform.with_named_sequence"] = ( + UnitAttr.get() + ) with InsertionPoint(schedule_module.body): - # CHECK-LABEL: transform.named_sequence @__transform_main named_seq = transform.named_sequence( - sym_name="__transform_main", + "__transform_main", input_types=[transform.any_op_t()], + result_types=[], arg_attrs=[{"transform.readonly": UnitAttr.get()}], ) @@ -71,22 +84,22 @@ def example_schedule() -> Module: with Context(), Location.unknown(): - payload_module = example_payload() + payload = example_payload() schedule_module = example_schedule() - schedule = schedule_module.body.operations[0] + schedule: transform.NamedSequenceOp = schedule_module.body.operations[0] print( - "NOTE: output of applying schedule to payload directly within Python process:" + "NOTE: result of applying schedule to payload directly within Python process:" ) - schedule.apply(payload_module) - print(payload_module) + schedule.apply(payload) + print(payload) - # Demonstrate apply a schedule from file to a payload from file + # Demonstrate applying a schedule from file to a payload from file with ( tempfile.NamedTemporaryFile("w", prefix="payload_") as payload_file, tempfile.NamedTemporaryFile("w", prefix="schedule_") as schedule_file, ): - print(payload_module, file=payload_file, flush=True) + print(payload, file=payload_file, flush=True) print("NOTE: Have dumped payload to temp file:", payload_file.name) print(schedule_module, file=schedule_file, flush=True) print("NOTE: Have dumped schedule to temp file:", schedule_file.name) diff --git a/python/lighthouse/schedule/__main__.py b/python/lighthouse/schedule/__main__.py index 96177dd..f17cd74 100644 --- a/python/lighthouse/schedule/__main__.py +++ b/python/lighthouse/schedule/__main__.py @@ -1,3 +1,7 @@ +# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + import argparse import sys @@ -7,8 +11,12 @@ if __name__ == "__main__": ArgParser = argparse.ArgumentParser(prog="lighthouse.transform") - ArgParser.add_argument("schedule", help="file which contains the MLIR schedule module") - ArgParser.add_argument("payload", help="file which contains the MLIR payload module") + ArgParser.add_argument( + "schedule", help="MLIR schedule module (path)" + ) + ArgParser.add_argument( + "payload", help="MLIR payload module (path)" + ) args = ArgParser.parse_args(sys.argv[1:]) with ir.Context(), ir.Location.unknown(): @@ -18,7 +26,11 @@ payload_module = ir.Module.parse(f.read()) schedule = schedule_module.body.operations[0] - assert(isinstance(schedule, transform.NamedSequenceOp)) + if not isinstance(schedule, transform.NamedSequenceOp): + sys.exit( + "The following op was expected to be a `transform.named_sequence`, instead got:\n" + + str(schedule) + ) schedule.apply(payload_module) print(payload_module) From e94f333527f7b655cac4b746279d6ff36916c2a8 Mon Sep 17 00:00:00 2001 From: Rolf Morel Date: Wed, 19 Nov 2025 05:09:41 -0800 Subject: [PATCH 5/8] Address Andrzej's feedback --- ...sform_a_payload_according_to_a_schedule.py | 36 ++++++++++++------- 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/python/examples/schedule/transform_a_payload_according_to_a_schedule.py b/python/examples/schedule/transform_a_payload_according_to_a_schedule.py index eee0e36..fc14a02 100644 --- a/python/examples/schedule/transform_a_payload_according_to_a_schedule.py +++ b/python/examples/schedule/transform_a_payload_according_to_a_schedule.py @@ -1,12 +1,15 @@ # RUN: %PYTHON %s | FileCheck %s -# A basic example of generating a payload, a schedule, and applying the latter -# to the former. Shows how to do it from Python and from the cmd given the -# payload and schedule are .mlir files. Run this file to see the concrete -# schedule IR, pre-transform payload IR and transformed payload IR. +# Simply demonstrates applying a schedule to a payload. +# +# To do so generates a basic payload and a basic schedule, purely as an example. +# Shows how to do it contained to a single Python process and how do invoke the +# same functionality from the cmdline. The latter to facilitate the case where +# the payload and schedule already exist as .mlir files. Run this file to see +# the concrete schedule IR, pre-transform payload IR and transformed payload IR. -import tempfile import subprocess +from tempfile import NamedTemporaryFile from mlir.ir import Context, Location, InsertionPoint, Operation, Module from mlir.ir import RankedTensorType, F32Type, FloatAttr, DenseElementsAttr, UnitAttr @@ -15,9 +18,15 @@ def example_payload() -> Module: - """Example payload where the results of two matmuls are summed together. - - Can be re-written so that the second matmul accumulates top of the the result of the first. + """IR for: + Zero = ... + X = matmul(..., C=Zero) + Y = matmul(..., C=Zero) + Res = add(X, Y) + + Can be re-written to: + X = matmul(..., C=Zero) + Res = matmul(..., C=X) """ print("NOTE: example payload module:") @@ -55,7 +64,8 @@ def fold_add_on_two_matmuls(matrixA, matrixB): def example_schedule() -> Module: - """Most basic schedule which doesn't just wrap a pass -- wraps a single rewrite pattern.""" + """Basic schedule wrapping a single rewrite pattern.""" + print("NOTE: example schedule module:") schedule_module = Module.create() schedule_module.operation.attributes["transform.with_named_sequence"] = ( @@ -86,6 +96,7 @@ def example_schedule() -> Module: with Context(), Location.unknown(): payload = example_payload() schedule_module = example_schedule() + # Demonstrate applying a schedule to a payload, both generated in-process: schedule: transform.NamedSequenceOp = schedule_module.body.operations[0] print( @@ -94,10 +105,11 @@ def example_schedule() -> Module: schedule.apply(payload) print(payload) - # Demonstrate applying a schedule from file to a payload from file + # Demonstrate applying a schedule to a payload, both as .mlir files, on the cmdline + # (to facilitate the same functionality for out-of-process generated schedules and payloads): with ( - tempfile.NamedTemporaryFile("w", prefix="payload_") as payload_file, - tempfile.NamedTemporaryFile("w", prefix="schedule_") as schedule_file, + NamedTemporaryFile("w", prefix="payload_", suffix=".mlir") as payload_file, + NamedTemporaryFile("w", prefix="schedule_", suffix=".mlir") as schedule_file, ): print(payload, file=payload_file, flush=True) print("NOTE: Have dumped payload to temp file:", payload_file.name) From 6f29ea983a2d8bba4f01252f0c167671c73e36e5 Mon Sep 17 00:00:00 2001 From: Rolf Morel Date: Wed, 19 Nov 2025 07:52:57 -0800 Subject: [PATCH 6/8] KISS. With an emphasis on S. --- .github/workflows/examples.yml | 6 +- python/examples/mlir/compile_and_run.py | 27 ++------ ...sform_a_payload_according_to_a_schedule.py | 67 +++---------------- python/lighthouse/schedule/__main__.py | 36 ---------- 4 files changed, 21 insertions(+), 115 deletions(-) delete mode 100644 python/lighthouse/schedule/__main__.py diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 9a9c63b..a745eb7 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -25,4 +25,8 @@ jobs: - name: Run MLP From Module run: |- - uv run python/examples/ingress/torch/mlp_from_model.py \ No newline at end of file + uv run python/examples/ingress/torch/mlp_from_model.py + + - name: Run schedule application example + run: |- + uv run python/examples/schedule/transform_a_payload_according_to_a_schedule.py diff --git a/python/examples/mlir/compile_and_run.py b/python/examples/mlir/compile_and_run.py index d0529f6..0d8b8f1 100644 --- a/python/examples/mlir/compile_and_run.py +++ b/python/examples/mlir/compile_and_run.py @@ -1,10 +1,11 @@ +# RUN: %PYTHON %s + import torch import argparse from mlir import ir from mlir.dialects import transform from mlir.dialects.transform import structured -from mlir.dialects.transform import interpreter from mlir.execution_engine import ExecutionEngine from mlir.passmanager import PassManager @@ -89,22 +90,6 @@ def create_schedule(ctx: ir.Context) -> ir.Module: return schedule -def apply_schedule(kernel: ir.Module, schedule: ir.Module) -> None: - """ - Apply transformation schedule to a kernel module. - The kernel is modified in-place. - - Args: - kernel: A module with payload function. - schedule: A module with transform schedule. - """ - interpreter.apply_named_sequence( - payload_root=kernel, - transform_root=schedule.body.operations[0], - transform_module=schedule, - ) - - def create_pass_pipeline(ctx: ir.Context) -> PassManager: """ Create an MLIR pass pipeline. @@ -141,9 +126,11 @@ def main(args): ctx = ir.Context() kernel = create_kernel(ctx) - # Create a transform schedule and apply initial lowering. - schedule = create_schedule(ctx) - apply_schedule(kernel, schedule) + # Create a transform schedule and apply initial lowering to kernel. + # The kernel is modified in-place. + schedule_module = create_schedule(ctx) + named_seq: transform.NamedSequenceOp = schedule_module.body.operations[0] + named_seq.apply(kernel) # Create a pass pipeline and lower the kernel to LLVM dialect. pm = create_pass_pipeline(ctx) diff --git a/python/examples/schedule/transform_a_payload_according_to_a_schedule.py b/python/examples/schedule/transform_a_payload_according_to_a_schedule.py index fc14a02..0cd1e2c 100644 --- a/python/examples/schedule/transform_a_payload_according_to_a_schedule.py +++ b/python/examples/schedule/transform_a_payload_according_to_a_schedule.py @@ -1,18 +1,8 @@ -# RUN: %PYTHON %s | FileCheck %s - # Simply demonstrates applying a schedule to a payload. -# # To do so generates a basic payload and a basic schedule, purely as an example. -# Shows how to do it contained to a single Python process and how do invoke the -# same functionality from the cmdline. The latter to facilitate the case where -# the payload and schedule already exist as .mlir files. Run this file to see -# the concrete schedule IR, pre-transform payload IR and transformed payload IR. - -import subprocess -from tempfile import NamedTemporaryFile from mlir.ir import Context, Location, InsertionPoint, Operation, Module -from mlir.ir import RankedTensorType, F32Type, FloatAttr, DenseElementsAttr, UnitAttr +from mlir.ir import RankedTensorType, F32Type, UnitAttr from mlir.dialects import arith, func, linalg, tensor, transform from mlir.dialects.transform import structured @@ -34,29 +24,16 @@ def example_payload() -> Module: with InsertionPoint(payload.body): matrixType = RankedTensorType.get([16, 16], F32Type.get()) - # NB: Do the CHECKing on the transformed output: - # CHECK-LABEL: result of applying schedule to payload - # CHECK: func.func @fold_add_on_two_matmuls - # CHECK-SAME: (%[[MATRIX_A:.*]]: {{.*}}, %[[MATRIX_B:.*]]: {{.*}}) - @func.func(matrixType, matrixType) - def fold_add_on_two_matmuls(matrixA, matrixB): - splat_float = FloatAttr.get(F32Type.get(), 1.111111) - splat_attr = DenseElementsAttr.get_splat(matrixType, splat_float) - # CHECK: %[[WEIGHTS:.*]] = arith.constant dense<1.11 - weights = arith.constant(matrixType, splat_attr) - c0 = arith.constant(F32Type.get(), 0.0) + @func.func(matrixType, matrixType, matrixType) + def fold_add_on_two_matmuls(matrixA, matrixB, weights): empty = tensor.empty(matrixType.shape, matrixType.element_type) - # CHECK: %[[ZERO_INIT:.*]] = linalg.fill + c0 = arith.constant(F32Type.get(), 0.0) zero_init = linalg.fill(c0, outs=[empty]) - # CHECK: %[[A_X_WEIGHTS:.*]] = linalg.matmul ins(%[[MATRIX_A]], %[[WEIGHTS]]{{.*}}) outs(%[[ZERO_INIT]] A_x_weights = linalg.matmul(matrixA, weights, outs=[zero_init]) empty2 = tensor.empty(matrixType.shape, matrixType.element_type) zero_init2 = linalg.fill(c0, outs=[empty2]) - # CHECK: %[[RES:.*]] = linalg.matmul ins(%[[MATRIX_B]], %[[WEIGHTS]]{{.*}}) outs(%[[A_X_WEIGHTS]] B_x_weights = linalg.matmul(matrixB, weights, outs=[zero_init2]) - # CHECK-NOT: linalg.add added = linalg.add(A_x_weights, B_x_weights, outs=[empty]) - # CHECK: return %[[RES]] return added print(payload) @@ -96,37 +73,11 @@ def example_schedule() -> Module: with Context(), Location.unknown(): payload = example_payload() schedule_module = example_schedule() - # Demonstrate applying a schedule to a payload, both generated in-process: + # Actual schedule is defined by the contained transfomr.named_sequence: schedule: transform.NamedSequenceOp = schedule_module.body.operations[0] + schedule_module.body - print( - "NOTE: result of applying schedule to payload directly within Python process:" - ) - schedule.apply(payload) - print(payload) + schedule.apply(payload) # The actual transformation happens here. - # Demonstrate applying a schedule to a payload, both as .mlir files, on the cmdline - # (to facilitate the same functionality for out-of-process generated schedules and payloads): - with ( - NamedTemporaryFile("w", prefix="payload_", suffix=".mlir") as payload_file, - NamedTemporaryFile("w", prefix="schedule_", suffix=".mlir") as schedule_file, - ): - print(payload, file=payload_file, flush=True) - print("NOTE: Have dumped payload to temp file:", payload_file.name) - print(schedule_module, file=schedule_file, flush=True) - print("NOTE: Have dumped schedule to temp file:", schedule_file.name) - - cmdline = [ - "python", - "-m", - "lighthouse.schedule", - schedule_file.name, - payload_file.name, - ] - print( - "NOTE: output of applying schedule to payload from commandline:", *cmdline - ) - print(subprocess.run(cmdline, capture_output=True).stdout.decode()) - print( - f"NOTE: cleaning-up temp files: {payload_file.name}, {schedule_file.name}" - ) + print("NOTE: result of applying schedule to payload:") + print(payload) diff --git a/python/lighthouse/schedule/__main__.py b/python/lighthouse/schedule/__main__.py deleted file mode 100644 index f17cd74..0000000 --- a/python/lighthouse/schedule/__main__.py +++ /dev/null @@ -1,36 +0,0 @@ -# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -# See https://llvm.org/LICENSE.txt for license information. -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import argparse -import sys - -from mlir import ir -from mlir.dialects import transform - - -if __name__ == "__main__": - ArgParser = argparse.ArgumentParser(prog="lighthouse.transform") - ArgParser.add_argument( - "schedule", help="MLIR schedule module (path)" - ) - ArgParser.add_argument( - "payload", help="MLIR payload module (path)" - ) - args = ArgParser.parse_args(sys.argv[1:]) - - with ir.Context(), ir.Location.unknown(): - with open(args.schedule) as f: - schedule_module = ir.Module.parse(f.read()) - with open(args.payload) as f: - payload_module = ir.Module.parse(f.read()) - - schedule = schedule_module.body.operations[0] - if not isinstance(schedule, transform.NamedSequenceOp): - sys.exit( - "The following op was expected to be a `transform.named_sequence`, instead got:\n" - + str(schedule) - ) - schedule.apply(payload_module) - - print(payload_module) From abc7403da68088e342bb0d264ab4b29e86f030e6 Mon Sep 17 00:00:00 2001 From: Rolf Morel Date: Wed, 19 Nov 2025 08:29:17 -0800 Subject: [PATCH 7/8] Remove line - thanks Adam! --- .../schedule/transform_a_payload_according_to_a_schedule.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/examples/schedule/transform_a_payload_according_to_a_schedule.py b/python/examples/schedule/transform_a_payload_according_to_a_schedule.py index 0cd1e2c..8a6855a 100644 --- a/python/examples/schedule/transform_a_payload_according_to_a_schedule.py +++ b/python/examples/schedule/transform_a_payload_according_to_a_schedule.py @@ -75,7 +75,6 @@ def example_schedule() -> Module: schedule_module = example_schedule() # Actual schedule is defined by the contained transfomr.named_sequence: schedule: transform.NamedSequenceOp = schedule_module.body.operations[0] - schedule_module.body schedule.apply(payload) # The actual transformation happens here. From fe032859b256dc4bb4961c45ea7b762009195194 Mon Sep 17 00:00:00 2001 From: Rolf Morel Date: Wed, 19 Nov 2025 08:51:50 -0800 Subject: [PATCH 8/8] Final cleanup --- .github/workflows/examples.yml | 2 +- python/examples/mlir/compile_and_run.py | 2 -- .../schedule/transform_a_payload_according_to_a_schedule.py | 2 +- 3 files changed, 2 insertions(+), 4 deletions(-) diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 4cbb876..991d947 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -33,4 +33,4 @@ jobs: - name: Run apply basic schedule to basic payload run: |- - uv run python/examples/schedule/transform_a_payload_according_to_a_schedule.py \ No newline at end of file + uv run python/examples/schedule/transform_a_payload_according_to_a_schedule.py diff --git a/python/examples/mlir/compile_and_run.py b/python/examples/mlir/compile_and_run.py index 0d8b8f1..009d9d5 100644 --- a/python/examples/mlir/compile_and_run.py +++ b/python/examples/mlir/compile_and_run.py @@ -1,5 +1,3 @@ -# RUN: %PYTHON %s - import torch import argparse diff --git a/python/examples/schedule/transform_a_payload_according_to_a_schedule.py b/python/examples/schedule/transform_a_payload_according_to_a_schedule.py index 8a6855a..12702ab 100644 --- a/python/examples/schedule/transform_a_payload_according_to_a_schedule.py +++ b/python/examples/schedule/transform_a_payload_according_to_a_schedule.py @@ -73,7 +73,7 @@ def example_schedule() -> Module: with Context(), Location.unknown(): payload = example_payload() schedule_module = example_schedule() - # Actual schedule is defined by the contained transfomr.named_sequence: + # Actual schedule is defined by the contained transform.named_sequence: schedule: transform.NamedSequenceOp = schedule_module.body.operations[0] schedule.apply(payload) # The actual transformation happens here.