Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Tensor Dialect canonicalizer FoldTensorCastProducerOp can produce invalid IR #91265

Closed
christopherbate opened this issue May 6, 2024 · 7 comments · Fixed by #92681
Closed

Tensor Dialect canonicalizer FoldTensorCastProducerOp can produce invalid IR #91265

christopherbate opened this issue May 6, 2024 · 7 comments · Fixed by #92681
Assignees
@matthias-springer
Labels
mlir:bufferization Bufferization infrastructure

Comments

@christopherbate
Copy link
Contributor

Reproducer:

// RUN: mlir-opt -canonicalize
func.func @materialize_in_destination_tensor_cast(%arg0: tensor<4xf32>, %arg1: index) -> tensor<?xf32> {
  %0 = bufferization.alloc_tensor(%arg1) : tensor<?xf32>
  %1 = tensor.cast %arg0 : tensor<4xf32> to tensor<?xf32>
  %2 = bufferization.materialize_in_destination %1 in %0 : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
  return %2 : tensor<?xf32>
}

FoldTensorCastProducerOp defined in lib/Dialect/Tensor/IR/TensorOps.cpp assumes that any operand of an op that implements DestinationStyleOpInterface can absorb a tensor.cast operation that is erasing information from the type. However, that assumption is not encoded anywhere in the spec of DestinationStyleOpInterface. For example, bufferization.materialize_in_destination` requires that the shapes of its DPS input and DPS output operands match, and therefore the canonicalizer will produce IR that fails verification :

within split at test/Dialect/Bufferization/canonicalize.mlir:392 offset :6:8: error: 'bufferization.materialize_in_destination' op failed to verify that all of {source, dest} have same shape
  %2 = bufferization.materialize_in_destination %1 in %0 : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
       ^
within split at test/Dialect/Bufferization/canonicalize.mlir:392 offset :6:8: note: see current operation: %1 = "bufferization.materialize_in_destination"(%arg0, %0) : (tensor<4xf32>, tensor<?xf32>) -> tensor<?xf32>
@christopherbate christopherbate mentioned this issue May 6, 2024
Closed
thexujie pushed a commit to thexujie/llvm-project that referenced this issue May 7, 2024
@christopherbate
[MLIR][Bufferization] Address tensor cast canonicalizer interaction w…
e597bde 
…ith `bufferization.materialize_in_destination`

Attempts to address a bug pointed out in llvm/llvm-project#91265
by relaxing the requirement for source/dest shapes to match in the
`bufferization.materialize_in_destination` operation. The relaxation
allows differences in static vs dynamic dims but still rejects cases
where the shapes are statically known to be different.
christopherbate added a commit to christopherbate/llvm-project that referenced this issue May 7, 2024
@christopherbate
[MLIR][Bufferization] Fix DPS op canonicalizer with tensor.cast`
3520470 
Attempts to address a bug pointed out in llvm#91265
by moving the FoldTensorCastProducerOp canonicalizer definition upward into
the MLIRDialectUtils library. Since the MLIRDialectUtils can't depend on any
dialect, the canonicalizer had to change slightly, and a templated version
is introduced.

Then, we need to add this canonicalization routine where it was used before,
except for places where it is incorrect as pointed out in the bug.
Based on cursory inspection of the TableGen definitions, only
`bufferization.materialize_in_destination` should *not* have the
canonicalizer, but existing tests passed if the canonicalizer as
only added for `tensor.pack|unpack|extract_slice` and the LinalgOp interface.
I went ahead and added tests where they were missing.
christopherbate added a commit to christopherbate/llvm-project that referenced this issue May 7, 2024
@christopherbate
[MLIR][Bufferization] Fix DPS op canonicalizer with tensor.cast`
7f87b72 
Attempts to address a bug pointed out in llvm#91265
by moving the FoldTensorCastProducerOp canonicalizer definition upward into
the MLIRDialectUtils library. Since the MLIRDialectUtils can't depend on any
dialect, the canonicalizer had to change slightly, and a templated version
is introduced.

Then, we need to add this canonicalization routine where it was used before,
except for places where it is incorrect as pointed out in the bug.
Based on cursory inspection of the TableGen definitions, only
`bufferization.materialize_in_destination` should *not* have the
canonicalizer, but existing tests passed if the canonicalizer as
only added for `tensor.pack|unpack|extract_slice` and the LinalgOp interface.
christopherbate added a commit to christopherbate/llvm-project that referenced this issue May 7, 2024
@christopherbate
[MLIR][Bufferization] Fix DPS op canonicalizer with tensor.cast`
809aa65 
Attempts to address a bug pointed out in llvm#91265
by moving the FoldTensorCastProducerOp canonicalizer definition upward into
the MLIRDialectUtils library. Since the MLIRDialectUtils can't depend on any
dialect, the canonicalizer had to change slightly, and a templated version
is introduced.

Then, we need to add this canonicalization routine where it was used before,
except for places where it is incorrect as pointed out in the bug.
Based on cursory inspection of the TableGen definitions, only
`bufferization.materialize_in_destination` should *not* have the
canonicalizer, but existing tests passed if the canonicalizer as
only added for `tensor.pack|unpack|extract_slice` and the LinalgOp interface.
@christopherbate christopherbate mentioned this issue May 7, 2024
Open
matthias-springer added a commit that referenced this issue May 19, 2024
@matthias-springer
[mlir][bufferization] Allow mixed static/dynamic shapes in `materiali…
150aa96 
…ze_in_destination` op

This commit relaxes the verifier of `bufferization.materialize_in_destination` such that mixed static/dynamic dimensions are allowed for the source and destination operands. E.g., `tensor<5xf32>` and `tensor<?xf32>` are now compatible, but it is assumed that the dynamic dimension is `5` at runtime.

This commit fixes #91265.
@matthias-springer matthias-springer mentioned this issue May 19, 2024
Merged
@matthias-springer
Copy link
Member

matthias-springer commented May 19, 2024 •
edited

#92681 fixes the issue for materialize_in_destination. A more general fix would make the canonicalization pattern opt-in for each op that supports it. There are already a few ops that are blacklisted in the canonicalization pattern (e.g., scf.forall) and it is not always possible/practical to add additional ops to that list.

(Generalizing materialize_in_destination as in #92681 still makes sense because memref.copy also supports mixed static/dynamic dimensions.)

@joker-eph
Copy link
Collaborator

This looks more like a workaround than a fix to me: a canonicalized shouldn't change the type of an SSA value without more information

@matthias-springer
Copy link
Member

I agree, it's likely still broken for other ops. I tried to make this pattern opt-in for each op that supports it. I.e., only ops that support the pattern should add it in their getCanonicalizationPatterns function; and rooting the pattern only the specified op name.

This requires a few more cleanups first, so I didn't get to it yet. E.g., tensor.pack, currently uses hasCanonicalizeMethod but we want hasCanonicalizer so that I can add the pattern. (I think the two cannot be used at the same time?)

@christopherbate
Copy link
Contributor Author

christopherbate commented May 20, 2024 •
edited

I agree, it's likely still broken for other ops. I tried to make this pattern opt-in for each op that supports it. I.e., only ops that support the pattern should add it in their getCanonicalizationPatterns function; and rooting the pattern only the specified op name.

@matthias-springer @joker-eph I opened PR #91382 two weeks ago to do just this. That PR makes the pattern opt-in.

@matthias-springer
Copy link
Member

Do you mean #91382?

I think there's a simpler solution (without that many templates): Add an additional StringRef opName constructor to FoldTensorCastProducerOp (can be passed to the Pattern constructor) such that the pattern applies only to the specified op name. (Same for op interface; in that case you need a TypeID.)

Then have a public function that populates the pattern for a given op name or interface type ID. The "populate function" can then be called from the getCanonicalizationPatterns functions of the ops that support this kind of transformation.

@matthias-springer
Copy link
Member

If you have time, I'd recommend two PRs: One that makes the canonicalizer of tensor.pack/unpack pattern-based (let hasCanonicalizer = 1). And another one that adds the opName/interfaceID constructor to FoldTensorCastProducerOp and adds the pattern to all necessary getCanonicalizationPatterns functions.

@christopherbate
Copy link
Contributor Author

Do you mean #91382?

Oops, yeah thanks.

If you have time, I'd recommend two PRs: One that makes the canonicalizer of tensor.pack/unpack pattern-based (let hasCanonicalizer = 1). And another one that adds the opName/interfaceID constructor to FoldTensorCastProducerOp and adds the pattern to all necessary getCanonicalizationPatterns functions.

Ok, should be able to work on it on Friday

matthias-springer added a commit that referenced this issue Jun 1, 2024
@matthias-springer
[mlir][bufferization] Allow mixed static/dynamic shapes in `materiali…
9d4b20a 
…ze_in_destination` op (#92681)

This commit relaxes the verifier of
`bufferization.materialize_in_destination` such that mixed
static/dynamic dimensions are allowed for the source and destination
operands. E.g., `tensor<5xf32>` and `tensor<?xf32>` are now compatible,
but it is assumed that the dynamic dimension is `5` at runtime.

This commit fixes #91265.
@EugeneZelenko EugeneZelenko added mlir:bufferization Bufferization infrastructure and removed mlir:tensor labels Jun 1, 2024
HendrikHuebner pushed a commit to HendrikHuebner/llvm-project that referenced this issue Jun 2, 2024
@matthias-springer @HendrikHuebner
[mlir][bufferization] Allow mixed static/dynamic shapes in `materiali…
519b2ca 
…ze_in_destination` op (llvm#92681)

This commit relaxes the verifier of
`bufferization.materialize_in_destination` such that mixed
static/dynamic dimensions are allowed for the source and destination
operands. E.g., `tensor<5xf32>` and `tensor<?xf32>` are now compatible,
but it is assumed that the dynamic dimension is `5` at runtime.

This commit fixes llvm#91265.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@matthias-springer matthias-springer

Labels
mlir:bufferization Bufferization infrastructure
Projects
None yet
Milestone
No milestone
4 participants
@matthias-springer @joker-eph @EugeneZelenko @christopherbate