[Mlir-commits] [mlir] [MLIR] Bug Fix: affine.prefetch replaceAffineOp invoked during canonicalization (PR #88346)

[Alexandre Eichenberger]

https://github.com/AlexandreEichenberger updated https://github.com/llvm/llvm-project/pull/88346

>From db2aa3dd7dbcea28bb06b4fa504da267c3b714f5 Mon Sep 17 00:00:00 2001
From: Alexandre Eichenberger <alexe at us.ibm.com>
Date: Wed, 10 Apr 2024 23:17:44 -0400
Subject: [PATCH 1/2] fix affine.prefetch replaceAffineOp invoked during
 canonicalization

Signed-off-by: Alexandre Eichenberger <alexe at us.ibm.com>
---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index c591e5056480ca..c9c0a7b4cc6860 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -1487,9 +1487,8 @@ void SimplifyAffineOp<AffinePrefetchOp>::replaceAffineOp(
     PatternRewriter &rewriter, AffinePrefetchOp prefetch, AffineMap map,
     ArrayRef<Value> mapOperands) const {
   rewriter.replaceOpWithNewOp<AffinePrefetchOp>(
-      prefetch, prefetch.getMemref(), map, mapOperands,
-      prefetch.getLocalityHint(), prefetch.getIsWrite(),
-      prefetch.getIsDataCache());
+      prefetch, prefetch.getMemref(), map, mapOperands, prefetch.getIsWrite(),
+      prefetch.getLocalityHint(), prefetch.getIsDataCache());
 }
 template <>
 void SimplifyAffineOp<AffineStoreOp>::replaceAffineOp(

>From 523a3b271992620d199d8edf4435ecb9e198ccd0 Mon Sep 17 00:00:00 2001
From: Alexandre Eichenberger <alexe at us.ibm.com>
Date: Thu, 11 Apr 2024 12:53:00 -0400
Subject: [PATCH 2/2] added test

Signed-off-by: Alexandre Eichenberger <alexe at us.ibm.com>
---
 mlir/test/Dialect/Affine/canonicalize.mlir | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 7c0930eedc8568..730808dbbdf180 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -1452,3 +1452,17 @@ func.func @mod_of_mod(%lb: index, %ub: index, %step: index) -> (index, index) {
   %1 = affine.apply affine_map<()[s0, s1, s2] -> ((s0 - ((s0 - s2) mod s1) - s2) mod s1)> ()[%ub, %step, %lb]
   return %0, %1 : index, index
 }
+
+// -----
+
+// CHECK-LABEL:  func.func @prefetch_canonicalize
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: memref<512xf32>) {
+func.func @prefetch_canonicalize(%arg0: memref<512xf32>) -> () {
+  // CHECK: affine.for [[I_0_:%.+]] = 0 to 8 {
+  affine.for %arg3 = 0 to 8  {
+    %1 = affine.apply affine_map<()[s0] -> (s0 * 64)>()[%arg3]
+    // CHECK: affine.prefetch [[PARAM_0_]][symbol([[I_0_]]) * 64], read, locality<3>, data : memref<512xf32>
+    affine.prefetch %arg0[%1], read, locality<3>, data : memref<512xf32>
+  }
+  return
+}