AffineIf fails with modular expression

MLIR
#1

An AffineExpr using an IntegerSet containing an expression of the form : (d0 mod N - K >= 0) where N > K + 1 fails when input into a AffineIf op. A compilation with the -debug flag reports : one or more local exprs do not have an explicit representation. Looking into the mlir code, it looks like it cannot represent an expression with a mod.

Caller: MLIR: lib/Analysis/AffineStructures.cpp Source File
Callee: MLIR: lib/Analysis/AffineStructures.cpp Source File

Is the callee function missing a call to detectAsMod()? Or can we just not have these types of expressions in an integer set?

#2

Can you please paste the minimal failing test case and what you ran?

#3 
const mlir::AffineExpr expr = rewriter.getAffineDimExpr(0) % 3 - 1;
const auto condition = mlir::IntegerSet::get(1, 0, llvm::ArrayRef(expr), llvm::ArrayRef(false);
const auto loop = rewriter.create<AffineForOp>(loc, 0, 10, 1);
rewriter.setInsertionPointToStart(loop.getBody());
{
  const auto branch = rewriter.create<mlir::AffineIfOp>(loc, condition, loop.getInductionVar(), false);
  rewriter.setInsertionPointToStart(branch.getThenBlock());
  {
    // do something here
  }
  rewriter.setInsertionPointAfter(branch);
}
rewriter.setInsertionPointAfter(loop);

This will compile but it will fail at runtime for the reason described earlier. The failure comes from the SimplifyAffineStructures pass.

#4

@bondhugula – Any ideas?

#5

Can you please just paste the corresponding test case (MLIR) and report what you ran and what you observed as output? From your original post, it isn’t clear whether you encountered an assertion or the input wasn’t simplified. If it’s an assertion, can you please file a bug on bugzilla if you have an account?

CC: @Groverkss

#6

This is a bug that I know of and already have a fix.

The issue here is that the inequality: d0 mod 3 - 1 >= 0 expands out to the set:

given inequality:     d0 - 3q0 - 1 >= 0
division upper bound: 3q0 - d0 + 2 >= 0
division lower bound: d0 - 3q0 >= 0

where q0 is a local variable

The problem is that when this is simplified, the given inequality makes the division lower bound redundant, and thus, the division lower bound is removed. The problem with this is, the division cannot be recovered anymore.

I already have a fix for this which allows the extraction of these kinds of divisions, which I plan to send after the current patch that works with divisions is landed (⚙ D112867 [MLIR][FlatAffineConstraints] Remove duplicate divisions while merging local ids).

#7

The assertion that I see is:
mlir/lib/Analysis/Utils.cpp:1076: mlir::IntegerSet mlir::simplifyIntegerSet(mlir::IntegerSet): Assertion simplifiedSet && "guaranteed to succeed while roundtripping"' failed.

and the error that I see is with -debug:
one or more local exprs do not have an explicit representation.

This is not a test case in MLIR, it is just some code I use to form an affine branch. @Groverkss seems to have found the issue. Thank you for the responses!

#8

@Groverkss – Do you have a rough ETA for this?

#9

Well, I have the patch fixing this ready and tested from my side. I’m just waiting for the previous patch on divisions to be accepted. The fix should be upstreamed maybe in 1-2 weeks. In case you need it earlier, you can apply this patch to your fork:

#10

Great, thanks

#11

Fixed by ⚙ D115096 [MLIR][FlatAffineConstraints] Add support for extracting divisions with tighter bounds

2 Likes