[mlir] detect integer overflow in debug mode

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

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D96530

Author: aartbik

mlir/lib/IR/BuiltinTypes.cpp

@@ -209,8 +209,10 @@ 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;
 }
 
@@ -801,10 +803,12 @@ AffineExpr mlir::makeCanonicalStridedLayoutExpr(ArrayRef<int64_t> sizes,
                             ? 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);
 }