I am not sure how this works. Assuming that I have a function like this:
func test_mul(%a:memref<10x10xf32>, %b:memref<10x10xf32>){
... ... ...
return
}
How to represent the cases that %a and %b might point to the same location on caller side; or they are never aliased from the caller.
This can be done via block argument attributes. The LLVM dialect for example has a noalias. I’m not aware of a well-defined encoding in mid-level dialects but attributes are a way to easily encode this information.
Can you please elaborate on how this can be done with block arguments? I’m new to MLIR and not sure where to look. Concretely, I’m trying to optimize the following function:
func.func @stencil1D(%input: memref<10xf32>{llvm.noalias}, %output: memref<10xf32>{llvm.noalias}) {
// Center of the stencil.
affine.for %i = 1 to 9 {
%0 = affine.load %input[%i] : memref<10xf32>
affine.store %0, %output[%i] : memref<10xf32>
}
// Left.
affine.for %i = 1 to 9 {
%0 = affine.load %input[%i - 1] : memref<10xf32>
%1 = affine.load %output[%i] : memref<10xf32>
%2 = arith.addf %0, %1 : f32
affine.store %2, %output[%i] : memref<10xf32>
}
// Right.
affine.for %i = 1 to 9 {
%0 = affine.load %input[%i + 1] : memref<10xf32>
%1 = affine.load %output[%i] : memref<10xf32>
%2 = arith.addf %0, %1 : f32
affine.store %2, %output[%i] : memref<10xf32>
}
return
}
I want to fuse these three loops and then remove redundant stores to turn this into a stencil with 3 loads and 1 store. With the affine loop fusion pass and scalrep pass, I can get something like this:
module {
func.func @stencil1D(%arg0: memref<10xf32> {llvm.noalias}, %arg1: memref<10xf32> {llvm.noalias}) {
affine.for %arg2 = 1 to 9 {
%0 = affine.load %arg0[%arg2] : memref<10xf32>
affine.store %0, %arg1[%arg2] : memref<10xf32>
%1 = affine.load %arg0[%arg2 - 1] : memref<10xf32>
%2 = arith.addf %1, %0 : f32
affine.store %2, %arg1[%arg2] : memref<10xf32>
%3 = affine.load %arg0[%arg2 + 1] : memref<10xf32>
%4 = arith.addf %3, %2 : f32
affine.store %4, %arg1[%arg2] : memref<10xf32>
}
return
}
}
However, the affine pass can’t remove the redundant stores because I’m guessing that the input and output memrefs may alias. If I rewrite the MLIR code to
func.func @stencil1D() -> memref<10xf32> {
%input = memref.alloc() : memref<10xf32>
%output = memref.alloc() : memref<10xf32>
// Center of the stencil.
affine.for %i = 1 to 9 {
%0 = affine.load %input[%i] : memref<10xf32>
affine.store %1, %output[%i] : memref<10xf32>
}
// Left.
affine.for %i = 1 to 9 {
%0 = affine.load %input[%i - 1] : memref<10xf32>
%1 = affine.load %output[%i] : memref<10xf32>
%2 = arith.addf %0, %1 : f32
affine.store %2, %output[%i] : memref<10xf32>
}
// Right.
affine.for %i = 1 to 9 {
%0 = affine.load %input[%i + 1] : memref<10xf32>
%1 = affine.load %output[%i] : memref<10xf32>
%2 = arith.addf %0, %1 : f32
affine.store %2, %output[%i] : memref<10xf32>
}
return %output : memref<10xf32>
}
then things get rewritten to what I would expect:
module {
func.func @stencil1D() -> memref<10xf32> {
%alloc = memref.alloc() : memref<10xf32>
%alloc_0 = memref.alloc() : memref<10xf32>
affine.for %arg0 = 1 to 9 {
%0 = affine.load %alloc[%arg0] : memref<10xf32>
%1 = affine.load %alloc[%arg0 - 1] : memref<10xf32>
%2 = arith.addf %1, %0 : f32
%3 = affine.load %alloc[%arg0 + 1] : memref<10xf32>
%4 = arith.addf %3, %2 : f32
affine.store %4, %alloc_0[%arg0] : memref<10xf32>
}
return %alloc_0 : memref<10xf32>
}
}
So I would like to be able to tell the affine passes that this aliasing isn’t present – adding {llvm.noalias} to the memref type doesn’t make that reasoning go through.
There is broader interest in the community to add attributes to funcOps, e.g. ⚙ D140357 [MLIR] Add readonly,writeonly,noalias to func arg attributes
Given that this hasn’t been solved yet, is there a way to emulate the behavior I want? Eventually the code I’m looking to generate will need to consume memory created outside the MLIR context, so I’ll need a way to thread these assumptions through somehow.
In your closed world compiler, you can add any attribute to the funcop. As long as your compiler understands them. Long term, readonly, writeonly, restrict, and more would be desirable.
I understand that I could define a dialect or something that has attributes, but there’s still the problem of eventually letting the affine passes understand the non-aliasing properties. I don’t want to re-invent the world by rewriting those passes to handle these new attributes. I’m new to MLIR, so I’m not sure what is the way to think about these kind of things within the framework.
Bumping this, is there a way forward for me to represent aliasing / lack of aliasing in memrefs without re-writing all of the affine loop passes that make assumptions about memref aliasing?