add noalias to all function (memref) arguments

nod-ai · Jan 13, 2025 · 7f2c00f · 7f2c00f
1 parent fa011a4
commit 7f2c00f
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 23 deletions.
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIELinalgFunctionOutlining.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIELinalgFunctionOutlining.cpp
@@ -22,7 +22,7 @@ namespace {
 static FailureOr<func::FuncOp> outline(IRRewriter &rewriter, ModuleOp moduleOp,
                                        linalg::LinalgOp computeOp,
                                        const std::string &funcName,
-                                       bool noAliasFinalArg) {
+                                       bool noAliasMemrefArgs) {
   // Form outlined FunctionType.
   for (const auto &operand : computeOp->getOperands()) {
     // Function signatures where the memrefs have layouts (strides / offsets)
@@ -64,16 +64,13 @@ static FailureOr<func::FuncOp> outline(IRRewriter &rewriter, ModuleOp moduleOp,
   // arguments.
   Operation *clonedComputeOp = rewriter.clone(*computeOp, operandMap);
 
-  if (noAliasFinalArg) {
+  if (noAliasMemrefArgs) {
+    StringRef noAliasAttrName = LLVM::LLVMDialect::getNoAliasAttrName();
     ArrayRef<BlockArgument> args = func.getArguments();
-    // Find the first MemRefType argument, starting at the end.
-    auto it = std::find_if(args.rbegin(), args.rend(), [](BlockArgument arg) {
-      return isa<MemRefType>(arg.getType());
-    });
-    if (it != args.rend()) {
-      int index = args.size() - std::distance(args.rbegin(), it) - 1;
-      StringRef noAliasAttrName = LLVM::LLVMDialect::getNoAliasAttrName();
-      func.setArgAttr(index, noAliasAttrName, rewriter.getUnitAttr());
+    for (auto iter : llvm::enumerate(args)) {
+      if (isa<MemRefType>(iter.value().getType())) {
+        func.setArgAttr(iter.index(), noAliasAttrName, rewriter.getUnitAttr());
+      }
     }
   }
 
@@ -155,7 +152,7 @@ class AMDAIELinalgFunctionOutliningPass
     }
 
     FailureOr<func::FuncOp> maybeFunc =
-        outline(rewriter, moduleOp, computeOp, funcName, noAliasFinalArg);
+        outline(rewriter, moduleOp, computeOp, funcName, noAliasMemrefArgs);
 
     if (succeeded(maybeFunc)) {
       computeOpToOutlinedFuncMap[computeOp] = maybeFunc.value();

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Passes.td b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Passes.td
@@ -291,7 +291,7 @@ def AMDAIELinalgFunctionOutlining :
       "i.e. it just returns. Useful for measuring the performance of data "
       "movement to/from the device -- by doing zero compute, all time is spent "
       "moving data to/from the AIE cores.">,
-    Option<"noAliasFinalArg", "no-alias-final-arg", "bool", /*default=*/"true",
+    Option<"noAliasMemrefArgs", "no-alias-memref-args", "bool", /*default=*/"true",
       "A developer only option. When 'true' (the default), "
       "the final memref argument of the outlined function "
       "will have the 'llvm.noalias' attribute attached to "

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/linalg_function_outlining.mlir b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/linalg_function_outlining.mlir
@@ -3,8 +3,8 @@
 // Test demonstrating multiple matmuls using different SSAs.
 
 // CHECK-LABEL: func.func private @generic_matmul_0_outlined
-// CHECK-SAME:  (%[[LHS:.*]]: memref<4x8xbf16>,
-// CHECK-SAME:   %[[RHS:.*]]: memref<8x4xbf16>,
+// CHECK-SAME:  (%[[LHS:.*]]: memref<4x8xbf16> {llvm.noalias},
+// CHECK-SAME:   %[[RHS:.*]]: memref<8x4xbf16> {llvm.noalias},
 // CHECK-SAME:   %[[OUT:.*]]: memref<4x4xf32> {llvm.noalias}) {
 // CHECK:           linalg.generic
 // CHECK-SAME:          ins(%[[LHS]], %[[RHS]] :
@@ -76,8 +76,8 @@ func.func @repeated_identical_matmul(%A: memref<4x8xbf16>, %B: memref<8x4xbf16>,
 // Test demonstrating different kind of matmul operations being mapped to a
 // unique corresponding outlined function.
 
-// CHECK-DAG:  func.func private @[[MATMUL_K6:.*]]({{.*}}memref<4x6xbf16>, {{.*}}memref<6x4xbf16>, {{.*}}memref<4x4xf32> {llvm.noalias})
-// CHECK-DAG:  func.func private @[[MATMUL_K4:.*]]({{.*}}memref<4x4xbf16>, {{.*}}memref<4x4xbf16>, {{.*}}memref<4x4xf32> {llvm.noalias})
+// CHECK-DAG:  func.func private @[[MATMUL_K6:.*]]({{.*}}memref<4x6xbf16> {llvm.noalias}, {{.*}}memref<6x4xbf16> {llvm.noalias}, {{.*}}memref<4x4xf32> {llvm.noalias})
+// CHECK-DAG:  func.func private @[[MATMUL_K4:.*]]({{.*}}memref<4x4xbf16> {llvm.noalias}, {{.*}}memref<4x4xbf16> {llvm.noalias}, {{.*}}memref<4x4xf32> {llvm.noalias})
 // CHECK-NOT:  func.func private
 // CHECK:      func.func @distinct_matmul_shapes(
 // CHECK-SAME:  %[[A0:.*]]: memref<4x4xbf16>, %[[B0:.*]]: memref<4x4xbf16>,
@@ -135,7 +135,7 @@ func.func @linalg_fill_copy(%A: memref<4xf32>, %B: memref<4xf32>) {
 // a matmul or elementwise operation. Specifically, one which has not been
 // 'blacklisted' like linalg.copy has (see test linalg_fill_copy above).
 // CHECK:       func.func private @generic_0_outlined
-// CHECK-SAME:    memref<4xbf16>,
+// CHECK-SAME:    memref<4xbf16>
 // CHECK-SAME:    memref<bf16>
 // CHECK:       linalg.generic
 // CHECK-SAME:    iterator_types = ["reduction"]

diff --git a/...lugins/target/AMD-AIE/iree-amd-aie/Transforms/test/linalg_function_outlining_options.mlir b/...lugins/target/AMD-AIE/iree-amd-aie/Transforms/test/linalg_function_outlining_options.mlir
@@ -1,25 +1,25 @@
 // Currently this is the default:
-// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(iree-amdaie-linalg-function-outlining{empty-functions=false no-alias-final-arg=true})" --verify-diagnostics --split-input-file  %s | FileCheck %s --check-prefix=CHECK_NOTEMPTY_NOALIAS
+// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(iree-amdaie-linalg-function-outlining{empty-functions=false no-alias-memref-args=true})" --verify-diagnostics --split-input-file  %s | FileCheck %s --check-prefix=CHECK_NOTEMPTY_NOALIAS
 
 // CHECK_NOTEMPTY_NOALIAS:       func.func private @
-// CHECK_NOTEMPTY_NOALIAS-SAME:    memref<4xbf16>,
+// CHECK_NOTEMPTY_NOALIAS-SAME:    memref<4xbf16> {llvm.noalias},
 // CHECK_NOTEMPTY_NOALIAS-SAME:    memref<bf16> {llvm.noalias}) {
 // CHECK_NOTEMPTY_NOALIAS:       linalg.generic
 // CHECK_NOTEMPTY_NOALIAS:       return
 // CHECK_NOTEMPTY_NOALIAS:       func.func @reduction
 
 // A run to check the option empty-functions=true:
-// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(iree-amdaie-linalg-function-outlining{empty-functions=true no-alias-final-arg=true})"  --verify-diagnostics --split-input-file  %s | FileCheck %s --check-prefix=CHECK_EMPTY_NOALIAS
+// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(iree-amdaie-linalg-function-outlining{empty-functions=true no-alias-memref-args=true})"  --verify-diagnostics --split-input-file  %s | FileCheck %s --check-prefix=CHECK_EMPTY_NOALIAS
 
 // CHECK_EMPTY_NOALIAS:       func.func private @
-// CHECK_EMPTY_NOALIAS-SAME:    memref<4xbf16>,
+// CHECK_EMPTY_NOALIAS-SAME:    memref<4xbf16> {llvm.noalias},
 // CHECK_EMPTY_NOALIAS-SAME:    memref<bf16> {llvm.noalias}) {
 // CHECK_EMPTY_NOALIAS-NOT:       linalg.generic
 // CHECK_EMPTY_NOALIAS:       return
 // CHECK_EMPTY_NOALIAS:       func.func @reduction
 
-// A run to check the option no-alias-final-arg=false:
-// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(iree-amdaie-linalg-function-outlining{empty-functions=false no-alias-final-arg=false})" --verify-diagnostics --split-input-file  %s | FileCheck %s --check-prefix=CHECK_NOTEMPTY_ALIAS
+// A run to check the option no-alias-memref-args=false:
+// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(iree-amdaie-linalg-function-outlining{empty-functions=false no-alias-memref-args=false})" --verify-diagnostics --split-input-file  %s | FileCheck %s --check-prefix=CHECK_NOTEMPTY_ALIAS
 
 // CHECK_NOTEMPTY_ALIAS:       func.func private @
 // CHECK_NOTEMPTY_ALIAS-SAME:    memref<4xbf16>,