[mlir] Avoid tensor canonicalizer crash on negative dimensions

[rikhuijzer]

Summary

Fixes #59703.

[ftynse]

It's strange that this is even allowed by the op verifier. Can we (also) change the verifier of tensor.generate to disallow this case?

[rikhuijzer]

[mlir][Tensor] Disallow negative dimensions in verifier of tensor.generate

Instead of refusing to transform negative dimensions, the compiler will now crash
which is probably the right thing to do. Thanks to Alex Zinenko for pointing this out.

[ftynse]

Thank you! Please run clang-format.

I'll mark this as "changes requested" so I don't forget to get back to this.

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

1114

It is usually undesirable to return SmallVector as it may copy a lot. Instead, you can pass a SmallVectorImpl<...> & into the function.

1115

Please declare auxiliary functions as static and document them.

1116

Nit: do not specify the number of stack-allocated vector elements unless there is a strong reason to pick a number, just SmallVector<Value> will do here.

1208–1213

Please expand auto to be the proper type unless the type is obvious from local context (e.g., the RHS has a cast) or difficult/impossible to spell (iterators, lambdas).

[rikhuijzer]

Incorporate reviewer comments

Thanks again for the review Alex. The comments were easy to understand and implement!

[ftynse]

Thanks again for the patch! Let me know if you need help committing this.

[rikhuijzer]

In D151611#4383856, @ftynse wrote:

Thanks again for the patch! Let me know if you need help committing this.

Yes. :) I don't have permissions.

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

for (auto dim : llvm::seq<int64_t>(0, getType().getRank())) {
     } else {
       reifiedReturnShapes[0][dim] =
           builder.getIndexAttr(getType().getDimSize(dim));
     }
   }
   return success();
 }
 
+/// Extract operands and shape from a tensor with dynamic extents.
+static void operandsAndShape(TensorType resultType,
+                             Operation::operand_range dynamicExtents,
+                             SmallVectorImpl<Value> &newOperands,
+                             SmallVectorImpl<int64_t> &newShape) {
+  auto operandsIt = dynamicExtents.begin();
+  for (int64_t dim : resultType.getShape()) {
+    if (!ShapedType::isDynamic(dim)) {
+      newShape.push_back(dim);
+      continue;
+    }
+    APInt index;
+    if (!matchPattern(*operandsIt, m_ConstantInt(&index))) {
+      newShape.push_back(ShapedType::kDynamic);
+      newOperands.push_back(*operandsIt++);
+      continue;
+    }
+    newShape.push_back(index.getSExtValue());
+    operandsIt++;
+  }
+}
+
 LogicalResult GenerateOp::verify() {
   // Ensure that the tensor type has as many dynamic dimensions as are
   // specified by the operands.
-  RankedTensorType resultTy = llvm::cast<RankedTensorType>(getType());
-  if (getNumOperands() != resultTy.getNumDynamicDims())
+  RankedTensorType resultType = llvm::cast<RankedTensorType>(getType());
+  if (getNumOperands() != resultType.getNumDynamicDims())
     return emitError("must have as many index operands as dynamic extents "
                      "in the result type");
-
+  // Ensure operands are non-negative.
+  SmallVector<Value> newOperands;
+  SmallVector<int64_t> newShape;
+  operandsAndShape(resultType, getDynamicExtents(), newOperands, newShape);
+  for (int64_t newdim : newShape) {
+    if (newdim < 0 && !ShapedType::isDynamic(newdim))
+      return emitError("tensor dimensions must be non-negative");
+  }
   return success();
 }
 
 LogicalResult GenerateOp::verifyRegions() {
   RankedTensorType resultTy = llvm::cast<RankedTensorType>(getType());
   // Ensure that region arguments span the index space.
   if (!llvm::all_of(getBody().getArgumentTypes(),
                     [](Type ty) { return ty.isIndex(); }))
struct StaticTensorGenerate : public OpRewritePattern<GenerateOp> {
   LogicalResult matchAndRewrite(GenerateOp tensorFromElements,
                                 PatternRewriter &rewriter) const final {
     auto resultType =
         llvm::cast<RankedTensorType>(tensorFromElements.getResult().getType());
 
     if (resultType.hasStaticShape())
       return failure();
 
-    SmallVector<Value, 4> newOperands;
-    SmallVector<int64_t, 4> newShape;
-    auto operandsIt = tensorFromElements.getDynamicExtents().begin();
-
-    for (int64_t dim : resultType.getShape()) {
-      if (!ShapedType::isDynamic(dim)) {
-        newShape.push_back(dim);
-        continue;
-      }
-      APInt index;
-      if (!matchPattern(*operandsIt, m_ConstantInt(&index))) {
-        newShape.push_back(ShapedType::kDynamic);
-        newOperands.push_back(*operandsIt++);
-        continue;
-      }
-      newShape.push_back(index.getSExtValue());
-      operandsIt++;
-    }
+    Operation::operand_range dynamicExtents =
+        tensorFromElements.getDynamicExtents();
+    SmallVector<Value> newOperands;
+    SmallVector<int64_t> newShape;
+    operandsAndShape(resultType, dynamicExtents, newOperands, newShape);
 
     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(),

mlir/test/Dialect/Tensor/invalid.mlir

func.func @tensor.generate(%m : index, %n : index)
   // expected-error @+1 {{body must be terminated with a `yield` operation of the tensor element type}}
   %tnsr = tensor.generate %m, %n {
     ^bb0(%i : index, %j : index, %k : index):
       %elem = arith.constant 8 : i32
       tensor.yield %elem : i32
   } : tensor<?x3x?xf32>
   return %tnsr : tensor<?x3x?xf32>
 }
+
+// -----
+
+func.func @generate_negative_size() -> tensor<?x8xi32> {
+  %cst = arith.constant 0 : i32
+  %size = index.constant -128
+  // 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>
+}
+
 // -----
 
 func.func @tensor.reshape_element_type_mismatch(
        %buf: tensor<*xf32>, %shape: tensor<1xi32>) {
   // expected-error @+1 {{element types of source and destination tensor types should be the same}}
   tensor.reshape %buf(%shape) : (tensor<*xf32>, tensor<1xi32>) -> tensor<?xi32>
 }