[Mlir-commits] [mlir] 4c654b7 - [MLIR][Python] Add missing peel_front argument to LoopPeelOp's extension class (#81424)

[llvmlistbot at llvm.org]

Author: Rolf Morel
Date: 2024-02-12T11:35:43-06:00
New Revision: 4c654b7b91aff61728619fc3cc955fa5169d17c6

URL: https://github.com/llvm/llvm-project/commit/4c654b7b91aff61728619fc3cc955fa5169d17c6
DIFF: https://github.com/llvm/llvm-project/commit/4c654b7b91aff61728619fc3cc955fa5169d17c6.diff

LOG: [MLIR][Python] Add missing peel_front argument to LoopPeelOp's extension class (#81424)

Added: 
    

Modified: 
    mlir/python/mlir/dialects/transform/loop.py
    mlir/test/python/dialects/transform_loop_ext.py

Removed: 
    


################################################################################
diff  --git a/mlir/python/mlir/dialects/transform/loop.py b/mlir/python/mlir/dialects/transform/loop.py
index 3bdd9ca3b22f07..c4770b1c4067e4 100644
--- a/mlir/python/mlir/dialects/transform/loop.py
+++ b/mlir/python/mlir/dialects/transform/loop.py
@@ -55,6 +55,7 @@ def __init__(
         remainder_loop_type: Type,
         target: Union[Operation, Value],
         *,
+        peel_front: Union[bool, BoolAttr] = False,
         fail_if_already_divisible: Union[bool, BoolAttr] = False,
         ip=None,
         loc=None,
@@ -63,6 +64,11 @@ def __init__(
             main_loop_type,
             remainder_loop_type,
             _get_op_result_or_value(target),
+            peel_front=(
+                peel_front
+                if isinstance(peel_front, BoolAttr)
+                else BoolAttr.get(peel_front)
+            ),
             fail_if_already_divisible=(
                 fail_if_already_divisible
                 if isinstance(fail_if_already_divisible, BoolAttr)

diff  --git a/mlir/test/python/dialects/transform_loop_ext.py b/mlir/test/python/dialects/transform_loop_ext.py
index 840e7a46e7ce09..430b33fba04c71 100644
--- a/mlir/test/python/dialects/transform_loop_ext.py
+++ b/mlir/test/python/dialects/transform_loop_ext.py
@@ -49,6 +49,24 @@ def loopPeel():
     # CHECK-LABEL: TEST: loopPeel
     # CHECK: = transform.loop.peel %
 
+ at run
+def loopPeel_peel_front():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.Propagate,
+        [],
+        transform.OperationType.get("scf.for"),
+    )
+    with InsertionPoint(sequence.body):
+        loop.LoopPeelOp(
+            transform.AnyOpType.get(),
+            transform.AnyOpType.get(),
+            sequence.bodyTarget,
+            peel_front=True,
+        )
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: loopPeel_peel_front
+    # CHECK: = transform.loop.peel %[[ARG0:.*]] {peel_front = true}
+
 
 @run
 def loopPipeline():