[MLIR][Tensor] Avoid crash on negative dimensions

[rikhuijzer]

Summary

In https://reviews.llvm.org/D151611, a check was added to the tensor verifier to
emit an error on negative tensor dimensions. This check allowed for dynamic
dimensions, hence negative dimensions were still able to get through the verifier.
This is a problem in situations such as #60558, where the dynamic dimension is
converted to a static (and possibly negative) dimension by another pass in the
compiler. This patch fixes that by doing another check during the
StaticTensorGenerate conversion, and return a failure if the dimension is
negative.

As a side-note, I have to admit that I do not know why returning a failure in
StaticTensorGenerate gives a nice "tensor dimensions must be non-negative"
error. I suspect that the verifier runs again when return failure() is called,
but I am not sure.

Fixes #60558.

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

if (resultType.hasStaticShape())
       return failure();
 
     Operation::operand_range dynamicExtents =
         tensorFromElements.getDynamicExtents();
     SmallVector<Value> newOperands;
     SmallVector<int64_t> newShape;
     operandsAndShape(resultType, dynamicExtents, newOperands, newShape);
 
+    for (int64_t newdim : newShape) {
+      // This check also occurs in the verifier, but we need it here too
+      // since intermediate passes may have some replaced dynamic dimensions
+      // by constants.
+      if (newdim < 0 && !ShapedType::isDynamic(newdim))
+        return failure();
+    }
+
     if (newOperands.size() == tensorFromElements.getDynamicExtents().size())
       return failure();
 
     auto loc = tensorFromElements.getLoc();
     auto newOp = rewriter.create<GenerateOp>(
         loc, RankedTensorType::get(newShape, resultType.getElementType()),
         newOperands);
     rewriter.inlineRegionBefore(tensorFromElements.getBody(), newOp.getBody(),
static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) {
   if (insertOp && insertOp.getSource().getType() == extractOp.getType() &&
       insertOp.isSameAs(extractOp, isSame))
     return insertOp.getSource();
 
   return {};
 }
 
 OpFoldResult ExtractSliceOp::fold(FoldAdaptor adaptor) {
-  if (auto splat = llvm::dyn_cast_if_present<SplatElementsAttr>(adaptor.getSource())) {
+  if (auto splat =
+          llvm::dyn_cast_if_present<SplatElementsAttr>(adaptor.getSource())) {
     auto resultType = llvm::cast<ShapedType>(getResult().getType());
     if (resultType.hasStaticShape())
       return splat.resizeSplat(resultType);
   }
   if (getSourceType() == getType() &&
       succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType())))
     return this->getSource();
   if (Value slice = foldExtractAfterInsertSlice(*this))

mlir/test/Dialect/Tensor/invalid-canonicalize.mlir

+// RUN: mlir-opt <%s -split-input-file -verify-diagnostics -canonicalize
+
+// -----
+
+func.func @indirectly_generate_negative_size() -> tensor<?x8xi32> {
+  %cst = arith.constant 0 : i32
+  %c0 = arith.constant 0 : index
+  %size = affine.max affine_map<(d0) -> (d0 mod 64 - 8)>(%c0)
+  // expected-error@+1 {{tensor dimensions must be non-negative}}
+  %tensor = tensor.generate %size {
+  ^bb0(%arg0: index, %arg1: index):
+    tensor.yield %cst : i32
+  } : tensor<?x8xi32>
+  return %tensor : tensor<?x8xi32>
+}