add in-depth documentation to insert_transform_script

Author: martin-luecke

mlir/python/mlir/dialects/transform/extras/__init__.py

@@ -82,17 +82,35 @@ def match_ops(
         return handle
 
 
-ValueT = TypeVar("ValueT", bound=Value)
-
-
 def insert_transform_script(
     module: ir.Module,
-    script: Callable[[ValueT], None],
+    script: Callable[[OpHandle], None],
     dump_script: bool = False,
 ) -> None:
-    """Inserts the transform script of the schedule into the module."""
+    """
+    Inserts the transform script of the schedule into the module. The script
+    should accept an instance of OpHandle as argument, which will be called with
+    the block arg of the newly created sequence op.
+
+    Example:
+    This python code
+    ```
+    module = ir.Module.create()
+    def test_match_ops_single(module: OpHandle):
+        module.match_ops(scf.ForOp)
+    insert_transform_script(module, script)
+    ```
+    generates the following IR:
+    ```
+    module {
+        transform.sequence failures(propagate) {
+        ^bb0(%arg0: !transform.any_op):
+            %0 = transform.structured.match ops{["scf.for"]} in %arg0 : (!transform.any_op) -> !transform.op<"scf.for">
+        }
+    }
+    ```
+    """
 
-    # Insert the script into the IR
     with module.context, ir.Location.unknown(module.context):
         with ir.InsertionPoint.at_block_begin(module.body):
             sequence_op = transform.SequenceOp(

mlir/test/python/dialects/transform_extras.py

@@ -1,15 +1,13 @@
 # RUN: %PYTHON %s | FileCheck %s
 
-from typing import Callable, TypeVar
+from typing import Callable
 from mlir import ir
 from mlir.dialects import scf
 from mlir.dialects.transform import structured
-from mlir.dialects.transform.extras import Value, OpHandle, insert_transform_script
+from mlir.dialects.transform.extras import OpHandle, insert_transform_script
 
-ValueT = TypeVar("ValueT", bound=Value)
 
-
-def build_transform_script(script: Callable[[ValueT], None]):
+def build_transform_script(script: Callable[[OpHandle], None]):
     print("\nTEST:", script.__name__)
     with ir.Context(), ir.Location.unknown():
         module = ir.Module.create()