[Mlir-commits] [mlir] 5f022ad - [mlir] detect integer overflow in debug mode

[Aart Bik]

Author: Aart Bik
Date: 2021-02-11T18:20:40-08:00
New Revision: 5f022ad6ed8dc23094bb50f83608d5c786a821a9

URL: https://github.com/llvm/llvm-project/commit/5f022ad6ed8dc23094bb50f83608d5c786a821a9
DIFF: https://github.com/llvm/llvm-project/commit/5f022ad6ed8dc23094bb50f83608d5c786a821a9.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/IR/BuiltinTypes.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/BuiltinTypes.cpp b/mlir/lib/IR/BuiltinTypes.cpp
index 289434ad52cd..1316bfd3bb52 100644
--- a/mlir/lib/IR/BuiltinTypes.cpp
+++ b/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);
 }