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

1196–1197

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();
 }
 
 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())
     return emitError("must have as many index operands as dynamic extents "
                      "in the result type");
 
   return success();
 }
 
for (int64_t dim : resultType.getShape()) {
       APInt index;
       if (!matchPattern(*operandsIt, m_ConstantInt(&index))) {
         newShape.push_back(ShapedType::kDynamic);
         newOperands.push_back(*operandsIt++);
         continue;
       }
       newShape.push_back(index.getSExtValue());
       operandsIt++;
     }
 
     if (newOperands.size() == tensorFromElements.getDynamicExtents().size())
       return failure();
 
+    for (int64_t newdim : newShape) {
+      if (newdim < 0 && !ShapedType::isDynamic(newdim))
+        return failure();
+    }
+
     auto loc = tensorFromElements.getLoc();
     auto newOp = rewriter.create<GenerateOp>(
         loc, RankedTensorType::get(newShape, resultType.getElementType()),
         newOperands);
     rewriter.inlineRegionBefore(tensorFromElements.getBody(), newOp.getBody(),
                                 newOp.getBody().begin());
     rewriter.replaceOpWithNewOp<tensor::CastOp>(tensorFromElements, resultType,
                                                 newOp);

mlir/test/Dialect/Tensor/canonicalize.mlir

func.func @static_tensor.generate(%size1: index, %size4: index) -> tensor<3x?x?x7x?xindex> {
   // CHECK: : tensor<3x?x5x7x?xindex>
   } : tensor<3x?x?x7x?xindex>
   // CHECK: tensor.cast %{{.*}} : tensor<3x?x5x7x?xindex> to tensor<3x?x?x7x?xindex>
   return %0 : tensor<3x?x?x7x?xindex>
 }
 
 // -----
 
+// CHECK-LABEL: @generate_negative_size
+func.func @generate_negative_size() -> tensor<?x8xi32> {
+  %cst = arith.constant 0 : i32
+  %size = index.constant -128
+  // CHECK: tensor.generate
+  %tensor = tensor.generate %size {
+  ^bb0(%arg0: index, %arg1: index):
+    tensor.yield %cst : i32
+  } : tensor<?x8xi32>
+  return %tensor : tensor<?x8xi32>
+}
+
+// -----
+
 // CHECK-LABEL: @from_elements.constant
 func.func @from_elements.constant() -> tensor<3xindex> {
   // CHECK: %[[CST:.*]] = arith.constant dense<[1, 2, 1]> : tensor<3xindex>
   // CHECK: return %[[CST]]
   %c1 = arith.constant 1 : index
   %c2 = arith.constant 2 : index
   %tensor = tensor.from_elements %c1, %c2, %c1 : tensor<3xindex>
   return %tensor : tensor<3xindex>