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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-ods

Author: None (MaheshRavishankar)

<details>
<summary>Changes</summary>

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

``````````

</details>


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