diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp --- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp +++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp @@ -99,7 +99,7 @@ } else if (!matchPattern(&op, m_Constant())) { // Register op in the set of ops that have users. opsWithUsers.insert(&op); - if (isa(op)) { + if (isa(op)) { auto read = dyn_cast(op); Value memref = read ? read.getMemRef() : cast(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 --- 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 @@ // 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