Introduce memref bound checking.

Introduce analysis to check memref accesses (in MLFunctions) for out of bound
ones. It works as follows:

$ mlir-opt -memref-bound-check test/Transforms/memref-bound-check.mlir

/tmp/single.mlir:10:12: error: 'load' op memref out of upper bound access along dimension tensorflow/mlir#1
      %x = load %A[%idxtensorflow/mlir#0, %idxtensorflow/mlir#1] : memref<9 x 9 x i32>
           ^
/tmp/single.mlir:10:12: error: 'load' op memref out of lower bound access along dimension tensorflow/mlir#1
      %x = load %A[%idxtensorflow/mlir#0, %idxtensorflow/mlir#1] : memref<9 x 9 x i32>
           ^
/tmp/single.mlir:10:12: error: 'load' op memref out of upper bound access along dimension tensorflow/mlir#2
      %x = load %A[%idxtensorflow/mlir#0, %idxtensorflow/mlir#1] : memref<9 x 9 x i32>
           ^
/tmp/single.mlir:10:12: error: 'load' op memref out of lower bound access along dimension tensorflow/mlir#2
      %x = load %A[%idxtensorflow/mlir#0, %idxtensorflow/mlir#1] : memref<9 x 9 x i32>
           ^
/tmp/single.mlir:12:12: error: 'load' op memref out of upper bound access along dimension tensorflow/mlir#1
      %y = load %B[%idy] : memref<128 x i32>
           ^
/tmp/single.mlir:12:12: error: 'load' op memref out of lower bound access along dimension tensorflow/mlir#1
      %y = load %B[%idy] : memref<128 x i32>
           ^
#map0 = (d0, d1) -> (d0, d1)
#map1 = (d0, d1) -> (d0 * 128 - d1)
mlfunc @test() {
  %0 = alloc() : memref<9x9xi32>
  %1 = alloc() : memref<128xi32>
  for %i0 = -1 to 9 {
    for %i1 = -1 to 9 {
      %2 = affine_apply #map0(%i0, %i1)
      %3 = load %0[%2tensorflow/mlir#0, %2tensorflow/mlir#1] : memref<9x9xi32>
      %4 = affine_apply #map1(%i0, %i1)
      %5 = load %1[%4] : memref<128xi32>
    }
  }
  return
}

- Improves productivity while manually / semi-automatically developing MLIR for
  testing / prototyping; also provides an indirect way to catch errors in
  transformations.

- This pass is an easy way to test the underlying affine analysis
  machinery including low level routines.

Some code (in getMemoryRegion()) borrowed from @andydavis cl/218263256.

While on this:

- create mlir/Analysis/Passes.h; move Pass.h up from mlir/Transforms/ to mlir/

- fix a bug in AffineAnalysis.cpp::toAffineExpr

TODO: extend to non-constant loop bounds (straightforward). Will transparently
work for all accesses once floordiv, mod, ceildiv are supported in the
AffineMap -> FlatAffineConstraints conversion.
PiperOrigin-RevId: 219397961

13 files changed, 205 insertions, 15 deletions

diff --git a/mlir/lib/Analysis/AffineAnalysis.cpp b/mlir/lib/Analysis/AffineAnalysis.cpp
index 8f792641dab..5cf8e6b13a6 100644
--- a/mlir/lib/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Analysis/AffineAnalysis.cpp
@@ -63,7 +63,7 @@ static AffineExpr toAffineExpr(ArrayRef<int64_t> eq, unsigned numDims,
   }
 
   // Constant term.
-  unsigned constTerm = eq[eq.size() - 1];
+  int64_t constTerm = eq[eq.size() - 1];
   if (constTerm != 0)
     expr = expr + constTerm;
   return expr;
@@ -299,7 +299,7 @@ AffineExpr mlir::simplifyAffineExpr(AffineExpr expr, unsigned numDims,
 }
 
 // Flattens 'expr' into 'flattenedExpr'. Returns true on success or false
-// if 'expr' was unable to be flattened (i.e. because it was not pur affine,
+// if 'expr' was unable to be flattened (i.e. because it was not pure affine,
 // or because it contained mod's and div's that could not be eliminated
 // without introducing local variables).
 bool mlir::getFlattenedAffineExpr(
diff --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp
index de7c5db3e07..52c15bb9237 100644
--- a/mlir/lib/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Analysis/AffineStructures.cpp
@@ -493,6 +493,22 @@ FlatAffineConstraints::FlatAffineConstraints(IntegerSet set)
   }
 }
 
+void FlatAffineConstraints::reset(unsigned numReservedInequalities,
+                                  unsigned numReservedEqualities,
+                                  unsigned newNumReservedCols,
+                                  unsigned newNumDims, unsigned newNumSymbols,
+                                  unsigned newNumLocals) {
+  assert(newNumReservedCols >= 1 && "minimum 1 column");
+  numReservedCols = newNumReservedCols;
+  numDims = newNumDims;
+  numSymbols = newNumSymbols;
+  numIds = numDims + numSymbols + newNumLocals;
+  equalities.clear();
+  inequalities.clear();
+  equalities.reserve(newNumReservedCols * numReservedEqualities);
+  inequalities.reserve(newNumReservedCols * numReservedInequalities);
+}
+
 /// Adds a dimensional identifier. The added column is initialized to
 /// zero.
 void FlatAffineConstraints::addDimId(unsigned pos) {
diff --git a/mlir/lib/Analysis/MemRefBoundCheck.cpp b/mlir/lib/Analysis/MemRefBoundCheck.cpp
new file mode 100644
index 00000000000..aa910f88092
--- /dev/null
+++ b/mlir/lib/Analysis/MemRefBoundCheck.cpp
@@ -0,0 +1,175 @@
+//===- MemRefBoundCheck.cpp - MLIR Affine Structures Class-----*- C++ -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file implements a pass to check memref accessses for out of bound
+// accesses.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/AffineAnalysis.h"
+#include "mlir/Analysis/AffineStructures.h"
+#include "mlir/Analysis/Passes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/StmtVisitor.h"
+#include "mlir/Pass.h"
+#include "mlir/StandardOps/StandardOps.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "memref-bound-check"
+
+using namespace mlir;
+
+namespace {
+
+/// Checks for out of bound memef access subscripts..
+struct MemRefBoundCheck : public FunctionPass, StmtWalker<MemRefBoundCheck> {
+  explicit MemRefBoundCheck() {}
+
+  PassResult runOnMLFunction(MLFunction *f) override;
+  // Not applicable to CFG functions.
+  PassResult runOnCFGFunction(CFGFunction *f) override { return success(); }
+
+  void visitOperationStmt(OperationStmt *opStmt);
+};
+
+} // end anonymous namespace
+
+FunctionPass *mlir::createMemRefBoundCheckPass() {
+  return new MemRefBoundCheck();
+}
+
+// Forward substitutes into 'valueMap' all AffineApplyOps reachable from the
+// operands of 'valueMap'.
+static void forwardSubstituteReachableOps(AffineValueMap *valueMap) {
+  // Gather AffineApplyOps reachable from 'indices'.
+  SmallVector<OperationStmt *, 4> affineApplyOps;
+  getReachableAffineApplyOps(valueMap->getOperands(), affineApplyOps);
+  // Compose AffineApplyOps in 'affineApplyOps'.
+  for (auto *opStmt : affineApplyOps) {
+    assert(opStmt->isa<AffineApplyOp>());
+    auto affineApplyOp = opStmt->dyn_cast<AffineApplyOp>();
+    // Forward substitute 'affineApplyOp' into 'valueMap'.
+    valueMap->forwardSubstitute(*affineApplyOp);
+  }
+}
+
+/// Returns the memory region accessed by this memref.
+// TODO(bondhugula): extend this to store's and other memref dereferencing ops.
+bool getMemoryRegion(OpPointer<LoadOp> loadOp, FlatAffineConstraints *region) {
+  OperationStmt *opStmt = dyn_cast<OperationStmt>(loadOp->getOperation());
+  // Only in MLFunctions.
+  if (!opStmt)
+    return false;
+
+  unsigned rank = loadOp->getMemRefType().getRank();
+  MLFuncBuilder b(opStmt);
+  auto idMap = b.getMultiDimIdentityMap(rank);
+
+  SmallVector<MLValue *, 4> indices;
+  for (auto *index : loadOp->getIndices()) {
+    indices.push_back(cast<MLValue>(index));
+  }
+
+  // Initialize 'srcValueMap' and compose with reachable AffineApplyOps.
+  AffineValueMap srcValueMap(idMap, indices);
+  forwardSubstituteReachableOps(&srcValueMap);
+  AffineMap srcMap = srcValueMap.getAffineMap();
+
+  region->reset(8, 8, srcMap.getNumInputs() + 1, srcMap.getNumDims(),
+                srcMap.getNumSymbols());
+
+  // Add equality constraints.
+  AffineMap map = srcValueMap.getAffineMap();
+  unsigned numDims = map.getNumDims();
+  unsigned numSymbols = map.getNumSymbols();
+  // Add inEqualties for loop lower/upper bounds.
+  for (unsigned i = 0; i < numDims + numSymbols; ++i) {
+    if (auto *loop = dyn_cast<ForStmt>(srcValueMap.getOperand(i))) {
+      if (!loop->hasConstantBounds())
+        return false;
+      // Add lower bound and upper bounds.
+      region->addConstantLowerBound(i, loop->getConstantLowerBound());
+      region->addConstantUpperBound(i, loop->getConstantUpperBound());
+    } else {
+      // Has to be a valid symbol.
+      auto *symbol = cast<MLValue>(srcValueMap.getOperand(i));
+      assert(symbol->isValidSymbol());
+      // Check if the symbols is a constant.
+      if (auto *opStmt = symbol->getDefiningStmt()) {
+        if (auto constOp = opStmt->dyn_cast<ConstantIndexOp>()) {
+          region->setIdToConstant(i, constOp->getValue());
+        }
+      }
+    }
+  }
+
+  // Add access function equalities to connect loop IVs to data dimensions.
+  region->addDimsForMap(0, srcValueMap.getAffineMap());
+
+  // Eliminate the loop IVs.
+  for (unsigned i = 0, e = srcValueMap.getNumOperands(); i < e; i++) {
+    region->FourierMotzkinEliminate(srcMap.getNumResults());
+  }
+  assert(region->getNumDimIds() == rank);
+  return true;
+}
+
+void MemRefBoundCheck::visitOperationStmt(OperationStmt *opStmt) {
+  // TODO(bondhugula): extend this to store's and other memref dereferencing
+  // op's.
+  if (auto loadOp = opStmt->dyn_cast<LoadOp>()) {
+    FlatAffineConstraints memoryRegion;
+    if (!getMemoryRegion(loadOp, &memoryRegion))
+      return;
+    unsigned rank = loadOp->getMemRefType().getRank();
+    // For each dimension, check for out of bounds.
+    for (unsigned r = 0; r < rank; r++) {
+      FlatAffineConstraints ucst(memoryRegion);
+      // Intersect memory region with constraint capturing out of bounds,
+      // and check if the constraint system is feasible. If it is, there is at
+      // least one point out of bounds.
+      SmallVector<int64_t, 4> ineq(rank + 1, 0);
+      // d_i >= memref dim size.
+      ucst.addConstantLowerBound(r, loadOp->getMemRefType().getDimSize(r));
+      LLVM_DEBUG(llvm::dbgs() << "System to check for overflow:\n");
+      LLVM_DEBUG(ucst.dump());
+      //
+      if (!ucst.isEmpty()) {
+        loadOp->emitOpError(
+            "memref out of upper bound access along dimension #" +
+            Twine(r + 1));
+      }
+      // Check for less than negative index.
+      FlatAffineConstraints lcst(memoryRegion);
+      std::fill(ineq.begin(), ineq.end(), 0);
+      // d_i <= -1;
+      lcst.addConstantUpperBound(r, -1);
+      LLVM_DEBUG(llvm::dbgs() << "System to check for underflow:\n");
+      LLVM_DEBUG(lcst.dump());
+      if (!lcst.isEmpty()) {
+        loadOp->emitOpError(
+            "memref out of lower bound access along dimension #" +
+            Twine(r + 1));
+      }
+    }
+  }
+}
+
+PassResult MemRefBoundCheck::runOnMLFunction(MLFunction *f) {
+  return walk(f), success();
+}
diff --git a/mlir/lib/Transforms/Utils/Pass.cpp b/mlir/lib/Analysis/Pass.cpp
index e4edee31900..1249c18c07e 100644
--- a/mlir/lib/Transforms/Utils/Pass.cpp
+++ b/mlir/lib/Analysis/Pass.cpp
@@ -19,7 +19,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Transforms/Pass.h"
+#include "mlir/Pass.h"
 #include "mlir/IR/CFGFunction.h"
 #include "mlir/IR/MLFunction.h"
 #include "mlir/IR/Module.h"
diff --git a/mlir/lib/Transforms/Canonicalizer.cpp b/mlir/lib/Transforms/Canonicalizer.cpp
index 4f308171467..f34118ce21a 100644
--- a/mlir/lib/Transforms/Canonicalizer.cpp
+++ b/mlir/lib/Transforms/Canonicalizer.cpp
@@ -22,7 +22,7 @@
 
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/PatternMatch.h"
-#include "mlir/Transforms/Pass.h"
+#include "mlir/Pass.h"
 #include "mlir/Transforms/Passes.h"
 using namespace mlir;
 
diff --git a/mlir/lib/Transforms/ComposeAffineMaps.cpp b/mlir/lib/Transforms/ComposeAffineMaps.cpp
index ced9cbcb86a..af4a5d11521 100644
--- a/mlir/lib/Transforms/ComposeAffineMaps.cpp
+++ b/mlir/lib/Transforms/ComposeAffineMaps.cpp
@@ -26,8 +26,8 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/StmtVisitor.h"
+#include "mlir/Pass.h"
 #include "mlir/StandardOps/StandardOps.h"
-#include "mlir/Transforms/Pass.h"
 #include "mlir/Transforms/Passes.h"
 #include "mlir/Transforms/Utils.h"
 #include "llvm/Support/CommandLine.h"
diff --git a/mlir/lib/Transforms/ConstantFold.cpp b/mlir/lib/Transforms/ConstantFold.cpp
index 15dd89bb758..411d1caae29 100644
--- a/mlir/lib/Transforms/ConstantFold.cpp
+++ b/mlir/lib/Transforms/ConstantFold.cpp
@@ -18,7 +18,7 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/CFGFunction.h"
 #include "mlir/IR/StmtVisitor.h"
-#include "mlir/Transforms/Pass.h"
+#include "mlir/Pass.h"
 #include "mlir/Transforms/Passes.h"
 #include "mlir/Transforms/Utils.h"
 
diff --git a/mlir/lib/Transforms/ConvertToCFG.cpp b/mlir/lib/Transforms/ConvertToCFG.cpp
index deeffb5bd9f..52687da65ba 100644
--- a/mlir/lib/Transforms/ConvertToCFG.cpp
+++ b/mlir/lib/Transforms/ConvertToCFG.cpp
@@ -23,7 +23,7 @@
 #include "mlir/IR/CFGFunction.h"
 #include "mlir/IR/MLFunction.h"
 #include "mlir/IR/Module.h"
-#include "mlir/Transforms/Pass.h"
+#include "mlir/Pass.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/DenseSet.h"
 using namespace mlir;
diff --git a/mlir/lib/Transforms/LoopUnroll.cpp b/mlir/lib/Transforms/LoopUnroll.cpp
index b8a5cf9d0a1..01cf61e5712 100644
--- a/mlir/lib/Transforms/LoopUnroll.cpp
+++ b/mlir/lib/Transforms/LoopUnroll.cpp
@@ -27,8 +27,8 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/StmtVisitor.h"
+#include "mlir/Pass.h"
 #include "mlir/Transforms/LoopUtils.h"
-#include "mlir/Transforms/Pass.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/CommandLine.h"
 
diff --git a/mlir/lib/Transforms/LoopUnrollAndJam.cpp b/mlir/lib/Transforms/LoopUnrollAndJam.cpp
index b05063232ed..f437b44ae26 100644
--- a/mlir/lib/Transforms/LoopUnrollAndJam.cpp
+++ b/mlir/lib/Transforms/LoopUnrollAndJam.cpp
@@ -50,8 +50,8 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/StmtVisitor.h"
+#include "mlir/Pass.h"
 #include "mlir/Transforms/LoopUtils.h"
-#include "mlir/Transforms/Pass.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/CommandLine.h"
 
diff --git a/mlir/lib/Transforms/PipelineDataTransfer.cpp b/mlir/lib/Transforms/PipelineDataTransfer.cpp
index 90421819d82..c59e007e543 100644
--- a/mlir/lib/Transforms/PipelineDataTransfer.cpp
+++ b/mlir/lib/Transforms/PipelineDataTransfer.cpp
@@ -26,9 +26,9 @@
 #include "mlir/Analysis/Utils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/StmtVisitor.h"
+#include "mlir/Pass.h"
 #include "mlir/StandardOps/StandardOps.h"
 #include "mlir/Transforms/LoopUtils.h"
-#include "mlir/Transforms/Pass.h"
 #include "mlir/Transforms/Utils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/Debug.h"
diff --git a/mlir/lib/Transforms/SimplifyAffineExpr.cpp b/mlir/lib/Transforms/SimplifyAffineExpr.cpp
index f4a23e6b153..a412a83f66c 100644
--- a/mlir/lib/Transforms/SimplifyAffineExpr.cpp
+++ b/mlir/lib/Transforms/SimplifyAffineExpr.cpp
@@ -20,14 +20,13 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/AffineStructures.h"
-#include "mlir/IR/AffineMap.h"
-#include "mlir/IR/Attributes.h"
 #include "mlir/IR/MLFunction.h"
-#include "mlir/IR/Statements.h"
 #include "mlir/IR/StmtVisitor.h"
-#include "mlir/Transforms/Pass.h"
+#include "mlir/Pass.h"
 #include "mlir/Transforms/Passes.h"
 
+#define DEBUG_TYPE "simplify-affine-structure"
+
 using namespace mlir;
 using llvm::report_fatal_error;
 
diff --git a/mlir/lib/Transforms/Vectorize.cpp b/mlir/lib/Transforms/Vectorize.cpp
index 511afa95993..a6c9681b6f5 100644
--- a/mlir/lib/Transforms/Vectorize.cpp
+++ b/mlir/lib/Transforms/Vectorize.cpp
@@ -29,9 +29,9 @@
 #include "mlir/IR/MLValue.h"
 #include "mlir/IR/SSAValue.h"
 #include "mlir/IR/Types.h"
+#include "mlir/Pass.h"
 #include "mlir/StandardOps/StandardOps.h"
 #include "mlir/Support/Functional.h"
-#include "mlir/Transforms/Pass.h"
 #include "mlir/Transforms/Passes.h"
 
 #include "llvm/ADT/DenseMap.h"