KISS. With an emphasis on S.

Author: rolfmorel

.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
+          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

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)

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)