[mlir][sparse] (re)introducing getRankedTensorType/getMemrefType

[wrengr]

Summary

The bulk of D142074 seems to have gotten overwritten due to some sort of merge conflict (afaict there's no record of it having been reverted intentionally). So this commit redoes those changes. In addition to the original changes, this commit also:

• moves the definition of getRankedTensorType (from Transforms/CodegenUtils.h to IR/SparseTensor.h), so that it can be used by IR/SparseTensorDialect.cpp.
• adds getMemRefType as another abbreviation.

[wrengr]

git-clang-format

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

 #define GET_OP_CLASSES
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOps.h.inc"
 
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOpsDialect.h.inc"
 
 namespace mlir {
 namespace sparse_tensor {
 
+/// Convenience method to abbreviate casting `getType()`.
+template <typename T>
+inline RankedTensorType getRankedTensorType(T t) {
+  return t.getType().template cast<RankedTensorType>();
+}
+
+/// Convenience method to abbreviate casting `getType()`.
+template <typename T>
+inline MemRefType getMemRefType(T t) {
+  return t.getType().template cast<MemRefType>();
+}
+
 /// Convenience method to get a sparse encoding attribute from a type.
 /// Returns null-attribute for any type without an encoding.
 SparseTensorEncodingAttr getSparseTensorEncoding(Type type);
 
 /// Returns true iff the given type is a type for a COO tensor with the last
 /// dimension level type being unique.
 bool isUniqueCOOType(RankedTensorType tp);
 

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

 
 #define GET_TYPEDEF_CLASSES
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorTypes.cpp.inc"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 //===----------------------------------------------------------------------===//
+// Additional convenience methods.
+//===----------------------------------------------------------------------===//
+
+template <typename T>
+static inline int64_t getTypeRank(T t) {
+  return getRankedTensorType(t).getRank();
+}
+
+//===----------------------------------------------------------------------===//
 // TensorDialect Attribute Methods.
 //===----------------------------------------------------------------------===//
 
 static bool acceptBitWidth(unsigned bitWidth) {
   switch (bitWidth) {
   case 0:
   case 8:
   case 16:
return getFieldType(kind, dim ? std::optional(dim.value().getZExtValue())
                                 : std::nullopt);
 }
 
 //===----------------------------------------------------------------------===//
 // SparseTensorDialect Operations.
 //===----------------------------------------------------------------------===//
 
 static LogicalResult isInBounds(uint64_t dim, Value tensor) {
-  return success(dim <
-                 (uint64_t)tensor.getType().cast<RankedTensorType>().getRank());
+  return success(dim < static_cast<uint64_t>(getTypeRank(tensor)));
 }
 
 static LogicalResult isMatchingWidth(Value result, unsigned width) {
-  const Type etp = result.getType().cast<MemRefType>().getElementType();
+  const Type etp = getMemRefType(result).getElementType();
   return success(width == 0 ? etp.isIndex() : etp.isInteger(width));
 }
 
 static LogicalResult verifySparsifierGetterSetter(
     StorageSpecifierKind mdKind, std::optional<APInt> dim,
     TypedValue<StorageSpecifierType> md, Operation *op) {
   if (mdKind == StorageSpecifierKind::ValMemSize && dim) {
     return op->emitError(
if (mdKind != StorageSpecifierKind::ValMemSize) {
     if (mdKind == StorageSpecifierKind::PtrMemSize && isSingletonDLT(dlts[d]))
       return op->emitError(
           "requested pointer memory size on a singleton level");
   }
   return success();
 }
 
 LogicalResult NewOp::verify() {
-  if (getExpandSymmetry() &&
-      getResult().getType().cast<RankedTensorType>().getRank() != 2)
+  if (getExpandSymmetry() && getTypeRank(getResult()) != 2)
     return emitOpError("expand_symmetry can only be used for 2D tensors");
   return success();
 }
 
 LogicalResult ConvertOp::verify() {
   if (auto tp1 = getSource().getType().dyn_cast<RankedTensorType>()) {
     if (auto tp2 = getDest().getType().dyn_cast<RankedTensorType>()) {
       if (tp1.getRank() != tp2.getRank())

 LogicalResult ToIndicesBufferOp::verify() {
   auto e = getSparseTensorEncoding(getTensor().getType());
   if (getCOOStart(e) >= e.getDimLevelType().size())
     return emitError("expected sparse tensor with a COO region");
   return success();
 }
 
 LogicalResult ToValuesOp::verify() {
-  RankedTensorType ttp = getTensor().getType().cast<RankedTensorType>();
-  MemRefType mtp = getResult().getType().cast<MemRefType>();
+  auto ttp = getRankedTensorType(getTensor());
+  auto mtp = getMemRefType(getResult());
   if (ttp.getElementType() != mtp.getElementType())
     return emitError("unexpected mismatch in element types");
   return success();
 }
 
 LogicalResult GetStorageSpecifierOp::verify() {
   RETURN_FAILURE_IF_FAILED(verifySparsifierGetterSetter(
       getSpecifierKind(), getDim(), getSpecifier(), getOperation()))
LogicalResult UnaryOp::verify() {
   if (!absent.empty()) {
     RETURN_FAILURE_IF_FAILED(
         verifyNumBlockArgs(this, absent, "absent", TypeRange{}, outputType))
   }
   return success();
 }
 
 LogicalResult ConcatenateOp::verify() {
-  auto dstTp = getType().cast<RankedTensorType>();
+  auto dstTp = getRankedTensorType(*this);
   uint64_t concatDim = getDimension().getZExtValue();
   unsigned rank = dstTp.getRank();
 
   if (getInputs().size() <= 1)
     return emitError("Need at least two tensors to concatenate.");
 
   for (auto type : getInputs().getTypes()) {
     auto shape = type.cast<RankedTensorType>().getShape();
     for (auto dim : shape) {
       if (ShapedType::isDynamic(dim))
         return emitError("Only statically-sized input tensors are supported.");
     }
   }
 
   if (concatDim >= rank)
     return emitError(llvm::formatv(
         "Failed to concatentate tensors with rank={0} on dimension={1}.", rank,
         concatDim));
 
   for (size_t i = 0, e = getInputs().size(); i < e; i++) {
-    Value input = getInputs()[i];
-    auto inputRank = input.getType().cast<RankedTensorType>().getRank();
+    const auto inputRank = getTypeRank(getInputs()[i]);
     if (inputRank != rank)
       return emitError(
           llvm::formatv("The input tensor ${0} has a different rank (rank={1}) "
                         "from the output tensor (rank={2}).",
                         i, inputRank, rank));
   }
 
   for (unsigned i = 0; i < rank; i++) {
-    auto dstDim = dstTp.getShape()[i];
+    const auto dstDim = dstTp.getShape()[i];
     if (i == concatDim) {
       if (!ShapedType::isDynamic(dstDim)) {
-        unsigned sumDim = 0;
-        for (auto src : getInputs()) {
-          // If we reach here, all inputs should have static shapes.
-          auto d = src.getType().cast<RankedTensorType>().getShape()[i];
-          sumDim += d;
-        }
+        unsigned sumDim = 0;
+        for (auto src : getInputs())
+          sumDim += getRankedTensorType(src).getShape()[i];
         // If all dimension are statically known, the sum of all the input
         // dimensions should be equal to the output dimension.
         if (sumDim != dstDim)
           return emitError(
               "The concatenation dimension of the output tensor should be the "
               "sum of all the concatenation dimensions of the input tensors.");
       }
     } else {
       int64_t prev = dstDim;
       for (auto src : getInputs()) {
-        auto d = src.getType().cast<RankedTensorType>().getShape()[i];
+        const auto d = getRankedTensorType(src).getShape()[i];
         if (!ShapedType::isDynamic(prev) && d != prev)
           return emitError("All dimensions (expect for the concatenating one) "
                            "should be equal.");
         prev = d;
       }
     }
   }
 
   return success();
 }
 
 LogicalResult InsertOp::verify() {
-  RankedTensorType ttp = getTensor().getType().cast<RankedTensorType>();
-  if (ttp.getRank() != static_cast<int64_t>(getIndices().size()))
+  if (getTypeRank(getTensor()) != static_cast<int64_t>(getIndices().size()))
     return emitOpError("incorrect number of indices");
   return success();
 }
 
 void PushBackOp::build(OpBuilder &builder, OperationState &result,
                        Value curSize, Value inBuffer, Value value) {
   build(builder, result, curSize, inBuffer, value, Value());
 }
 
 LogicalResult PushBackOp::verify() {
   if (Value n = getN()) {
     auto nValue = dyn_cast_or_null<arith::ConstantIndexOp>(n.getDefiningOp());
     if (nValue && nValue.value() < 1)
       return emitOpError("n must be not less than 1");
   }
   return success();
 }
 
 LogicalResult CompressOp::verify() {
-  RankedTensorType ttp = getTensor().getType().cast<RankedTensorType>();
-  if (ttp.getRank() != 1 + static_cast<int64_t>(getIndices().size()))
+  if (getTypeRank(getTensor()) != 1 + static_cast<int64_t>(getIndices().size()))
     return emitOpError("incorrect number of indices");
   return success();
 }
 
 void ForeachOp::build(
     OpBuilder &builder, OperationState &result, Value tensor,
     function_ref<void(OpBuilder &, Location, ValueRange, Value, ValueRange)>
         bodyBuilder) {
   build(builder, result, tensor, std::nullopt, bodyBuilder);
 }
 
 void ForeachOp::build(
     OpBuilder &builder, OperationState &result, Value tensor,
     ValueRange initArgs,
     function_ref<void(OpBuilder &, Location, ValueRange, Value, ValueRange)>
         bodyBuilder) {
   build(builder, result, initArgs.getTypes(), tensor, initArgs);
   // Builds foreach body.
   if (!bodyBuilder)
     return;
-  auto rtp = tensor.getType().cast<RankedTensorType>();
+  auto rtp = getRankedTensorType(tensor);
   int64_t rank = rtp.getRank();
 
   SmallVector<Type> blockArgTypes;
   // Starts with n index.
   std::fill_n(std::back_inserter(blockArgTypes), rank, builder.getIndexType());
   // Followed by one value.
   blockArgTypes.push_back(rtp.getElementType());
   // Followed by reduction variable.
Block *bodyBlock =
       builder.createBlock(&region, region.end(), blockArgTypes, blockArgLocs);
   bodyBuilder(builder, result.location,
               bodyBlock->getArguments().slice(0, rank),
               bodyBlock->getArguments()[rank],
               bodyBlock->getArguments().drop_front(rank + 1));
 }
 
 LogicalResult ForeachOp::verify() {
-  auto t = getTensor().getType().cast<RankedTensorType>();
+  auto t = getRankedTensorType(getTensor());
   auto args = getBody()->getArguments();
 
   if (static_cast<size_t>(t.getRank()) + 1 + getInitArgs().size() !=
       args.size())
     return emitError("Unmatched number of arguments in the block");
 
   if (getNumResults() != getInitArgs().size())
     return emitError("Mismatch in number of init arguments and results");

 }
 
 LogicalResult SortOp::verify() {
   if (getXs().empty())
     return emitError("need at least one xs buffer.");
 
   auto n = getN().getDefiningOp<arith::ConstantIndexOp>();
 
-  Type xtp = getXs().front().getType().cast<MemRefType>().getElementType();
+  Type xtp = getMemRefType(getXs().front()).getElementType();
   auto checkTypes = [&](ValueRange operands,
                         bool checkEleType = true) -> LogicalResult {
     for (Value opnd : operands) {
-      MemRefType mtp = opnd.getType().cast<MemRefType>();
+      auto mtp = getMemRefType(opnd);
       int64_t dim = mtp.getShape()[0];
       // We can't check the size of dynamic dimension at compile-time, but all
       // xs and ys should have a dimension not less than n at runtime.
       if (n && !ShapedType::isDynamic(dim) && dim < n.value())
         return emitError(llvm::formatv("xs and ys need to have a dimension >= n"
                                        ": {0} < {1}",
                                        dim, n.value()));
 
if (nx < 1)
       emitError(llvm::formatv("Expected nx > 1, got {0}", nx));
   }
   uint64_t ny = 0;
   if (auto nyAttr = getNyAttr()) {
     ny = nyAttr.getInt();
   }
 
   auto checkDim = [&](Value v, uint64_t min, const char *message) {
-    MemRefType tp = v.getType().cast<MemRefType>();
+    auto tp = getMemRefType(v);
     int64_t dim = tp.getShape()[0];
     if (!ShapedType::isDynamic(dim) && dim < (int64_t)min) {
       emitError(llvm::formatv("{0} got {1} < {2}", message, dim, min));
     }
   };
 
   checkDim(getXy(), n * (nx + ny), "Expected dimension(xy) >= n * (nx + ny)");
 

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

 
 /// Converts a primary storage type to its function-name suffix.
 StringRef primaryTypeFunctionSuffix(Type elemTp);
 
 //===----------------------------------------------------------------------===//
 // Misc code generators and utilities.
 //===----------------------------------------------------------------------===//
 
-template <typename T>
-inline RankedTensorType getRankedTensorType(T t) {
-  return t.getType().template cast<RankedTensorType>();
-}
-
 /// Generates a 1-valued attribute of the given type.  This supports
 /// all the same types as `getZeroAttr`; however, unlike `getZeroAttr`,
 /// for unsupported types we raise `llvm_unreachable` rather than
 /// returning a null attribute.
 Attribute getOneAttr(Builder &builder, Type tp);
 
 /// Generates the comparison `v != 0` where `v` is of numeric type.
 /// For floating types, we use the "unordered" comparator (i.e., returns

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

Value sparse_tensor::reshapeValuesToLevels(
   for (unsigned i = 0; i < dimSizes.size(); i++)
     lvlSizes.push_back(dimSizes[toOrigDim(enc, i)]);
   storeIndices(builder, loc, rank, idxBuffer, lvlSizes);
   // The memref ReshapeOp requires the sizes buffer to have a static
   // shape.
   idxBuffer = builder.create<memref::CastOp>(
       loc, MemRefType::get({rank}, builder.getIndexType()), idxBuffer);
   SmallVector<int64_t> shape(rank, ShapedType::kDynamic);
-  Type elemTp = valuesBuffer.getType().cast<MemRefType>().getElementType();
+  Type elemTp = getMemRefType(valuesBuffer).getElementType();
   return builder.create<memref::ReshapeOp>(loc, MemRefType::get(shape, elemTp),
                                            valuesBuffer, idxBuffer);
 }
 
 Value sparse_tensor::genToPointers(OpBuilder &builder, Location loc,
                                    Value tensor, uint64_t d) {
   RankedTensorType srcTp = getRankedTensorType(tensor);
   SparseTensorEncodingAttr encSrc = getSparseTensorEncoding(srcTp);

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp

 
 /// Constructs a function name with this format to facilitate quick sort:
 ///   <namePrefix><nx>_<x type>_<y0 type>..._<yn type> for sort
 ///   <namePrefix><nx>_<x type>_coo_<ny>_<y0 type>..._<yn type> for sort_coo
 static void getMangledSortHelperFuncName(llvm::raw_svector_ostream &nameOstream,
                                          StringRef namePrefix, uint64_t nx,
                                          uint64_t ny, bool isCoo,
                                          ValueRange operands) {
-  nameOstream
-      << namePrefix << nx << "_"
-      << operands[xStartIdx].getType().cast<MemRefType>().getElementType();
+  nameOstream << namePrefix << nx << "_"
+              << getMemRefType(operands[xStartIdx]).getElementType();
 
   if (isCoo)
     nameOstream << "_coo_" << ny;
 
   uint64_t yBufferOffset = isCoo ? 1 : nx;
   for (Value v : operands.drop_front(xStartIdx + yBufferOffset))
-    nameOstream << "_" << v.getType().cast<MemRefType>().getElementType();
+    nameOstream << "_" << getMemRefType(v).getElementType();
 }
 
 /// Looks up a function that is appropriate for the given operands being
 /// sorted, and creates such a function if it doesn't exist yet. The
 /// parameters `nx` and `ny` tell the number of x and y values provided
 /// by the buffer in xStartIdx, and `isCoo` indicates whether the instruction
 /// being processed is a sparse_tensor.sort or sparse_tensor.sort_coo.
 static FlatSymbolRefAttr

 LogicalResult matchAndRewriteSortOp(OpTy op, ValueRange xys, uint64_t nx,
                                     uint64_t ny, bool isCoo,
                                     PatternRewriter &rewriter) {
   Location loc = op.getLoc();
   SmallVector<Value> operands{constantIndex(rewriter, loc, 0), op.getN()};
 
   // Convert `values` to have dynamic shape and append them to `operands`.
   for (Value v : xys) {
-    auto mtp = v.getType().cast<MemRefType>();
+    auto mtp = getMemRefType(v);
     if (!mtp.isDynamicDim(0)) {
       auto newMtp =
           MemRefType::get({ShapedType::kDynamic}, mtp.getElementType());
       v = rewriter.create<memref::CastOp>(loc, newMtp, v);
     }
     operands.push_back(v);
   }
   auto insertPoint = op->template getParentOfType<func::FuncOp>();

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

 ///     coo->add(reshape(elem.indices), elem.value)
 ///   }
 ///   s = newSparseTensor(coo)
 template <typename ReshapeOp>
 static LogicalResult
 genSparse2SparseReshape(ReshapeOp op, typename ReshapeOp::Adaptor adaptor,
                         ConversionPatternRewriter &rewriter) {
   Location loc = op.getLoc();
-  auto srcTp = op.getSrc().getType().template cast<RankedTensorType>();
-  auto dstTp = op.getResult().getType().template cast<RankedTensorType>();
+  auto srcTp = getRankedTensorType(op.getSrc());
+  auto dstTp = getRankedTensorType(op.getResult());
   auto encSrc = getSparseTensorEncoding(srcTp);
   auto encDst = getSparseTensorEncoding(dstTp);
   if (!encDst || !encSrc)
     return failure();
   Type elemTp = srcTp.getElementType();
   assert(elemTp == dstTp.getElementType() &&
          "reshape should not change element type");
   // Start an iterator over the source tensor (in original index order).
public:
   SparseTensorConvertConverter(TypeConverter &typeConv, MLIRContext *context,
                                SparseTensorConversionOptions o)
       : OpConversionPattern<ConvertOp>(typeConv, context), options(o) {}
 
   LogicalResult
   matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
-    Type resType = op.getType();
-    Type srcType = op.getSource().getType();
+    auto resType = getRankedTensorType(op);
+    auto srcType = getRankedTensorType(op.getSource());
     auto encDst = getSparseTensorEncoding(resType);
     auto encSrc = getSparseTensorEncoding(srcType);
     Value src = adaptor.getOperands()[0];
     if (encDst && encSrc) {
       // This is a sparse => sparse conversion, which is handled as follows:
       //   t = src->toCOO();         ; src to COO in dst order
       //   dst = newSparseTensor(t)
       // Using the coordinate scheme as an intermediate does not always
matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
     if (!encDst && encSrc) {
       // This is sparse => dense conversion, which is handled as follows:
       //   dst = new Tensor(0);
       //   iter = new SparseTensorIterator(src);
       //   while (elem = iter->getNext()) {
       //     dst[elem.indices] = elem.value;
       //   }
       //   delete iter;
-      RankedTensorType dstTensorTp = resType.cast<RankedTensorType>();
-      RankedTensorType srcTensorTp = srcType.cast<RankedTensorType>();
-      unsigned rank = dstTensorTp.getRank();
-      Type elemTp = dstTensorTp.getElementType();
+      const unsigned rank = resType.getRank();
+      const Type elemTp = resType.getElementType();
       // Fabricate a no-permutation encoding for NewCallParams
       // The pointer/index types must be those of `src`.
       // The dimLevelTypes aren't actually used by Action::kToIterator.
       encDst = SparseTensorEncodingAttr::get(
           op->getContext(),
           SmallVector<DimLevelType>(rank, DimLevelType::Dense), AffineMap(),
           AffineMap(), encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
       SmallVector<Value> dimSizes =
-          getDimSizes(rewriter, loc, encSrc, srcTensorTp, src);
+          getDimSizes(rewriter, loc, encSrc, srcType, src);
       Value iter = NewCallParams(rewriter, loc)
-                       .genBuffers(encDst, dimSizes, dstTensorTp)
+                       .genBuffers(encDst, dimSizes, resType)
                        .genNewCall(Action::kToIterator, src);
       Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
       Value elemPtr = genAllocaScalar(rewriter, loc, elemTp);
       Block *insertionBlock = rewriter.getInsertionBlock();
       // TODO: Dense buffers should be allocated/deallocated via the callback
       // in BufferizationOptions.
-      Value dst = allocDenseTensor(rewriter, loc, dstTensorTp, dimSizes);
+      Value dst = allocDenseTensor(rewriter, loc, resType, dimSizes);
       SmallVector<Value> noArgs;
       SmallVector<Type> noTypes;
       auto whileOp = rewriter.create<scf::WhileOp>(loc, noTypes, noArgs);
       Block *before = rewriter.createBlock(&whileOp.getBefore(), {}, noTypes);
       rewriter.setInsertionPointToEnd(before);
       Value cond = genGetNextCall(rewriter, loc, iter, ind, elemPtr);
       rewriter.create<scf::ConditionOp>(loc, cond, before->getArguments());
       Block *after = rewriter.createBlock(&whileOp.getAfter(), {}, noTypes);
public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(InsertOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Note that the current regime only allows for strict lexicographic
     // index order. All values are passed by reference through stack
     // allocated memrefs.
     Location loc = op->getLoc();
-    auto tp = op.getTensor().getType().cast<RankedTensorType>();
+    auto tp = getRankedTensorType(op.getTensor());
     auto elemTp = tp.getElementType();
     unsigned rank = tp.getRank();
     auto mref = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
     auto vref = genAllocaScalar(rewriter, loc, elemTp);
     for (unsigned i = 0; i < rank; i++)
       rewriter.create<memref::StoreOp>(loc, adaptor.getIndices()[i], mref,
                                        constantIndex(rewriter, loc, i));
     rewriter.create<memref::StoreOp>(loc, adaptor.getValue(), vref);
     SmallString<12> name{"lexInsert", primaryTypeFunctionSuffix(elemTp)};
     createFuncCall(rewriter, loc, name, {}, {adaptor.getTensor(), mref, vref},
                    EmitCInterface::On);
     rewriter.replaceOp(op, adaptor.getTensor());
     return success();
   }
 };
 
 /// Sparse conversion rule for the expand operator.
 class SparseTensorExpandConverter : public OpConversionPattern<ExpandOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ExpandOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
-    RankedTensorType srcType =
-        op.getTensor().getType().cast<RankedTensorType>();
+    auto srcType = getRankedTensorType(op.getTensor());
     Type eltType = srcType.getElementType();
     Type boolType = rewriter.getIntegerType(1);
     Type idxType = rewriter.getIndexType();
     // All initialization should be done on entry of the loop nest.
     rewriter.setInsertionPointAfter(op.getTensor().getDefiningOp());
     // Get the cardinality of valid coordinates for the innermost level.
     auto srcEnc = getSparseTensorEncoding(srcType);
     unsigned lvlRank =
matchAndRewrite(CompressOp op, OpAdaptor adaptor,
     // all-zero/false by only iterating over the set elements, so the
     // complexity remains proportional to the sparsity of the expanded
     // access pattern.
     Value values = adaptor.getValues();
     Value filled = adaptor.getFilled();
     Value added = adaptor.getAdded();
     Value count = adaptor.getCount();
     Value tensor = adaptor.getTensor();
-    auto tp = op.getTensor().getType().cast<RankedTensorType>();
+    auto tp = getRankedTensorType(op.getTensor());
     Type elemTp = tp.getElementType();
     unsigned rank = tp.getRank();
     auto mref = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
     for (unsigned i = 0; i < rank - 1; i++)
       rewriter.create<memref::StoreOp>(loc, adaptor.getIndices()[i], mref,
                                        constantIndex(rewriter, loc, i));
     SmallString<12> name{"expInsert", primaryTypeFunctionSuffix(elemTp)};
     createFuncCall(rewriter, loc, name, {},
matchAndRewrite(ConcatenateOp op, OpAdaptor adaptor,
     //    a = malloc(shapeOf(a)) or newSparseAllDense(shapeOf(a))
     //    for i, j, k // dense input
     //      a[ adjustForOffset(i,j,k) ] = b[i,j,k]
     //
     //    for elem in sparse_input
     //      a[ adjustForOffset(elem.indices) ] = elem.value
     //    return a
     Location loc = op.getLoc();
-    auto dstTp = op.getType().cast<RankedTensorType>();
+    auto dstTp = getRankedTensorType(op);
     auto encDst = getSparseTensorEncoding(dstTp);
     Type elemTp = dstTp.getElementType();
     uint64_t concatDim = op.getDimension().getZExtValue();
     unsigned rank = dstTp.getRank();
 
     Value dst;     // destination tensor
     Value dstPerm; // destination tensor permutation (if sparse out)
     // A pointer to the value being inserted (if dense => sparse)
auto dimIdx2LvlIdx = [&](ValueRange dIdx) -> SmallVector<Value> {
       SmallVector<Value> lIdx;
       for (unsigned i = 0; i < dIdx.size(); i++)
         lIdx.push_back(dIdx[toOrigDim(encDst, i)]);
       return lIdx;
     };
     for (auto it : llvm::zip(op.getInputs(), adaptor.getInputs())) {
       Value orignalOp = std::get<0>(it); // Input (with encoding) from Op
       Value adaptedOp = std::get<1>(it); // Input (type converted) from adaptor
-      RankedTensorType srcTp = orignalOp.getType().cast<RankedTensorType>();
+      auto srcTp = getRankedTensorType(orignalOp);
       auto encSrc = getSparseTensorEncoding(srcTp);
       if (encSrc) {
         genSparseCOOIterationLoop(
             rewriter, loc, adaptedOp, srcTp,
             [&](OpBuilder &builder, Location loc, Value idx,
                 Value elemPtr) -> void {
               SmallVector<Value> dimInd =
                   loadIndices(builder, loc, rank, idx, concatDim, offset);

mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp

 /// Constructs vector type for element type.
 static VectorType vectorType(VL vl, Type etp) {
   unsigned numScalableDims = vl.enableVLAVectorization;
   return VectorType::get(vl.vectorLength, etp, numScalableDims);
 }
 
 /// Constructs vector type from pointer.
 static VectorType vectorType(VL vl, Value ptr) {
-  return vectorType(vl, ptr.getType().cast<MemRefType>().getElementType());
+  return vectorType(vl, getMemRefType(ptr).getElementType());
 }
 
 /// Constructs vector iteration mask.
 static Value genVectorMask(PatternRewriter &rewriter, Location loc, VL vl,
                            Value iv, Value lo, Value hi, Value step) {
   VectorType mtp = vectorType(vl, rewriter.getI1Type());
   // Special case if the vector length evenly divides the trip count (for
   // example, "for i = 0, 128, 16"). A constant all-true mask is generated