update normalizeMemRef utility; handle missing failure check + add more tests

- take care of symbolic operands with alloc
- add missing check for compose map failure and a test case
- add test cases on strides
- drop incorrect check for one-to-one'ness

Signed-off-by: Uday Bondhugula <uday@polymagelabs.com>

Closes tensorflow/mlir#132

COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/mlir/pull/132 from bondhugula:normalize-memrefs 8aebf285fb0d7c19269d85255aed644657e327b7
PiperOrigin-RevId: 269105947

1 files changed, 9 insertions, 10 deletions

diff --git a/mlir/lib/Transforms/Utils/Utils.cpp b/mlir/lib/Transforms/Utils/Utils.cpp
index 5f1bf93f0c0..e57d40e5a1c 100644
--- a/mlir/lib/Transforms/Utils/Utils.cpp
+++ b/mlir/lib/Transforms/Utils/Utils.cpp
@@ -389,7 +389,7 @@ void mlir::createAffineComputationSlice(
   }
 }
 
-// TODO: Currently works for static memrefs with single non-identity layout map.
+// TODO: Currently works for static memrefs with a single layout map.
 LogicalResult mlir::normalizeMemRef(AllocOp allocOp) {
   MemRefType memrefType = allocOp.getType();
   unsigned rank = memrefType.getRank();
@@ -403,16 +403,12 @@ LogicalResult mlir::normalizeMemRef(AllocOp allocOp) {
 
   AffineMap layoutMap = layoutMaps.front();
 
+  // Nothing to do for identity layout maps.
   if (layoutMap == b.getMultiDimIdentityMap(rank))
     return success();
 
-  if (layoutMap.getNumResults() < rank)
-    // This is a sufficient condition for not being one-to-one; the map is thus
-    // invalid. Leave it alone. (Undefined behavior?)
-    return failure();
-
-  // We don't do any more non-trivial checks for one-to-one'ness; we
-  // assume that it is one-to-one.
+  // We don't do any checks for one-to-one'ness; we assume that it is
+  // one-to-one.
 
   // TODO: Only for static memref's for now.
   if (memrefType.getNumDynamicDims() > 0)
@@ -421,7 +417,7 @@ LogicalResult mlir::normalizeMemRef(AllocOp allocOp) {
   // We have a single map that is not an identity map. Create a new memref with
   // the right shape and an identity layout map.
   auto shape = memrefType.getShape();
-  FlatAffineConstraints fac(rank, 0);
+  FlatAffineConstraints fac(rank, allocOp.getNumSymbolicOperands());
   for (unsigned d = 0; d < rank; ++d) {
     fac.addConstantLowerBound(d, 0);
     fac.addConstantUpperBound(d, shape[d] - 1);
@@ -430,7 +426,10 @@ LogicalResult mlir::normalizeMemRef(AllocOp allocOp) {
   // We compose this map with the original index (logical) space to derive the
   // upper bounds for the new index space.
   unsigned newRank = layoutMap.getNumResults();
-  fac.composeMatchingMap(layoutMap);
+  if (failed(fac.composeMatchingMap(layoutMap)))
+    // TODO: semi-affine maps.
+    return failure();
+
   // Project out the old data dimensions.
   fac.projectOut(newRank, fac.getNumIds() - newRank - fac.getNumLocalIds());
   SmallVector<int64_t, 4> newShape(newRank);