[mlir][affine] fix affine LICM pass for has effected memory's user

[lipracer]

Summary

When the memory is written by dma, its user is moved

[lipracer]

bug report https://github.com/llvm/llvm-project/issues/59838

[bondhugula]

LGTM - thanks for fixing this.

[lipracer]

This comment was removed by lipracer.

[lipracer]

In D141106#4036569, @bondhugula wrote:

LGTM - thanks for fixing this.

Can you help me to land it？

[Mogball]

landed

[lipracer]

In D141106#4100415, @Mogball wrote:

landed

Thanks.

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

if (auto ifOp = dyn_cast<AffineIfOp>(op)) {
     opsWithUsers.insert(&op);
     if (isa<AffineMapAccessInterface>(op)) {
       auto read = dyn_cast<AffineReadOpInterface>(op);
       Value memref = read ? read.getMemRef()
                           : cast<AffineWriteOpInterface>(op).getMemRef();
       for (auto *user : memref.getUsers()) {
         // If this memref has a user that is a DMA, give up because these
         // operations write to this memref.
-        if (isa<AffineDmaStartOp, AffineDmaWaitOp>(op))
+        if (isa<AffineDmaStartOp, AffineDmaWaitOp>(user))
           return false;
         // If the memref used by the load/store is used in a store elsewhere in
         // the loop nest, we do not hoist. Similarly, if the memref used in a
         // load is also being stored too, we do not hoist the load.
         if (isa<AffineWriteOpInterface>(user) ||
             (isa<AffineReadOpInterface>(user) &&
              isa<AffineWriteOpInterface>(op))) {
           if (&op != user) {

mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir

func.func @affine_parallel(%memref_8: memref<4090x2040xf32>, %x: index) {
   // CHECK:       affine.for
   // CHECK-NEXT:    affine.for
   // CHECK-NEXT:      affine.parallel
   // CHECK-NEXT:        scf.for
   // CHECK-NEXT:          affine.store
 
   return
 }
+
+// -----
+
+// CHECK-LABEL: func.func @affine_invariant_use_after_dma
+#map = affine_map<(d0) -> (d0 * 163840)>
+func.func @affine_invariant_use_after_dma(%arg0: memref<10485760xi32>, %arg1: memref<1xi32>, %arg2: memref<10485760xi32>) {
+  %c320 = arith.constant 320 : index
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %alloc = memref.alloc() {alignment = 16 : i64} : memref<0xi32, 2>
+  %alloc_0 = memref.alloc() : memref<1xi32, 2>
+  affine.for %arg3 = 0 to 64 {
+    %0 = affine.apply #map(%arg3)
+    %alloc_1 = memref.alloc() {alignment = 16 : i64} : memref<0xi32, 2>
+    %alloc_2 = memref.alloc() : memref<320xi32, 2>
+    affine.dma_start %arg0[%0], %alloc_2[%c0], %alloc_1[%c0], %c320 : memref<10485760xi32>, memref<320xi32, 2>, memref<0xi32, 2>
+    affine.dma_start %arg1[%c0], %alloc_0[%c0], %alloc[%c0], %c1 : memref<1xi32>, memref<1xi32, 2>, memref<0xi32, 2>
+    affine.dma_wait %alloc_1[%c0], %c320 : memref<0xi32, 2>
+    affine.dma_wait %alloc[%c0], %c1 : memref<0xi32, 2>
+    %1 = affine.apply #map(%arg3)
+    %alloc_3 = memref.alloc() {alignment = 16 : i64} : memref<0xi32, 2>
+    %alloc_4 = memref.alloc() : memref<320xi32, 2>
+    affine.for %arg4 = 0 to 320 {
+      %2 = affine.load %alloc_2[%arg4] : memref<320xi32, 2>
+      %3 = affine.load %alloc_0[0] : memref<1xi32, 2>
+      %4 = arith.addi %2, %3 : i32
+      %5 = arith.addi %4, %2 : i32
+      affine.store %5, %alloc_4[%arg4] : memref<320xi32, 2>
+    }
+    affine.dma_start %alloc_4[%c0], %arg2[%1], %alloc_3[%c0], %c320 : memref<320xi32, 2>, memref<10485760xi32>, memref<0xi32, 2>
+    affine.dma_wait %alloc_3[%c0], %c320 : memref<0xi32, 2>
+  }
+  return
+}
+// CHECK: %[[zero:.*]] = arith.constant 0 : index
+// CHECK: %[[scalar_mem:.*]] = memref.alloc() : memref<1xi32, 2>
+// CHECK: affine.dma_start %arg1[%[[zero]]], %alloc_0[%[[zero]]], %alloc[%[[zero]]], %c1
+// CHECK: affine.load %[[scalar_mem]][0]