Allow multiple devices in one design and add device configuration op (aiex.configure) (#2532)

Author: andrej

docs/buildHostLin.md

@@ -403,7 +403,7 @@ For your design of interest, for instance from [programming_examples](../program
 1. Run (program arguments are just an example for add_one design)
     ```bash
     cd Release
-    .\<testName>.exe -x ..\..\build\final.xclbin -k MLIR_AIE -i ..\..\build\insts.txt -v 1
+    .\<testName>.exe -x ..\..\build\main.xclbin -k MLIR_AIE -i ..\..\build\main_sequence.bin -v 1
     ```
 
 # Troubleshooting

include/aie-c/Translation.h

@@ -21,20 +21,23 @@ extern "C" {
 MLIR_CAPI_EXPORTED MlirStringRef aieTranslateAIEVecToCpp(MlirOperation op,
                                                          bool aie2);
 MLIR_CAPI_EXPORTED MlirStringRef aieTranslateModuleToLLVMIR(MlirOperation op);
-MLIR_CAPI_EXPORTED MlirStringRef aieTranslateNpuToBinary(MlirOperation op,
-                                                         MlirStringRef name);
+MLIR_CAPI_EXPORTED MlirStringRef aieTranslateNpuToBinary(
+    MlirOperation op, MlirStringRef deviceName, MlirStringRef sequenceName);
 MLIR_CAPI_EXPORTED MlirStringRef
-aieTranslateControlPacketsToUI32Vec(MlirOperation op);
-MLIR_CAPI_EXPORTED MlirStringRef aieTranslateToXAIEV2(MlirOperation op);
-MLIR_CAPI_EXPORTED MlirStringRef aieTranslateToHSA(MlirOperation op);
+aieTranslateControlPacketsToUI32Vec(MlirOperation op, MlirStringRef deviceName);
+MLIR_CAPI_EXPORTED MlirStringRef aieTranslateToXAIEV2(MlirOperation op,
+                                                      MlirStringRef deviceName);
+MLIR_CAPI_EXPORTED MlirStringRef aieTranslateToHSA(MlirOperation op,
+                                                   MlirStringRef deviceName);
 MLIR_CAPI_EXPORTED MlirStringRef aieTranslateToBCF(MlirOperation op, int col,
-                                                   int row);
+                                                   int row,
+                                                   MlirStringRef deviceName);
 MLIR_CAPI_EXPORTED MlirStringRef aieLLVMLink(MlirStringRef *modules,
                                              int nModules);
-MLIR_CAPI_EXPORTED MlirLogicalResult
-aieTranslateToCDODirect(MlirOperation moduleOp, MlirStringRef workDirPath,
-                        bool bigEndian, bool emitUnified, bool cdoDebug,
-                        bool aieSim, bool xaieDebug, bool enableCores);
+MLIR_CAPI_EXPORTED MlirLogicalResult aieTranslateToCDODirect(
+    MlirOperation moduleOp, MlirStringRef workDirPath, MlirStringRef deviceName,
+    bool bigEndian, bool emitUnified, bool cdoDebug, bool aieSim,
+    bool xaieDebug, bool enableCores);
 MLIR_CAPI_EXPORTED MlirOperation aieTranslateBinaryToTxn(MlirContext ctx,
                                                          MlirStringRef binary);
 

include/aie/Conversion/Passes.td

@@ -72,6 +72,8 @@ def ConvertAIEToTransaction : Pass<"convert-aie-to-transaction",
   let options = [
       Option<"clElfDir", "elf-dir", "std::string", /*default=*/"",
              "Where to find ELF files">,
+      Option<"clDeviceName", "device-name", "std::string", /*default=*/"",
+             "Which device to generate code for; if none specified, process all devices">,
   ];
 }
 
@@ -94,6 +96,8 @@ def ConvertAIEToControlPackets : Pass<"convert-aie-to-control-packets",
   let options = [
       Option<"clElfDir", "elf-dir", "std::string", /*default=*/"",
              "Where to find ELF files">,
+      Option<"clDeviceName", "device-name", "std::string", /*default=*/"",
+             "Which device to generate code for; if none specified, process all devices">,
   ];
 }
 

include/aie/Dialect/AIE/IR/AIEDialect.h

@@ -22,6 +22,8 @@
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 
+#include "llvm/ADT/StringRef.h"
+
 namespace xilinx::AIE {
 
 // Check that the given DMA-like op (e.g. MemOp, ShimDMAOp)

include/aie/Dialect/AIE/IR/AIEOps.td

@@ -27,8 +27,12 @@ class AIE_Op<string mnemonic, list<Trait> traits = []> :
 
 
 def AIE_DeviceOp: AIE_Op<"device", [
-    AIETarget, HasParent<"mlir::ModuleOp">,
-    SymbolTable, SingleBlockImplicitTerminator<"EndOp">, IsolatedFromAbove
+    Symbol,
+    AIETarget, 
+    HasParent<"mlir::ModuleOp">,
+    SymbolTable, 
+    SingleBlockImplicitTerminator<"EndOp">, 
+    IsolatedFromAbove
   ]> {
   let summary = "Define an AIE design targetting a complete device";
   let description = [{
@@ -55,13 +59,19 @@ def AIE_DeviceOp: AIE_Op<"device", [
     ```
   }];
 
-  let arguments = (ins AIEDevice:$device);
+  let arguments = (
+    ins AIEDevice:$device,
+        DefaultValuedAttr<SymbolNameAttr, "getDefaultDeviceName()">:$sym_name
+  );
   let regions = (region AnyRegion:$body_region);
   let assemblyFormat = [{
-    `(` $device `)` regions attr-dict
+    `(` $device `)` ($sym_name^)? regions attr-dict
   }];
   let extraClassDeclaration = [{
     const xilinx::AIE::AIETargetModel &getTargetModel();
+    static xilinx::AIE::DeviceOp getForSymbolInModule(mlir::ModuleOp module, llvm::StringRef symbol);
+    static xilinx::AIE::DeviceOp getForSymbolInModuleOrError(mlir::ModuleOp module, llvm::StringRef symbol);
+    static llvm::StringRef getDefaultDeviceName() { return "main"; }
   }];
 }
 
@@ -325,6 +335,11 @@ def AIE_CoreOp: AIE_Op<"core", [
     for each core.  The name of this file can be be specified using the `elf_file`
     attribute.
 
+    If the `elf_file` attribute is present, no MLIR besides a terminator may be
+    present in the core; in that case, the binary file linked dictates what 
+    will run in the core. The path specified should is relative to the MLIR
+    file.
+
     This op has an optional `stackSize` attribute, to control the amount of memory (in bytes)
     reserved for the stack.  The default value is 1024.  The stack (and other data allocations)
     are always stored in the local core memory, to avoid conflicts with static data allocations
@@ -359,6 +374,7 @@ def AIE_CoreOp: AIE_Op<"core", [
     int colIndex();
     int rowIndex();
     bool isMemWest() { return ((rowIndex() % 2) == 0); };
+    bool isEmpty();
     TileOp getTileOp();
     using ::xilinx::AIE::TileElement::Trait<CoreOp>::getAsmResultNames;
   }];

include/aie/Dialect/AIE/Transforms/AIEPasses.h

@@ -63,16 +63,12 @@ std::unique_ptr<mlir::OperationPass<DeviceOp>> createAIEAssignTileCtrlIDsPass();
 /// 5. rewrite stream-switches (within a bounding box) back to flows
 struct AIEPathfinderPass : AIERoutePathfinderFlowsBase<AIEPathfinderPass> {
 
-  DynamicTileAnalysis analyzer;
-  mlir::DenseMap<TileID, mlir::Operation *> tiles;
-
   AIEPathfinderPass() = default;
-  AIEPathfinderPass(DynamicTileAnalysis analyzer)
-      : analyzer(std::move(analyzer)) {}
 
   void runOnOperation() override;
-  void runOnFlow(DeviceOp d);
-  void runOnPacketFlow(DeviceOp d, mlir::OpBuilder &builder);
+  void runOnFlow(DeviceOp d, DynamicTileAnalysis &analyzer);
+  void runOnPacketFlow(DeviceOp d, mlir::OpBuilder &builder,
+                       DynamicTileAnalysis &analyzer);
 
   typedef std::pair<mlir::Operation *, Port> PhysPort;
 

include/aie/Dialect/AIE/Transforms/AIEPasses.td

@@ -66,6 +66,7 @@ def AIECoreToStandard : Pass<"aie-standard-lowering", "mlir::ModuleOp"> {
 
   }];
   let options = [
+    Option<"deviceName", "device", "std::string", /*default=*/"\"\"", "Device to generate code for">,
     Option<"tileCol", "tilecol", "unsigned",
            /*default=*/"-1", "X coordinate of tile to generate code for">,
     Option<"tileRow", "tilerow", "unsigned",

include/aie/Dialect/AIE/Transforms/AIEPathFinder.h

@@ -245,6 +245,8 @@ class DynamicTileAnalysis {
 
   DynamicTileAnalysis() : pathfinder(std::make_shared<Pathfinder>()) {}
   DynamicTileAnalysis(std::shared_ptr<Router> p) : pathfinder(std::move(p)) {}
+  DynamicTileAnalysis(mlir::Operation *op)
+      : pathfinder(std::make_shared<Pathfinder>()) {}
 
   mlir::LogicalResult runAnalysis(DeviceOp &device);
 

include/aie/Dialect/AIEX/AIEUtils.h

@@ -0,0 +1,24 @@
+//===- AIEUtils.h -----------------------------------------------*- C++ -*-===//
+//
+// 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
+//
+// (c) Copyright 2025 Advanced Micro Devices, Inc.
+//
+//===----------------------------------------------------------------------===//
+
+#include "aie/Dialect/AIE/IR/AIEDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+
+using namespace mlir;
+
+namespace xilinx {
+namespace AIEX {
+
+memref::GlobalOp getOrCreateDataMemref(OpBuilder &builder, AIE::DeviceOp dev,
+                                       mlir::Location loc,
+                                       ArrayRef<uint32_t> words);
+
+}
+} // namespace xilinx

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

@@ -17,6 +17,7 @@ include "aie/Dialect/AIE/IR/AIEInterfaces.td"
 include "mlir/IR/OpBase.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/EnumAttr.td"
+include "mlir/IR/BuiltinAttributes.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"
@@ -527,7 +528,11 @@ def AIE_SelectOp: AIEX_Op<"select", []>, Results<(outs Index)> {
   ];
 }
 
-def AIE_RuntimeSequenceOp : AIEX_Op<"runtime_sequence", [NoTerminator, HasParent<"AIE::DeviceOp">]> {
+def AIE_RuntimeSequenceOp : AIEX_Op<"runtime_sequence", [
+  Symbol,
+  NoTerminator, 
+  HasParent<"AIE::DeviceOp">, 
+]> {
   let summary = "Program the configuration co-processor of the AI Engine array";
   let description = [{
     Instructions in this operation allow for runtime (re-)configuration of the AI Engine array, such as configuring data movement buffer descriptors.
@@ -537,13 +542,69 @@ def AIE_RuntimeSequenceOp : AIEX_Op<"runtime_sequence", [NoTerminator, HasParent
     The input arguments are arguments passed in from the host at kernel invocation time. This may include buffers on the host.
   }];
   let arguments = (
-    ins OptionalAttr<SymbolNameAttr>:$sym_name
+    ins DefaultValuedAttr<SymbolNameAttr, "getDefaultRuntimeSequenceName()">:$sym_name
   );
   let regions = (region
     AnyRegion:$body
   );
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
+  let extraClassDeclaration = [{
+    static llvm::StringRef getDefaultRuntimeSequenceName() { return "sequence"; }
+    static RuntimeSequenceOp getForSymbolInDevice(AIE::DeviceOp module, llvm::StringRef symbol);
+    static RuntimeSequenceOp getForSymbolInDeviceOrError(AIE::DeviceOp module, llvm::StringRef symbol);
+  }];
+  let extraClassDefinition = [{
+  }];
+}
+
+def AIE_ConfigureOp: AIEX_Op<"configure", [
+  HasParent<"RuntimeSequenceOp">,
+  NoTerminator
+]>,
+  Results<(outs Index:$result)>
+ {
+  let summary = "Set up a configuration (program memories, stream switches, etc.) on the NPU device.";
+  let arguments = (
+    ins FlatSymbolRefAttr:$symbol
+  );
+
+  let assemblyFormat = [{
+     $symbol regions attr-dict
+  }];
+
+  let extraClassDeclaration = [{
+    AIE::DeviceOp getReferencedDeviceOp();
+  }];
+
+  let regions = (region
+    AnyRegion:$body
+  );
+
+  let hasVerifier = 1;
+}
+
+def AIE_RunOp: AIEX_Op<"run", [HasParent<"ConfigureOp">]> {
+  let arguments = (
+    ins FlatSymbolRefAttr:$runtime_sequence_symbol,
+        Variadic<AnyType>:$args
+  );
+
+  let assemblyFormat = [{
+     $runtime_sequence_symbol `(` $args `)` `:` `(` type($args) `)` attr-dict
+  }];
+
+  let extraClassDeclaration = [{
+    AIE::DeviceOp getCalleeDeviceOp();
+    RuntimeSequenceOp getCalleeRuntimeSequenceOp();
+  }];
+
+  let hasVerifier = 1;
+  
+  let summary = "Execute a runtime sequence";
+  let description = [{
+    Executes an `aiex.runtime_sequence` with the given name and arguments by inlining its instructions at the call site.
+  }];
 }
 
 def AIE_NpuDmaMemcpyNdOp: AIEX_Op<"npu.dma_memcpy_nd", [
@@ -772,6 +833,9 @@ def AIE_NpuWrite32Op: AIEX_Op<"npu.write32", []> {
     If 'buffer' is not present and 'column' and 'row' are not present then
     'address' is interpreted as a full 32-bit address in the AIE array.
   }];
+  let extraClassDeclaration = [{
+    std::optional<uint32_t> getAbsoluteAddress();
+  }];
 }
 
 // MASKWRITE
@@ -798,6 +862,9 @@ def AIE_NpuMaskWrite32Op: AIEX_Op<"npu.maskwrite32", []> {
     If 'buffer' is not present and 'column' and 'row' are not present then
     'address' is interpreted as a full 32-bit address in the AIE array.
   }];
+  let extraClassDeclaration = [{
+    std::optional<uint32_t> getAbsoluteAddress();
+  }];
 }
 
 // BLOCKWRITE
@@ -823,6 +890,10 @@ def AIE_NpuBlockWriteOp: AIEX_Op<"npu.blockwrite", []> {
     If 'buffer' is not present and 'column' and 'row' are not present then
     'address' is interpreted as a full 32-bit address in the AIE array.
   }];
+  let extraClassDeclaration = [{
+    std::optional<uint32_t> getAbsoluteAddress();
+    mlir::DenseIntElementsAttr getDataWords(); 
+  }];
 }
 
 // OP_SYNC