Bitwise or on consecutive ctrl pkts operating on the same address (#…

…1778)

lib/Conversion/AIEToConfiguration/AIEToConfiguration.cpp

@@ -297,6 +297,48 @@ emitControlPacketOps(OpBuilder &builder,
   return success();
 }
 
+// Perform bitwise or on consecutive control packets operating on the same
+// address, to resolve the lack of mask write in control packets.
+LogicalResult orConsecutiveWritesOnSameAddr(Block *body) {
+  SmallVector<AIEX::NpuControlPacketOp> ctrlPktOps;
+  body->walk(
+      [&](AIEX::NpuControlPacketOp cpOp) { ctrlPktOps.push_back(cpOp); });
+  if (ctrlPktOps.empty())
+    return success();
+
+  SmallVector<Operation *> erased;
+  int addrBuffer = ctrlPktOps[0].getAddress();
+  AIEX::NpuControlPacketOp ctrlPktBuffer = ctrlPktOps[0];
+  for (size_t i = 1; i < ctrlPktOps.size(); i++) {
+    int currentAddrBuffer = ctrlPktOps[i].getAddress();
+    if (addrBuffer != currentAddrBuffer) {
+      addrBuffer = currentAddrBuffer;
+      ctrlPktBuffer = ctrlPktOps[i];
+      continue;
+    }
+    auto bufferedData = ctrlPktBuffer.getData().value();
+    auto currentData = ctrlPktOps[i].getData().value();
+    SmallVector<int> newData;
+    for (unsigned j = 0; j < std::max(bufferedData.size(), currentData.size());
+         j++) {
+      if (j < std::min(bufferedData.size(), currentData.size())) {
+        newData.push_back(bufferedData[j] | currentData[j]);
+        continue;
+      }
+      newData.push_back(j < bufferedData.size() ? bufferedData[j]
+                                                : currentData[j]);
+    }
+    ctrlPktBuffer.getProperties().data = DenseI32ArrayAttr::get(
+        ctrlPktBuffer->getContext(), ArrayRef<int>{newData});
+    erased.push_back(ctrlPktOps[i]);
+  }
+
+  for (auto e : erased)
+    e->erase();
+
+  return success();
+}
+
 // an enum to represent the output type of the transaction binary
 enum OutputType {
   Transaction,
@@ -350,6 +392,9 @@ static LogicalResult convertTransactionOpsToMLIR(
   } else if (outputType == OutputType::ControlPacket) {
     if (failed(emitControlPacketOps(builder, operations, global_data)))
       return failure();
+    // resolve mask writes; control packet doesn't natively support mask write.
+    if (failed(orConsecutiveWritesOnSameAddr(&seq.getBody().front())))
+      return failure();
   } else {
     llvm_unreachable("bad output type");
   }

test/Conversion/AIEToConfiguration/convert_aie_to_ctrl_pkts.mlir

@@ -0,0 +1,28 @@
+//===- convert_aie_to_ctrl_pkts.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
+//
+// (c) Copyright 2024 Advanced Micro Devices, Inc.
+//
+//===----------------------------------------------------------------------===//
+
+// RUN: aie-opt -convert-aie-to-control-packets="elf-dir=%S/convert_aie_to_ctrl_pkts_elfs/" %s | FileCheck %s
+
+// CHECK-label: aiex.runtime_sequence @configure
+// CHECK-COUNT-5: aiex.control_packet {address = {{.*}} : ui32, data = array<i32: {{.*}}>
+// CHECK-COUNT-2: aiex.control_packet {address = {{.*}} : ui32, data = array<i32: {{.*}}, {{.*}}, {{.*}}, {{.*}}>
+// CHECK: aiex.control_packet {address = {{.*}} : ui32, data = array<i32: {{.*}}>
+// CHECK-COUNT-36: aiex.control_packet {address = {{.*}} : ui32, data = array<i32: {{.*}}, {{.*}}, {{.*}}, {{.*}}>
+// CHECK-COUNT-22: aiex.control_packet {address = {{.*}} : ui32, data = array<i32: {{.*}}>
+aie.device(npu1_1col) {
+  %12 = aie.tile(0, 2)
+  %buf = aie.buffer(%12) : memref<256xi32>
+  %4 = aie.core(%12)  {
+    %0 = arith.constant 0 : i32
+    %1 = arith.constant 0 : index
+    memref.store %0, %buf[%1] : memref<256xi32>
+    aie.end
+  }
+}