[Mlir-commits] [mlir] 46ca425 - [MLIR][Affine] Fix crash in normalizeMemRefType

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2022-10-07T16:12:29+05:30
New Revision: 46ca4259c5d066ccee797fea7233cf21eb0e6070

URL: https://github.com/llvm/llvm-project/commit/46ca4259c5d066ccee797fea7233cf21eb0e6070
DIFF: https://github.com/llvm/llvm-project/commit/46ca4259c5d066ccee797fea7233cf21eb0e6070.diff

LOG: [MLIR][Affine] Fix crash in normalizeMemRefType

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

Reviewed By: dcaballe

Differential Revision: https://reviews.llvm.org/D135424

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/Utils/Utils.cpp
    mlir/test/Transforms/normalize-memrefs.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index eaacce294bf13..2f077cc3bae0c 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1792,20 +1792,24 @@ MemRefType mlir::normalizeMemRefType(MemRefType memrefType,
   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;
     }

diff  --git a/mlir/test/Transforms/normalize-memrefs.mlir b/mlir/test/Transforms/normalize-memrefs.mlir
index f3d0a12ac72df..892d5e53f690c 100644
--- a/mlir/test/Transforms/normalize-memrefs.mlir
+++ b/mlir/test/Transforms/normalize-memrefs.mlir
@@ -332,3 +332,23 @@ func.func @affine_parallel_norm() ->  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>
+}