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

examples/schedule/transform_a_payload_according_to_a_schedule.py

@@ -2,6 +2,15 @@
 
 # Simply demonstrates applying a schedule to a payload.
 # To do so generates a basic payload and a basic schedule, purely as an example.
+# The following demonstrates doing so from the cmdline:
+
+# RUN: mkdir -p %t
+# RUN: %PYTHON %S/transform_a_payload_according_to_a_schedule.py payload > %t/payload.mlir
+# RUN: %PYTHON %S/transform_a_payload_according_to_a_schedule.py schedule > %t/schedule.mlir
+# RUN: lh-transform %t/schedule.mlir %t/payload.mlir | FileCheck %S/transform_a_payload_according_to_a_schedule.py
+
+
+import sys
 
 from mlir.ir import Context, Location, InsertionPoint, Operation, Module
 from mlir.ir import RankedTensorType, F32Type, UnitAttr
@@ -21,13 +30,11 @@ def example_payload() -> Module:
     Res = matmul(..., C=X)
     """
 
-    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: result of applying schedule to payload
         # CHECK: func.func @fold_add_on_two_matmuls
         # CHECK-SAME:      (%[[MATRIX_A:.*]]: {{.*}}, %[[MATRIX_B:.*]]: {{.*}}, %[[WEIGHTS:.*]]: {{.*}})
         @func.func(matrixType, matrixType, matrixType)
@@ -45,14 +52,12 @@ def fold_add_on_two_matmuls(matrixA, matrixB, weights):
             # CHECK: return %[[RES]]
             return added
 
-    print(payload)
     return payload
 
 
 def example_schedule() -> Module:
     """Basic schedule wrapping a single rewrite pattern."""
 
-    print("NOTE: example schedule module:")
     schedule_module = Module.create()
     schedule_module.operation.attributes["transform.with_named_sequence"] = (
         UnitAttr.get()
@@ -75,17 +80,27 @@ def example_schedule() -> Module:
                 )  # TODO: expose dedicated builder upstream
             transform.yield_([])
 
-    print(schedule_module)
     return schedule_module
 
 
 with Context(), Location.unknown():
     payload = example_payload()
     schedule_module = example_schedule()
-    # 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.
+    if len(sys.argv) > 1 and sys.argv[1] == "schedule":
+        print(schedule_module)
+    elif len(sys.argv) > 1 and sys.argv[1] == "payload":
+        print(payload)
+    else:
+        print("// NOTE: example payload module:")
+        print(payload)
+        print("// NOTE: example schedule module:")
+        print(schedule_module)
+
+        # 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.
 
-    print("NOTE: result of applying schedule to payload:")
-    print(payload)
+        print("// NOTE: result of applying schedule to payload:")
+        print(payload)

lit.cfg.py

@@ -51,3 +51,8 @@ def find_filecheck() -> str:
 torch_kernels_dir = project_root + "/third_party/KernelBench/KernelBench"
 if os.path.isdir(torch_kernels_dir):
     config.available_features.add("kernel_bench")
+
+for tool in os.listdir(project_root + "/tools"):
+    tool_path = os.path.join(project_root, "tools", tool)
+    if os.access(tool_path, os.X_OK):
+        config.substitutions.append((tool, tool_path))

tools/lh-transform

@@ -0,0 +1,39 @@
+#!/usr/bin/env python
+
+#  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__":
+    arg_parser = argparse.ArgumentParser(
+        prog="lh-transform", description="Apply a schedule to a payload module"
+    )
+    arg_parser.add_argument(
+        "schedule", help="path to file containing MLIR schedule module"
+    )
+    arg_parser.add_argument(
+        "payload", help="path to file containing MLIR payload module"
+    )
+    args = arg_parser.parse_args(sys.argv[1:])
+
+    with ir.Context(), ir.Location.unknown():
+        with open(args.schedule) as sched_file, open(args.payload) as payload_file:
+            schedule_module = ir.Module.parse(sched_file.read())
+            payload_module = ir.Module.parse(payload_file.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)