[MLIR] ShapedType accessor minor fixes + add isDynamicDim accessor

[bondhugula]

Summary

Minor fixes for ShapedType accessors, use ShapedType::kDynamicSize, add ShapedType::isDynamicDim.

[ftynse]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1071

Can we initialize it to kDynamicSize instead of hardcoded -1?

[bondhugula]

Address review comment.

[nicolasvasilache]

Can we initialize it to kDynamicSize instead of hardcoded -1?

Orthogonally to this revision, note that this is a particular pain point: there are a lot of -1 scattered around the codebase both in MLIR and things that depend on MLIR.
-1 is not good enough for representing dynamic offset and strides.
I had started an effort to clean that up a few months ago but did not see the light at the end of the tunnel..

mlir/include/mlir/IR/StandardTypes.h

public:
   bool hasStaticShape(ArrayRef<int64_t> shape) const;
 
   /// If this is a ranked type, return the number of dimensions with dynamic
   /// size. Otherwise, abort.
   int64_t getNumDynamicDims() const;
 
   /// If this is ranked type, return the size of the specified dimension.
   /// Otherwise, abort.
-  int64_t getDimSize(int64_t i) const;
+  int64_t getDimSize(unsigned idx) const;
+
+  /// Returns true if this dimension has a dynamic size (for ranked types);
+  /// aborts for unranked types.
+  bool isDynamicDim(unsigned idx) const;
 
   /// Returns the position of the dynamic dimension relative to just the dynamic
   /// dimensions, given its `index` within the shape.
   unsigned getDynamicDimIndex(unsigned index) const;
 
   /// Get the total amount of bits occupied by a value of this type.  This does
   /// not take into account any memory layout or widening constraints, e.g. a
   /// vector<3xi57> is reported to occupy 3x57=171 bit, even though in practice
   /// it will likely be stored as in a 4xi64 vector register.  Fail an assertion
   /// if the size cannot be computed statically, i.e. if the type has a dynamic
   /// shape or if its elemental type does not have a known bit width.
   int64_t getSizeInBits() const;
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(Type type) {
     return type.getKind() == StandardTypes::Vector ||
            type.getKind() == StandardTypes::RankedTensor ||
            type.getKind() == StandardTypes::UnrankedTensor ||
            type.getKind() == StandardTypes::UnrankedMemRef ||
            type.getKind() == StandardTypes::MemRef;
   }
 
   /// Whether the given dimension size indicates a dynamic dimension.
-  static constexpr bool isDynamic(int64_t dSize) { return dSize < 0; }
+  static constexpr bool isDynamic(int64_t dSize) {
+    return dSize == kDynamicSize;
+  }
   static constexpr bool isDynamicStrideOrOffset(int64_t dStrideOrOffset) {
     return dStrideOrOffset == kDynamicStrideOrOffset;
   }
 };
 
 /// Vector types represent multi-dimensional SIMD vectors, and have a fixed
 /// known constant shape with one or more dimension.
 class VectorType

mlir/lib/Analysis/Utils.cpp

LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
   assert(cst.getNumDimIds() == rank && "unexpected MemRefRegion format");
 
   // Add upper/lower bounds for each memref dimension with static size
   // to guard against potential over-approximation from projection.
   // TODO(andydavis) Support dynamic memref dimensions.
   if (addMemRefDimBounds) {
     auto memRefType = memref.getType().cast<MemRefType>();
     for (unsigned r = 0; r < rank; r++) {
-      cst.addConstantLowerBound(r, 0);
-      int64_t dimSize = memRefType.getDimSize(r);
-      if (ShapedType::isDynamic(dimSize))
+      cst.addConstantLowerBound(/*pos=*/r, /*lb=*/0);
+      if (memRefType.isDynamicDim(r))
         continue;
-      cst.addConstantUpperBound(r, dimSize - 1);
+      cst.addConstantUpperBound(/*pos=*/r, memRefType.getDimSize(r) - 1);
     }
   }
   cst.removeTrivialRedundancy();
 
   LLVM_DEBUG(llvm::dbgs() << "Memory region:\n");
   LLVM_DEBUG(cst.dump());
   return success();
 }

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

struct DimOpLowering : public ConvertOpToLLVMPattern<DimOp> {
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto dimOp = cast<DimOp>(op);
     OperandAdaptor<DimOp> transformed(operands);
     MemRefType type = dimOp.getOperand().getType().cast<MemRefType>();
 
-    auto shape = type.getShape();
     int64_t index = dimOp.getIndex();
     // Extract dynamic size from the memref descriptor.
-    if (ShapedType::isDynamic(shape[index]))
+    if (type.isDynamicDim(index))
       rewriter.replaceOp(op, {MemRefDescriptor(transformed.memrefOrTensor())
                                   .size(rewriter, op->getLoc(), index)});
     else
       // Use constant for static size.
-      rewriter.replaceOp(
-          op, createIndexConstant(rewriter, op->getLoc(), shape[index]));
+      rewriter.replaceOp(op, createIndexConstant(rewriter, op->getLoc(),
+                                                 type.getDimSize(index)));
     return success();
   }
 };
 
 // Common base for load and store operations on MemRefs.  Restricts the match
 // to supported MemRef types.  Provides functionality to emit code accessing a
 // specific element of the underlying data buffer.
 template <typename Derived>

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 
 /// Returns true if the 'index' dimension of the `memref` defined by
 /// `memrefDefOp` is a statically  shaped one or defined using a valid symbol.
 template <typename AnyMemRefDefOp>
 static bool isMemRefSizeValidSymbol(AnyMemRefDefOp memrefDefOp,
                                     unsigned index) {
   auto memRefType = memrefDefOp.getType();
   // Statically shaped.
-  if (!ShapedType::isDynamic(memRefType.getDimSize(index)))
+  if (!memRefType.isDynamicDim(index))
     return true;
   // Get the position of the dimension among dynamic dimensions;
   unsigned dynamicDimPos = memRefType.getDynamicDimIndex(index);
   return isValidSymbol(
       *(memrefDefOp.getDynamicSizes().begin() + dynamicDimPos));
 }
 
 /// Returns true if the result of the dim op is a valid symbol.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

static LogicalResult verify(DimOp op) {
   }
 
   return success();
 }
 
 OpFoldResult DimOp::fold(ArrayRef<Attribute> operands) {
   // Constant fold dim when the size along the index referred to is a constant.
   auto opType = memrefOrTensor().getType();
-  int64_t indexSize = ShapedType::kDynamicSize;
+  int64_t dimSize = ShapedType::kDynamicSize;
   if (auto tensorType = opType.dyn_cast<RankedTensorType>())
-    indexSize = tensorType.getShape()[getIndex()];
+    dimSize = tensorType.getShape()[getIndex()];
   else if (auto memrefType = opType.dyn_cast<MemRefType>())
-    indexSize = memrefType.getShape()[getIndex()];
+    dimSize = memrefType.getShape()[getIndex()];
 
-  if (!ShapedType::isDynamic(indexSize))
-    return IntegerAttr::get(IndexType::get(getContext()), indexSize);
+  if (!ShapedType::isDynamic(dimSize))
+    return IntegerAttr::get(IndexType::get(getContext()), dimSize);
 
   // Fold dim to the size argument for an AllocOp/ViewOp/SubViewOp.
   auto memrefType = opType.dyn_cast<MemRefType>();
   if (!memrefType)
     return {};
 
   // The size at getIndex() is now a dynamic size of a memref.
   auto memref = memrefOrTensor().getDefiningOp();
Value ViewOp::getDynamicOffset() {
   assert(result);
   if (result && offset == MemRefType::getDynamicStrideOrOffset())
     return getOperand(1);
   return nullptr;
 }
 
 static LogicalResult verifyDynamicStrides(MemRefType memrefType,
                                           ArrayRef<int64_t> strides) {
-  ArrayRef<int64_t> shape = memrefType.getShape();
   unsigned rank = memrefType.getRank();
   assert(rank == strides.size());
   bool dynamicStrides = false;
   for (int i = rank - 2; i >= 0; --i) {
     // If size at dim 'i + 1' is dynamic, set the 'dynamicStrides' flag.
-    if (ShapedType::isDynamic(shape[i + 1]))
+    if (memrefType.isDynamicDim(i + 1))
       dynamicStrides = true;
     // If stride at dim 'i' is not dynamic, return error.
     if (dynamicStrides && strides[i] != MemRefType::getDynamicStrideOrOffset())
       return failure();
   }
   return success();
 }
 

mlir/lib/IR/StandardTypes.cpp

for (auto dim : shape)
     num *= dim;
   return num;
 }
 
 int64_t ShapedType::getRank() const { return getShape().size(); }
 
 bool ShapedType::hasRank() const { return !isa<UnrankedTensorType>(); }
 
-int64_t ShapedType::getDimSize(int64_t i) const {
-  assert(i >= 0 && i < getRank() && "invalid index for shaped type");
-  return getShape()[i];
+int64_t ShapedType::getDimSize(unsigned idx) const {
+  assert(idx < getRank() && "invalid index for shaped type");
+  return getShape()[idx];
+}
+
+bool ShapedType::isDynamicDim(unsigned idx) const {
+  assert(idx < getRank() && "invalid index for shaped type");
+  return isDynamic(getShape()[idx]);
 }
 
 unsigned ShapedType::getDynamicDimIndex(unsigned index) const {
   assert(index < getRank() && "invalid index");
   assert(ShapedType::isDynamic(getDimSize(index)) && "invalid index");
   return llvm::count_if(getShape().take_front(index), ShapedType::isDynamic);
 }