Hit an assertion in Affine patterns when running canonicalizer

[hanhanW]

I hit an assertion when running canonicalizer. It's broken in one of affine.apply canonicalization patterns. It shouldn't crash given it's a canonicalization pattern.

To repro: mlir-opt --canonicalize a.mlir

#map0 = affine_map<(d0) -> (d0 mod 2)>
#map1 = affine_map<(d0, d1)[s0] -> ((d0 floordiv 2) * s0 + d1 floordiv 2)>
#map2 = affine_map<()[s0, s1, s2] -> (s0 * s1 + s2)>
#map3 = affine_map<()[s0, s1] -> (s0 floordiv s1)>
#map4 = affine_map<()[s0, s1] -> (s0 + s1)>
module {
  func.func @dynamic_unpack_simple_dispatch_0(%arg0: memref<?xi32>, %arg1: memref<?xi32>, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index) {
    %c0 = arith.constant 0 : index
    %c4 = arith.constant 4 : index
    %c2 = arith.constant 2 : index
    %c1 = arith.constant 1 : index
    scf.for %arg7 = %c0 to %arg5 step %c1 {
      scf.for %arg8 = %c0 to %arg6 step %c1 {
        %0 = affine.apply #map0(%arg7)
        %1 = affine.apply #map0(%arg8)
        %2 = affine.apply #map1(%arg7, %arg8)[%arg3]
        %3 = affine.apply #map2()[%2, %c2, %0]
        %4 = affine.apply #map2()[%3, %c2, %1]
        %5 = affine.apply #map3()[%arg2, %c4]
        %6 = affine.apply #map4()[%4, %5]
        %7 = memref.load %arg0[%6] : memref<?xi32>
        %8 = affine.apply #map2()[%arg7, %arg6, %arg8]
        memref.store %7, %arg1[%8] : memref<?xi32>
      }
    }
    return
  }
}

Error logs:

llvm-project/mlir/lib/IR/AffineExpr.cpp:978: auto getSemiAffineExprFromFlatForm(ArrayRef<int64_t>, unsigned int, unsigned int, ArrayRef<mlir::AffineExpr>, mlir::MLIRContext *)::(anonymous class)::operator()(std::pair<unsigned int, int>, int64_t, mlir::AffineExpr) const: Assertion `!llvm::is_contained(indices, index) && "Key is already present in indices vector and overwriting will " "happen in `indexToExprMap` and `coefficients`!"' failed.

[llvmbot]

@llvm/issue-subscribers-mlir-affine

[llvmbot]

@llvm/issue-subscribers-mlir

[matthias-springer]

Minimized test case:

#map2 = affine_map<()[s0] -> (((s0 floordiv 2) * s0) + (s0 mod 2))>
module {
  func.func @dynamic_unpack_simple_dispatch_0(%arg3: index) {
    %3 = affine.apply #map2()[%arg3] {__d__}
    vector.print %3 : index
    return
  }
}

[matthias-springer]

This is a bug in getSemiAffineExprFromFlatForm (introduced with https://reviews.llvm.org/D112808).

expression: (s0 floordiv 2) * s0 + s0 mod 2
numSyms = 1
numDims = 0
numLocals = 2
flat expression: [1, -2, 1, 0]

getAffineExprFromFlatForm can handle this expression, but getSemiAffineExprFromFlatForm fails.

getSemiAffineExprFromFlatForm computes an "index" for every part of the flat expression, so that the AffineExpr can be built in a certain order (as described in the comment of the function). I think the problem is that there are two subexprs with the same index.

index(s0) = (0, 1)
index(s0 floordiv 2) = (0, 1)

I don't know how to fix this index computation. @bondhugula @arnab-oss Any ideas?

[bondhugula]

This is a bug in getSemiAffineExprFromFlatForm (introduced with https://reviews.llvm.org/D112808).

expression: (s0 floordiv 2) * s0 + s0 mod 2
numSyms = 1
numDims = 0
numLocals = 2
flat expression: [1, -2, 1, 0]

getAffineExprFromFlatForm can handle this expression, but getSemiAffineExprFromFlatForm fails.

getSemiAffineExprFromFlatForm computes an "index" for every part of the flat expression, so that the AffineExpr can be built in a certain order (as described in the comment of the function). I think the problem is that there are two subexprs with the same index.

index(s0) = (0, 1)
index(s0 floordiv 2) = (0, 1)

I don't know how to fix this index computation. @bondhugula @arnab-oss Any ideas?

Thank you for reporting this -- this is clearly a bug. @arnab-oss will fix this.

[bjacob]

Thanks to everyone here --- very much interested in a resolution to unblock downstream https://github.com/iree-org/iree/issues/10796!

[arnab-oss]

This patch fixes the duplicate index issue.

[hanhanW]

Hi @arnab-oss thanks for the patch. I verified that it fixes the minimized test case, but not the original repro. Can you take a look at that case?

#map0 = affine_map<(d0) -> (d0 mod 2)>
#map1 = affine_map<(d0, d1)[s0] -> ((d0 floordiv 2) * s0 + d1 floordiv 2)>
#map2 = affine_map<()[s0, s1, s2] -> (s0 * s1 + s2)>
#map3 = affine_map<()[s0, s1] -> (s0 floordiv s1)>
#map4 = affine_map<()[s0, s1] -> (s0 + s1)>
module {
  func.func @dynamic_unpack_simple_dispatch_0(%arg0: memref<?xi32>, %arg1: memref<?xi32>, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index) {
    %c0 = arith.constant 0 : index
    %c4 = arith.constant 4 : index
    %c2 = arith.constant 2 : index
    %c1 = arith.constant 1 : index
    scf.for %arg7 = %c0 to %arg5 step %c1 {
      scf.for %arg8 = %c0 to %arg6 step %c1 {
        %0 = affine.apply #map0(%arg7)
        %1 = affine.apply #map0(%arg8)
        %2 = affine.apply #map1(%arg7, %arg8)[%arg3]
        %3 = affine.apply #map2()[%2, %c2, %0]
        %4 = affine.apply #map2()[%3, %c2, %1]
        %5 = affine.apply #map3()[%arg2, %c4]
        %6 = affine.apply #map4()[%4, %5]
        %7 = memref.load %arg0[%6] : memref<?xi32>
        %8 = affine.apply #map2()[%arg7, %arg6, %arg8]
        memref.store %7, %arg1[%8] : memref<?xi32>
      }
    }
    return
  }
}

[hanhanW]

I simplified that test case to

#map = affine_map<(d0) -> (d0 mod 2)>
#map1 = affine_map<(d0, d1)[s0] -> ((d0 floordiv 2) * s0 + d1 floordiv 2)>
#map2 = affine_map<()[s0, s1, s2] -> (s0 * s1 + s2)>
module {
  func.func @foo(%arg0: memref<?xi32>, %arg1: memref<?xi32>, %arg2: index, %arg3: index, %arg4: index) {
    %c0 = arith.constant 0 : index
    %c2 = arith.constant 2 : index
    %c1 = arith.constant 1 : index
    scf.for %arg5 = %c0 to %arg3 step %c1 {
      %0 = affine.apply #map(%arg5)
      %1 = affine.apply #map1(%arg5, %arg4)[%arg2]
      %2 = affine.apply #map2()[%1, %c2, %0]
      %3 = memref.load %arg0[%2] : memref<?xi32>
      vector.print %3 : i32
    }
    return
  }
}

[arnab-oss]

Hi @hanhanW , I'm not sure what is the issue here. All the expected semi affine expressions after simplification looks ok to me.

[matthias-springer]

getSemiAffineExprFromFlatForm crashes in the same location as before.

assertion failed at third_party/llvm/llvm-project/mlir/lib/IR/AffineExpr.cpp:979 in auto getSemiAffineExprFromFlatForm(ArrayRef<int64_t>, unsigned int, unsigned int, ArrayRef<AffineExpr>, MLIRContext *)::(anonymous class)::operator()(std::pair<unsigned int, int>, int64_t, AffineExpr) const: !llvm::is_contained(indices, index) && "Key is already present in indices vector and overwriting will " "happen in `indexToExprMap` and `coefficients`!"

The problematic AffineExpr is:

expression: ((d0 floordiv 2) * s0 + s1 floordiv 2) * 2 + d0 mod 2
numDims = 1
numSymbols = 2
numLocalExprs = 3
flat expression: [1, 0, 0, -2, 2, 2, 0]

[arnab-oss]

This patch resolves the bug. I've checked that it works for the original repro.

[hanhanW]

thanks for the quick fix!

[hanhanW]

I confirmed that it fixes the issue we've hit in IREE, closing the issue.