From f4fe32b3e8e116b748542b57b76f55f511e09055 Mon Sep 17 00:00:00 2001
From: erwei-xilinx <erwei.wang@amd.com>
Date: Mon, 24 Feb 2025 15:57:41 -0800
Subject: [PATCH] Modify the --lower-linalg-to-func option in aircc to take in
 the obj file name as str

---
 python/air/backend/xrt.py                 | 3 ++-
 python/air/compiler/aircc/cl_arguments.py | 6 +++---
 python/air/compiler/aircc/main.py         | 6 +++---
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/python/air/backend/xrt.py b/python/air/backend/xrt.py
index 5d2cd0e4d..57b8986ea 100644
--- a/python/air/backend/xrt.py
+++ b/python/air/backend/xrt.py
@@ -48,7 +48,7 @@ def __init__(
         verbose: bool = False,
         omit_while_true_loop: bool = False,
         omit_pingpong: bool = False,
-        lower_linalg_to_func: bool = False,
+        lower_linalg_to_func: str = None,
         air_loop_fusion: bool = False,
         runtime_loop_tiling_sizes: list[int] = [4, 4],
         omit_auto_broadcast: bool = False,
@@ -137,6 +137,7 @@ def compile(
 
             if self.lower_linalg_to_func:
                 aircc_options += ["--lower-linalg-to-func"]
+                aircc_options += [self.lower_linalg_to_func]
 
             if self.air_loop_fusion:
                 aircc_options += ["--air-loop-fusion"]
diff --git a/python/air/compiler/aircc/cl_arguments.py b/python/air/compiler/aircc/cl_arguments.py
index 55013dbe0..f5bf422ef 100644
--- a/python/air/compiler/aircc/cl_arguments.py
+++ b/python/air/compiler/aircc/cl_arguments.py
@@ -125,10 +125,10 @@ def parse_args(args=None):
     )
     parser.add_argument(
         "--lower-linalg-to-func",
+        type=str,
         dest="lower_linalg_to_func",
-        default=False,
-        action="store_true",
-        help="Whether to run pass which lowers linalg.generic ops to function calls. If False, then they lower to loops.",
+        default=None,
+        help="Whether to run pass which lowers linalg.generic ops to function calls. If a string is passed in, then register the string value as the object file name to link with.",
     )
     parser.add_argument(
         "--air-loop-fusion",
diff --git a/python/air/compiler/aircc/main.py b/python/air/compiler/aircc/main.py
index 2e4feeb9b..fd1bc043f 100644
--- a/python/air/compiler/aircc/main.py
+++ b/python/air/compiler/aircc/main.py
@@ -40,7 +40,7 @@ def get_L2_splitting_analysis_pass():
 def get_air_optimization_pass(
     device,
     omit_pingpong=False,
-    lower_linalg_to_func=False,
+    lower_linalg_to_func=None,
     air_loop_fusion=False,
     omit_auto_broadcast=False,
     channel_multiplexing=[],
@@ -96,9 +96,9 @@ def get_air_optimization_pass(
             "canonicalize",
             "cse",
         ]
-    if lower_linalg_to_func:
+    if lower_linalg_to_func != None:
         OPTIMIZATION_PASSES += [
-            "air-linalg-to-func",
+            "air-linalg-to-func{link-with=" + f"{lower_linalg_to_func}" + "}",
         ]
     else:
         OPTIMIZATION_PASSES += [