AffineMap is not simplified in scf.if regions

I am expecting %b in the code snippet to be reduced to arith.constant 1 : index, but it is not simplified under MLIR. Any ideas why?

Code snippet:

#map1 = affine_map<(a)[b] -> (a - b + 1)>
func.func @test(%a : index, %a0 : index, %a1 : index, %cond : i1) -> (index){
        %out = scf.if %cond -> (index){
                %k = arith.addi %a1, %a0 : index
                %b = affine.apply #map1(%k)[%k]
                scf.yield %b : index

        } else {
          scf.yield %a : index
        return %out: index

LLVM/MLIR version: 14.0.0

The canonicalization treats affine dimensions and symbols differently and, in this case, cannot promote the dimension a to be a symbol because it is defined in an operation that is not an affine scope. We could potentially add the simplification that works with a flat operand list without considering the symbols, but we need to be careful not to break not to break the semantics in the case the result of the operation is used in another affine operation (this leads to maps being composed before application, and dimensions compose differently than symbols).

1 Like

Thank you for the reply. However, when I remove the %k = arith.addi %a1, %a0 : index and change to %b = affine.apply #map1(%a0)[%a0], the reduction works perfectly well. Perhaps, it might be a special case in canonicalization.

%a0 is defined in an operation that is an affine scope (it’s the argument of the func.func entry block), so it can be promoted to a symbol. The canonicalization rewrites the map to be affine_map<()[a,b] -> (a - b + 1)> and then “merges” the two symbols locally because they are passed the same value.

I see. Thank you for clarifying!

Two things here:

  1. scf.if (and scf.for) are supposed to be marked with an AffineScope trait. That will make %k a symbol (and rightly so). The desired simplification will then happen. Otherwise, as @ftynse mentions, the map won’t be simplified. In the latter cases, the desired simplification is expected to happen on “application” (de-abstraction/lowering) of the affine.apply.
  2. Separately, if you lower the affine operations away (with -lower-affine), the simplification should happen (the subi should fold away or some algebraic simplification downstream should take care of it) – but I can understand we expect better/earlier simplification.

@lituzou - can you add the AffineScope trait to the op definition of scf IfOp and try?

@ftynse Note that the fact that %k currently not being treated as a symbol but as a dimension is a problem with the current codebase. %k can’t be a dimensional identifier! (because it’s not an affine function of loop IVs) – it is and can only be a symbol. The issue is that certain ops of interest scf.if/scf.for don’t have the AffineScope trait attached. Also, my revision (⚙ D112891 [MLIR][Affine] Replace AffineScope by its complement trait) completely fixes this problem in the process of changing the default trait – %k would be correctly inferred as a symbol over there. That would effectively be the same as adding the AffineScope trait on the scf IfOp and that would still be a valid/acceptable change if someone wants to commit it.

Adding AffineScope to the op definition did work.

I am currently considering another revision, but I am not very sure about it:
If a dim operand and a symbol refer to the same ssa value (like %k in the first example), can we automatically promote the dim operand to a symbol?
#map(%k)[%k] -> #map()[%k,%k]

You shouldn’t – canonicalization already does this for you. Just run -canonicalize, or, if needed in your pass, add the canonicalization patterns for relevant affine ops.

I’m merely describing the current state. The ball is on your side with the revision AFAICS.