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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-affine

Author: Alexandre Eichenberger (AlexandreEichenberger)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/89675.diff


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/Affine/IR/AffineOps.td (+1-1) 
- (modified) mlir/test/Dialect/MemRef/normalize-memrefs-ops.mlir (+16) 


``````````diff
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
+}

``````````

</details>


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