some progress with peano, no progress with chess

Author: hunhoffe

lib/Dialect/AIE/Transforms/AIEExternalMangle.cpp

@@ -27,7 +27,6 @@ struct AIEExternalManglePass
     : public AIEExternalMangleBase<AIEExternalManglePass> {
   void runOnOperation() override {
     ModuleOp module = getOperation();
-    SymbolTable symbolTable(module);
 
     // Helper to mangle a function based on an object file
     auto mangleFunction = [&](func::FuncOp func,
@@ -82,12 +81,14 @@ struct AIEExternalManglePass
             return func;
           }
 
+          SymbolTable parentSymbolTable(func->getParentOp());
+
           // Found a match! Check if we need to rename/clone.
           std::string newName = mangledName.str();
           int suffix = 0;
-          while (symbolTable.lookup(newName)) {
+          while (parentSymbolTable.lookup(newName)) {
             // If the existing symbol is the function itself, we are good.
-            auto existingOp = symbolTable.lookup(newName);
+            auto existingOp = parentSymbolTable.lookup(newName);
             if (existingOp == func)
               break;
 
@@ -106,12 +107,12 @@ struct AIEExternalManglePass
 
           if (newName != functionName) {
             // If 'func' already has the correct name and link_with, return it.
-            if (auto existingOp = symbolTable.lookup(newName)) {
+            if (auto existingOp = parentSymbolTable.lookup(newName)) {
               return cast<func::FuncOp>(existingOp);
             }
 
             // Create a new function declaration
-            OpBuilder builder(module.getBodyRegion());
+            OpBuilder builder(func);
             auto newFunc = func::FuncOp::create(builder, func.getLoc(), newName,
                                                 func.getFunctionType());
             newFunc.setPrivate();
@@ -128,7 +129,9 @@ struct AIEExternalManglePass
                                StringAttr::get(func.getContext(), mangledName));
             }
 
-            symbolTable.insert(newFunc);
+            // Insert the new function in the same symbol table as the original
+            // function.
+            parentSymbolTable.insert(newFunc, func->getIterator());
             return newFunc;
           }
           return func;
@@ -146,7 +149,8 @@ struct AIEExternalManglePass
         StringRef objectFileName = linkWithAttr.getValue();
 
         core.walk([&](func::CallOp call) {
-          auto callee = symbolTable.lookup<func::FuncOp>(call.getCallee());
+          auto callee = SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(
+              call, call.getCalleeAttr());
           if (callee && callee.isExternal()) {
             // Mangle/Clone the callee for this core
             auto newCallee = mangleFunction(callee, objectFileName);
@@ -166,10 +170,11 @@ struct AIEExternalManglePass
         if (newFunc != func) {
           // If newFunc is a different op, replace uses and erase original.
           if (failed(SymbolTable::replaceAllSymbolUses(
-                  func, newFunc.getNameAttr(), module))) {
+                  func, newFunc.getNameAttr(), func->getParentOp()))) {
             func.emitError("failed to replace symbol uses");
             return;
           }
+
           func.erase();
         }
       }

lib/Targets/AIETargetLdScript.cpp

@@ -11,6 +11,11 @@
 #include "aie/Dialect/AIE/IR/AIEDialect.h"
 #include "aie/Targets/AIETargets.h"
 
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/SymbolTable.h"
+
+#include <set>
+
 using namespace mlir;
 using namespace xilinx;
 using namespace xilinx::AIE;
@@ -177,8 +182,23 @@ SECTIONS
       output << "  .bss : { *(.bss*) } > data\n";
       output << "}\n";
       if (auto coreOp = tile.getCoreOp()) {
+        std::set<std::string> linkedFiles;
         if (auto fileAttr = coreOp.getLinkWith())
-          output << "INPUT(" << fileAttr.value().str() << ")\n";
+          linkedFiles.insert(fileAttr.value().str());
+
+        coreOp.walk([&](func::CallOp call) {
+          auto callee = call.getCallee();
+          if (auto func = SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(
+                  call, StringAttr::get(call.getContext(), callee))) {
+            if (auto linkWith = func.getOperation()->getAttrOfType<StringAttr>(
+                    "link_with")) {
+              linkedFiles.insert(linkWith.getValue().str());
+            }
+          }
+        });
+
+        for (const auto &file : linkedFiles)
+          output << "INPUT(" << file << ")\n";
 
         output << "PROVIDE(main = core_" << tile.getCol() << "_"
                << tile.getRow() << ");\n";

python/compiler/aiecc/main.py

@@ -699,6 +699,18 @@ def __init__(self, mlir_module_str, opts, tmpdirname):
         self.peano_clang_path = os.path.join(opts.peano_install_dir, "bin", "clang")
         self.peano_opt_path = os.path.join(opts.peano_install_dir, "bin", "opt")
         self.peano_llc_path = os.path.join(opts.peano_install_dir, "bin", "llc")
+        self.peano_objcopy_path = os.path.join(
+            opts.peano_install_dir, "bin", "llvm-objcopy"
+        )
+        if not os.path.exists(self.peano_objcopy_path):
+            if shutil.which("llvm-objcopy"):
+                self.peano_objcopy_path = "llvm-objcopy"
+            elif shutil.which("llvm-objcopy-18"):
+                self.peano_objcopy_path = "llvm-objcopy-18"
+            elif shutil.which("llvm-objcopy-20"):
+                self.peano_objcopy_path = "llvm-objcopy-20"
+            else:
+                self.peano_objcopy_path = "objcopy"
         self.repeater_output_dir = opts.repeater_output_dir or tempfile.gettempdir()
 
     def prepend_tmp(self, x):
@@ -792,7 +804,7 @@ async def do_call(self, task_id, command, force=False):
             ret = proc.returncode
             if self.opts.verbose and stdout:
                 print(f"{stdout.decode()}")
-            if ret != 0 and stderr:
+            if (self.opts.verbose or ret != 0) and stderr:
                 print(f"{stderr.decode()}", file=sys.stderr)
         else:
             ret = 0
@@ -828,6 +840,7 @@ async def chesshack(self, task, llvmir, aie_target):
 
         llvmir_ir = await read_file_async(llvmir)
         llvmir_hacked_ir = downgrade_ir_for_chess(llvmir_ir)
+
         await write_file_async(llvmir_hacked_ir, llvmir_chesshack)
 
         if aie_target.casefold() == "AIE2".casefold():
@@ -954,46 +967,47 @@ async def process_cores(
         return elf_paths
 
     async def handle_mangled_collisions(self, module):
-        # Find all func.func ops with link_name attribute
-        funcs = find_ops(
-            module.operation,
-            lambda o: isinstance(o.operation.opview, funcdialect.FuncOp)
-            and "link_name" in o.attributes,
-        )
+        with module.context:
+            # Find all func.func ops with link_name attribute
+            funcs = find_ops(
+                module.operation,
+                lambda o: isinstance(o.operation.opview, funcdialect.FuncOp)
+                and "link_name" in o.attributes,
+            )
 
-        for func in funcs:
-            link_name = func.attributes["link_name"].value
-            if "link_with" not in func.attributes:
-                print(
-                    f"Warning: Function '{func.sym_name.value}' has 'link_name' but no 'link_with' attribute. Skipping symbol renaming."
-                )
-                continue
-            link_with = func.attributes["link_with"].value
-            sym_name = func.sym_name.value
-
-            # If sym_name is different from link_name, we need to rename the symbol in the object file
-            if sym_name != link_name:
-                # Create a new object file name
-                new_obj_file = self.prepend_tmp(f"{sym_name}.o")
-
-                # Run llvm-objcopy to rename the symbol
-                cmd = [
-                    "llvm-objcopy",
-                    "--redefine-sym",
-                    f"{link_name}={sym_name}",
-                    link_with,
-                    new_obj_file,
-                ]
-
-                if self.opts.verbose:
+            for func in funcs:
+                link_name = func.attributes["link_name"].value
+                if "link_with" not in func.attributes:
                     print(
-                        f"Renaming symbol {link_name} to {sym_name} in {link_with} -> {new_obj_file}"
+                        f"Warning: Function '{func.sym_name.value}' has 'link_name' but no 'link_with' attribute. Skipping symbol renaming."
                     )
+                    continue
+                link_with = func.attributes["link_with"].value
+                sym_name = func.sym_name.value
+
+                # If sym_name is different from link_name, we need to rename the symbol in the object file
+                if sym_name != link_name:
+                    # Create a new object file name
+                    new_obj_file = self.prepend_tmp(f"{sym_name}.o")
+
+                    # Run llvm-objcopy to rename the symbol
+                    cmd = [
+                        self.peano_objcopy_path,
+                        "--redefine-sym",
+                        f"{link_name}={sym_name}",
+                        link_with,
+                        new_obj_file,
+                    ]
+
+                    if self.opts.verbose:
+                        print(
+                            f"Renaming symbol {link_name} to {sym_name} in {link_with} -> {new_obj_file}"
+                        )
 
-                await self.do_call(None, cmd)
+                    await self.do_call(None, cmd)
 
-                # Update link_with attribute
-                func.attributes["link_with"] = StringAttr.get(new_obj_file)
+                    # Update link_with attribute
+                    func.attributes["link_with"] = StringAttr.get(new_obj_file)
 
     async def process_core(
         self,
@@ -1053,6 +1067,11 @@ async def process_core(
             if opts.compile and opts.xchesscc:
                 if not opts.unified:
                     file_core_llvmir_chesslinked = await self.chesshack(task, file_core_llvmir, aie_target)
+                    
+                    # Hack: rename symbol in LLVM IR
+                    cmd = ["sed", "-i", "s/_Z6kernelPii_1/_Z8kernel_1Pii/g", file_core_llvmir_chesslinked]
+                    await self.do_call(task, cmd)
+
                     if self.opts.link and self.opts.xbridge:
                         link_with_obj = await extract_input_files(file_core_bcf)
                         await self.do_call(task, ["xchesscc_wrapper", aie_target.lower(), "+w", self.prepend_tmp("work"), "-d", "+Wclang,-xir", "-f", file_core_llvmir_chesslinked, link_with_obj, "+l", file_core_bcf, "-o", file_core_elf])

test/npu-xrt/external_mangle_collision/aie.mlir

@@ -1,8 +1,8 @@
 module {
-  func.func private @kernel_add(memref<8xi32>, i32) attributes { link_with="kernel1.o", link_symbol="kernel(int*, int)" }
-  func.func private @kernel_sub(memref<8xi32>, i32) attributes { link_with="kernel2.o", link_symbol="kernel(int*, int)" }
-
   aie.device(npu1_1col) {
+    func.func private @kernel_add(memref<8xi32>, i32) attributes { link_with="kernel1.o", link_symbol="kernel(int*, int)" }
+    func.func private @kernel_sub(memref<8xi32>, i32) attributes { link_with="kernel2.o", link_symbol="kernel(int*, int)" }
+
     %t00 = aie.tile(0, 0)
     %t01 = aie.tile(0, 1)
     %t02 = aie.tile(0, 2)