Sketch out affine analysis structures: AffineValueMap, IntegerValueSet,
FlatAffineConstraints, and MutableAffineMap.

All four classes introduced reside in lib/Analysis and are not meant to be
used in the IR (from lib/IR or lib/Parser/). They are all mutable, alloc'ed,
dealloc'ed - although with their fields pointing to immutable affine
expressions (AffineExpr *).

While on this, update simplifyMod to fold mod to a zero when possible.

PiperOrigin-RevId: 209618437
diff --git a/lib/Analysis/AffineStructures.cpp b/lib/Analysis/AffineStructures.cpp
new file mode 100644
index 0000000..295c76a
--- /dev/null
+++ b/lib/Analysis/AffineStructures.cpp
@@ -0,0 +1,61 @@
+//===- AffineStructures.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.
+// =============================================================================
+//
+// Structures for affine/polyhedral analysis of MLIR functions.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/AffineStructures.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/IntegerSet.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/StandardOps.h"
+
+namespace mlir {
+
+MutableAffineMap::MutableAffineMap(AffineMap *map) {
+  for (auto *result : map->getResults())
+    results.push_back(result);
+  for (auto *rangeSize : map->getRangeSizes())
+    results.push_back(rangeSize);
+}
+
+MutableIntegerSet::MutableIntegerSet(IntegerSet *set)
+    : numDims(set->getNumDims()), numSymbols(set->getNumSymbols()) {
+  // TODO(bondhugula)
+}
+
+AffineValueMap::AffineValueMap(const AffineApplyOp &op)
+    : map(op.getAffineMap()) {
+  // TODO: pull operands and results in.
+}
+
+bool AffineValueMap::isMultipleOf(unsigned idx, int64_t factor) const {
+  /* Check if the (first result expr) % factor becomes 0. */
+  if (auto *expr = dyn_cast<AffineConstantExpr>(AffineBinaryOpExpr::get(
+          AffineExpr::Kind::Mod, map.getResult(idx),
+          AffineConstantExpr::get(factor, context), context)))
+    return expr->getValue() == 0;
+
+  // TODO(bondhugula): use FlatAffineConstraints to complete this.
+  assert(0 && "isMultipleOf implementation incomplete");
+}
+
+AffineValueMap::~AffineValueMap() {}
+
+} // end namespace mlir
diff --git a/lib/IR/AffineMap.cpp b/lib/IR/AffineMap.cpp
index 3f640ba..95cd8ff 100644
--- a/lib/IR/AffineMap.cpp
+++ b/lib/IR/AffineMap.cpp
@@ -205,20 +205,33 @@
   }
 
   return nullptr;
-  // TODO(someone): implement more simplification along the lines described in
-  // simplifyMod TODO. For eg: 128*N ceildiv 128 is N.
 }
 
 AffineExpr *AffineBinaryOpExpr::simplifyMod(AffineExpr *lhs, AffineExpr *rhs,
                                             MLIRContext *context) {
-  if (auto *l = dyn_cast<AffineConstantExpr>(lhs))
-    if (auto *r = dyn_cast<AffineConstantExpr>(rhs))
-      return AffineConstantExpr::get(l->getValue() % r->getValue(), context);
+  auto *lhsConst = dyn_cast<AffineConstantExpr>(lhs);
+  auto *rhsConst = dyn_cast<AffineConstantExpr>(rhs);
+
+  if (lhsConst && rhsConst)
+    return AffineConstantExpr::get(lhsConst->getValue() % rhsConst->getValue(),
+                                   context);
+
+  // Fold modulo of a multiply with a constant that is a multiple of the
+  // modulo factor to zero. Eg: (i * 128) mod 64 = 0.
+  if (rhsConst) {
+    auto *lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
+    if (lBin && lBin->getKind() == Kind::Mul) {
+      if (auto *lrhs = dyn_cast<AffineConstantExpr>(lBin->getRHS())) {
+        // rhsConst is known to be positive if a constant.
+        if (lrhs->getValue() % rhsConst->getValue() == 0)
+          return AffineConstantExpr::get(0, context);
+      }
+    }
+  }
 
   return nullptr;
-  // TODO(someone): implement more simplification; for eg: 2*x mod 2 is 0; (2*x
-  // + 1) mod 2 is 1. In general, this can be simplified by using the GCD test
-  // iteratively if the RHS of the mod is a small number, or in general using
-  // quantifier elimination (add two new variables q and r, and eliminate all
-  // variables from the linear system other than r.
+  // TODO(bondhugula): In general, this can be simplified more by using the GCD
+  // test, or in general using quantifier elimination (add two new variables q
+  // and r, and eliminate all variables from the linear system other than r. All
+  // of this can be done through mlir/Analysis/'s FlatAffineConstraints.
 }