[mlir] detect integer overflow in debug mode

[aartbik]

Summary

Rationale:
This computation failed ASAN for the following input
(integer overflow during 4032000000000000000 * 100):

tensor<100x200x300x400x500x600x700x800xf32>

This change adds a simple overflow detection during
debug mode (which we run more regularly than ASAN).
Arguably this is an unrealistic tensor input, but
in the context of sparse tensors, we may start to
see cases like this.

Bug:
https://bugs.llvm.org/show_bug.cgi?id=49136

[mehdi_amini]

Would we catch the issue from yesterday if you moved this assertions in ShapedType::getNumElements()?
That would be more centralized and can cover more cases ideally.

[aartbik]

In D96530#2557732, @mehdi_amini wrote:

Would we catch the issue from yesterday if you moved this assertions in ShapedType::getNumElements()?
That would be more centralized and can cover more cases ideally.

That wouldn't have caught this particular issue, but looking at the logic in that method, happy to add that there too.

Please also give some thought on how to detect overflow (I present two methods in the bug as alternatives).
Testing > 0 is by far the easiest, but technically not a portable way, since the behavior is undefined (perhaps
a good reason to switch to unsigned :-)?

[mehdi_amini]

In D96530#2557752, @aartbik wrote:

In D96530#2557732, @mehdi_amini wrote:

Would we catch the issue from yesterday if you moved this assertions in ShapedType::getNumElements()?
That would be more centralized and can cover more cases ideally.

That wouldn't have caught this particular issue, but looking at the logic in that method, happy to add that there too.

OK, if you can add it there as well that'd be great!

Please also give some thought on how to detect overflow (I present two methods in the bug as alternatives).
Testing > 0 is by far the easiest, but technically not a portable way, since the behavior is undefined

Ah I missed the comment in the bug, answered there: I'm fine with 2) since we're just talking about an assertions for a poor-man UBSAN.

(perhaps good reason to switch to unsigned :-)?

Actually UBSAN catching the bug in the first place is the perfect justification why I never use unsigned :)
(unless I really intended to wrap around intentionally).

[aartbik]

similar quick overflow test in number of elements
(note that here we need >= since total size can be zero)

mlir/lib/IR/BuiltinTypes.cpp

 unsigned ShapedType::getElementTypeBitWidth() const {
   return getElementType().getIntOrFloatBitWidth();
 }
 
 int64_t ShapedType::getNumElements() const {
   assert(hasStaticShape() && "cannot get element count of dynamic shaped type");
   auto shape = getShape();
   int64_t num = 1;
-  for (auto dim : shape)
+  for (auto dim : shape) {
     num *= dim;
+    assert(num >= 0 && "integer overflow in element count computation");
+  }
   return num;
 }
 
 int64_t ShapedType::getRank() const {
   assert(hasRank() && "cannot query rank of unranked shaped type");
   return getShape().size();
 }
 
for (auto en : llvm::zip(llvm::reverse(exprs), llvm::reverse(sizes))) {
     // Degenerate case, no size =-> no stride
     if (size == 0)
       continue;
     AffineExpr dimExpr = std::get<0>(en);
     AffineExpr stride = dynamicPoisonBit
                             ? getAffineSymbolExpr(nSymbols++, context)
                             : getAffineConstantExpr(runningSize, context);
     expr = expr ? expr + dimExpr * stride : dimExpr * stride;
-    if (size > 0)
+    if (size > 0) {
       runningSize *= size;
-    else
+      assert(runningSize > 0 && "integer overflow in size computation");
+    } else {
       dynamicPoisonBit = true;
     }
+  }
   return simplifyAffineExpr(expr, numDims, nSymbols);
 }
 
 /// Return a version of `t` with a layout that has all dynamic offset and
 /// strides. This is used to erase the static layout.
 MemRefType mlir::eraseStridedLayout(MemRefType t) {
   auto val = ShapedType::kDynamicStrideOrOffset;
   return MemRefType::Builder(t).setAffineMaps(makeStridedLinearLayoutMap(