Add aiex.npu.blockwrite operation (Xilinx#1638)

Author: fifield

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

@@ -683,6 +683,22 @@ def AIE_NpuWrite32Op: AIEX_Op<"npu.write32", []> {
   }];
 }
 
+// BLOCKWRITE
+def AIE_NpuBlockWriteOp: AIEX_Op<"npu.blockwrite", []> {
+  let summary = "blockwrite operator";
+  let arguments = (
+    ins AnyMemRef:$data,
+        UI32Attr:$address
+  );
+  let results = (outs );
+  let assemblyFormat = [{
+    `(` $data `)` attr-dict `:` type($data)
+  }];
+  let description = [{
+    blockwrite operator
+  }];
+}
+
 // OP_SYNC
 def AIE_NpuSyncOp: AIEX_Op<"npu.sync", []> {
   let summary = "sync operator";

lib/Dialect/AIEX/Transforms/AIEDmaToNpu.cpp

@@ -70,10 +70,10 @@ struct ShimDMAllocationGetter {
 };
 } // namespace
 
-struct RtpToNpuPattern : OpConversionPattern<NpuWriteRTPOp> {
+struct RtpToWrite32Pattern : OpConversionPattern<NpuWriteRTPOp> {
   using OpConversionPattern::OpConversionPattern;
 
-  RtpToNpuPattern(MLIRContext *context, PatternBenefit benefit = 1)
+  RtpToWrite32Pattern(MLIRContext *context, PatternBenefit benefit = 1)
       : OpConversionPattern(context, benefit) {}
 
   LogicalResult
@@ -118,12 +118,12 @@ struct RtpToNpuPattern : OpConversionPattern<NpuWriteRTPOp> {
   }
 };
 
-struct PushToNpuPattern : OpConversionPattern<NpuPushQueueOp> {
+struct PushQueuetoWrite32Pattern : OpConversionPattern<NpuPushQueueOp> {
 
 public:
   using OpConversionPattern::OpConversionPattern;
 
-  PushToNpuPattern(MLIRContext *context, PatternBenefit benefit = 1)
+  PushQueuetoWrite32Pattern(MLIRContext *context, PatternBenefit benefit = 1)
       : OpConversionPattern(context, benefit) {}
 
   LogicalResult
@@ -343,16 +343,16 @@ struct DmaToNpuPattern : OpConversionPattern<NpuDmaMemcpyNdOp> {
 /// Convert NpuDmaWaitOp into NpuSyncOp by retrieving the necessary
 /// information from the ShimDMAAllocationOp referenced through the
 /// symbol argument of this op.
-struct DmaWaitToNpuPattern : OpConversionPattern<NpuDmaWaitOp> {
+struct DmaWaitToSyncPattern : OpConversionPattern<NpuDmaWaitOp> {
 
 private:
   ShimDMAllocationGetter &allocGetter;
 
 public:
   using OpConversionPattern::OpConversionPattern;
 
-  DmaWaitToNpuPattern(MLIRContext *context, ShimDMAllocationGetter &getter,
-                      PatternBenefit benefit = 1)
+  DmaWaitToSyncPattern(MLIRContext *context, ShimDMAllocationGetter &getter,
+                       PatternBenefit benefit = 1)
       : OpConversionPattern(context, benefit), allocGetter(getter) {}
 
   LogicalResult
@@ -374,11 +374,103 @@ struct DmaWaitToNpuPattern : OpConversionPattern<NpuDmaWaitOp> {
         op, shimDmaAllocOp->getCol(), /* row */ 0,
         static_cast<uint32_t>(shimDmaAllocOp->getChannelDir()),
         shimDmaAllocOp->getChannelIndex(), 1, 1);
+
+    return success();
+  }
+};
+
+struct WriteBdToBlockWritePattern : OpConversionPattern<NpuWriteBdOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  WriteBdToBlockWritePattern(MLIRContext *context, PatternBenefit benefit = 1)
+      : OpConversionPattern(context, benefit) {}
+
+  LogicalResult
+  matchAndRewrite(NpuWriteBdOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    AIE::DeviceOp dev = op->getParentOfType<AIE::DeviceOp>();
+    const AIE::AIETargetModel &tm = dev.getTargetModel();
+
+    auto bd_id = op.getBdId();
+    uint32_t bd_addr = (op.getColumn() << tm.getColumnShift()) |
+                       (op.getRow() << tm.getRowShift()) |
+                       (0x1D000 + bd_id * 0x20);
+
+    std::vector<uint32_t> words(8, 0);
+
+    // DMA_BDX_0
+    words[0] = op.getBufferLength();
+
+    // DMA_BDX_1
+    words[1] = op.getBufferOffset();
+
+    // DMA_BDX_2
+    // En Packet , OoO BD ID , Packet ID , Packet Type
+    words[2] |= (op.getEnablePacket() & 0x1) << 30;
+    words[2] |= (op.getOutOfOrderId() & 0x3f) << 24;
+    words[2] |= (op.getPacketId() & 0x1f) << 19;
+    words[2] |= (op.getPacketType() & 0x7) << 16;
+
+    // DMA_BDX_3
+    // TODO: Secure Access
+    words[3] |= (op.getD0Size() & 0x3ff) << 20;
+    words[3] |= op.getD0Stride() & 0xfffff;
+
+    // DMA_BDX_4
+    words[4] = 0x80000000; // burst length;
+    words[4] |= (op.getD1Size() & 0x3ff) << 20;
+    words[4] |= op.getD1Stride() & 0xfffff;
+
+    // DMA_BDX_5
+    // TODO: SIMID, AxCache, AXQoS
+    words[5] = op.getD2Stride() & 0xfffff;
+
+    // DMA_BDX_6
+    words[6] |= (op.getIterationCurrent() & 0x3f) << 26;
+    words[6] |= (op.getIterationSize() & 0x3f) << 20;
+    words[6] |= op.getIterationStride() & 0xfffff;
+
+    // DMA_BDX_7
+    // TODO: TLAST Suppress
+    words[7] |= (op.getNextBd() & 0xf) << 27;
+    words[7] |= (op.getUseNextBd() & 0x1) << 26;
+    words[7] |= (op.getValidBd() & 0x1) << 25;
+    words[7] |= (op.getLockRelVal() & 0xef) << 18;
+    words[7] |= (op.getLockRelId() & 0xf) << 13;
+    words[7] |= (op.getLockAcqEnable() & 0x1) << 12;
+    words[7] |= (op.getLockAcqVal() & 0xef) << 5;
+    words[7] |= op.getLockAcqId() & 0xf;
+
+    MemRefType memrefType = MemRefType::get({8}, rewriter.getI32Type());
+    TensorType tensorType = RankedTensorType::get({8}, rewriter.getI32Type());
+    memref::GlobalOp global = nullptr;
+    {
+      OpBuilder::InsertionGuard guard(rewriter);
+      std::string name = "blockwrite_data_";
+      rewriter.setInsertionPoint(op->getParentOfType<func::FuncOp>());
+      int id = 0;
+      while (dev.lookupSymbol(name + std::to_string(id)))
+        id++;
+      name += std::to_string(id);
+      global = rewriter.create<memref::GlobalOp>(
+          op->getLoc(), name, rewriter.getStringAttr("private"), memrefType,
+          DenseElementsAttr::get<uint32_t>(tensorType, words), true, nullptr);
+    }
+    auto memref = rewriter.create<memref::GetGlobalOp>(op->getLoc(), memrefType,
+                                                       global.getName());
+    (void)rewriter.replaceOpWithNewOp<NpuBlockWriteOp>(
+        op, memref.getResult(), rewriter.getUI32IntegerAttr(bd_addr));
     return success();
   }
 };
 
 struct AIEDmaToNpuPass : AIEDmaToNpuBase<AIEDmaToNpuPass> {
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<memref::MemRefDialect>();
+  }
+
   void runOnOperation() override {
 
     ShimDMAllocationGetter cachingGetter;
@@ -387,18 +479,22 @@ struct AIEDmaToNpuPass : AIEDmaToNpuBase<AIEDmaToNpuPass> {
 
     ConversionTarget target(getContext());
     target.addLegalDialect<AIEXDialect>();
+    target.addLegalDialect<memref::MemRefDialect>();
     target.addLegalOp<AIE::BufferOp>();
     target.addLegalOp<AIE::ShimDMAAllocationOp>();
-    target.addIllegalOp<NpuWriteRTPOp>();
+
     target.addIllegalOp<NpuDmaMemcpyNdOp>();
     target.addIllegalOp<NpuDmaWaitOp>();
     target.addIllegalOp<NpuPushQueueOp>();
+    target.addIllegalOp<NpuWriteRTPOp>();
+    target.addIllegalOp<NpuWriteBdOp>();
 
     RewritePatternSet patterns(&getContext());
     patterns.insert<DmaToNpuPattern>(&getContext(), cachingGetter);
-    patterns.insert<DmaWaitToNpuPattern>(&getContext(), cachingGetter);
-    patterns.insert<PushToNpuPattern>(&getContext());
-    patterns.insert<RtpToNpuPattern>(&getContext());
+    patterns.insert<DmaWaitToSyncPattern>(&getContext(), cachingGetter);
+    patterns.insert<PushQueuetoWrite32Pattern>(&getContext());
+    patterns.insert<RtpToWrite32Pattern>(&getContext());
+    patterns.insert<WriteBdToBlockWritePattern>(&getContext());
 
     if (failed(applyPartialConversion(device, target, std::move(patterns))))
       signalPassFailure();

lib/Targets/AIETargetNPU.cpp

@@ -104,66 +104,50 @@ void appendAddressPatch(std::vector<uint32_t> &instructions,
   words[11] = 0;
 }
 
-void appendWriteBdShimTile(std::vector<uint32_t> &instructions,
-                           NpuWriteBdOp op) {
+void appendBlockWrite(std::vector<uint32_t> &instructions, NpuBlockWriteOp op) {
 
-  auto words = reserveAndGetTail(instructions, 12);
-  const AIETargetModel &tm = op->getParentOfType<DeviceOp>().getTargetModel();
+  Value memref = op.getData();
+  int64_t width = cast<MemRefType>(memref.getType()).getElementTypeBitWidth();
+  if (width != 32) {
+    op.emitWarning("Only 32-bit data type is supported for now");
+    return;
+  }
+
+  memref::GetGlobalOp getGlobal = memref.getDefiningOp<memref::GetGlobalOp>();
+  if (!getGlobal) {
+    op.emitError("Only MemRefs from memref.get_global are supported");
+    return;
+  }
+
+  auto global = dyn_cast_if_present<memref::GlobalOp>(
+      op->getParentOfType<AIE::DeviceOp>().lookupSymbol(getGlobal.getName()));
+  if (!global) {
+    op.emitError("Global symbol not found");
+    return;
+  }
+
+  auto initVal = global.getInitialValue();
+  if (!initVal) {
+    op.emitError("Global symbol has no initial value");
+    return;
+  }
+
+  auto data = dyn_cast<DenseIntElementsAttr>(*initVal);
+  if (!data) {
+    op.emitError("Global symbol initial value is not a dense int array");
+    return;
+  }
 
   // XAIE_IO_BLOCKWRITE
+  auto words = reserveAndGetTail(instructions, data.size() + 4);
   words[0] = TXN_OPC_BLOCKWRITE;
   words[1] = 0;
-
-  // RegOff
-  auto bd_id = op.getBdId();
-  uint32_t bd_addr = (op.getColumn() << tm.getColumnShift()) |
-                     (op.getRow() << tm.getRowShift()) |
-                     (0x1D000 + bd_id * 0x20);
-  words[2] = bd_addr;                         // ADDR
+  words[2] = op.getAddress();
   words[3] = words.size() * sizeof(uint32_t); // Operation Size
 
-  // DMA_BDX_0
-  words[4] = op.getBufferLength();
-
-  // DMA_BDX_1
-  words[5] = op.getBufferOffset();
-
-  // DMA_BDX_2
-  // En Packet , OoO BD ID , Packet ID , Packet Type
-  words[6] |= (op.getEnablePacket() & 0x1) << 30;
-  words[6] |= (op.getOutOfOrderId() & 0x3f) << 24;
-  words[6] |= (op.getPacketId() & 0x1f) << 19;
-  words[6] |= (op.getPacketType() & 0x7) << 16;
-
-  // DMA_BDX_3
-  // TODO: Secure Access
-  words[7] |= (op.getD0Size() & 0x3ff) << 20;
-  words[7] |= op.getD0Stride() & 0xfffff;
-
-  // DMA_BDX_4
-  words[8] = 0x80000000; // burst length;
-  words[8] |= (op.getD1Size() & 0x3ff) << 20;
-  words[8] |= op.getD1Stride() & 0xfffff;
-
-  // DMA_BDX_5
-  // TODO: SIMID, AxCache, AXQoS
-  words[9] = op.getD2Stride() & 0xfffff;
-
-  // DMA_BDX_6
-  words[10] |= (op.getIterationCurrent() & 0x3f) << 26;
-  words[10] |= (op.getIterationSize() & 0x3f) << 20;
-  words[10] |= op.getIterationStride() & 0xfffff;
-
-  // DMA_BDX_7
-  // TODO: TLAST Suppress
-  words[11] |= (op.getNextBd() & 0xf) << 27;
-  words[11] |= (op.getUseNextBd() & 0x1) << 26;
-  words[11] |= (op.getValidBd() & 0x1) << 25;
-  words[11] |= (op.getLockRelVal() & 0xef) << 18;
-  words[11] |= (op.getLockRelId() & 0xf) << 13;
-  words[11] |= (op.getLockAcqEnable() & 0x1) << 12;
-  words[11] |= (op.getLockAcqVal() & 0xef) << 5;
-  words[11] |= op.getLockAcqId() & 0xf;
+  unsigned i = 4;
+  for (auto d : data)
+    words[i++] = d.getZExtValue();
 }
 
 } // namespace
@@ -195,13 +179,13 @@ std::vector<uint32_t> xilinx::AIE::AIETranslateToNPU(ModuleOp module) {
             count++;
             appendWrite32(instructions, op);
           })
-          .Case<NpuAddressPatchOp>([&](auto op) {
+          .Case<NpuBlockWriteOp>([&](auto op) {
             count++;
-            appendAddressPatch(instructions, op);
+            appendBlockWrite(instructions, op);
           })
-          .Case<NpuWriteBdOp>([&](auto op) {
+          .Case<NpuAddressPatchOp>([&](auto op) {
             count++;
-            appendWriteBdShimTile(instructions, op);
+            appendAddressPatch(instructions, op);
           });
     }
   }

test/Conversion/DmaToNpu/aiert_insts.mlir

@@ -7,9 +7,9 @@
 //===----------------------------------------------------------------------===//
 
 // RUN: aie-opt --aie-dma-to-npu %s | FileCheck %s
-// CHECK: aiex.npu.writebd {bd_id = 1 : i32, buffer_length = 32 : i32, buffer_offset = 0 : i32, column = 0 : i32, d0_size = 0 : i32, d0_stride = 0 : i32, d1_size = 0 : i32, d1_stride = 0 : i32, d2_stride = 0 : i32, enable_packet = 0 : i32, iteration_current = 0 : i32, iteration_size = 0 : i32, iteration_stride = 0 : i32, lock_acq_enable = 0 : i32, lock_acq_id = 0 : i32, lock_acq_val = 0 : i32, lock_rel_id = 0 : i32, lock_rel_val = 0 : i32, next_bd = 0 : i32, out_of_order_id = 0 : i32, packet_id = 0 : i32, packet_type = 0 : i32, row = 0 : i32, use_next_bd = 0 : i32, valid_bd = 1 : i32}
+// CHECK: aiex.npu.blockwrite(%{{.*}}) {address = 118816 : ui32} : memref<8xi32>
 // CHECK: aiex.npu.write32 {address = 119300 : ui32, column = 0 : i32, row = 0 : i32, value = 2147483649 : ui32}
-// CHECK: aiex.npu.writebd {bd_id = 0 : i32, buffer_length = 32 : i32, buffer_offset = 128 : i32, column = 0 : i32, d0_size = 8 : i32, d0_stride = 0 : i32, d1_size = 2 : i32, d1_stride = 7 : i32, d2_stride = 15 : i32, enable_packet = 0 : i32, iteration_current = 0 : i32, iteration_size = 0 : i32, iteration_stride = 0 : i32, lock_acq_enable = 0 : i32, lock_acq_id = 0 : i32, lock_acq_val = 0 : i32, lock_rel_id = 0 : i32, lock_rel_val = 0 : i32, next_bd = 0 : i32, out_of_order_id = 0 : i32, packet_id = 0 : i32, packet_type = 0 : i32, row = 0 : i32, use_next_bd = 0 : i32, valid_bd = 1 : i32}
+// CHECK: aiex.npu.blockwrite(%{{.*}}) {address = 118784 : ui32} : memref<8xi32>
 // CHECK: aiex.npu.write32 {address = 119316 : ui32, column = 0 : i32, row = 0 : i32, value = 0 : ui32}
 
 module {
@@ -24,8 +24,8 @@ module {
       %c8 = arith.constant 8 : i64
       %c16 = arith.constant 16 : i64
       %c32 = arith.constant 32 : i64
-      aiex.npu.dma_memcpy_nd (0, 0, %out[%c0,%c0,%c0,%c0][%c1,%c1,%c1,%c32][%c0,%c0,%c0, %c1]) { metadata = @of_toMem, id = 1 : i64 } : memref<64xi32>
-      aiex.npu.dma_memcpy_nd (0, 0, %in[%c0,%c2,%c0,%c0][%c1,%c2,%c2,%c8][%c0,%c16,%c8, %c1]) { metadata = @of_fromMem, id = 0 : i64 } : memref<4x2x8xi32>
+      aiex.npu.dma_memcpy_nd (0, 0, %out[%c0,%c0,%c0,%c0][%c1,%c1,%c1,%c32][%c0,%c0,%c0, %c1]) { metadata = @of_toMem, id = 1 : i64, issue_token = true } : memref<64xi32>
+      aiex.npu.dma_memcpy_nd (0, 0, %in[%c0,%c2,%c0,%c0][%c1,%c2,%c2,%c8][%c0,%c16,%c8, %c1]) { metadata = @of_fromMem, id = 0 : i64, issue_token = false } : memref<4x2x8xi32>
       return
     }
     aie.shim_dma_allocation @of_fromMem (MM2S, 0, 0)