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

[llvmlistbot at llvm.org]

Author: Alexandre Eichenberger
Date: 2024-04-30T10:54:15-04:00
New Revision: a7b968a57834a0e522505b56fab0ca4b979cb68f

URL: https://github.com/llvm/llvm-project/commit/a7b968a57834a0e522505b56fab0ca4b979cb68f
DIFF: https://github.com/llvm/llvm-project/commit/a7b968a57834a0e522505b56fab0ca4b979cb68f.diff

LOG: Adding memref normalization of affine.prefetch (#89675)

Added support for memref-normalization for prefetch.

Signed-off-by: Alexandre Eichenberger <alexe at us.ibm.com>

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
    mlir/test/Dialect/MemRef/normalize-memrefs-ops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index edcfcfd830c443..3640055ea8da8f 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -742,7 +742,8 @@ 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
+}