[mlir] Use RankedTensorType when rank is required

[springerm]

Summary

RankedTensorOf and TensorRankOf (in Tablegen files) now generate code that uses RankedTensorType instead of TensorType. This gives us more accurate type information (e.g., when calling op.getType()).

Also use restrict tensor.expand_shape/tensor.collapse_shape/tensor.pad to ranked tensors. Only cast ops should deal with unranked tensors.

Also improves a few places in the code base (e.g., Toy tutorial) where a ranked tensor is assumed (e.g., because getRank is called) but a TensorType is currently used: cast to RankedTensorType directly, so that the assertion is triggered directly at the cast.

[rkayaith]

mlir/include/mlir/IR/OpBase.td

796–797

can this use RankedTensorOf now?

[springerm]

address comments

mlir/examples/toy/Ch2/mlir/Dialect.cpp

mlir::LogicalResult ConstantOp::verify() {
   // If the return type of the constant is not an unranked tensor, the shape
   // must match the shape of the attribute holding the data.
   auto resultType = getResult().getType().dyn_cast<mlir::RankedTensorType>();
   if (!resultType)
     return success();
 
   // Check that the rank of the attribute type matches the rank of the constant
   // result type.
-  auto attrType = getValue().getType().cast<mlir::TensorType>();
+  auto attrType = getValue().getType().cast<mlir::RankedTensorType>();
   if (attrType.getRank() != resultType.getRank()) {
     return emitOpError("return type must match the one of the attached value "
                        "attribute: ")
            << attrType.getRank() << " != " << resultType.getRank();
   }
 
   // Check that each of the dimensions match between the two types.
   for (int dim = 0, dimE = attrType.getRank(); dim < dimE; ++dim) {

mlir/examples/toy/Ch3/mlir/Dialect.cpp

mlir::LogicalResult ConstantOp::verify() {
   // If the return type of the constant is not an unranked tensor, the shape
   // must match the shape of the attribute holding the data.
   auto resultType = getResult().getType().dyn_cast<mlir::RankedTensorType>();
   if (!resultType)
     return success();
 
   // Check that the rank of the attribute type matches the rank of the constant
   // result type.
-  auto attrType = getValue().getType().cast<mlir::TensorType>();
+  auto attrType = getValue().getType().cast<mlir::RankedTensorType>();
   if (attrType.getRank() != resultType.getRank()) {
     return emitOpError("return type must match the one of the attached value "
                        "attribute: ")
            << attrType.getRank() << " != " << resultType.getRank();
   }
 
   // Check that each of the dimensions match between the two types.
   for (int dim = 0, dimE = attrType.getRank(); dim < dimE; ++dim) {

mlir/examples/toy/Ch4/mlir/Dialect.cpp

mlir::LogicalResult ConstantOp::verify() {
   // If the return type of the constant is not an unranked tensor, the shape
   // must match the shape of the attribute holding the data.
   auto resultType = getResult().getType().dyn_cast<mlir::RankedTensorType>();
   if (!resultType)
     return success();
 
   // Check that the rank of the attribute type matches the rank of the constant
   // result type.
-  auto attrType = getValue().getType().cast<mlir::TensorType>();
+  auto attrType = getValue().getType().cast<mlir::RankedTensorType>();
   if (attrType.getRank() != resultType.getRank()) {
     return emitOpError("return type must match the one of the attached value "
                        "attribute: ")
            << attrType.getRank() << " != " << resultType.getRank();
   }
 
   // Check that each of the dimensions match between the two types.
   for (int dim = 0, dimE = attrType.getRank(); dim < dimE; ++dim) {

mlir/examples/toy/Ch5/mlir/Dialect.cpp

mlir::LogicalResult ConstantOp::verify() {
   // If the return type of the constant is not an unranked tensor, the shape
   // must match the shape of the attribute holding the data.
   auto resultType = getResult().getType().dyn_cast<mlir::RankedTensorType>();
   if (!resultType)
     return success();
 
   // Check that the rank of the attribute type matches the rank of the constant
   // result type.
-  auto attrType = getValue().getType().cast<mlir::TensorType>();
+  auto attrType = getValue().getType().cast<mlir::RankedTensorType>();
   if (attrType.getRank() != resultType.getRank()) {
     return emitOpError("return type must match the one of the attached value "
                        "attribute: ")
            << attrType.getRank() << " != " << resultType.getRank();
   }
 
   // Check that each of the dimensions match between the two types.
   for (int dim = 0, dimE = attrType.getRank(); dim < dimE; ++dim) {

mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp

 #include "llvm/ADT/Sequence.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // ToyToAffine RewritePatterns
 //===----------------------------------------------------------------------===//
 
-/// Convert the given TensorType into the corresponding MemRefType.
-static MemRefType convertTensorToMemRef(TensorType type) {
-  assert(type.hasRank() && "expected only ranked shapes");
+/// Convert the given RankedTensorType into the corresponding MemRefType.
+static MemRefType convertTensorToMemRef(RankedTensorType type) {
   return MemRefType::get(type.getShape(), type.getElementType());
 }
 
 /// Insert an allocation and deallocation for the given MemRefType.
 static Value insertAllocAndDealloc(MemRefType type, Location loc,
                                    PatternRewriter &rewriter) {
   auto alloc = rewriter.create<memref::AllocOp>(loc, type);
 

 /// induction variables for the iteration. It returns a value to store at the
 /// current index of the iteration.
 using LoopIterationFn = function_ref<Value(
     OpBuilder &rewriter, ValueRange memRefOperands, ValueRange loopIvs)>;
 
 static void lowerOpToLoops(Operation *op, ValueRange operands,
                            PatternRewriter &rewriter,
                            LoopIterationFn processIteration) {
-  auto tensorType = (*op->result_type_begin()).cast<TensorType>();
+  auto tensorType = (*op->result_type_begin()).cast<RankedTensorType>();
   auto loc = op->getLoc();
 
   // Insert an allocation and deallocation for the result of this operation.
   auto memRefType = convertTensorToMemRef(tensorType);
   auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
   // Create a nest of affine loops, with one loop per dimension of the shape.
   // The buildAffineLoopNest function takes a callback that is used to construct
struct ConstantOpLowering : public OpRewritePattern<toy::ConstantOp> {
 
   LogicalResult matchAndRewrite(toy::ConstantOp op,
                                 PatternRewriter &rewriter) const final {
     DenseElementsAttr constantValue = op.getValue();
     Location loc = op.getLoc();
 
     // When lowering the constant operation, we allocate and assign the constant
     // values to a corresponding memref allocation.
-    auto tensorType = op.getType().cast<TensorType>();
+    auto tensorType = op.getType().cast<RankedTensorType>();
     auto memRefType = convertTensorToMemRef(tensorType);
     auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
     // We will be generating constant indices up-to the largest dimension.
     // Create these constants up-front to avoid large amounts of redundant
     // operations.
     auto valueShape = memRefType.getShape();
     SmallVector<Value, 8> constantIndices;

mlir/examples/toy/Ch6/mlir/Dialect.cpp

mlir::LogicalResult ConstantOp::verify() {
   // If the return type of the constant is not an unranked tensor, the shape
   // must match the shape of the attribute holding the data.
   auto resultType = getResult().getType().dyn_cast<mlir::RankedTensorType>();
   if (!resultType)
     return success();
 
   // Check that the rank of the attribute type matches the rank of the constant
   // result type.
-  auto attrType = getValue().getType().cast<mlir::TensorType>();
+  auto attrType = getValue().getType().cast<mlir::RankedTensorType>();
   if (attrType.getRank() != resultType.getRank()) {
     return emitOpError("return type must match the one of the attached value "
                        "attribute: ")
            << attrType.getRank() << " != " << resultType.getRank();
   }
 
   // Check that each of the dimensions match between the two types.
   for (int dim = 0, dimE = attrType.getRank(); dim < dimE; ++dim) {

mlir/examples/toy/Ch6/mlir/LowerToAffineLoops.cpp

 #include "llvm/ADT/Sequence.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // ToyToAffine RewritePatterns
 //===----------------------------------------------------------------------===//
 
-/// Convert the given TensorType into the corresponding MemRefType.
-static MemRefType convertTensorToMemRef(TensorType type) {
-  assert(type.hasRank() && "expected only ranked shapes");
+/// Convert the given RankedTensorType into the corresponding MemRefType.
+static MemRefType convertTensorToMemRef(RankedTensorType type) {
   return MemRefType::get(type.getShape(), type.getElementType());
 }
 
 /// Insert an allocation and deallocation for the given MemRefType.
 static Value insertAllocAndDealloc(MemRefType type, Location loc,
                                    PatternRewriter &rewriter) {
   auto alloc = rewriter.create<memref::AllocOp>(loc, type);
 

 /// induction variables for the iteration. It returns a value to store at the
 /// current index of the iteration.
 using LoopIterationFn = function_ref<Value(
     OpBuilder &rewriter, ValueRange memRefOperands, ValueRange loopIvs)>;
 
 static void lowerOpToLoops(Operation *op, ValueRange operands,
                            PatternRewriter &rewriter,
                            LoopIterationFn processIteration) {
-  auto tensorType = (*op->result_type_begin()).cast<TensorType>();
+  auto tensorType = (*op->result_type_begin()).cast<RankedTensorType>();
   auto loc = op->getLoc();
 
   // Insert an allocation and deallocation for the result of this operation.
   auto memRefType = convertTensorToMemRef(tensorType);
   auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
   // Create a nest of affine loops, with one loop per dimension of the shape.
   // The buildAffineLoopNest function takes a callback that is used to construct
struct ConstantOpLowering : public OpRewritePattern<toy::ConstantOp> {
 
   LogicalResult matchAndRewrite(toy::ConstantOp op,
                                 PatternRewriter &rewriter) const final {
     DenseElementsAttr constantValue = op.getValue();
     Location loc = op.getLoc();
 
     // When lowering the constant operation, we allocate and assign the constant
     // values to a corresponding memref allocation.
-    auto tensorType = op.getType().cast<TensorType>();
+    auto tensorType = op.getType().cast<RankedTensorType>();
     auto memRefType = convertTensorToMemRef(tensorType);
     auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
     // We will be generating constant indices up-to the largest dimension.
     // Create these constants up-front to avoid large amounts of redundant
     // operations.
     auto valueShape = memRefType.getShape();
     SmallVector<Value, 8> constantIndices;

mlir/examples/toy/Ch7/mlir/Dialect.cpp

if (type.isa<mlir::TensorType>()) {
     // If the return type of the constant is not an unranked tensor, the shape
     // must match the shape of the attribute holding the data.
     auto resultType = type.dyn_cast<mlir::RankedTensorType>();
     if (!resultType)
       return success();
 
     // Check that the rank of the attribute type matches the rank of the
     // constant result type.
-    auto attrType = attrValue.getType().cast<mlir::TensorType>();
+    auto attrType = attrValue.getType().cast<mlir::RankedTensorType>();
     if (attrType.getRank() != resultType.getRank()) {
       return op->emitOpError("return type must match the one of the attached "
                              "value attribute: ")
              << attrType.getRank() << " != " << resultType.getRank();
     }
 
     // Check that each of the dimensions match between the two types.
     for (int dim = 0, dimE = attrType.getRank(); dim < dimE; ++dim) {

mlir/examples/toy/Ch7/mlir/LowerToAffineLoops.cpp

 #include "llvm/ADT/Sequence.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // ToyToAffine RewritePatterns
 //===----------------------------------------------------------------------===//
 
-/// Convert the given TensorType into the corresponding MemRefType.
-static MemRefType convertTensorToMemRef(TensorType type) {
-  assert(type.hasRank() && "expected only ranked shapes");
+/// Convert the given RankedTensorType into the corresponding MemRefType.
+static MemRefType convertTensorToMemRef(RankedTensorType type) {
   return MemRefType::get(type.getShape(), type.getElementType());
 }
 
 /// Insert an allocation and deallocation for the given MemRefType.
 static Value insertAllocAndDealloc(MemRefType type, Location loc,
                                    PatternRewriter &rewriter) {
   auto alloc = rewriter.create<memref::AllocOp>(loc, type);
 

 /// induction variables for the iteration. It returns a value to store at the
 /// current index of the iteration.
 using LoopIterationFn = function_ref<Value(
     OpBuilder &rewriter, ValueRange memRefOperands, ValueRange loopIvs)>;
 
 static void lowerOpToLoops(Operation *op, ValueRange operands,
                            PatternRewriter &rewriter,
                            LoopIterationFn processIteration) {
-  auto tensorType = (*op->result_type_begin()).cast<TensorType>();
+  auto tensorType = (*op->result_type_begin()).cast<RankedTensorType>();
   auto loc = op->getLoc();
 
   // Insert an allocation and deallocation for the result of this operation.
   auto memRefType = convertTensorToMemRef(tensorType);
   auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
   // Create a nest of affine loops, with one loop per dimension of the shape.
   // The buildAffineLoopNest function takes a callback that is used to construct
struct ConstantOpLowering : public OpRewritePattern<toy::ConstantOp> {
 
   LogicalResult matchAndRewrite(toy::ConstantOp op,
                                 PatternRewriter &rewriter) const final {
     DenseElementsAttr constantValue = op.getValue();
     Location loc = op.getLoc();
 
     // When lowering the constant operation, we allocate and assign the constant
     // values to a corresponding memref allocation.
-    auto tensorType = op.getType().cast<TensorType>();
+    auto tensorType = op.getType().cast<RankedTensorType>();
     auto memRefType = convertTensorToMemRef(tensorType);
     auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
     // We will be generating constant indices up-to the largest dimension.
     // Create these constants up-front to avoid large amounts of redundant
     // operations.
     auto valueShape = memRefType.getShape();
     SmallVector<Value, 8> constantIndices;

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

 //===----------------------------------------------------------------------===//
 // ExpandShapeOp / CollapseShapeOp
 //===----------------------------------------------------------------------===//
 
 class Tensor_ReassociativeReshapeOp<string mnemonic, list<Trait> traits = []> :
     Tensor_Op<mnemonic, !listconcat(traits, [
       DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
       Pure])>,
-    Arguments<(ins AnyTensor:$src, IndexListArrayAttr:$reassociation)>,
-    Results<(outs AnyTensor:$result)> {
+    Arguments<(ins AnyRankedTensor:$src, IndexListArrayAttr:$reassociation)>,
+    Results<(outs AnyRankedTensor:$result)> {
 
   code commonExtraClassDeclaration = [{
     static StringRef getReassociationAttrStrName() { return "reassociation"; }
     SmallVector<AffineMap, 4> getReassociationMaps();
     SmallVector<ReassociationExprs, 4> getReassociationExprs();
     SmallVector<ReassociationIndices, 4> getReassociationIndices() {
       SmallVector<ReassociationIndices, 4> reassociationIndices;
       for (auto attr : getReassociation())
```mlir
       %0 = tensor.pad %arg0 nofold low[0, 0] high[0, 0] {
       ^bb0(%arg1: index, %arg2: index):
         tensor.yield %pad_value : f32
       } : tensor<2x3xf32> to tensor<2x3xf32>
     ```
   }];
 
   let arguments = (ins
-    AnyTensor:$source,
+    AnyRankedTensor:$source,
     Variadic<Index>:$low,
     Variadic<Index>:$high,
     DenseI64ArrayAttr:$static_low,
     DenseI64ArrayAttr:$static_high,
     UnitAttr:$nofold);
 
   let regions = (region SizedRegion<1>:$region);
 
-  let results = (outs AnyTensor:$result);
+  let results = (outs AnyRankedTensor:$result);
 
   // TODO: Remove custom<InferType> when AllTypesMatch supports opt. operands.
   let assemblyFormat = [{
     $source
     (`nofold` $nofold^)?
     `low` `` custom<DynamicIndexList>($low, $static_low)
     `high` `` custom<DynamicIndexList>($high, $static_high)
     $region attr-dict `:` type($source) `to` type($result)
class Tensor_RelayoutOp<string mnemonic, list<Trait> traits = []> :
         DestinationStyleOpInterface,
         ConditionallySpeculatable, NoMemoryEffect,
         DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
         TypesMatchWith<"result type matches type of dest",
                    "dest", "result",
                    "$_self">])> {
 
   code commonExtraClassDeclaration = [{
-    size_t getSourceRank() { return getSource().getType().getRank(); };
-    size_t getDestRank() { return getDest().getType().getRank(); };
+    size_t getSourceRank() { return getSourceType().getRank(); };
+    size_t getDestRank() { return getDestType().getRank(); };
     RankedTensorType getSourceType() {
       return getSource().getType().cast<RankedTensorType>(); };
     RankedTensorType getDestType() {
       return getDest().getType().cast<RankedTensorType>(); };
 
     /// Return position for init operand. Init operand is `dest`.
     std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
       return {1, 2}; // `dest` operand

mlir/include/mlir/IR/OpBase.td

 // Whether a type is an UnrankedMemRefType
 def IsUnrankedMemRefTypePred
         : CPred<"$_self.isa<::mlir::UnrankedMemRefType>()">;
 
 // Whether a type is an UnrankedTensorType
 def IsUnrankedTensorTypePred
         : CPred<"$_self.isa<::mlir::UnrankedTensorType>()">;
 
+// Whether a type is a RankedTensorType
+def IsRankedTensorTypePred
+        : CPred<"$_self.isa<::mlir::RankedTensorType>()">;
+
 // Whether a type is a BaseMemRefType
 def IsBaseMemRefTypePred
         : CPred<"$_self.isa<::mlir::BaseMemRefType>()">;
 
 // Whether a type is a ShapedType.
 def IsShapedTypePred : CPred<"$_self.isa<::mlir::ShapedType>()">;
 
 // For a ShapedType, verify that it has a static shape.

 // Shaped types.
 
 def AnyShaped: ShapedContainerType<[AnyType], IsShapedTypePred, "shaped",
                                    "::mlir::ShapedType">;
 
 //===----------------------------------------------------------------------===//
 // Tensor types.
 
-// Unranked tensor type whose element type is from the given
-// `allowedTypes` list.
-class UnrankedTensorOf<list<Type> allowedTypes>
-  : ShapedContainerType<allowedTypes, IsUnrankedTensorTypePred,
-      "unranked.tensor", "::mlir::UnrankedTensorType">;
+// Unranked tensor type whose element type is from the given `allowedTypes`
+// list, and which additionally satisfies an optional list of predicates.
+class UnrankedTensorOf<list<Type> allowedTypes, list<Pred> preds = [],
+                       string summary = "unranked tensor">
+  : ShapedContainerType<
+      allowedTypes, And<!listconcat([IsUnrankedTensorTypePred], preds)>,
+      summary, "::mlir::UnrankedTensorType">;
+
+// Ranked tensor type whose element type is from the given `allowedTypes` list,
+// and which additionally satisfies an optional list of predicates.
+class RankedTensorOf<list<Type> allowedTypes, list<Pred> preds = [],
+                     string summary = "ranked tensor">
+  : ShapedContainerType<
+      allowedTypes, And<!listconcat([IsRankedTensorTypePred], preds)>,
+      summary, "::mlir::RankedTensorType">;
 
 // Any tensor type whose element type is from the given `allowedTypes`
 // list, and which additionally satisfies an optional list of predicates.
 //
 // TODO: use `Constraint` instead of `Pred`, so we can generate a better
 // default summary (a la `ConfinedAttr`).
 class TensorOf<
     list<Type> allowedTypes,

 def I64Tensor  : TensorOf<[I64]>;
 def IndexTensor: TensorOf<[Index]>;
 
 def BF16Tensor : TensorOf<[BF16]>;
 def F16Tensor  : TensorOf<[F16]>;
 def F32Tensor  : TensorOf<[F32]>;
 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]>;
+def AnyNon0RankedOrUnrankedTensor
+  : AnyTypeOf<[AnyUnrankedTensor, AnyNon0RankedTensor],
+              "non-0-ranked or unranked tensor", "::mlir::TensorType">;
 
 // Ranked tensor type with one of the specified types and ranks.
 class TensorRankOf<list<Type> allowedTypes, list<int> ranks>
-  : TensorOf<allowedTypes,
+  : RankedTensorOf<allowedTypes,
       [HasAnyRankOfPred<ranks>],
       !interleave(!foreach(rank, ranks, rank # "D"), "/") # " tensor">;
 
 class 0DTensorOf<list<Type> allowedTypes> : TensorRankOf<allowedTypes, [0]>;
 class 1DTensorOf<list<Type> allowedTypes> : TensorRankOf<allowedTypes, [1]>;
 class 2DTensorOf<list<Type> allowedTypes> : TensorRankOf<allowedTypes, [2]>;
 class 3DTensorOf<list<Type> allowedTypes> : TensorRankOf<allowedTypes, [3]>;
 class 4DTensorOf<list<Type> allowedTypes> : TensorRankOf<allowedTypes, [4]>;
 
 class StaticShapeTensorOf<list<Type> allowedTypes>
-  : TensorOf<allowedTypes, [HasStaticShapePred], "statically shaped tensor">;
+  : RankedTensorOf<allowedTypes, [HasStaticShapePred],
+                   "statically shaped tensor">;
 
 def AnyStaticShapeTensor : StaticShapeTensorOf<[AnyType]>;
 
 //===----------------------------------------------------------------------===//
 // Memref type.
 
 // Unranked Memref type
 class UnrankedMemRefOf<list<Type> allowedTypes> :

mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp

public:
   TensorExtractPattern(TypeConverter &typeConverter, MLIRContext *context,
                        int64_t threshold, PatternBenefit benefit = 1)
       : OpConversionPattern(typeConverter, context, benefit),
         byteCountThreshold(threshold) {}
 
   LogicalResult
   matchAndRewrite(tensor::ExtractOp extractOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    TensorType tensorType = extractOp.getTensor().getType().cast<TensorType>();
+    auto tensorType = extractOp.getTensor().getType().cast<RankedTensorType>();
 
     if (!tensorType.hasStaticShape())
       return rewriter.notifyMatchFailure(extractOp, "non-static tensor");
 
     if (tensorType.getNumElements() * tensorType.getElementTypeBitWidth() >
         byteCountThreshold * 8)
       return rewriter.notifyMatchFailure(extractOp,
                                          "exceeding byte count threshold");

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

 struct TensorCastExtractSlice : public OpRewritePattern<CastOp> {
   using OpRewritePattern<CastOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(CastOp tensorCast,
                                 PatternRewriter &rewriter) const final {
     auto extractOperand =
         tensorCast.getOperand().getDefiningOp<ExtractSliceOp>();
 
+    // Cannot fold cast to unranked tensor.
+    auto rankedResultType = tensorCast.getType().dyn_cast<RankedTensorType>();
+    if (!rankedResultType)
+      return failure();
+
     if (!extractOperand || !canFoldIntoProducerOp(tensorCast) ||
-        tensorCast.getType().getShape() == tensorCast.getSource()
-                                               .getType()
-                                               .cast<RankedTensorType>()
-                                               .getShape())
+        rankedResultType.getShape() == tensorCast.getSource()
       return failure();
 
     SmallVector<OpFoldResult, 4> sizes = extractOperand.getMixedSizes();
     auto dimMask = computeRankReductionMask(
         extractOperand.getStaticSizes(), extractOperand.getType().getShape());
     size_t dimIndex = 0;
     for (size_t i = 0, e = sizes.size(); i < e; i++) {
       if (dimMask && dimMask->count(i))
         continue;
-      int64_t dim = tensorCast.getType().getShape()[dimIndex++];
+      int64_t dim = rankedResultType.getShape()[dimIndex++];
       if (ShapedType::isDynamic(dim))
         continue;
       sizes[i] = rewriter.getIndexAttr(dim);
     }
 
     rewriter.replaceOpWithNewOp<ExtractSliceOp>(
-        tensorCast, tensorCast.getType().cast<RankedTensorType>(),
-        extractOperand.getSource(), extractOperand.getMixedOffsets(), sizes,
+        tensorCast, rankedResultType, extractOperand.getSource(),
+        extractOperand.getMixedOffsets(), sizes,
         extractOperand.getMixedStrides());
     return success();
   }
 };
 
 } // namespace
 
 void CastOp::getCanonicalizationPatterns(RewritePatternSet &results,
if (!expandShapeOp)
       return failure();
 
     // Only constant dimension values are supported.
     std::optional<int64_t> dim = dimOp.getConstantIndex();
     if (!dim.has_value())
       return failure();
 
     // Skip static dims. These are folded to constant ops.
-    TensorType resultType = expandShapeOp.getResultType();
+    RankedTensorType resultType = expandShapeOp.getResultType();
     if (!resultType.isDynamicDim(*dim))
       return failure();
 
     // Find reassociation group that contains this result dimension.
     int64_t srcDim = expandShapeOp.getCorrespondingSourceDim(*dim);
 
     // `dim` is the only dynamic dimension in `group`. (Otherwise, the
     // ExpandShapeOp would be ambiguous.)
if (!collapseShapeOp)
       return failure();
 
     // Only constant dimension values are supported.
     std::optional<int64_t> dim = dimOp.getConstantIndex();
     if (!dim.has_value())
       return failure();
 
     // Skip static dims. These are folded to constant ops.
-    TensorType resultType = collapseShapeOp.getResultType();
+    RankedTensorType resultType = collapseShapeOp.getResultType();
     if (!resultType.isDynamicDim(*dim))
       return failure();
 
     // Get reassociation group of the result dimension.
     ReassociationIndices group =
         collapseShapeOp.getReassociationIndices()[*dim];
 
     // result dim size = product(dims in reassoc group)

mlir/test/Dialect/Tensor/invalid.mlir

 // RUN: mlir-opt <%s -split-input-file -verify-diagnostics
 
 // 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>'}}
+  tensor.dim %arg, %arg1 : tensor<f32> // expected-error {{'tensor.dim' op operand #0 must be non-0-ranked or unranked tensor, 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>
func.func @insert_too_many_indices(%arg0: f32, %arg1: tensor<?xf32>) {
   // expected-error@+1 {{incorrect number of indices}}
   %0 = tensor.insert %arg0 into %arg1[] : tensor<?xf32>
   return
 }
 
 // -----
 
 func.func @tensor.from_elements_wrong_result_type() {
-  // expected-error@+2 {{'result' must be statically shaped tensor of any type values, but got 'tensor<*xi32>'}}
+  // expected-error@+2 {{'tensor.from_elements' invalid kind of type specified}}
   %c0 = arith.constant 0 : i32
   %0 = tensor.from_elements %c0 : tensor<*xi32>
   return
 }
 
 // -----
 
 func.func @tensor.from_elements_wrong_elements_count() {