Update aiex.runtime_sequence to add optional sym_name, remove unused …

…args (#1650)
Xilinx · Jul 30, 2024 · b9cdb91 · b9cdb91
1 parent 142fb0c
commit b9cdb91
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 4 deletions.
diff --git a/include/aie/Dialect/AIEX/IR/AIEX.td b/include/aie/Dialect/AIEX/IR/AIEX.td
@@ -473,8 +473,8 @@ def AIE_RuntimeSequenceOp : AIEX_Op<"runtime_sequence", [NoTerminator, HasParent
     Typically, these instructions include configuring the data transfers between host and AIE array on the shims.
     The input arguments are arguments passed in from the host at kernel invocation time. This may include buffers on the host.
   }];
-  let arguments = (ins
-    Variadic<AnyType>:$args
+  let arguments = (
+    ins OptionalAttr<SymbolNameAttr>:$sym_name
   );
   let regions = (region
     AnyRegion:$body

diff --git a/lib/Dialect/AIEX/IR/AIEXDialect.cpp b/lib/Dialect/AIEX/IR/AIEXDialect.cpp
@@ -1,4 +1,4 @@
-//===- AIEDialect.cpp -------------------------------------------*- C++ -*-===//
+//===- AIEXDialect.cpp ------------------------------------------*- C++ -*-===//
 //
 // This file is licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -284,6 +284,11 @@ LogicalResult AIEX::NpuWriteBdOp::verify() {
 
 ParseResult AIEX::RuntimeSequenceOp::parse(OpAsmParser &parser,
                                            OperationState &result) {
+
+  StringAttr nameAttr;
+  (void)parser.parseOptionalSymbolName(
+      nameAttr, mlir::SymbolTable::getSymbolAttrName(), result.attributes);
+
   SmallVector<OpAsmParser::Argument> entryArgs;
 
   // Entry arguments,  e.g. (%addr: memref<1xi32>)
@@ -313,6 +318,13 @@ ParseResult AIEX::RuntimeSequenceOp::parse(OpAsmParser &parser,
 void AIEX::RuntimeSequenceOp::print(OpAsmPrinter &printer) {
   Region &body = getRegion();
 
+  auto nameAttr = (*this)->getAttrOfType<StringAttr>(
+      mlir::SymbolTable::getSymbolAttrName());
+  if (nameAttr) {
+    printer << ' ';
+    printer.printSymbolName(nameAttr);
+  }
+
   printer << '(';
   for (unsigned i = 0, n = body.getNumArguments(); i < n; i++) {
     if (i > 0) {

diff --git a/python/dialects/aiex.py b/python/dialects/aiex.py
@@ -768,7 +768,7 @@ def _find_next_channel(used_channels):
 
 def runtime_sequence(*inputs: Type):
     def decorator(f):
-        seq_op = RuntimeSequenceOp(args=[])
+        seq_op = RuntimeSequenceOp()
         entry_block = seq_op.body.blocks.append(*inputs)
         args = entry_block.arguments
         with InsertionPoint(entry_block):

diff --git a/test/dialect/AIEX/roundtrip.mlir b/test/dialect/AIEX/roundtrip.mlir
@@ -38,4 +38,15 @@ aie.device(npu1_4col) {
   aiex.runtime_sequence() {
     aiex.npu.address_patch {addr = 123 : ui32, arg_idx = 3 : i32, arg_plus = 0 : i32}
   }
+}
+
+// -----
+
+// CHECK: aie.device
+// CHECK: runtime_sequence @seq(%arg0: memref<1xi32>)
+// CHECK: aiex.npu.write32 {address = 432 : ui32, value = 1 : ui32}
+aie.device(npu1_4col) {
+  aiex.runtime_sequence @seq(%arg0 : memref<1xi32>) {
+    aiex.npu.write32 {address = 432 : ui32, value = 1 : ui32}
+  }
 }