[Mlir-commits] [mlir] Allow MemRefLayoutInterface to return null AffineMaps (PR #162757)

[llvmlistbot at llvm.org]

https://github.com/MaheshRavishankar created https://github.com/llvm/llvm-project/pull/162757

The current requirement of `MemRefLayoutAttrInterface` seems to be that any attribute that implements the interface needs to return an `AffineMap`. This seems like too strict a requirement for an interface. It could be upto the implementation to either return an `AffineMap` if possible, or NULL if the strides cannot be represented by an `AffineMap`.  Any method that uses this interface and gets the layout as `AffineMap` should handle the case where such a representation is not supported by the underlying interface.

>From d2db8418adcc4374be5247401bb69f0b13b14754 Mon Sep 17 00:00:00 2001
From: MaheshRavishankar <mahesh.ravishankar at gmail.com>
Date: Tue, 30 Sep 2025 09:55:01 -0700
Subject: [PATCH] Allow MemRefLayoutInterface to return null AffineMaps

Signed-off-by: MaheshRavishankar <mahesh.ravishankar at gmail.com>
---
 mlir/include/mlir/IR/BuiltinAttributeInterfaces.td | 5 +++--
 mlir/lib/IR/BuiltinAttributeInterfaces.cpp         | 6 +++++-
 2 files changed, 8 insertions(+), 3 deletions(-)

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();