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

mlir/include/mlir/Analysis/AffineStructures.h

@@ -0,0 +1,268 @@
+//===- AffineStructures.h - 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 ML functions.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_AFFINE_STRUCTURES_H
+#define MLIR_ANALYSIS_AFFINE_STRUCTURES_H
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace mlir {
+
+class AffineExpr;
+class AffineApplyOp;
+class AffineBound;
+class AffineCondition;
+class AffineMap;
+class IntegerSet;
+class MLIRContext;
+class MLValue;
+
+/// A mutable affine map. Its affine expressions are however unique.
+struct MutableAffineMap {
+public:
+  explicit MutableAffineMap(AffineMap *map);
+
+  AffineExpr *getResult(unsigned idx) const { return results[idx]; }
+  unsigned getNumResults() const { return results.size(); }
+
+private:
+  SmallVector<AffineExpr *, 8> results;
+  SmallVector<AffineExpr *, 8> rangeSizes;
+};
+
+/// A mutable integer set. Its affine expressions are however unique.
+struct MutableIntegerSet {
+public:
+  explicit MutableIntegerSet(IntegerSet *set);
+
+  unsigned getNumDims() const { return numDims; }
+  unsigned getNumSymbols() const { return numSymbols; }
+  unsigned getNumConstraints() const { return constraints.size(); }
+
+private:
+  unsigned numDims;
+  unsigned numSymbols;
+
+  SmallVector<AffineExpr *, 8> constraints;
+  SmallVector<bool, 8> eqFlags;
+};
+
+/// An AffineValueMap is an affine map plus its ML value operands and
+/// results for analysis purposes. The structure is still a tree form that is
+/// same as that of an affine map or an AffineApplyOp. However, its operands,
+/// results, and its map can themselves change  as a result of
+/// substitutions, simplifications, and other analysis.
+// An affine value map can readily be constructed from an AffineApplyOp, or an
+// AffineBound of a ForStmt. It can be further transformed, substituted into,
+// or simplified. Unlike AffineMap's, AffineValueMap's are created and destroyed
+// during analysis. Only the AffineMap expressions that are pointed by them are
+// unique'd.
+// TODO(bondhugula): Some of these classes could go into separate files.
+class AffineValueMap {
+public:
+  explicit AffineValueMap(const AffineApplyOp &op);
+  explicit AffineValueMap(const AffineBound &bound);
+  explicit AffineValueMap(AffineMap *map);
+
+  ~AffineValueMap();
+
+  /// Substitute the results of inputMap into the operands of this map.
+  // The new list of operands will be a union of this map's and that of the map
+  // we are substituting from.
+  // Example usage scenario: a subscript operand for a 'load' is forward
+  // substituted into the memref's access map. The subscript operand itself is
+  // then substituted by its defining affine_apply op instructions and
+  // successively by a loop IV remap expression, eventually resulting in an
+  // affine value map that has only the loop IVs and symbols as its operands.
+  // Hence, the access pattern can then be analyzed for example.
+  // TODO(bondhugula)
+  void fwdSubstitute(const AffineValueMap &inputMap);
+  void fwdSubstitute(const AffineApplyOp &inputOp);
+
+  /// Return true if the idx^th result can be proved to be a multiple of
+  /// 'factor', false otherwise.
+  bool isMultipleOf(unsigned idx, int64_t factor) const;
+
+  /// Return true if the result at 'idx' is a constant, false
+  /// otherwise.
+  bool isConstant(unsigned idx) const;
+
+  /// Return true if this is an identity map.
+  bool isIdentity() const;
+
+private:
+  // A mutable affine map.
+  MutableAffineMap map;
+
+  // TODO: make these trailing objects?
+  /// The SSA operands binding to the dim's and symbols of 'map'.
+  SmallVector<MLValue *, 4> operands;
+  /// The SSA results binding to the results of 'map'.
+  SmallVector<MLValue *, 4> results;
+  /// A pointer to the IR's context to store all newly created AffineExpr's.
+  MLIRContext *context;
+};
+
+/// An IntegerValueSet is an integer set plus its operands.
+// Both, the integer set being pointed to and the operands can change during
+// analysis, simplification, and transformation.
+class IntegerValueSet {
+  // Constructs an integer value set map from an IntegerSet and operands.
+  explicit IntegerValueSet(const AffineCondition &cond);
+
+  /// Constructs an integer value set from an affine value map.
+  // This will lead to a single equality in 'set'.
+  explicit IntegerValueSet(const AffineValueMap &avm);
+
+  /// Returns true if this integer set is empty.
+  bool isEmpty() const;
+
+  bool getNumDims() const { return set.getNumDims(); }
+  bool getNumSymbols() const { return set.getNumSymbols(); }
+
+private:
+  // The set pointed to may itself change unlike in IR structures like
+  // 'AffineCondition'.
+  MutableIntegerSet set;
+  /// The SSA operands binding to the dim's and symbols of 'set'.
+  SmallVector<MLValue *, 4> operands;
+  /// A pointer to the IR's context to store all newly created AffineExpr's.
+  MLIRContext *context;
+};
+
+/// A flat list of affine equalities and inequalities in the form.
+/// Inequality: c_0*x_0 + c_1*x_1 + .... + c_{n-1}*x_{n-1} == 0
+/// Equality: c_0*x_0 + c_1*x_1 + .... + c_{n-1}*x_{n-1} >= 0
+///
+/// The coefficients are stored. x_0, x_1, ... appear in the order: dimensional
+/// identifiers, symbolic identifiers, and local identifiers.  / The local
+/// identifiers correspond to local/internal variables created / temporarily and
+/// are needed to increase representational power. Local identifiers / are
+/// typically obtained when eliminating % and div constraints.
+//  Usage scenario:
+//  For a register tiling or unroll-jam, for example, if we need to check loop
+//  bounds are a multiple of say the tile size say 4:
+//
+// %lb = affine_apply #map1 (%s, %i0)
+// %ub = affine_apply #map2 (%N, %i0)
+// for %i1 = %lb to %ub
+//   ...
+//
+// Create AffineValueMap's that have result %lb, %ub (successively fwd
+// substituting all affine_apply that lead to the %lb, %ub). Create another
+// AffineValueMap: %trip_count = (%ub - %lb + 1). Create a
+// FlatAffineConstraints set using all these. Add %trip_count % 4 = 0 to this,
+// and check for feasibility.
+class FlatAffineConstraints {
+public:
+  enum IdKind { Dimension, Symbol, Local };
+
+  /// Construct a constraint system reserving memory for the specified number of
+  /// constraints and identifiers..
+  FlatAffineConstraints(unsigned numReservedInequalities,
+                        unsigned numReservedEqualities, unsigned numReservedIds)
+      : numEqualities(0), numInequalities(0),
+        numReservedEqualities(numReservedEqualities),
+        numReservedInequalities(numReservedInequalities),
+        numReservedIds(numReservedIds) {
+    equalities.reserve(numReservedIds * numReservedEqualities);
+    inequalities.reserve(numReservedIds * numReservedInequalities);
+  }
+
+  /// Create a flat affine constraint system from an AffineValueMap or a list of
+  /// these. The constructed system will only include equalities.
+  // TODO(bondhugula)
+  explicit FlatAffineConstraints(const AffineValueMap &avm);
+  explicit FlatAffineConstraints(ArrayRef<const AffineValueMap *> avmRef);
+
+  /// Create an affine constraint system from an IntegerValueSet.
+  // TODO(bondhugula)
+  FlatAffineConstraints(const IntegerValueSet &set);
+
+  FlatAffineConstraints(ArrayRef<const AffineValueMap *> avmRef,
+                        const IntegerSet &set);
+
+  ~FlatAffineConstraints() {}
+
+  inline int64_t atEq(unsigned i, unsigned j) {
+    return equalities[i * numIds + j];
+  }
+
+  inline int64_t atIneq(unsigned i, unsigned j) {
+    return inequalities[i * numIds + j];
+  }
+
+  unsigned getNumEqualities() const { return equalities.size(); }
+  unsigned getNumInequalities() const { return inequalities.size(); }
+
+  void addInequality(ArrayRef<int64_t> inEq);
+  void addEquality(ArrayRef<int64_t> eq);
+
+  void addId(IdKind idKind, unsigned pos);
+  void addDimId(unsigned pos);
+  void addSymbolId(unsigned pos);
+  void addLocalId(unsigned pos);
+
+  void removeId(IdKind idKind, unsigned pos);
+
+  void removeEquality(unsigned pos);
+  void removeInequality(unsigned pos);
+
+  unsigned getNumConstraints() const { return numEqualities + numInequalities; }
+  unsigned getNumIds() const { return numIds; }
+  unsigned getNumDimIds() const { return numDims; }
+  unsigned getNumSymbolIds() const { return numSymbols; }
+  unsigned getNumLocalIds() const { return numDims - numSymbols - numDims; }
+
+private:
+  /// Coefficients of affine equalities (in == 0 form).
+  SmallVector<int64_t, 64> equalities;
+
+  /// Coefficients of affine inequalities (in >= 0 form).
+  SmallVector<int64_t, 64> inequalities;
+
+  /// Number of equalities in this system.
+  unsigned numEqualities;
+
+  /// Number of inequalities in this system.
+  unsigned numInequalities;
+
+  // Pre-allocated space.
+  unsigned numReservedEqualities;
+  unsigned numReservedInequalities;
+  unsigned numReservedIds;
+
+  /// Total number of identifiers.
+  unsigned numIds;
+
+  /// Number of identifiers corresponding to real dimensions.
+  unsigned numDims;
+
+  /// Number of identifiers corresponding to symbols (unknown but constant for
+  /// analysis).
+  unsigned numSymbols;
+};
+
+} // end namespace mlir.
+
+#endif // MLIR_ANALYSIS_AFFINE_STRUCTURES_H

mlir/include/mlir/IR/Statements.h

@@ -238,6 +238,8 @@ public:
   }
 
 private:
+  // TODO(shpeisman): please name the ForStmt's bounds encapsulating
+  // an affinemap and its operands as AffineBound.
   AffineConstantExpr *lowerBound;
   AffineConstantExpr *upperBound;
   AffineConstantExpr *step;
@@ -292,6 +294,8 @@ public:
 private:
   IfClause *thenClause;
   IfClause *elseClause;
+  // TODO(shpeisman): please name the ifStmt's conditional encapsulating
+  // IntegerSet + its operands as AffineCondition.
   // The integer set capturing the conditional guard.
   IntegerSet *condition;
   // TODO: arguments to integer set

mlir/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

mlir/lib/IR/AffineMap.cpp

@@ -205,20 +205,33 @@ AffineExpr *AffineBinaryOpExpr::simplifyCeilDiv(AffineExpr *lhs,
   }
 
   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.
 }

mlir/test/IR/affine-map.mlir

@@ -146,8 +146,11 @@
 // CHECK: #map{{[0-9]+}} = (d0, d1, d2)[s0] -> (d0 + d1 + d2 + 1, d2 + d1, (d0 * s0) * 8)
 #map45 = (i, j, k) [N] -> (1 + i + 3 + j - 3 + k, k + 5 + j - 5, 2*i*4*N)
 
-// CHECK: #map{{[0-9]+}} = (d0, d1, d2) -> (0, d0 * 2, 0, d0, d0 * 4)
-#map46 = (i, j, k) -> (i*0, i * 128 floordiv 64, j * 0 floordiv 64, i * 64 ceildiv 64, i * 512 ceildiv 128)
+// CHECK: #map{{[0-9]+}} = (d0, d1, d2) -> (0, d1, d0 * 2, 0)
+#map46 = (i, j, k) -> (i*0, 1*j, i * 128 floordiv 64, j * 0 floordiv 64)
+
+// CHECK: #map{{[0-9]+}} = (d0, d1, d2) -> (d0, d0 * 4, 0, 0)
+#map47 = (i, j, k) -> (i * 64 ceildiv 64, i * 512 ceildiv 128, 4 * j mod 4, 4*j*4 mod 8)
 
 // CHECK: extfunc @f0(memref<2x4xi8, #map{{[0-9]+}}, 1>)
 extfunc @f0(memref<2x4xi8, #map0, 1>)
@@ -302,3 +305,6 @@ extfunc @f45(memref<100x100x100xi8, #map45>)
 
 // CHECK: extfunc @f46(memref<100x100x100xi8, #map{{[0-9]+}}>)
 extfunc @f46(memref<100x100x100xi8, #map46>)
+
+// CHECK: extfunc @f47(memref<100x100x100xi8, #map{{[0-9]+}}>)
+extfunc @f47(memref<100x100x100xi8, #map47>)