[mlir][bufferization][transform][python] Add enums to bindings & mixins.

[ingomueller-net]

Summary

This patch uses the new enum binding generation to add the enums of the
dialect to the Python bindings and uses them in the mix-in class where
it was still missing (namely, the LayoutMapOption for the
function_boundary_type_conversion of the OneShotBufferizeOp.

The patch also piggy-backs a few smaller clean-ups:

• Order the keyword-only arguments alphabetically.
• Add the keyword-only arguments to an overload where they were left out by accident.
• Change some of the attribute values used in the tests to non-default values such that they show up in the output IR and check for that output.

[ftynse]

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

782

I don't think this line is necessary here.

[ingomueller-net]

• Rebase to main. This should fix the build failure since the dependent patch has now landed.
• Remove unnecessary line from BUILD file.

mlir/python/CMakeLists.txt

declare_mlir_dialect_extension_python_bindings(
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/BufferizationTransformOps.td
   SOURCES
     dialects/_bufferization_transform_ops_ext.py
     dialects/transform/bufferization.py
   DIALECT_NAME transform
   EXTENSION_NAME bufferization_transform)
 
+set(LLVM_TARGET_DEFINITIONS "${CMAKE_CURRENT_SOURCE_DIR}/mlir/dialects/BufferizationTransformOps.td")
+mlir_tablegen("dialects/_bufferization_transform_enum_gen.py" -gen-python-enum-bindings)
+add_public_tablegen_target(MLIRBufferizationTransformDialectPyEnumGen)
+declare_mlir_python_sources(
+  MLIRPythonSources.Dialects.bufferization_transform.enum_gen
+  ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
+  ADD_TO_PARENT MLIRPythonSources.Dialects.bufferization_transform
+  SOURCES "dialects/_bufferization_transform_enum_gen.py")
+
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/GPUTransformOps.td
   SOURCES
     dialects/_gpu_transform_ops_ext.py
     dialects/transform/gpu.py
   DIALECT_NAME transform

mlir/python/mlir/dialects/_bufferization_transform_ops_ext.py

 #  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
 
 try:
     from ..ir import *
     from ..dialects import transform
 except ImportError as e:
     raise RuntimeError("Error loading imports from extension module") from e
 
+from enum import Enum
 from typing import Optional, overload, Union
 
 
 class EmptyTensorToAllocTensorOp:
     """Specialization for EmptyTensorToAllocTensorOp class."""
 
     @overload
     def __init__(
def __init__(
         self,
         transformed_type: Type,
         target: Union[Operation, OpView, Value],
         *,
         allow_return_allocs: Optional[bool] = None,
         allow_unknown_ops: Optional[bool] = None,
         bufferize_function_boundaries: Optional[bool] = None,
         create_deallocs: Optional[bool] = None,
-        test_analysis_only: Optional[bool] = None,
-        print_conflicts: Optional[bool] = None,
+        function_boundary_type_conversion: Optional[Enum] = None,
         memcpy_op: Optional[str] = None,
+        print_conflicts: Optional[bool] = None,
+        test_analysis_only: Optional[bool] = None,
         loc=None,
         ip=None
     ):
         ...
 
     @overload
-    def __init__(self, target: Union[Operation, OpView, Value], *, loc=None, ip=None):
+    def __init__(
+        self,
+        target: Union[Operation, OpView, Value],
+        *,
+        allow_return_allocs: Optional[bool] = None,
+        allow_unknown_ops: Optional[bool] = None,
+        bufferize_function_boundaries: Optional[bool] = None,
+        create_deallocs: Optional[bool] = None,
+        function_boundary_type_conversion: Optional[Enum] = None,
+        memcpy_op: Optional[str] = None,
+        print_conflicts: Optional[bool] = None,
+        test_analysis_only: Optional[bool] = None,
+        loc=None,
+        ip=None
+    ):
         ...
 
     def __init__(
         self,
         transformed_type_or_target: Type,
         target_or_none: Optional[Union[Operation, OpView, Value]] = None,
         *,
         allow_return_allocs: Optional[bool] = None,
         allow_unknown_ops: Optional[bool] = None,
         bufferize_function_boundaries: Optional[bool] = None,
         create_deallocs: Optional[bool] = None,
-        test_analysis_only: Optional[bool] = None,
-        print_conflicts: Optional[bool] = None,
+        function_boundary_type_conversion: Optional[Enum] = None,
         memcpy_op: Optional[str] = None,
+        print_conflicts: Optional[bool] = None,
+        test_analysis_only: Optional[bool] = None,
         loc=None,
         ip=None
     ):
         if isinstance(transformed_type_or_target, Type):
             transformed_type = transformed_type_or_target
             target = target_or_none
         else:
             transformed_type = transform.AnyOpType.get()
             target = transformed_type_or_target
 
         super().__init__(
             transformed_type,
             target,
             allow_return_allocs=allow_return_allocs,
             allow_unknown_ops=allow_unknown_ops,
             bufferize_function_boundaries=bufferize_function_boundaries,
             create_deallocs=create_deallocs,
-            test_analysis_only=test_analysis_only,
-            print_conflicts=print_conflicts,
+            function_boundary_type_conversion=function_boundary_type_conversion,
             memcpy_op=memcpy_op,
+            print_conflicts=print_conflicts,
+            test_analysis_only=test_analysis_only,
             loc=loc,
             ip=ip,
         )

mlir/python/mlir/dialects/transform/bufferization.py

 #  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
 
+from .._bufferization_transform_enum_gen import *
 from .._bufferization_transform_ops_gen import *

mlir/test/python/dialects/transform_bufferization_ext.py

sequence = transform.SequenceOp(
         transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get()
     )
     with InsertionPoint(sequence.body):
         bufferization.OneShotBufferizeOp(
             sequence.bodyTarget,
             allow_return_allocs=True,
             allow_unknown_ops=True,
             bufferize_function_boundaries=True,
-            create_deallocs=True,
-            test_analysis_only=True,
+            create_deallocs=False,
+            function_boundary_type_conversion=bufferization.LayoutMapOption.IDENTITY_LAYOUT_MAP,
+            memcpy_op="linalg.copy",
             print_conflicts=True,
-            memcpy_op="memref.copy",
+            test_analysis_only=True,
         )
         transform.YieldOp()
     # CHECK-LABEL: TEST: testOneShotBufferizeOpAttributes
     # CHECK: = transform.bufferization.one_shot_bufferize
+    # CHECK-SAME: layout{IdentityLayoutMap}
     # CHECK-SAME: allow_return_allocs = true
     # CHECK-SAME: allow_unknown_ops = true
     # CHECK-SAME: bufferize_function_boundaries = true
+    # CHECK-SAME: create_deallocs = false
+    # CHECK-SAME: memcpy_op = "linalg.copy"
     # CHECK-SAME: print_conflicts = true
     # CHECK-SAME: test_analysis_only = true
     # CHECK-SAME: (!transform.any_op) -> !transform.any_op

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

gentbl_filegroup(
     ],
 )
 
 gentbl_filegroup(
     name = "BufferizationTransformOpsPyGen",
     tbl_outs = [
         (
             [
+                "-gen-python-enum-bindings",
+                "-bind-dialect=transform",
+            ],
+            "mlir/dialects/_bufferization_transform_enums_gen.py",
+        ),
+        (
+            [
                 "-gen-python-op-bindings",
                 "-bind-dialect=transform",
                 "-dialect-extension=bufferization_transform",
             ],
             "mlir/dialects/_bufferization_transform_ops_gen.py",
         ),
     ],
     tblgen = "//mlir:mlir-tblgen",