[mlir][memref][transform][python] Create mix-in for MemRefMultiBufferOp.

[ingomueller-net]

Summary

Create a mix-in class with an overloaded constructor that makes the
return type optional.

[ingomueller-net]

Delcaring D156536 as dependent patch.

[ftynse]

In D156561#4543580, @ingomueller-net wrote:

Delcaring D156536 as dependent patch.

You can click "edit related revisions" on the right to add parent revisions.

[ftynse]

Did you forget to git add?

[ingomueller-net]

• Add main *_ext.py file, which I had forgotten to add.
• Rebase to main.

mlir/python/CMakeLists.txt

declare_mlir_dialect_extension_python_bindings(
   DIALECT_NAME transform
   EXTENSION_NAME loop_transform)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/MemRefTransformOps.td
   SOURCES
+    dialects/_memref_transform_ops_ext.py
     dialects/transform/memref.py
   DIALECT_NAME transform
   EXTENSION_NAME memref_transform)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/LinalgStructuredTransformOps.td

mlir/test/python/dialects/transform_memref_ext.py

+# RUN: %PYTHON %s | FileCheck %s
+
+
+from mlir.ir import *
+from mlir.dialects import transform
+from mlir.dialects.transform import memref
+
+
+def run(f):
+    with Context(), Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            print("\nTEST:", f.__name__)
+            f()
+        print(module)
+    return f
+
+
+@run
+def testMemRefMultiBufferOpCompact():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.PROPAGATE,
+        [],
+        transform.OperationType.get("memref.alloc"),
+    )
+    with InsertionPoint(sequence.body):
+        memref.MemRefMultiBufferOp(sequence.bodyTarget, 4)
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testMemRefMultiBufferOpCompact
+    # CHECK: = transform.memref.multibuffer
+    # CHECK-SAME: factor = 4 : i64
+    # CHECK-SAME: (!transform.op<"memref.alloc">) -> !transform.any_op
+
+
+@run
+def testMemRefMultiBufferOpTyped():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.PROPAGATE,
+        [],
+        transform.OperationType.get("memref.alloc"),
+    )
+    with InsertionPoint(sequence.body):
+        memref.MemRefMultiBufferOp(
+            transform.OperationType.get("memref.alloc"), sequence.bodyTarget, 4
+        )
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testMemRefMultiBufferOpTyped
+    # CHECK: = transform.memref.multibuffer
+    # CHECK-SAME: factor = 4 : i64
+    # CHECK-SAME: (!transform.op<"memref.alloc">) -> !transform.op<"memref.alloc">
+
+
+@run
+def testMemRefMultiBufferOpAttributes():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.PROPAGATE,
+        [],
+        transform.OperationType.get("memref.alloc"),
+    )
+    with InsertionPoint(sequence.body):
+        memref.MemRefMultiBufferOp(sequence.bodyTarget, 4, skip_analysis=True)
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testMemRefMultiBufferOpAttributes
+    # CHECK: = transform.memref.multibuffer
+    # CHECK-SAME: factor = 4 : i64
+    # CHECK-SAME: skip_analysis
+    # CHECK-SAME: (!transform.op<"memref.alloc">) -> !transform.any_op

utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel

 )
 
 filegroup(
     name = "TransformOpsPyFiles",
     srcs = [
         "mlir/dialects/_bufferization_transform_ops_ext.py",
         "mlir/dialects/_gpu_transform_ops_ext.py",
         "mlir/dialects/_loop_transform_ops_ext.py",
+        "mlir/dialects/_memref_transform_ops_ext.py",
         "mlir/dialects/_structured_transform_ops_ext.py",
         "mlir/dialects/_transform_ops_ext.py",
         "mlir/dialects/_transform_pdl_extension_ops_ext.py",
         ":BufferizationTransformOpsPyGen",
         ":GPUTransformOpsPyGen",
         ":LoopTransformOpsPyGen",
         ":MemRefTransformOpsPyGen",
         ":PDLTransformOpsPyGen",