[mlir][IR] RankedShapedType::getNumElements supports non-static shapes

[springerm]

Summary

Return kDynamic in case of non-static shapes. This addresses a comment in D148453.

Depends On: D148453

[mehdi_amini]

Why is this a better API than the assertion?

[rkayaith]

In D148605#4276618, @mehdi_amini wrote:

Why is this a better API than the assertion?

It seems more consistent with the other methods, e.g. getDimSize(dim) doesn't assert that the dim size is known. Also a RankedShapedType without static shape is still a valid RankedShapeType, so it seems odd for this to be a pre-condition on a method (this seems somewhat similar to the motivation for changing the ShapedType methods that required the type be ranked).

On the other hand, since kDynamic is a sentinel value, it could be easy to miss that it needs to be handled. But I think that's an issue with how kDynamic is used elsewhere as well...

[mehdi_amini]

In D148605#4277871, @rkayaith wrote:

In D148605#4276618, @mehdi_amini wrote:

Why is this a better API than the assertion?

It seems more consistent with the other methods, e.g. getDimSize(dim) doesn't assert that the dim size is known.

That does not seem comparable to me: it getDimSize(dim) returns always a valid dimension, if not known it return kDynamic, which is a valid value for a "dimension".
I don't see how getNumElements() can return anything "valid" when shapes are dynamic.

Also a RankedShapedType without static shape is still a valid RankedShapeType, so it seems odd for this to be a pre-condition on a method

I don't understand the connection you're making here?
There are many objects with methods you cannot call because of preconditions, even if the object is valid. Here a valid RankedShapeType object does not necessarily have a statically computable "number of elements".

[rkayaith]

In D148605#4277956, @mehdi_amini wrote:

In D148605#4277871, @rkayaith wrote:

In D148605#4276618, @mehdi_amini wrote:

Why is this a better API than the assertion?

It seems more consistent with the other methods, e.g. getDimSize(dim) doesn't assert that the dim size is known.

That does not seem comparable to me: it getDimSize(dim) returns always a valid dimension, if not known it return kDynamic, which is a valid value for a "dimension".
I don't see how getNumElements() can return anything "valid" when shapes are dynamic.

The interpretation would be that kDynamic is a valid value for a "shape element count" (maybe this isn't useful though, and actually I would've preferred returning something like std::optional<int64_t> instead, I just suggested kDynamic initially since it seemed like there was precedent).

Also a RankedShapedType without static shape is still a valid RankedShapeType, so it seems odd for this to be a pre-condition on a method

I don't understand the connection you're making here?
There are many objects with methods you cannot call because of preconditions, even if the object is valid.

Fair, though to make the API harder to misuse I think it'd be desirable to remove the precondition if we can find a sensible behaviour (i.e. somehow indicate that there's no statically known element count, rather than just having undefined behaviour)

Anyways I'm not convinced this should be changed either, just thought it was worth considering :)

mlir/include/mlir/IR/BuiltinTypeInterfaces.td

let extraClassDeclaration = [{
     }
 
     /// Return "true" if 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. Asserts that
-    /// all dimensions are static.
+    /// Return the number of elements present in the given shape. Returns
+    /// kDynamic if at least one dimension is dynamic.
     static int64_t getNumElements(ArrayRef<int64_t> shape);
   }];
 
   let extraSharedClassDeclaration = [{
     /// Return the rank of this type.
     int64_t getRank() const {
       return $_type.getShape().size();
     }
 
-    /// If it has static shape, return the number of elements. Otherwise, abort.
+    /// If it has static shape, return the number of elements. Returns kDynamic
+    /// if at least one dimension is dynamic.
     int64_t getNumElements() const {
-      assert(hasStaticShape()
-             && "cannot get element count of dynamically shaped type");
       return ::mlir::RankedShapedType::getNumElements($_type.getShape());
     }
 
     /// Return "true" if this dimension has a dynamic size.
     bool isDynamicDim(unsigned idx) const {
       assert(idx < getRank() && "invalid index for shaped type");
       return ::mlir::RankedShapedType::isDynamic($_type.getShape()[idx]);
     }

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

 // ReshapeOp
 //===----------------------------------------------------------------------===//
 
 void ReshapeOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getResult(), "reshape");
 }
 
-static int64_t getNumElements(RankedShapedType type) {
-  int64_t numElements = 1;
-  for (auto dim : type.getShape())
-    numElements *= dim;
-  return numElements;
-}
-
 LogicalResult ReshapeOp::verify() {
   TensorType operandType = getSource().getType().cast<TensorType>();
   TensorType resultType = getResult().getType().cast<TensorType>();
 
   if (operandType.getElementType() != resultType.getElementType())
     return emitOpError("element types of source and destination tensor "
                        "types should be the same");
 
   int64_t shapeSize =
       getShape().getType().cast<RankedTensorType>().getDimSize(0);
   auto resultRankedType = resultType.dyn_cast<RankedTensorType>();
   auto operandRankedType = operandType.dyn_cast<RankedTensorType>();
 
   if (resultRankedType) {
     if (operandRankedType && resultRankedType.hasStaticShape() &&
         operandRankedType.hasStaticShape()) {
-      if (getNumElements(operandRankedType) != getNumElements(resultRankedType))
+      if (operandRankedType.getNumElements() !=
+          resultRankedType.getNumElements())
         return emitOpError("source and destination tensor should have the "
                            "same number of elements");
     }
     if (RankedShapedType::isDynamic(shapeSize))
       return emitOpError("cannot use shape operand with dynamic length to "
                          "reshape to statically-ranked tensor type");
     if (shapeSize != resultRankedType.getRank())
       return emitOpError(

mlir/lib/IR/BuiltinTypeInterfaces.cpp

 // RankedShapedType
 //===----------------------------------------------------------------------===//
 
 constexpr int64_t RankedShapedType::kDynamic;
 
 int64_t RankedShapedType::getNumElements(ArrayRef<int64_t> shape) {
   int64_t num = 1;
   for (int64_t dim : shape) {
-    assert(!RankedShapedType::isDynamic(dim) && "expected only static dims");
+    if (RankedShapedType::isDynamic(dim))
+      return RankedShapedType::kDynamic;
     num *= dim;
     assert(num >= 0 && "integer overflow in element count computation");
   }
   return num;
 }