diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -178,6 +178,7 @@ 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) diff --git a/mlir/python/mlir/dialects/_memref_transform_ops_ext.py b/mlir/python/mlir/dialects/_memref_transform_ops_ext.py new file mode 100644 --- /dev/null +++ b/mlir/python/mlir/dialects/_memref_transform_ops_ext.py @@ -0,0 +1,68 @@ +# 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 typing import Optional, overload, Union + + +class MemRefMultiBufferOp: + """Specialization for MemRefMultiBufferOp class.""" + + @overload + def __init__( + self, + transformed_type: Type, + target: Union[Operation, OpView, Value], + factor: Union[int, IntegerAttr], + *, + skip_analysis: Optional[bool] = None, + loc=None, + ip=None + ): + ... + + @overload + def __init__( + self, + target: Union[Operation, OpView, Value], + factor: Union[int, IntegerAttr], + *, + skip_analysis: Optional[bool] = None, + loc=None, + ip=None + ): + ... + + def __init__( + self, + transformed_type_or_target: Type, + target_or_factor: Union[int, IntegerAttr, Operation, OpView, Value] = None, + factor_or_none: Optional[Union[int, IntegerAttr]] = None, + *, + skip_analysis: Optional[bool] = None, + loc=None, + ip=None + ): + if isinstance(transformed_type_or_target, Type): + transformed_type = transformed_type_or_target + target = target_or_factor + factor = factor_or_none + else: + transformed_type = transform.AnyOpType.get() + target = transformed_type_or_target + factor = target_or_factor + + super().__init__( + transformed_type, + target, + factor, + skip_analysis=skip_analysis, + loc=loc, + ip=ip, + ) diff --git a/mlir/test/python/dialects/transform_memref_ext.py b/mlir/test/python/dialects/transform_memref_ext.py new file mode 100644 --- /dev/null +++ b/mlir/test/python/dialects/transform_memref_ext.py @@ -0,0 +1,67 @@ +# 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 diff --git a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel --- a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel @@ -947,6 +947,7 @@ "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",