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

[Rolf Morel]

https://github.com/rolfmorel updated https://github.com/llvm/llvm-project/pull/81424

>From 33a0c28b1447537e51aabe2afd334bb203b4039d Mon Sep 17 00:00:00 2001
From: Rolf Morel <rolf.morel at huawei.com>
Date: Sun, 11 Feb 2024 17:43:45 +0000
Subject: [PATCH] [MLIR][Python] Add missing peel_front argument to
 LoopPeelOp's extension class

While PR 74015 added support for peeling an iteration from the front of a loop,
it did not update the (non-automatically generated) part of the Python
interface that exposes the op.
---
 mlir/python/mlir/dialects/transform/loop.py    |  6 ++++++
 .../test/python/dialects/transform_loop_ext.py | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+)

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():