[mlir][MemRef|Tensor] Fix the handling of DimOp

[qcolombet]

Summary

Although specifying an index that is out of bounds for both memref.dim
and tensor.dim produces an undefined behavior, this is still valid IR.
In particular, we could expose an out of bound index because of some
optimizations, for instance as demonstrated with
https://github.com/llvm/llvm-project/issues/60295, and this shouldn't
cause the compiler to abort.

This patch removes the overzealous verifier checks and properly handles
out of bound indices (as in it doesn't crash the compiler, but still
produces UB).

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

Note: That shape.dim has a similar problem but we're not supposed to
produce UB in this case. Instead we're supposed to propagate an error in
the resulting value and I don't know how to do that at the moment. Hence I
left this part out of the patch.

[qcolombet]

Filed https://github.com/llvm/llvm-project/issues/60808 for the remaining shape.dim issue.

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

let extraClassDeclaration = [{
     OpFoldResult getDimension() { return getIndex(); }
 
     /// Interface method for ConditionallySpeculatable.
     Speculation::Speculatability getSpeculatability();
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
-  let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DmaStartOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_DmaStartOp : MemRef_Op<"dma_start"> {
   let summary = "non-blocking DMA operation that starts a transfer";

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

let extraClassDeclaration = [{
     OpFoldResult getDimension() { return getIndex(); }
 
     /// Interface method for ConditionallySpeculatable.
     Speculation::Speculatability getSpeculatability();
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
-  let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // EmptyOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_EmptyOp : Tensor_Op<"empty",
     [Pure,

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

Value extractSizeOfRankedMemRef(Type operandType, memref::DimOp dimOp,
                                   OpAdaptor adaptor,
                                   ConversionPatternRewriter &rewriter) const {
     Location loc = dimOp.getLoc();
 
     // Take advantage if index is constant.
     MemRefType memRefType = operandType.cast<MemRefType>();
     if (std::optional<int64_t> index = getConstantDimIndex(dimOp)) {
       int64_t i = *index;
+      if (i >= 0 && i < memRefType.getRank()) {
         if (memRefType.isDynamicDim(i)) {
           // extract dynamic size from the memref descriptor.
           MemRefDescriptor descriptor(adaptor.getSource());
           return descriptor.size(rewriter, loc, i);
         }
         // Use constant for static size.
         int64_t dimSize = memRefType.getDimSize(i);
         return createIndexConstant(rewriter, loc, dimSize);
       }
+    }
     Value index = adaptor.getIndex();
     int64_t rank = memRefType.getRank();
     MemRefDescriptor memrefDescriptor(adaptor.getSource());
     return memrefDescriptor.size(rewriter, loc, index, rank);
   }
 };
 
 /// Common base for load and store operations on MemRefs. Restricts the match

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

Speculation::Speculatability DimOp::getSpeculatability() {
   if (!rankedSourceType)
     return Speculation::NotSpeculatable;
 
   // The verifier rejects operations that violate this assertion.
   assert(constantIndex < rankedSourceType.getRank());
   return Speculation::Speculatable;
 }
 
-LogicalResult DimOp::verify() {
-  // Assume unknown index to be in range.
-  std::optional<int64_t> index = getConstantIndex();
-  if (!index)
-    return success();
-
-  // 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())
-      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();
-}
-
 /// Return a map with key being elements in `vals` and data being number of
 /// occurences of it. Use std::map, since the `vals` here are strides and the
 /// dynamic stride value is the same as the tombstone value for
 /// `DenseMap<int64_t>`.
 static std::map<int64_t, unsigned> getNumOccurences(ArrayRef<int64_t> vals) {
   std::map<int64_t, unsigned> numOccurences;
   for (auto val : vals)
     numOccurences[val]++;
OpFoldResult DimOp::fold(FoldAdaptor adaptor) {
   if (!index)
     return {};
 
   // Folding for unranked types (UnrankedMemRefType) is not supported.
   auto memrefType = getSource().getType().dyn_cast<MemRefType>();
   if (!memrefType)
     return {};
 
+  // Out of bound indices produce undefined behavior but are still valid IR.
+  // Don't choke on them.
+  int64_t indexVal = index.getInt();
+  if (indexVal < 0 || indexVal >= memrefType.getRank())
+    return {};
+
   // Fold if the shape extent along the given index is known.
   if (!memrefType.isDynamicDim(index.getInt())) {
     Builder builder(getContext());
     return builder.getIndexAttr(memrefType.getShape()[index.getInt()]);
   }
 
   // The size at the given index is now known to be a dynamic size.
   unsigned unsignedIndex = index.getValue().getZExtValue();

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

Speculation::Speculatability DimOp::getSpeculatability() {
   if (!rankedSourceType)
     return Speculation::NotSpeculatable;
 
   // The verifier rejects operations that violate this assertion.
   assert(constantIndex < rankedSourceType.getRank());
   return Speculation::Speculatable;
 }
 
-LogicalResult DimOp::verify() {
-  // Assume unknown index to be in range.
-  std::optional<int64_t> index = getConstantIndex();
-  if (!index)
-    return success();
-
-  // 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())
-      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();
-}
-
 OpFoldResult DimOp::fold(FoldAdaptor adaptor) {
   // All forms of folding require a known index.
   auto index = adaptor.getIndex().dyn_cast_or_null<IntegerAttr>();
   if (!index)
     return {};
 
   // Folding for unranked types (UnrankedTensorType) is not supported.
   auto tensorType = getSource().getType().dyn_cast<RankedTensorType>();
   if (!tensorType)
     return {};
 
+  // Out of bound indices produce undefined behavior but are still valid IR.
+  // Don't choke on them.
+  int64_t indexVal = index.getInt();
+  if (indexVal < 0 || indexVal >= tensorType.getRank())
+    return {};
+
   // Fold if the shape extent along the given index is known.
   if (!tensorType.isDynamicDim(index.getInt())) {
     Builder builder(getContext());
     return builder.getIndexAttr(tensorType.getShape()[index.getInt()]);
   }
 
   Operation *definingOp = getSource().getDefiningOp();
 

mlir/test/Conversion/MemRefToLLVM/convert-static-memref-ops.mlir

 // CHECK:  llvm.mlir.constant(27 : index) : i64
   %c4 = arith.constant 4 : index
   %4 = memref.dim %static, %c4 : memref<42x32x15x13x27xf32>
   return
 }
 
 // -----
 
+// CHECK-LABEL: func @static_out_of_bound_memref_dim
+func.func @static_out_of_bound_memref_dim(%static : memref<42x32x15x13x27xf32>) -> index {
+// CHECK: %[[C_MINUS_7:.*]] = arith.constant -7 : index
+// CHECK: %[[C_MINUS_7_I64:.*]] = builtin.unrealized_conversion_cast %[[C_MINUS_7]] : index to i64
+// CHECK: %[[UB_IDX:.*]] = llvm.getelementptr %{{.*}}[0, %[[C_MINUS_7_I64]]] : (!llvm.ptr, i64) -> !llvm.ptr
+// CHECK: %[[UB_DIM_I64:.*]] = llvm.load %[[UB_IDX]] : !llvm.ptr
+// CHECK: %[[UB_DIM:.*]] = builtin.unrealized_conversion_cast %[[UB_DIM_I64]] : i64 to index
+// CHECK: return %[[UB_DIM]] : index
+  %c-7 = arith.constant -7 : index
+  %1 = memref.dim %static, %c-7 : memref<42x32x15x13x27xf32>
+  return %1 : index
+}
+// -----
+
 // Check that consistent types are emitted in address arithemic in presence of
 // a data layout specification.
 module attributes { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<index, 32>> } {
   func.func @address() {
     %c1 = arith.constant 1 : index
     %0 = memref.alloc(%c1) : memref<? x vector<2xf32>>
     // CHECK: %[[CST_S:.*]] = arith.constant 1 : index
     // CHECK: %[[CST:.*]] = builtin.unrealized_conversion_cast

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
-}
-
-// -----
-
 // Asking the dimension of a 0-D shape doesn't make sense.
 func.func @dim_0_ranked(%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
 }
 
 // -----