[MLIR][Affine] Fix crash in normalizeMemRefType

[bondhugula]

Summary

Fix crash in normalizeMemRefType. Correctly handle scenario and replace
assertion with a failure.

[dcaballe]

Thanks! LGTM

mlir/test/Transforms/normalize-memrefs.mlir

336

do we have a test for the negative case? Could you please add one?

[bondhugula]

Add test for the negative case. Rebase.

mlir/test/Transforms/normalize-memrefs.mlir

336

Good catch. Sure, added it now.

mlir/lib/Dialect/Affine/Utils/Utils.cpp

MemRefType mlir::normalizeMemRefType(MemRefType memrefType,
   if (failed(fac.composeMatchingMap(layoutMap)))
     return memrefType;
   // TODO: Handle semi-affine maps.
   // Project out the old data dimensions.
   fac.projectOut(newRank, fac.getNumVars() - newRank - fac.getNumLocalVars());
   SmallVector<int64_t, 4> newShape(newRank);
   MLIRContext *context = memrefType.getContext();
   for (unsigned d = 0; d < newRank; ++d) {
-    // Check if each dimension of normalized memrefType is dynamic.
+    // Check if this dimension is dynamic.
     bool isDynDim =
         isNormalizedMemRefDynamicDim(d, layoutMap, memrefTypeDynDims, context);
     if (isDynDim) {
       newShape[d] = -1;
     } else {
       // The lower bound for the shape is always zero.
-      auto ubConst = fac.getConstantBound64(IntegerPolyhedron::UB, d);
+      Optional<int64_t> ubConst =
+          fac.getConstantBound64(IntegerPolyhedron::UB, d);
       // For a static memref and an affine map with no symbols, this is
-      // always bounded.
-      assert(ubConst && "should always have an upper bound");
-      if (ubConst.value() < 0)
-        // This is due to an invalid map that maps to a negative space.
+      // always bounded. However, when we have symbols, we may not be able to
+      // obtain a constant upper bound. Also, mapping to a negative space is
+      // invalid for normalization.
+      if (!ubConst.has_value() || ubConst.value() < 0) {
+        LLVM_DEBUG(llvm::dbgs()
+                   << "can't normalize map due to unknown/invalid upper bound");
         return memrefType;
+      }
       // If dimension of new memrefType is dynamic, the value is -1.
       newShape[d] = ubConst.value() + 1;
     }
   }
 
   // Create the new memref type after trivializing the old layout map.
   MemRefType newMemRefType =
       MemRefType::Builder(memrefType)

mlir/test/Transforms/normalize-memrefs.mlir

func.func @affine_parallel_norm() ->  memref<8xf32, #tile> {
   // CHECK: affine.parallel (%{{.*}}) = (0) to (8) reduce ("assign") -> (memref<2x4xf32>)
   %1 = affine.parallel (%i) = (0) to (8) reduce ("assign") ->  memref<8xf32, #tile> {
     affine.store %c, %a[%i] : memref<8xf32, #tile>
     // CHECK: affine.yield %{{.*}} : memref<2x4xf32>
     affine.yield %a : memref<8xf32, #tile>
   }
   return %1 : memref<8xf32, #tile>
 }
+
+#map = affine_map<(d0, d1)[s0] -> (d0 * 3 + s0 + d1)>
+// CHECK-LABEL: func.func @map_symbol
+func.func @map_symbol() -> memref<2x3xf32, #map> {
+  %c1 = arith.constant 1 : index
+  // The constant isn't propagated here and the utility can't compute a constant
+  // upper bound for the memref dimension in the absence of that.
+  // CHECK: memref.alloc()[%{{.*}}]
+  %0 = memref.alloc()[%c1] : memref<2x3xf32, #map>
+  return %0 : memref<2x3xf32, #map>
+}
+
+#neg = affine_map<(d0, d1) -> (d0, d1 - 100)>
+// CHECK-LABEL: func.func @neg_map
+func.func @neg_map() -> memref<2x3xf32, #neg> {
+  // This isn't a valid map for normalization.
+  // CHECK: memref.alloc() : memref<2x3xf32, #{{.*}}>
+  %0 = memref.alloc() : memref<2x3xf32, #neg>
+  return %0 : memref<2x3xf32, #neg>
+}