[Mlir-commits] [mlir] Fix memref load vector element crash (PR #224834)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-affine

Author: purnima shrivastava (purnima-nlp)

<details>
<summary>Changes</summary>

`ConvertMemRefLoad` assumed memref element types were always integer or floating-point and called `getIntOrFloatBitWidth()` unconditionally.

For vector element types such as `vector<1xi1>`, this caused an assertion. This change makes the pattern fail gracefully instead, and adds a regression test.

Fixes #<!-- -->224811.



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


4 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Analysis/Utils.cpp (+14-17) 
- (modified) mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp (+2) 
- (modified) mlir/test/Dialect/Affine/loop-fusion.mlir (+17) 
- (added) mlir/test/Dialect/MemRef/emulate-narrow-type-invalid.mlir (+7) 


``````````diff
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 67885513dcd3b..2b0338276024c 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -243,17 +243,15 @@ addNodeToMDG(Operation *nodeOp, MemRefDependenceGraph &mdg,
   return &node;
 }
 
-/// Returns the memref being read/written by a memref/affine load/store op.
-static Value getMemRef(Operation *memOp) {
-  if (auto memrefLoad = dyn_cast<memref::LoadOp>(memOp))
-    return memrefLoad.getMemRef();
-  if (auto affineLoad = dyn_cast<AffineReadOpInterface>(memOp))
-    return affineLoad.getMemRef();
-  if (auto memrefStore = dyn_cast<memref::StoreOp>(memOp))
-    return memrefStore.getMemRef();
-  if (auto affineStore = dyn_cast<AffineWriteOpInterface>(memOp))
-    return affineStore.getMemRef();
-  llvm_unreachable("unexpected op");
+/// Returns true if `op` may read from or write to `memref`.
+static bool mayAccessMemRef(Operation *op, Value memref) {
+  SmallVector<Value> effectedValues;
+  getEffectedValues<MemoryEffects::Read, MemoryEffects::Write>(op,
+                                                               effectedValues);
+  for (Value effectedValue : effectedValues)
+    if (effectedValue == memref)
+      return true;
+  return false;
 }
 
 /// Returns true if there may be a dependence on `memref` from srcNode's
@@ -272,16 +270,15 @@ static bool mayDependence(const Node &srcNode, const Node &dstNode,
   // true if there exists a conflicting read/write access involving such.
 
   // Check whether there is a dependence from a source read/write op to a
-  // destination read/write one; all expected to be memref/affine load/store.
+  // destination read/write op on `memref`.
   auto hasNonAffineDep = [&](ArrayRef<Operation *> srcMemOps,
                              ArrayRef<Operation *> dstMemOps) {
     return llvm::any_of(srcMemOps, [&](Operation *srcOp) {
-      Value srcMemref = getMemRef(srcOp);
-      if (srcMemref != memref)
+      if (!mayAccessMemRef(srcOp, memref))
         return false;
-      return llvm::find_if(dstMemOps, [&](Operation *dstOp) {
-               return srcMemref == getMemRef(dstOp);
-             }) != dstMemOps.end();
+      return llvm::any_of(dstMemOps, [&](Operation *dstOp) {
+        return mayAccessMemRef(dstOp, memref);
+      });
     });
   };
 
diff --git a/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp b/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
index a11e14faa5475..2a62fe7a7ecfe 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
@@ -331,6 +331,8 @@ struct ConvertMemRefLoad final : OpConversionPattern<memref::LoadOp> {
     auto convertedType = cast<MemRefType>(adaptor.getMemref().getType());
     auto convertedElementType = convertedType.getElementType();
     auto oldElementType = op.getMemRefType().getElementType();
+    if (!oldElementType.isIntOrFloat() || !convertedElementType.isIntOrFloat())
+      return failure();
     int srcBits = oldElementType.getIntOrFloatBitWidth();
     int dstBits = convertedElementType.getIntOrFloatBitWidth();
     if (dstBits % srcBits != 0) {
diff --git a/mlir/test/Dialect/Affine/loop-fusion.mlir b/mlir/test/Dialect/Affine/loop-fusion.mlir
index 1ea42517988c3..4f72b2deab50c 100644
--- a/mlir/test/Dialect/Affine/loop-fusion.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion.mlir
@@ -1576,3 +1576,20 @@ func.func @producer_consumer_with_outmost_user(%arg0 : f16) {
 
 // Add further tests in mlir/test/Transforms/loop-fusion-4.mlir
 
+// -----
+
+// CHECK-LABEL: func @vector_transfer_read
+func.func @vector_transfer_read(%arg0: memref<4xf32>) {
+  %c0 = arith.constant 0.0 : f32
+
+  affine.for %i = 0 to 4 {
+    affine.store %c0, %arg0[%i] : memref<4xf32>
+  }
+
+  affine.for %i = 0 to 4 {
+    %0 = vector.transfer_read %arg0[%i], %c0
+        : memref<4xf32>, vector<4xf32>
+  }
+
+  return
+}
diff --git a/mlir/test/Dialect/MemRef/emulate-narrow-type-invalid.mlir b/mlir/test/Dialect/MemRef/emulate-narrow-type-invalid.mlir
new file mode 100644
index 0000000000000..07e2c89e6b096
--- /dev/null
+++ b/mlir/test/Dialect/MemRef/emulate-narrow-type-invalid.mlir
@@ -0,0 +1,7 @@
+// RUN: mlir-opt --test-emulate-narrow-int --verify-diagnostics %s
+
+func.func @memref_load_vector_element(%arg0: memref<4xvector<1xi1>>, %idx: index) {
+  // expected-error @+1 {{failed to legalize operation 'memref.load'}}
+  %0 = memref.load %arg0[%idx] : memref<4xvector<1xi1>>
+  return
+}

``````````

</details>


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