[Mlir-commits] [mlir] Adding memref normalization of affine.prefetch (PR #89675)

[Alexandre Eichenberger]

https://github.com/AlexandreEichenberger created https://github.com/llvm/llvm-project/pull/89675

As of now, `affine.prefetch` cannot be memref normalized; namely the memory reference being prefetch has to be a simple constant-offset based indexing.

For example:
```mlir
func.func @prefetch_normalize(%arg0: memref<512xf32, affine_map<(d0) -> (d0 floordiv 32, d0 mod 32)>>) -> () {
  affine.for %arg3 = 0 to 8  {
    affine.prefetch %arg0[%arg3], read, locality<3>, data : memref<512xf32, affine_map<(d0) -> (d0 floordiv 32, d0 mod 32)>>
  }
  return
}
```
 
cannot be normalized.

Using this PR (which simply add the `MemRefsNormalizable` property to the `affine.prefetch` operation) proper code is being generated, as shown below:

```mlir
func.func @prefetch_normalize(%arg0: memref<16x32xf32>) {
    affine.for %arg1 = 0 to 8 {
      affine.prefetch %arg0[%arg1 floordiv 32, %arg1 mod 32], read, locality<3>, data : memref<16x32xf32>
    }
    return
}
```

PR include a new lit test.

>From 2a017beaba1fd4b49bee40a2928b77eedbaffd4e Mon Sep 17 00:00:00 2001
From: Alexandre Eichenberger <alexe at us.ibm.com>
Date: Mon, 22 Apr 2024 17:36:56 -0400
Subject: [PATCH] Adding memref normalization of affine.prefetch

Signed-off-by: Alexandre Eichenberger <alexe at us.ibm.com>
---
 mlir/include/mlir/Dialect/Affine/IR/AffineOps.td |  2 +-
 .../Dialect/MemRef/normalize-memrefs-ops.mlir    | 16 ++++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index edcfcfd830c443..b9decdf1e4ff25 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -742,7 +742,7 @@ def AffineParallelOp : Affine_Op<"parallel",
 }
 
 def AffinePrefetchOp : Affine_Op<"prefetch",
-  [DeclareOpInterfaceMethods<AffineMapAccessInterface>]> {
+  [DeclareOpInterfaceMethods<AffineMapAccessInterface>, MemRefsNormalizable]> {
   let summary = "affine prefetch operation";
   let description = [{
     The `affine.prefetch` op prefetches data from a memref location described
diff --git a/mlir/test/Dialect/MemRef/normalize-memrefs-ops.mlir b/mlir/test/Dialect/MemRef/normalize-memrefs-ops.mlir
index 34420c50a51ab4..3bede131325a7f 100644
--- a/mlir/test/Dialect/MemRef/normalize-memrefs-ops.mlir
+++ b/mlir/test/Dialect/MemRef/normalize-memrefs-ops.mlir
@@ -149,3 +149,19 @@ func.func @test_norm_reinterpret_cast(%arg0 : memref<3xf32, #map_1d_tile>) -> (m
     // CHECK: memref.reinterpret_cast %[[v0]] to offset: [0], sizes: [3, 1, 1], strides: [1, 1, 1] : memref<3xf32> to memref<3x1x1xf32>
     return %1 : memref<3x1x1xf32>
 }
+
+
+// -----
+
+// Test normalization of memrefs for prefetch.affine
+
+// CHECK-LABEL: func.func @prefetch_normalize
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: memref<16x32xf32>) {
+func.func @prefetch_normalize(%arg0: memref<512xf32, affine_map<(d0) -> (d0 floordiv 32, d0 mod 32)>>) -> () {
+  // CHECK: affine.for [[I_0_:%.+]] = 0 to 8 {
+  affine.for %arg3 = 0 to 8  {
+    // CHECK: affine.prefetch [[PARAM_0_]]{{.}}[[I_0_]] floordiv 32, [[I_0_]] mod 32], read, locality<3>, data : memref<16x32xf32>
+    affine.prefetch %arg0[%arg3], read, locality<3>, data : memref<512xf32, affine_map<(d0) -> (d0 floordiv 32, d0 mod 32)>>
+  }
+  return
+}