[mlir][affine] Fix crash on -affine-loop-invariant-code-motion pass with affine.prefe...

Affine Prefetch Op impelements AffineMapAccessInterface but does not implement
AffineReadOpInterface or AffineWriteOpInterface. Prefetch Op was cast to
AffineWriteOpinterface causing the crash.

Reviewed By: bondhugula

Differential Revision: https://reviews.llvm.org/D146836

diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
index 320e184..3d8ea2e 100644 (file)
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
@@ -99,7 +99,7 @@ bool isOpLoopInvariant(Operation &op, Value indVar, ValueRange iterArgs,
   } else if (!matchPattern(&op, m_Constant())) {
     // Register op in the set of ops that have users.
     opsWithUsers.insert(&op);
-    if (isa<AffineMapAccessInterface>(op)) {
+    if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op)) {
       auto read = dyn_cast<AffineReadOpInterface>(op);
       Value memref = read ? read.getMemRef()
                           : cast<AffineWriteOpInterface>(op).getMemRef();

diff --git a/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir b/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
index 923cead..50bc938 100644 (file)
--- a/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
+++ b/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
@@ -846,3 +846,25 @@ func.func @affine_invariant_use_after_dma(%arg0: memref<10485760xi32>, %arg1: me
 // 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]
+
+// -----
+
+// CHECK-LABEL: func @affine_prefetch_invariant
+func.func @affine_prefetch_invariant() {
+  %0 = memref.alloc() : memref<10x10xf32>
+  affine.for %i0 = 0 to 10 {
+    affine.for %i1 = 0 to 10 {
+      %1 = affine.load %0[%i0, %i1] : memref<10x10xf32>
+      affine.prefetch %0[%i0, %i0], write, locality<0>, data : memref<10x10xf32>
+    }
+  }
+
+  // CHECK:      memref.alloc() : memref<10x10xf32>
+  // CHECK-NEXT: affine.for %{{.*}} = 0 to 10 {
+  // CHECK-NEXT:   affine.prefetch
+  // CHECK-NEXT:   affine.for %{{.*}} = 0 to 10 {
+  // CHECK-NEXT:     %{{.*}}  = affine.load %{{.*}}[%{{.*}}  : memref<10x10xf32>
+  // CHECK-NEXT:   }
+  // CHECK-NEXT: }
+  return
+}
\ No newline at end of file