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

[qcolombet]

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.

[springerm]

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]

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]

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

perhaps assign this to

uint64_t rank = ....

and then compare for readability

mlir/test/Dialect/MemRef/invalid.mlir

1064

yeah, thanks for rejecting this!

[mehdi_amini]

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?

[qcolombet]

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.

[mehdi_amini]

In D142445#4077807, @qcolombet wrote:

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...

[qcolombet]

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.

[mehdi_amini]

In D142445#4079181, @qcolombet wrote:

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)

[mehdi_amini]

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]

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?

[mehdi_amini]

https://github.com/llvm/llvm-project/issues/60295

[qcolombet]

• Keep only the part of the change that disallow dim up with 0-ranked shapes. (I.e., remove the check with negative indices)

[qcolombet]

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

[mehdi_amini]

In D142445#4081964, @qcolombet wrote:

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!

[qcolombet]

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.

[qcolombet]

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.

[mehdi_amini]

In D142445#4085082, @qcolombet wrote:

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

[qcolombet]

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.

[qcolombet]

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.

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let description = [{
     %y = memref.dim %A, %c1 : memref<4 x ? x f32>
 
     // Equivalent generic form:
     %x = "memref.dim"(%A, %c0) : (memref<4 x ? x f32>, index) -> index
     %y = "memref.dim"(%A, %c1) : (memref<4 x ? x f32>, index) -> index
     ```
   }];
 
-  let arguments = (ins AnyRankedOrUnrankedMemRef:$source,
+  let arguments = (ins AnyNon0RankedOrUnrankedMemRef:$source,
                        Index:$index);
   let results = (outs Index:$result);
 
   let assemblyFormat = [{
     attr-dict $source `,` $index `:` type($source)
   }];
 
   let builders = [

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

let description = [{
     %y = tensor.dim %A, %c1 : memref<4x?xf32>
 
     // Equivalent generic form:
     %x = "tensor.dim"(%A, %c0) : (memref<4x?xf32>, index) -> index
     %y = "tensor.dim"(%A, %c1) : (memref<4x?xf32>, index) -> index
     ```
   }];
 
-  let arguments = (ins AnyTensor:$source,
+  let arguments = (ins AnyNon0RankedOrUnrankedTensor:$source,
                        Index:$index);
   let results = (outs Index:$result);
 
   let assemblyFormat = [{
     attr-dict $source `,` $index `:` type($source)
   }];
 
   let builders = [

mlir/include/mlir/IR/OpBase.td

 // Whether a shaped type has one of the specified ranks.
 class HasAnyRankOfPred<list<int> ranks> : And<[
     HasRankPred,
     Or<!foreach(rank, ranks,
                 CPred<[{$_self.cast<::mlir::ShapedType>().getRank()
                          == }]
                       # rank>)>]>;
 
+// Whether a shaped type has a rank greater than or equal of the specified rank.
+class HasRankGreaterOrEqualPred<int rank> : And<[
+    HasRankPred,
+    CPred<[{$_self.cast<::mlir::ShapedType>().getRank() >= }] # rank>
+]>;
+
 // Vector types.
 
 class VectorOf<list<Type> allowedTypes> :
   ShapedContainerType<allowedTypes, IsVectorTypePred, "vector",
                       "::mlir::VectorType">;
 
 // Temporary vector type clone that allows gradual transition to 0-D vectors.
 // TODO: Remove this when all ops support 0-D vectors.

 def F64Tensor  : TensorOf<[F64]>;
 
 class RankedTensorOf<
     list<Type> allowedTypes,
     list<Pred> preds = [],
     string summary = "ranked tensor">
   : TensorOf<allowedTypes, !listconcat([HasRankPred], preds), summary>;
 
+class Non0RankedTensorOf<list<Type> allowedTypes>
+  : TensorOf<allowedTypes, [HasRankGreaterOrEqualPred<1>],
+      "non-0-ranked.tensor">;
+
 def AnyRankedTensor : RankedTensorOf<[AnyType]>;
+def AnyNon0RankedTensor  : Non0RankedTensorOf<[AnyType]>;
+def AnyUnrankedTensor  : UnrankedTensorOf<[AnyType]>;
+
+def AnyNon0RankedOrUnrankedTensor:
+    AnyTypeOf<[AnyUnrankedTensor, AnyNon0RankedTensor]>;
 
 // Ranked tensor type with one of the specified types and ranks.
 class TensorRankOf<list<Type> allowedTypes, list<int> ranks>
   : TensorOf<allowedTypes,
       [HasAnyRankOfPred<ranks>],
       !interleave(!foreach(rank, ranks, rank # "D"), "/") # " tensor">;
 
 class 0DTensorOf<list<Type> allowedTypes> : TensorRankOf<allowedTypes, [0]>;
ShapedContainerType<allowedTypes,
                         "::mlir::UnrankedMemRefType">;
 
 def AnyUnrankedMemRef : UnrankedMemRefOf<[AnyType]>;
 
 // Memrefs are blocks of data with fixed type and rank.
 class MemRefOf<list<Type> allowedTypes> :
     ShapedContainerType<allowedTypes, IsMemRefTypePred, "memref",
                         "::mlir::MemRefType">;
+class Non0RankedMemRefOf<list<Type> allowedTypes> :
+    ConfinedType<MemRefOf<allowedTypes>, [HasRankGreaterOrEqualPred<1>],
+         "non-0-ranked." # MemRefOf<allowedTypes>.summary,
+         "::mlir::MemRefType">;
 
 def AnyMemRef : MemRefOf<[AnyType]>;
+def AnyNon0RankedMemRef : Non0RankedMemRefOf<[AnyType]>;
 
 class RankedOrUnrankedMemRefOf<list<Type> allowedTypes>:
     AnyTypeOf<[UnrankedMemRefOf<allowedTypes>, MemRefOf<allowedTypes>]>;
 
 def AnyRankedOrUnrankedMemRef: AnyTypeOf<[AnyUnrankedMemRef, AnyMemRef]>;
+def AnyNon0RankedOrUnrankedMemRef:
+    AnyTypeOf<[AnyUnrankedMemRef, AnyNon0RankedMemRef]>;
 
 // Memref declarations handle any memref, independent of rank, size, (static or
 // dynamic), layout, or memory space.
 def I1MemRef  : MemRefOf<[I1]>;
 def I8MemRef  : MemRefOf<[I8]>;
 def I16MemRef : MemRefOf<[I16]>;
 def I32MemRef : MemRefOf<[I32]>;
 def I64MemRef : MemRefOf<[I64]>;

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

 }
 
 LogicalResult DimOp::verify() {
   // Assume unknown index to be in range.
   std::optional<int64_t> index = getConstantIndex();
   if (!index)
     return success();
 
+  // Convert index to unsigned since negative indices don't make sense.
+  // In other words, if we get a negative index assume we want a big value here.
+  uint64_t uindex = *index;
+
   // Check that constant index is not knowingly out of range.
   auto type = getSource().getType();
   if (auto memrefType = type.dyn_cast<MemRefType>()) {
-    if (*index >= memrefType.getRank())
+    if (uindex >= (uint64_t)memrefType.getRank())
       return emitOpError("index is out of range");
   } else if (type.isa<UnrankedMemRefType>()) {
     // Assume index to be in range.
   } else {
     llvm_unreachable("expected operand with memref type");
   }
   return success();
 }

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

 }
 
 LogicalResult DimOp::verify() {
   // Assume unknown index to be in range.
   std::optional<int64_t> index = getConstantIndex();
   if (!index)
     return success();
 
+  // Convert index to unsigned since negative indices don't make sense.
+  // In other words, if we get a negative index assume we want a big value here.
+  uint64_t uindex = *index;
+
   // Check that constant index is not knowingly out of range.
   auto type = getSource().getType();
   if (auto tensorType = type.dyn_cast<RankedTensorType>()) {
-    if (*index >= tensorType.getRank())
+    if (uindex >= (uint64_t)tensorType.getRank())
       return emitOpError("index is out of range");
   } else if (type.isa<UnrankedTensorType>()) {
     // Assume index to be in range.
   } else {
     llvm_unreachable("expected operand with tensor type");
   }
   return success();
 }

mlir/test/Dialect/MemRef/invalid.mlir

 
 // -----
 
 func.func @memref_realloc_type(%src : memref<256xf32>) -> memref<?xi32>{
   // expected-error@+1 {{different element types}}
   %0 = memref.realloc %src : memref<256xf32> to memref<?xi32>
   return %0 : memref<?xi32>
 }
+
+// -----
+
+func.func @memref_dim(%arg : memref<1x?xf32>) {
+  %c2 = arith.constant 2 : index
+  memref.dim %arg, %c2 : memref<1x?xf32> // expected-error {{'memref.dim' op index is out of range}}
+  return
+}
+
+// -----
+
+func.func @dim_negative_idx(%arg : memref<1x?xf32>) {
+  %c_minus_2 = arith.constant -2 : index
+  memref.dim %arg, %c_minus_2 : memref<1x?xf32> // expected-error {{'memref.dim' op index is out of range}}
+  return
+}
+
+// -----
+
+// Asking the dimension of a 0-D shape doesn't make sense.
+func.func @dim_negative_idx(%arg : memref<f32>, %arg1 : index) {
+  memref.dim %arg, %arg1 : memref<f32> // expected-error {{'memref.dim' op operand #0 must be unranked.memref of any type values or non-0-ranked.memref of any type values, but got 'memref<f32>'}}
+  return
+}

mlir/test/Dialect/Tensor/bufferize.mlir

 // RUN: mlir-opt %s -tensor-bufferize -cse -split-input-file -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL:   func @dim(
-// CHECK-SAME:              %[[TENSOR:.*]]: tensor<f32>,
+// CHECK-SAME:              %[[TENSOR:.*]]: tensor<*xf32>,
 // CHECK-SAME:              %[[INDEX:.*]]: index) -> index {
-// CHECK:           %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<f32>
-// CHECK:           %[[EXTENT:.*]] = memref.dim %[[MEMREF]], %[[INDEX]] : memref<f32>
+// CHECK:           %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<*xf32>
+// CHECK:           %[[EXTENT:.*]] = memref.dim %[[MEMREF]], %[[INDEX]] : memref<*xf32>
 // CHECK:           return %[[EXTENT]] : index
-func.func @dim(%arg0: tensor<f32>, %arg1: index) -> index {
-  %0 = tensor.dim %arg0, %arg1 : tensor<f32>
+func.func @dim(%arg0: tensor<*xf32>, %arg1: index) -> index {
+  %0 = tensor.dim %arg0, %arg1 : tensor<*xf32>
   return %0 : index
 }
 
 // -----
 
 // CHECK-LABEL: func @rank(
 // CHECK-SAME:    %[[TENSOR:.*]]: tensor<*xf32>) -> index {
 // CHECK:           %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]]

mlir/test/Dialect/Tensor/invalid.mlir

 // RUN: mlir-opt <%s -split-input-file -verify-diagnostics
 
 func.func @dim(%arg : tensor<1x?xf32>) {
   %c2 = arith.constant 2 : index
   tensor.dim %arg, %c2 : tensor<1x?xf32> // expected-error {{'tensor.dim' op index is out of range}}
   return
 }
 
 // -----
 
+func.func @dim_negative_idx(%arg : tensor<1x?xf32>) {
+  %c_minus_2 = arith.constant -2 : index
+  tensor.dim %arg, %c_minus_2 : tensor<1x?xf32> // expected-error {{'tensor.dim' op index is out of range}}
+  return
+}
+
+// -----
+
+// Asking the dimension of a 0-D shape doesn't make sense.
+func.func @dim_negative_idx(%arg : tensor<f32>, %arg1 : index) {
+  tensor.dim %arg, %arg1 : tensor<f32> // expected-error {{'tensor.dim' op operand #0 must be unranked.tensor of any type values or non-0-ranked.tensor of any type values, but got 'tensor<f32>'}}
+  return
+}
+
+// -----
+
 func.func @tensor.cast_mismatching_constants(%arg0: tensor<1xf32>) {
   // expected-error@+1 {{operand type 'tensor<1xf32>' and result type 'tensor<2xf32>' are cast incompatible}}
   %0 = tensor.cast %arg0 : tensor<1xf32> to tensor<2xf32>
   return
 }
 
 // -----