This is an archive of the discontinued LLVM Phabricator instance.

[mlir][tensor|memref] Harden the checks on dim op
ClosedPublic

Authored by qcolombet on Jan 24 2023, 2:55 AM.

Details

Summary

Prior to this patch it was possible to use the dim operation on a 0-D
memref/tensor.
Unless we want to change the semantic of a 0-D shape, this doesn't make
sense because, paraphrasing the dim op semantic, this is guaranteed to
produce something that is undefined. (The requested index is guaranteed
to be equal to or greater than the rank.)

Harden the type requirements for the dim op by disallowing 0-D shaped
types.

This "fixes" llvm.org/PR60195 by rejecting dim op on 0-D shapes instead of
crashing during LLVM conversion.

Diff Detail

Event Timeline

qcolombet created this revision.Jan 24 2023, 2:55 AM
Herald added a project: Restricted Project. · View Herald Transcript
qcolombet requested review of this revision.Jan 24 2023, 2:55 AM
springerm added inline comments.Jan 24 2023, 3:29 AM
mlir/include/mlir/IR/OpBase.td
757–766

Side note: Adding a manual check to the op verifier would just be two lines. Are there any other ops that could benefit from Non0RankedTensorOf etc.?

qcolombet added inline comments.Jan 24 2023, 3:38 AM
mlir/include/mlir/IR/OpBase.td
757–766

Not that I know of on top of my head.

For some reasons, I feel that having the check in the td file is easier to discover (the "documentation" is all in one place), but I'm happy to move it to pure C++ in the verifier if that's what people prefer.

aartbik accepted this revision.Jan 24 2023, 8:50 AM
aartbik added inline comments.
mlir/include/mlir/IR/OpBase.td
757–766

Isn't the declarative way always preferred over the imperative C++ way? This looks just fine to me ;-)

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
957 ↗(On Diff #491681)

perhaps assign this to

uint64_t rank = ....

and then compare for readability

mlir/test/Dialect/MemRef/invalid.mlir
1064

yeah, thanks for rejecting this!

This revision is now accepted and ready to land.Jan 24 2023, 8:50 AM

I'm always a bit worried when we turn something from "UB" into "invalid IR": the issue being "action at distance". That is, are there patterns, fold, etc. that reduce the rank of a memref and build invalid IR unexpectedly?

I'm always a bit worried when we turn something from "UB" into "invalid IR": the issue being "action at distance". That is, are there patterns, fold, etc. that reduce the rank of a memref and build invalid IR unexpectedly?

I can see that, but at the same time, if that happens the generated code would be silently broken. If there were practical reasons to support that, I could live with it but I don't know if it is a pattern that happens in practice.
Anyhow, to be more concrete here, do you think we should support (as in don't crash the compiler) dim op of 0-D shapes?

This patch basically changes a crash on clearly UB IR (I guess we could come up with a semantic for dim of 0-D though) to the compiler emitting an error.

I'm always a bit worried when we turn something from "UB" into "invalid IR": the issue being "action at distance". That is, are there patterns, fold, etc. that reduce the rank of a memref and build invalid IR unexpectedly?

I can see that, but at the same time, if that happens the generated code would be silently broken.

That is the concept of UB though :)

If there were practical reasons to support that, I could live with it but I don't know if it is a pattern that happens in practice.
Anyhow, to be more concrete here, do you think we should support (as in don't crash the compiler) dim op of 0-D shapes?

We should not crash the compiler, we can generate ud2 / llvm.unreachable or anything instead :)

This patch basically changes a crash on clearly UB IR (I guess we could come up with a semantic for dim of 0-D though) to the compiler emitting an error.

Yes but you can't prove this path will be executed at runtime...

Yes but you can't prove this path will be executed at runtime...

Although true, I feel that if anything this change makes the verifier more consistent.
Right now, the verifier will reject dim ops that are known to be out-of-bound for non-0-ranked shapes. This patch extends the check to 0-ranked shapes.

If we follow the UB argument, we should not reject out-of-bound accesses for the non-0-ranked shapes either and instead generate unreachable.

What I am saying is I feel the UB semantic is usually our way out of things we can't check statically, but we actually don't stick to that when we can catch it in the verifier.

The bottom line is for me it feels arbitrary to decide that 0-ranked should produce unreachable whereas non-0-ranked don't.

Yes but you can't prove this path will be executed at runtime...

Although true, I feel that if anything this change makes the verifier more consistent.
Right now, the verifier will reject dim ops that are known to be out-of-bound for non-0-ranked shapes. This patch extends the check to 0-ranked shapes.

While I'm not sure about your patch (that is: it may very well be reasonable!), I am quite concerned about the out-of-bound check: any folding can make it so that we go from a dynamic offset to a constant one which would trigger a verifier failure (I need to craft a test-case...)

If we follow the UB argument, we should not reject out-of-bound accesses for the non-0-ranked shapes either and instead generate unreachable.

Yes absolutely :)

What I am saying is I feel the UB semantic is usually our way out of things we can't check statically

I think this is too limited of a view: this is a problem of scaling and "action at a distance": invalid IR differs from UB to me in terms of construction, and in particular the "dynamically reachable" is a critical aspect, take for example:

if (cond()) *(int*)nullptr;

We can very well statically check and forbid this, however it is possible that cond() is always false. While you can design a language where "by construction" this can't happen, you would still need to think what kind of program transformation can lead there.

The bottom line is for me it feels arbitrary to decide that 0-ranked should produce unreachable whereas non-0-ranked don't.

I agree, and as mentioned above, the 0-ranked is much more defendable than the other one to me! (That is because I'm not entirely sure about rank-reduction operation or how a transformation would expose the pattern in a way that justifies the original IR)

Here is some IR that passes the verifier, but crashes -canonicalize:

func.func @dim(%arg : tensor<1x?xf32>) -> index {
  %c2 = arith.constant 2 : index
  %add = arith.addi %c2, %c2 : index
  %dim = tensor.dim %arg, %add : tensor<1x?xf32>
  return %dim : index
}

The issue is that folding arith.addi is creating invalid IR: this is something that should never be possible, it breaks a strong invariant of the compiler.

The crash happens because the tensor.dim folder expects only valid IR (this is expected), but somehow we defined that tensor.dim with an out-of-bound constant access was invalid (hence my argument that this not a sound design for the verifier).

qcolombet abandoned this revision.Jan 25 2023, 8:38 AM

hence my argument that this not a sound design for the verifier.

Agree!
I was going for consistency but looks like what is already here is not the right thing to do.

I'll stir the change in the other direction.
Abandoning this PR.

Thanks @mehdi_amini!

Could you file an issue with your example?

qcolombet edited the summary of this revision. (Show Details)
  • Keep only the part of the change that disallow dim up with 0-ranked shapes. (I.e., remove the check with negative indices)
This revision is now accepted and ready to land.Jan 26 2023, 12:16 AM

Hi @mehdi_amini,

After talking with @aartbik and sleeping on it, I think making 0-ranked dim op invalid makes sense. I don't believe we can produce them with optimizations.

Regarding the out-of-bound accesses for the non-0-ranked shapes, I will do a separate patch to remove the verifier checks and do something sensible (unreachable, ..., anything but crash).
Thanks for filing the issue BTW.

What do you think?

Cheers,
-Quentin

Hi @mehdi_amini,

After talking with @aartbik and sleeping on it, I think making 0-ranked dim op invalid makes sense. I don't believe we can produce them with optimizations.

Seems reasonable for now: producing this would involve a type change, which isn't a "normal" safe thing to do (which is why I wasn't sure about this).

The part that is a bit speculative that worried me a bit is how it creates an "edge case" make it not uniform to handle. For example take some pseudo-IR like this:

// return the first dimension or 0 if rank == 0.
func @unranked(%arg0 : memref<*xf32>) : index {
  %rank = memref.rank %arg0  : memref<*xf32>
  %zero = arith.constant 0 : index
  %rank_not_zero = arith.icmpi ne, %rank, %zero : (index, index) -> i1
  %res = scf.if (%rank) {
     %dim = memref.dim %arg0[0] : memref<*xf32>
     scf.yield %dim
  } else {
     scf.yield %zero
  }
  return %res
}

This is fairly useless code, but just to illustrate my point: assume this is a generic / reusable routine part of a library.
Now when we integrate this in a larger program we could try to do some inlining and/or function specialization, to turn the unranked into ranked. So if you have a call-site with a rank-1 memref you would specialize as such:

func @unranked1(%arg0 : memref<?xf32>) : index {
  %rank = memref.rank %arg0  : memref<?xf32>
  %zero = arith.constant 0 : index
  %rank_not_zero = arith.icmpi ne, %rank, %zero : (index, index) -> i1
  %res = scf.if (%rank) {
     %dim = memref.dim %arg0[0] : memref<?xf32>
     scf.yield %dim
  } else {
     scf.yield %zero
  }
  return %res
}

which then can fold (memref.rank folds to 1):

func @unranked1(%arg0 : memref<?xf32>) : index {
   %dim = memref.dim %arg0[0] : memref<?xf32>
  return %res
}

But the same logic trying to specialize for rank 0 is impossible:

func @unranked0(%arg0 : memref<f32>) : index {
  %rank = memref.rank %arg0  : memref<f32>
  %zero = arith.constant 0 : index
  %rank_not_zero = arith.icmpi ne, %rank, %zero : (index, index) -> i1
  %res = scf.if (%rank) {
     %dim = memref.dim %arg0[0] : memref<f32> // <= verifier error
     scf.yield %dim
  } else {
     scf.yield %zero
  }
  return %res
}

Even though the folding would yield:

func @unranked0(%arg0 : memref<f32>) : index {
  %zero = arith.constant 0 : index
  return %zero
}

Regarding the out-of-bound accesses for the non-0-ranked shapes, I will do a separate patch to remove the verifier checks and do something sensible (unreachable, ..., anything but crash).

Thanks!

producing this would involve a type change, which isn't a "normal" safe thing to do (which is why I wasn't sure about this).

My thinking exactly.

I am going to go ahead with the patch as is for now, but I found the exact same issue with shape.dim.
I don't know what the fix should look like though because the description says:

Gets the extent indexed by `dim` from the shape of the `value` operand. If
the index is error or out-of-bound then it returns an invalid size if the
return type carries error information else the behavior is undefined.

What is this type carries error thing and how do you propagate the information there?
Right now we would just bail out in the verifier.

What is this type carries error thing and how do you propagate the information there?

The shape dialect has a special type to model a "shape" and this type has an invalid state: https://mlir.llvm.org/docs/Dialects/ShapeDialect/#shapetype

The shape dialect has a special type to model a "shape" and this type has an invalid state: https://mlir.llvm.org/docs/Dialects/ShapeDialect/#shapetype

Thanks!

I don't see how it is used though. Meaning, I don't see any actual implementation using/supporting that.

E.g., in the shape-to-standard conversion pass I see that pretty much all operations have something resembling:

// For now, only error-free types are supported by this lowering.
if (op.getType().isa<xxxType>())
  return failure();

I don't see any test with an invalid input or an invalid output.

This revision was automatically updated to reflect the committed changes.

I landed the patch as is like we discussed.
I have a fix for https://github.com/llvm/llvm-project/issues/60295 that I need to send for review but it doesn't include the shape dialect.

I'll file a different issue for that part for now, while we figure out what it means to return invalid in the shape dialect.