[mlir][IR] Improve `clone` function return type of shaped types

[springerm]

Summary

There are clone overloads that take a shape as a parameter. These overloads are guaranteed to return a ranked shaped type.

TensorType::clone/BaseMemRefType::clone used to always return a TensorType/BaseMemRefType. The variants that take a shape parameter now return a RankedTensorType/MemRefType. Better static type information can make extra casts at the call site obsolete.

E.g.:

{TensorType/RankedTensorType} t;
t.clone({1, 2})  // now returns RankedTensorType instead of TensorType

Also improve documentation for clone.

mlir/include/mlir/IR/BuiltinTypeInterfaces.td

let description = [{
     A shape is a list of sizes corresponding to the dimensions of the container.
     If the number of dimensions in the shape is unknown, the shape is "unranked".
     If the number of dimensions is known, the shape "ranked". The sizes of the
     dimensions of the shape must be positive, or kDynamic (in which case the
     size of the dimension is dynamic, or not statically known).
   }];
   let methods = [
     InterfaceMethod<[{
-      Returns a clone of this type with the given shape and element
-      type. If a shape is not provided, the current shape of the type is used.
+      Returns a clone of this type with the given shape and element type.
+
+      If no shape is provided, the shape of this type is used. In that case, if
+      this type is unranked, so is the resulting type.
+
+      If a shape is provided, the resulting type is always ranked, even if this
+      type is unranked.
     }],
     "::mlir::ShapedType", "cloneWith", (ins
       "::std::optional<::llvm::ArrayRef<int64_t>>":$shape,
       "::mlir::Type":$elementType
     )>,
 
     InterfaceMethod<[{
       Returns the element type of this shaped type.
def ShapedTypeInterface : TypeInterface<"ShapedType"> {
   ];
 
   let extraClassDeclaration = [{
     static constexpr int64_t kDynamic =
         std::numeric_limits<int64_t>::min();
 
     /// Whether the given dimension size indicates a dynamic dimension.
     static constexpr bool isDynamic(int64_t dValue) {
-  return dValue == kDynamic;
+      return dValue == kDynamic;
     }
 
     /// Whether the given shape has any size that indicates a dynamic dimension.
     static bool isDynamicShape(ArrayRef<int64_t> dSizes) {
       return any_of(dSizes, [](int64_t dSize) { return isDynamic(dSize); });
     }
 
     /// Return the number of elements present in the given shape.
     static int64_t getNumElements(ArrayRef<int64_t> shape);
-  }];
 
-  let extraSharedClassDeclaration = [{
     /// Return a clone of this type with the given new shape and element type.
+    /// The returned type is ranked, even if this type is unranked.
     auto clone(::llvm::ArrayRef<int64_t> shape, Type elementType) {
-      return $_type.cloneWith(shape, elementType);
+      return cloneWith(shape, elementType);
     }
-    /// Return a clone of this type with the given new shape.
+
+    /// Return a clone of this type with the given new shape. The returned type
+    /// is ranked, even if this type is unranked.
     auto clone(::llvm::ArrayRef<int64_t> shape) {
-      return $_type.cloneWith(shape, $_type.getElementType());
+      return cloneWith(shape, getElementType());
     }
-    /// Return a clone of this type with the given new element type.
+  }];
+
+  let extraSharedClassDeclaration = [{
+    /// Return a clone of this type with the given new element type. The
+    /// returned type is ranked if and only if this type is ranked. In that
+    /// case, the returned type has the same shape as this type.
     auto clone(::mlir::Type elementType) {
       return $_type.cloneWith(/*shape=*/std::nullopt, elementType);
     }
 
     /// If an element type is an integer or a float, return its width. Otherwise,
     /// abort.
     unsigned getElementTypeBitWidth() const {
       return $_type.getElementType().getIntOrFloatBitWidth();

mlir/include/mlir/IR/BuiltinTypes.h

 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 class AffineExpr;
 class AffineMap;
 class FloatType;
 class IndexType;
 class IntegerType;
+class MemRefType;
+class RankedTensorType;
 class StringAttr;
 class TypeRange;
 
 //===----------------------------------------------------------------------===//
 // FloatType
 //===----------------------------------------------------------------------===//
 
 class FloatType : public Type {
public:
   /// Returns the shape of this tensor type.
   ArrayRef<int64_t> getShape() const;
 
   /// Clone this type with the given shape and element type. If the
   /// provided shape is `std::nullopt`, the current shape of the type is used.
   TensorType cloneWith(std::optional<ArrayRef<int64_t>> shape,
                        Type elementType) const;
 
+  // Make sure that base class overloads are visible.
+  using ShapedType::Trait<TensorType>::clone;
+
+  /// Return a clone of this type with the given new shape and element type.
+  /// The returned type is ranked, even if this type is unranked.
+  RankedTensorType clone(ArrayRef<int64_t> shape, Type elementType) const;
+
+  /// Return a clone of this type with the given new shape. The returned type
+  /// is ranked, even if this type is unranked.
+  RankedTensorType clone(ArrayRef<int64_t> shape) const;
+
   /// Return true if the specified element type is ok in a tensor.
   static bool isValidElementType(Type type);
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(Type type);
 
   /// Allow implicit conversion to ShapedType.
   operator ShapedType() const { return cast<ShapedType>(); }
public:
   /// Returns the shape of this memref type.
   ArrayRef<int64_t> getShape() const;
 
   /// Clone this type with the given shape and element type. If the
   /// provided shape is `std::nullopt`, the current shape of the type is used.
   BaseMemRefType cloneWith(std::optional<ArrayRef<int64_t>> shape,
                            Type elementType) const;
 
+  // Make sure that base class overloads are visible.
+  using ShapedType::Trait<BaseMemRefType>::clone;
+
+  /// Return a clone of this type with the given new shape and element type.
+  /// The returned type is ranked, even if this type is unranked.
+  MemRefType clone(ArrayRef<int64_t> shape, Type elementType) const;
+
+  /// Return a clone of this type with the given new shape. The returned type
+  /// is ranked, even if this type is unranked.
+  MemRefType clone(ArrayRef<int64_t> shape) const;
+
   /// Return true if the specified element type is ok in a memref.
   static bool isValidElementType(Type type);
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(Type type);
 
   /// Returns the memory space in which data referred to by this memref resides.
   Attribute getMemorySpace() const;

mlir/include/mlir/IR/BuiltinTypes.td

TypeBuilderWithInferredContext<(ins
       CArg<"Attribute", "{}">:$memorySpace)>,
     /// [deprecated] `Attribute`-based form should be used instead.
     TypeBuilderWithInferredContext<(ins
       "ArrayRef<int64_t>":$shape, "Type":$elementType,
       "AffineMap":$map,
       "unsigned":$memorySpaceInd)>
   ];
   let extraClassDeclaration = [{
-    using ShapedType::Trait<MemRefType>::clone;
+    using BaseMemRefType::clone;
     using ShapedType::Trait<MemRefType>::getElementTypeBitWidth;
     using ShapedType::Trait<MemRefType>::getRank;
     using ShapedType::Trait<MemRefType>::getNumElements;
     using ShapedType::Trait<MemRefType>::isDynamicDim;
     using ShapedType::Trait<MemRefType>::hasStaticShape;
     using ShapedType::Trait<MemRefType>::getNumDynamicDims;
     using ShapedType::Trait<MemRefType>::getDimSize;
     using ShapedType::Trait<MemRefType>::getDynamicDimIndex;
TypeBuilderWithInferredContext<(ins
       "ArrayRef<int64_t>":$shape,
       "Type":$elementType,
       CArg<"Attribute", "{}">:$encoding
     ), [{
       return $_get(elementType.getContext(), shape, elementType, encoding);
     }]>
   ];
   let extraClassDeclaration = [{
-    using ShapedType::Trait<RankedTensorType>::clone;
+    using TensorType::clone;
     using ShapedType::Trait<RankedTensorType>::getElementTypeBitWidth;
     using ShapedType::Trait<RankedTensorType>::getRank;
     using ShapedType::Trait<RankedTensorType>::getNumElements;
     using ShapedType::Trait<RankedTensorType>::isDynamicDim;
     using ShapedType::Trait<RankedTensorType>::hasStaticShape;
     using ShapedType::Trait<RankedTensorType>::getNumDynamicDims;
     using ShapedType::Trait<RankedTensorType>::getDimSize;
     using ShapedType::Trait<RankedTensorType>::getDynamicDimIndex;
 
     /// This is a builder type that keeps local references to arguments.
     /// Arguments that are passed into the builder must outlive the builder.
     class Builder;
+
+    /// Return a clone of this type with the given new element type and the same
+    /// shape as this type.
+    RankedTensorType clone(::mlir::Type elementType) {
+      return ::llvm::cast<RankedTensorType>(cloneWith(getShape(), elementType));
+    }
   }];
   let skipDefaultBuilders = 1;
   let genVerifyDecl = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TupleType
 //===----------------------------------------------------------------------===//
TypeBuilderWithInferredContext<(ins "Type":$elementType,
                                         "unsigned":$memorySpace), [{
       // Convert deprecated integer-like memory space to Attribute.
       Attribute memorySpaceAttr =
           wrapIntegerMemorySpace(memorySpace, elementType.getContext());
       return UnrankedMemRefType::get(elementType, memorySpaceAttr);
     }]>
   ];
   let extraClassDeclaration = [{
-    using ShapedType::Trait<UnrankedMemRefType>::clone;
+    using BaseMemRefType::clone;
     using ShapedType::Trait<UnrankedMemRefType>::getElementTypeBitWidth;
     using ShapedType::Trait<UnrankedMemRefType>::getRank;
     using ShapedType::Trait<UnrankedMemRefType>::getNumElements;
     using ShapedType::Trait<UnrankedMemRefType>::isDynamicDim;
     using ShapedType::Trait<UnrankedMemRefType>::hasStaticShape;
     using ShapedType::Trait<UnrankedMemRefType>::getNumDynamicDims;
     using ShapedType::Trait<UnrankedMemRefType>::getDimSize;
     using ShapedType::Trait<UnrankedMemRefType>::getDynamicDimIndex;
 
     ArrayRef<int64_t> getShape() const { return std::nullopt; }
 
     /// [deprecated] Returns the memory space in old raw integer representation.
     /// New `Attribute getMemorySpace()` method should be used instead.
     unsigned getMemorySpaceAsInt() const;
+
+    /// Return a clone of this type with the given new element type and the same
+    /// shape as this type.
+    MemRefType clone(::mlir::Type elementType) {
+      return ::llvm::cast<MemRefType>(cloneWith(getShape(), elementType));
+    }
   }];
   let skipDefaultBuilders = 1;
   let genVerifyDecl = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // UnrankedTensorType
 //===----------------------------------------------------------------------===//
def Builtin_UnrankedTensor : Builtin_Type<"UnrankedTensor", [
   let parameters = (ins "Type":$elementType);
 
   let builders = [
     TypeBuilderWithInferredContext<(ins "Type":$elementType), [{
       return $_get(elementType.getContext(), elementType);
     }]>
   ];
   let extraClassDeclaration = [{
-    using ShapedType::Trait<UnrankedTensorType>::clone;
+    using TensorType::clone;
     using ShapedType::Trait<UnrankedTensorType>::getElementTypeBitWidth;
     using ShapedType::Trait<UnrankedTensorType>::getRank;
     using ShapedType::Trait<UnrankedTensorType>::getNumElements;
     using ShapedType::Trait<UnrankedTensorType>::isDynamicDim;
     using ShapedType::Trait<UnrankedTensorType>::hasStaticShape;
     using ShapedType::Trait<UnrankedTensorType>::getNumDynamicDims;
     using ShapedType::Trait<UnrankedTensorType>::getDimSize;
     using ShapedType::Trait<UnrankedTensorType>::getDynamicDimIndex;

mlir/lib/IR/BuiltinTypes.cpp

TensorType TensorType::cloneWith(std::optional<ArrayRef<int64_t>> shape,
   auto rankedTy = cast<RankedTensorType>();
   if (!shape)
     return RankedTensorType::get(rankedTy.getShape(), elementType,
                                  rankedTy.getEncoding());
   return RankedTensorType::get(shape.value_or(rankedTy.getShape()), elementType,
                                rankedTy.getEncoding());
 }
 
+RankedTensorType TensorType::clone(::llvm::ArrayRef<int64_t> shape,
+                                   Type elementType) const {
+  return ::llvm::cast<RankedTensorType>(cloneWith(shape, elementType));
+}
+
+RankedTensorType TensorType::clone(::llvm::ArrayRef<int64_t> shape) const {
+  return ::llvm::cast<RankedTensorType>(cloneWith(shape, getElementType()));
+}
+
 // Check if "elementType" can be an element type of a tensor.
 static LogicalResult
 checkTensorElementType(function_ref<InFlightDiagnostic()> emitError,
                        Type elementType) {
   if (!TensorType::isValidElementType(elementType))
     return emitError() << "invalid tensor element type: " << elementType;
   return success();
 }
BaseMemRefType BaseMemRefType::cloneWith(std::optional<ArrayRef<int64_t>> shape,
 
   MemRefType::Builder builder(cast<MemRefType>());
   if (shape)
     builder.setShape(*shape);
   builder.setElementType(elementType);
   return builder;
 }
 
+MemRefType BaseMemRefType::clone(::llvm::ArrayRef<int64_t> shape,
+                                 Type elementType) const {
+  return ::llvm::cast<MemRefType>(cloneWith(shape, elementType));
+}
+
+MemRefType BaseMemRefType::clone(::llvm::ArrayRef<int64_t> shape) const {
+  return ::llvm::cast<MemRefType>(cloneWith(shape, getElementType()));
+}
+
 Attribute BaseMemRefType::getMemorySpace() const {
   if (auto rankedMemRefTy = dyn_cast<MemRefType>())
     return rankedMemRefTy.getMemorySpace();
   return cast<UnrankedMemRefType>().getMemorySpace();
 }
 
 unsigned BaseMemRefType::getMemorySpaceAsInt() const {
   if (auto rankedMemRefTy = dyn_cast<MemRefType>())