Page MenuHomePhabricator

Please use GitHub pull requests for new patches. Avoid migrating existing patches. Phabricator shutdown timeline

[mlir][Linalg] Fix crash in `dropUnitDims`
Needs RevisionPublic

Authored by springerm on Aug 25 2023, 3:00 AM.



dropUnitDims is a transformation that rank-reduces the shape of operands of Linalg ops: dimensions of size 1 are removed. The transformation used to crash or produce invalid IR when dropping a dynamic unit dimensions. I.e., the dimension was statically detected to have size 1 because the dim size of another operand (mapped to the same loop) is a static 1.

This fixes #64977.

Diff Detail

Event Timeline

springerm created this revision.Aug 25 2023, 3:00 AM
springerm requested review of this revision.Aug 25 2023, 3:00 AM
mravishankar requested changes to this revision.Aug 25 2023, 11:11 PM

Could you post some more info on the crash? Crash is not good, so that is worth fixing, but IMO if the two dimensions are equal then one of them shouldn't be dynamic. I would expect the canonicalization to make both dimensions equal before calling these patterns. Having every pattern account for such inconsistencies is going to be very onerous. But it really depends on the crash. So some more info on the crash would help in review. Thanks!

This revision now requires changes to proceed.Aug 25 2023, 11:11 PM

There are two cases.

  1. When lowering to tensor.collapse_shape:
error: 'memref.collapse_shape' op collapsed dim (0) must be dynamic if and only if reassociation group is dynamic
%1 = "memref.collapse_shape"(%arg0) <{reassociation = [[0, 1], [2]]}> : (memref<?x28x4xf16>) -> memref<28x4xf16>

We could also fix this by making the result type memref<?x4xf16>, but we would loose some static type information.

  1. When lowering to tensor.extract_slice:
assert.h assertion failed at third_party/llvm/llvm-project/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp:293 in Value collapseValue(RewriterBase &, Location, Value, ArrayRef<int64_t>, ArrayRef<ReassociationIndices>, ControlDropUnitDims::RankReductionStrategy): succeeded(rankReducingExtract) && "not a unit-extent collapse"

Same thing, we know statically that the dimension is 1, but it is actually dynamic. It may be as simple as removing the assertion.

The canonicalizer does currently not promote the ? to 1, but I think it would make sense to do that (and ignore operands with ? dims in this transformation).