diff --git a/mlir/include/mlir/IR/BuiltinAttributeInterfaces.td b/mlir/include/mlir/IR/BuiltinAttributeInterfaces.td
index 7bc7fbe8c50f2..beccdf847383b 100644
--- a/mlir/include/mlir/IR/BuiltinAttributeInterfaces.td
+++ b/mlir/include/mlir/IR/BuiltinAttributeInterfaces.td
@@ -486,7 +486,7 @@ def MemRefLayoutAttrInterface : AttrInterface<"MemRefLayoutAttrInterface"> {
 
   let methods = [
     InterfaceMethod<
-      "Get the MemRef layout as an AffineMap, the method must not return NULL",
+      "Get the MemRef layout as an AffineMap, returns NULL if cannot be expressed as affine map",
       "::mlir::AffineMap", "getAffineMap", (ins)
     >,
 
@@ -495,7 +495,8 @@ def MemRefLayoutAttrInterface : AttrInterface<"MemRefLayoutAttrInterface"> {
       "bool", "isIdentity", (ins),
       [{}],
       [{
-        return $_attr.getAffineMap().isIdentity();
+        AffineMap map = $_attr.getAffineMap();
+        return map && $_attr.getAffineMap().isIdentity();
       }]
     >,
 
diff --git a/mlir/lib/IR/BuiltinAttributeInterfaces.cpp b/mlir/lib/IR/BuiltinAttributeInterfaces.cpp
index 9e8ce4ca3a902..d490ad54ab78d 100644
--- a/mlir/lib/IR/BuiltinAttributeInterfaces.cpp
+++ b/mlir/lib/IR/BuiltinAttributeInterfaces.cpp
@@ -77,6 +77,9 @@ uint64_t ElementsAttr::getFlattenedIndex(Type type, ArrayRef<uint64_t> index) {
 LogicalResult mlir::detail::verifyAffineMapAsLayout(
     AffineMap m, ArrayRef<int64_t> shape,
     function_ref<InFlightDiagnostic()> emitError) {
+  if (!m) {
+    return success();
+  }
   if (m.getNumDims() != shape.size())
     return emitError() << "memref layout mismatch between rank and affine map: "
                        << shape.size() << " != " << m.getNumDims();
@@ -204,7 +207,8 @@ LogicalResult mlir::detail::getAffineMapStridesAndOffset(
     int64_t &offset) {
   AffineExpr offsetExpr;
   SmallVector<AffineExpr, 4> strideExprs;
-  if (failed(::getStridesAndOffset(map, shape, strideExprs, offsetExpr)))
+  if (!map ||
+      failed(::getStridesAndOffset(map, shape, strideExprs, offsetExpr)))
     return failure();
   if (auto cst = llvm::dyn_cast<AffineConstantExpr>(offsetExpr))
     offset = cst.getValue();