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

[anoopjs]

Summary

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

[bondhugula]

LGTM - thanks.

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

102

Use variadic isa.

[anoopjs]

Use variadic isa

[anoopjs]

Thanks @bondhugula. I don’t have commit access, can you land this patch for me? Please use “[e-mail]” to commit the change.

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

if (auto ifOp = dyn_cast<AffineIfOp>(op)) {
     // instead rely on side effects.
     return false;
   } else if (op.getNumRegions() > 0) {
     // We can't handle region-holding ops we don't know about.
     return false;
   } 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();
       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>(user))
           return false;

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

affine.for %arg3 = 0 to 64 {
     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]
+
+// -----
+
+// 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