WIP: [mlir][transform][python] Test that mix-ins stay in sync with tablegen.

[ingomueller-net]

Summary

Mix-in classes of the Python bindings have the risk of running out of
sync with respect to the tablegen definition of their underlying op. If
that op gets new parameters or required parameters are renamed, existing
test cases usually don't break. This patch creates a helper that allows
to print the parameters of the generated base classes such that they can
be CHECKed in the corresponding tests.

I have only currently added such a test for one op to gather feedback
and plan to add more tests once I see people like the idea.

[ingomueller-net]

• Add file I forgot to add in the initial revision.

mlir/test/python/dialects/ext_test_helper.py

+import inspect
+
+
+def print_gen_init_arg_names(cls: type):
+    """
+    Prints the argument names of the __init__ function of the generated base
+    class of a mix-in extension.
+
+    This is useful for ensuring that mix-in classes remain in sync with their
+    underlying op. If operands are added or existing ones are renamed in that
+    op, the same change needs to occur in the mix-in class. However, in many
+    cases, these changes do not otherwise break the tests of the mix-ins, such
+    that the mix-in classes become out of sync. This function allows to print
+    the argument names of the generated base class, which can then be checked
+    for in a `CHECK` statement.
+
+    If you are modifying a tablegen definition and that modification breaks a
+    test that uses this function, please update the corresponding mix-in class
+    in the `*_ext.py` of the dialect.
+    """
+    mro = cls.__mro__[2]
+    assert mro.__module__.endswith("_gen")
+    print(list(inspect.signature(mro.__init__).parameters.keys()))

mlir/test/python/dialects/transform_gpu_ext.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 from mlir.dialects import transform
 from mlir.dialects.transform import gpu
 
+from ext_test_helper import print_gen_init_arg_names
+
 
 def run(f):
+    print("\nTEST:", f.__name__)
     with Context(), Location.unknown():
         module = Module.create()
         with InsertionPoint(module.body):
             sequence = transform.SequenceOp(
                 transform.FailurePropagationMode.Propagate,
                 [],
                 transform.AnyOpType.get(),
             )
             with InsertionPoint(sequence.body):
                 f(sequence.bodyTarget)
                 transform.YieldOp()
-        print("\nTEST:", f.__name__)
         print(module)
     return f
 
 
 @run
 def testMapForallToBlocksCompact(target):
     gpu.MapForallToBlocks(target)
     # CHECK-LABEL: TEST: testMapForallToBlocksCompact
def testMapNestedForallToThreadsAttributes(target):
     gpu.MapNestedForallToThreads(
         target, block_dims=[4, 2], warp_size=64, sync_after_distribute=False
     )
     # CHECK-LABEL: TEST: testMapNestedForallToThreadsAttributes
     # CHECK: transform.gpu.map_nested_forall_to_threads
     # CHECK-SAME: block_dims = [4, 2]
     # CHECK-SAME: sync_after_distribute = false
     # CHECK-SAME: warp_size = 64
+
+
+@run
+def testMapNestedForallToThreadsGenArgs(_):
+    print_gen_init_arg_names(gpu.MapNestedForallToThreads)
+    # CHECK-LABEL: TEST: testMapNestedForallToThreadsGenArgs
+    # CHECK: ['self', 'result', 'target', 'block_dims', 'sync_after_distribute', 'warp_size', 'loc', 'ip']