Allow MemRefLayoutInterface to return null AffineMaps

[MaheshRavishankar]

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.

[llvmbot]

@llvm/pr-subscribers-mlir-core
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-ods

Author: None (MaheshRavishankar)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/162757.diff

2 Files Affected:

• (modified) mlir/include/mlir/IR/BuiltinAttributeInterfaces.td (+3-2)
• (modified) mlir/lib/IR/BuiltinAttributeInterfaces.cpp (+5-1)

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

[krzysz00]

At the very least, this needs documentation updates

I also don't know if this is the only place in the code where there's a relyance on the convertability of memref layouts to affine maps.

[ftynse]

I'd recommend a forum post. This is a builtin interface and we want to be a bit mindful of downstream failure. I remember some discussions about simplifying/relaxing restrictions on layouts when we switched to strided as first class (as opposed to being extracted from an affine map).

[MaheshRavishankar]

Posted on discourse here https://discourse.llvm.org/t/drop-the-requirement-of-having-to-return-affinemap-s-from-memreflayoutinterface/88564 .

@krzysz00 what changes to the documentation (on top of changing the method description) are you refering to.