Adding memref normalization of affine.prefetch #89675
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: Alexandre Eichenberger <[email protected]>
|
@llvm/pr-subscribers-mlir-memref @llvm/pr-subscribers-mlir-affine Author: Alexandre Eichenberger (AlexandreEichenberger) ChangesAs of now, For example: 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 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:
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
+}
|
|
@llvm/pr-subscribers-mlir Author: Alexandre Eichenberger (AlexandreEichenberger) ChangesAs of now, For example: 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 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:
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
+}
|
Signed-off-by: Alexandre Eichenberger <[email protected]>
|
Not an expert, the log indicates so I assume that restarting the test should handle this. |
|
@bondhugula any objections? Since load/store ops are already normalizable, feels like having prefetch too is not a big deal? We encounter memory layouts with floordiv/ceildiv/mods (for tiling). |
|
@bondhugula @ftynse no review for about a week, do you think someone could look at it? |
ftynse
approved these changes
Apr 29, 2024
|
@ftynse Used the feedback you gave me last time for another PR for a smooth sailing this time. Thanks for the review. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
As of now,
affine.prefetchcannot be memref normalized; namely the memory reference being prefetch has to be a simple constant-offset based indexing.For example:
cannot be normalized.
Using this PR (which simply add the
MemRefsNormalizableproperty to theaffine.prefetchoperation) proper code is being generated, as shown below:PR include a new lit test.