Fixup device symbol ref with runtime sequence materialization (#2849)

Co-authored-by: Joseph Melber <jgmelber@gmail.com>

Author: erwei-xilinx

include/aie/Dialect/AIEX/IR/AIEX.td

@@ -1014,6 +1014,7 @@ def AIE_NpuLoadPdiOp: AIEX_Op<"npu.load_pdi", []> {
 
     If a symbol reference is provided, the compiler driver (aiecc.py) will match it to a device symbol name and assign the PDI ID field based on it.
   }];
+  let hasCanonicalizeMethod = 1;
 }
 
 def AIE_DMAConfigureTaskOp : AIEX_Op<"dma_configure_task", [HasParent<"AIE::RuntimeSequenceOp">, TileElement]>, Results<(outs Index:$result)> {

lib/Dialect/AIEX/IR/AIEXDialect.cpp

@@ -1010,3 +1010,32 @@ LogicalResult AIEX::RunOp::verify() {
 
   return success();
 }
+
+//===----------------------------------------------------------------------===//
+// NpuLoadPdiOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult AIEX::NpuLoadPdiOp::canonicalize(AIEX::NpuLoadPdiOp op,
+                                               PatternRewriter &rewriter) {
+  // Check for back-to-back identical load_pdi ops and remove duplicates
+  Operation *nextOp = op->getNextNode();
+  if (!nextOp)
+    return failure();
+
+  // Check if next op is also a NpuLoadPdiOp
+  auto nextLoadPdi = dyn_cast<AIEX::NpuLoadPdiOp>(nextOp);
+  if (!nextLoadPdi)
+    return failure();
+
+  // Check if they are identical (all attributes match)
+  if (op.getDeviceRefAttr() == nextLoadPdi.getDeviceRefAttr() &&
+      op.getId() == nextLoadPdi.getId() &&
+      op.getSize() == nextLoadPdi.getSize() &&
+      op.getAddress() == nextLoadPdi.getAddress()) {
+    // Erase the first one, keeping the second
+    rewriter.eraseOp(op);
+    return success();
+  }
+
+  return failure();
+}

lib/Dialect/AIEX/Transforms/AIEMaterializeRuntimeSequences.cpp

@@ -129,6 +129,7 @@ struct InsertLoadPdiForConfigurePattern : RewritePattern {
     } else {
       configureBlock = &configureOp.getBody().front();
     }
+
     rewriter.setInsertionPointToStart(configureBlock);
     AIEX::NpuLoadPdiOp::create(
         rewriter, configureOp.getLoc(),
@@ -344,12 +345,26 @@ static LogicalResult inlineReferencedSymbolDefinitions(
         previouslyInlinedSymbolMap[oldSymbolRef] = newSymbolRef;
 
         // Add the new symbol definition
+        // First try to look up from the lookupFrom operation (e.g., within the
+        // callee device). If not found, try looking up from the module level
+        // (for cross-device references).
         Operation *symbolDefOp =
             SymbolTable::lookupNearestSymbolFrom(lookupFrom, oldSymbolRef);
+        if (!symbolDefOp) {
+          if (ModuleOp moduleOp = lookupFrom->getParentOfType<ModuleOp>()) {
+            symbolDefOp = SymbolTable::lookupSymbolIn(moduleOp, oldSymbolRef);
+          }
+        }
         if (!symbolDefOp) {
           return std::make_pair(newSymbolRef, WalkResult::interrupt());
         }
 
+        // If the symbol is a device, don't clone it - keep the original
+        // reference. Device ops must stay at module level.
+        if (llvm::isa<AIE::DeviceOp>(symbolDefOp)) {
+          return std::make_pair(oldSymbolRef, WalkResult::advance());
+        }
+
         // Collect SSA values referenced by the symbol definition operation
         llvm::SetVector<Value> symbolReferencedValues;
         collectReferencedSSAValues(symbolDefOp, argMap, symbolReferencedValues);
@@ -561,6 +576,15 @@ struct AIEMaterializeRuntimeSequencesPass
       patterns_1.insert<InsertLoadPdiForConfigurePattern>(ctx);
       walkAndApplyPatterns(deviceOp, std::move(patterns_1));
 
+      // Canonicalize to remove duplicate back-to-back load_pdi ops
+      RewritePatternSet canonicalize_patterns(ctx);
+      AIEX::NpuLoadPdiOp::getCanonicalizationPatterns(canonicalize_patterns,
+                                                      ctx);
+      if (failed(applyPatternsGreedily(
+              deviceOp, std::move(canonicalize_patterns), rewriter_config))) {
+        return signalPassFailure();
+      }
+
       // Flatten the IR: hoist all operations inside aiex.configure to be direct
       // children of the runtime sequence, preserving order
       for (AIE::RuntimeSequenceOp runtimeSequenceOp :

test/Passes/materialize-runtime-sequences/cross_device_inline_load_pdi.mlir

@@ -0,0 +1,120 @@
+//===- cross_device_inline_load_pdi.mlir -----------------------*- MLIR -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// Copyright (C) 2026, Advanced Micro Devices, Inc.
+//
+//===----------------------------------------------------------------------===//
+
+// RUN: aie-opt --aie-materialize-runtime-sequences --split-input-file %s | FileCheck %s
+
+// Test cross-device inlining of runtime sequences:
+// 1. Inlined load_pdi operations are preserved with original device_ref
+// 2. InsertLoadPdiForConfigurePattern adds load_pdi at start only if needed
+// 3. When callee starts with load_pdi, no duplicate is added
+
+//===----------------------------------------------------------------------===//
+// TEST 1: Callee does NOT start with load_pdi - one should be added at start
+//===----------------------------------------------------------------------===//
+
+module {
+  // The outer/caller device - anonymous (no symbol name)
+  // CHECK: aie.device(npu2) {
+  aie.device(npu2) {
+    %tile00 = aie.tile(0, 0)
+    
+    // CHECK: aie.runtime_sequence @caller_seq
+    aie.runtime_sequence @caller_seq(%arg0: memref<16xi32>) {
+      // After inlining, we should have:
+      // 1. A load_pdi added by InsertLoadPdiForConfigurePattern (since the first inlined op is write32, not load_pdi)
+      // 2. All operations from the callee sequence inlined
+      // 3. The inlined load_pdi operations preserved with their original device_ref
+      
+      // CHECK: aiex.npu.load_pdi {device_ref = @callee_device}
+      // CHECK-NEXT: aiex.npu.write32 {address = 100 : ui32
+      // CHECK: aiex.npu.load_pdi {device_ref = @callee_device}
+      // CHECK-NEXT: aiex.npu.write32 {address = 200 : ui32
+      // CHECK: aiex.npu.load_pdi {device_ref = @callee_device}
+      // CHECK-NEXT: aiex.npu.write32 {address = 300 : ui32
+      aiex.configure @callee_device {
+        aiex.run @callee_seq(%arg0) : (memref<16xi32>)
+      }
+    }
+  }
+  
+  // The callee device that contains the runtime sequence to be inlined
+  // CHECK: aie.device(npu2) @callee_device
+  aie.device(npu2) @callee_device {
+    %tile10 = aie.tile(1, 0)
+    
+    // This sequence has load_pdi operations embedded in it for reconfiguration
+    // between iterations. When inlined, these should be preserved.
+    aie.runtime_sequence @callee_seq(%arg0: memref<16xi32>) {
+      // First iteration
+      aiex.npu.write32 {address = 100 : ui32, column = 0 : i32, row = 0 : i32, value = 1 : ui32}
+      
+      // Reconfigure for second iteration
+      aiex.npu.load_pdi {device_ref = @callee_device}
+      aiex.npu.write32 {address = 200 : ui32, column = 0 : i32, row = 0 : i32, value = 2 : ui32}
+      
+      // Reconfigure for third iteration
+      aiex.npu.load_pdi {device_ref = @callee_device}
+      aiex.npu.write32 {address = 300 : ui32, column = 0 : i32, row = 0 : i32, value = 3 : ui32}
+    }
+  }
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// TEST 2: Callee STARTS with load_pdi - no duplicate should be added
+//===----------------------------------------------------------------------===//
+
+module {
+  // CHECK: aie.device(npu2) {
+  aie.device(npu2) {
+    %tile00 = aie.tile(0, 0)
+    
+    // CHECK: aie.runtime_sequence @caller_seq2
+    aie.runtime_sequence @caller_seq2(%arg0: memref<16xi32>) {
+      // After inlining, the callee's load_pdi is at the start of the configure block.
+      // InsertLoadPdiForConfigurePattern should detect this and NOT add another one.
+      // We should see exactly 3 load_pdi operations (from the callee), not 4.
+      
+      // CHECK: aiex.npu.load_pdi {device_ref = @callee_device2}
+      // CHECK-NEXT: aiex.npu.write32 {address = 100 : ui32
+      // CHECK: aiex.npu.load_pdi {device_ref = @callee_device2}
+      // CHECK-NEXT: aiex.npu.write32 {address = 200 : ui32
+      // CHECK: aiex.npu.load_pdi {device_ref = @callee_device2}
+      // CHECK-NEXT: aiex.npu.write32 {address = 300 : ui32
+      aiex.configure @callee_device2 {
+        aiex.run @callee_seq2(%arg0) : (memref<16xi32>)
+      }
+    }
+  }
+  
+  // The callee device - its runtime sequence starts with load_pdi
+  // CHECK: aie.device(npu2) @callee_device2
+  aie.device(npu2) @callee_device2 {
+    %tile10 = aie.tile(1, 0)
+    
+    // This sequence STARTS with a load_pdi operation.
+    // When inlined into a configure block, InsertLoadPdiForConfigurePattern
+    // should NOT add another load_pdi at the start.
+    aie.runtime_sequence @callee_seq2(%arg0: memref<16xi32>) {
+      // First load_pdi at the very start
+      aiex.npu.load_pdi {device_ref = @callee_device2}
+      aiex.npu.write32 {address = 100 : ui32, column = 0 : i32, row = 0 : i32, value = 1 : ui32}
+      
+      // Second iteration
+      aiex.npu.load_pdi {device_ref = @callee_device2}
+      aiex.npu.write32 {address = 200 : ui32, column = 0 : i32, row = 0 : i32, value = 2 : ui32}
+      
+      // Third iteration
+      aiex.npu.load_pdi {device_ref = @callee_device2}
+      aiex.npu.write32 {address = 300 : ui32, column = 0 : i32, row = 0 : i32, value = 3 : ui32}
+    }
+  }
+}