Skip to content

Commit 1335a05

Browse files
authored
[MLIR][Python] Fix AffineIfOp insertion point (#171957)
This bug was introduced by #108323, where the loc and ip were not properly set. It may lead to errors when the operations are not linearly asserted to the IR.
1 parent ecaf673 commit 1335a05

File tree

2 files changed

+29
-1
lines changed

2 files changed

+29
-1
lines changed

mlir/python/mlir/dialects/affine.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -198,7 +198,7 @@ def __init__(
198198
results = []
199199
results.extend(results_)
200200

201-
super().__init__(results, cond_operands, cond)
201+
super().__init__(results, cond_operands, cond, loc=loc, ip=ip)
202202
self.regions[0].blocks.append(*[])
203203
if has_else:
204204
self.regions[1].blocks.append(*[])

mlir/test/python/dialects/affine.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -333,3 +333,31 @@ def simple_affine_if_else(cond_operands):
333333
affine.AffineYieldOp([x_false, y_false])
334334
add = arith.AddIOp(if_op.results[0], if_op.results[1])
335335
return
336+
337+
338+
# CHECK-LABEL: TEST: testAffineIfOpInsertionPoint
339+
@constructAndPrintInModule
340+
def testAffineIfOpInsertionPoint():
341+
index = IndexType.get()
342+
i32 = IntegerType.get_signless(32)
343+
d0 = AffineDimExpr.get(0)
344+
cond = IntegerSet.get(1, 0, [d0 - 5], [False])
345+
346+
# CHECK: func.func @affine_if_insertion_point_test
347+
@func.FuncOp.from_py_func(index)
348+
def affine_if_insertion_point_test(cond_operands):
349+
first_const = arith.ConstantOp(i32, 42)
350+
ip_before = InsertionPoint(first_const)
351+
second_const = arith.ConstantOp(i32, 100)
352+
353+
# Now create AffineIfOp with ip pointing BEFORE first_const
354+
# If ip works correctly, the if_op should appear BEFORE first_const
355+
if_op = affine.AffineIfOp(cond, cond_operands=[cond_operands], ip=ip_before)
356+
with InsertionPoint(if_op.then_block):
357+
affine.AffineYieldOp([])
358+
359+
# CHECK: affine.if
360+
# CHECK-NEXT: }
361+
# CHECK-NEXT: %c42_i32 = arith.constant 42 : i32
362+
# CHECK-NEXT: %c100_i32 = arith.constant 100 : i32
363+
return

0 commit comments

Comments
 (0)