Add NpuSyncOp generation to AIEDmaToNpu

lib/Dialect/AIEX/Transforms/AIEDmaToIpu.cpp

@@ -332,6 +332,32 @@ struct DmaToIpuPattern : OpConversionPattern<IpuDmaMemcpyNdOp> {
   }
 };
 
+void insertIpuSyncOpForResults(AIE::DeviceOp device) {
+  device.walk([&](mlir::func::FuncOp f) {
+    SmallVector<AIEX::IpuDmaMemcpyNdOp> dmas;
+    Operation *returnOp = nullptr;
+    f.walk([&](mlir::func::ReturnOp op) { returnOp = op.getOperation(); });
+    f.walk([&](AIEX::IpuDmaMemcpyNdOp dma) { dmas.push_back(dma); });
+    for (auto dma : dmas) {
+      if (auto infoOp = getAllocOpForSymbol(device, dma.getMetadata())) {
+        if (infoOp->getChannelDir() == AIE::DMAChannelDir::S2MM) {
+          // Found dma op copying results to host
+          OpBuilder builder(dma);
+          auto col = builder.getI32IntegerAttr(infoOp->getCol());
+          auto row = builder.getI32IntegerAttr(0);
+          auto dir = builder.getI32IntegerAttr(0);
+          auto chan = builder.getI32IntegerAttr(infoOp->getChannelIndex());
+          auto col_num = builder.getI32IntegerAttr(1);
+          auto row_num = builder.getI32IntegerAttr(1);
+          builder.setInsertionPoint(returnOp);
+          builder.create<AIEX::IpuSyncOp>(dma->getLoc(), col, row, dir, chan,
+                                          col_num, row_num);
+        }
+      }
+    }
+  });
+}
+
 struct AIEDmaToIpuPass : AIEDmaToIpuBase<AIEDmaToIpuPass> {
   void runOnOperation() override {
 
@@ -350,6 +376,9 @@ struct AIEDmaToIpuPass : AIEDmaToIpuBase<AIEDmaToIpuPass> {
     patterns.insert<PushToIpuPattern>(&getContext());
     patterns.insert<RtpToIpuPattern>(&getContext());
 
+    // Insert sync op after copying data out to host
+    insertIpuSyncOpForResults(device);
+
     if (failed(applyPartialConversion(device, target, std::move(patterns))))
       signalPassFailure();
   }

reference_designs/ipu-xrt/add_one_objFifo/aie2.py

@@ -67,7 +67,6 @@ def sequence(inTensor, notUsed, outTensor):
                 ipu_dma_memcpy_nd(
                     metadata="in0", bd_id=1, mem=inTensor, sizes=[1, 1, 1, 64]
                 )
-                ipu_sync(column=0, row=0, direction=0, channel=0)
 
     print(ctx.module)
 

reference_designs/ipu-xrt/log_hello_world/hello_world.py

@@ -56,7 +56,6 @@ def sequence(in_mem, out_mem, logout):
                 ipu_dma_memcpy_nd(
                     metadata="logoutOF", bd_id=2, mem=logout, sizes=[1, 1, 1, N]
                 )
-                ipu_sync(column=0, row=0, direction=0, channel=0)
 
     print(ctx.module)
 

reference_designs/ipu-xrt/matrix_multiplication/aie2.py

@@ -297,7 +297,6 @@ def sequence(A, B, C):
                             strides=[n_in_i32s, k_x_N_in_i32s, N_in_i32s],
                         )
 
-                    ipu_sync(column=0, row=0, direction=0, channel=0)
 
     print(ctx.module)
 

reference_designs/ipu-xrt/matrix_multiplication_array/aie2.py

@@ -354,8 +354,6 @@ def sequence(A, B, C):
                                 sizes=[N_div_n_div_n_cols, K_div_k, k, n_in_i32s],
                                 strides=[n_x_n_cols_in_i32s, k_x_N_in_i32s, N_in_i32s],
                             )
-                    for i in range(n_cols):
-                        ipu_sync(column=i, row=0, direction=0, channel=0)
 
     # print(ctx.module.operation.verify())
     print(ctx.module)

reference_designs/ipu-xrt/matrix_multiplication_column/aie2.py

@@ -246,7 +246,6 @@ def sequence(A, B, C):
                             strides=[n_in_i32s, k_x_N_in_i32s, N_in_i32s],
                         )
 
-                    ipu_sync(column=0, row=0, direction=0, channel=0)
 
     print(ctx.module)
 

reference_designs/ipu-xrt/matrix_vector_multiplication/aie2.py

@@ -204,8 +204,6 @@ def sequence(A, B, C):
                         strides=[0, 0, 0],
                     )
 
-                for i in range(n_cores):
-                    ipu_sync(column=i, row=0, direction=0, channel=0)
 
     print(ctx.module)
 

reference_designs/ipu-xrt/passthrough_hardware/aie2.py

@@ -52,7 +52,6 @@ def core_body():
             def sequence(A, B, C):
                 ipu_dma_memcpy_nd(metadata="out", bd_id=0, mem=C, sizes=[1, 1, 1, N])
                 ipu_dma_memcpy_nd(metadata="in", bd_id=1, mem=A, sizes=[1, 1, 1, N])
-                ipu_sync(column=0, row=0, direction=0, channel=0)
 
     print(ctx.module)
 

reference_designs/ipu-xrt/vector_scalar/aie2.py

@@ -61,7 +61,6 @@ def core_body():
             def sequence(A, B, C):
                 ipu_dma_memcpy_nd(metadata="out", bd_id=0, mem=C, sizes=[1, 1, 1, N])
                 ipu_dma_memcpy_nd(metadata="in", bd_id=1, mem=A, sizes=[1, 1, 1, N])
-                ipu_sync(column=0, row=0, direction=0, channel=0)
 
     print(ctx.module)
 

reference_designs/ipu-xrt/vision_pipelines/color_detect/aie2_colorDetect.py

@@ -266,7 +266,6 @@ def sequence(I, B, O):
                     mem=O,
                     sizes=[1, 1, 1, height * lineWidthInInt32s],
                 )
-                ipu_sync(column=0, row=0, direction=0, channel=0)
 
     print(ctx.module)
 

reference_designs/ipu-xrt/vision_pipelines/color_threshold/aie2_colorThreshold.py

@@ -284,7 +284,6 @@ def sequence(inTensor, notUsed, outTensor):
                     mem=outTensor,
                     sizes=[1, 1, 1, tensorSizeInInt32s],
                 )
-                ipu_sync(column=0, row=0, direction=0, channel=0)
 
     # print(ctx.module.operation.verify())
     print(ctx.module)

reference_designs/ipu-xrt/vision_pipelines/edge_detect/aie2_edgeDetect.py

@@ -312,7 +312,6 @@ def sequence(I, B, O):
                     mem=I,
                     sizes=[1, 1, 1, tensorSizeInInt32s],
                 )
-                ipu_sync(column=0, row=0, direction=0, channel=0)
 
     #    print(ctx.module.operation.verify())
     print(ctx.module)

reference_designs/ipu-xrt/vision_pipelines/passthrough/aie2.py

@@ -165,7 +165,6 @@ def sequence(inTensor, notUsed, outTensor):
                     mem=outTensor,
                     sizes=[1, 1, 1, tensorSizeInInt32s],
                 )
-                ipu_sync(column=0, row=0, direction=0, channel=0)
 
     print(ctx.module)
 

test/ipu-xrt/add_256_using_dma_op_no_double_buffering/aie.mlir

@@ -102,7 +102,6 @@ module {
     func.func @bobsyouruncle(%arg0: memref<64xi32>, %arg1: memref<32xi32>, %arg2: memref<64xi32>) {
       aiex.ipu.dma_memcpy_nd(0, 0, %arg0[0, 0, 0, 0][1, 1, 1, 64][0, 0, 0]) {id = 0 : i64, metadata = @this_just_creates_a_symbol_and_the_type_means_nothing_in} : memref<64xi32>
       aiex.ipu.dma_memcpy_nd(0, 0, %arg2[0, 0, 0, 0][1, 1, 1, 64][0, 0, 0]) {id = 1 : i64, metadata = @this_just_creates_a_symbol_and_the_type_means_nothing_out} : memref<64xi32>
-      aiex.ipu.sync {channel = 0 : i32, column = 0 : i32, column_num = 1 : i32, direction = 0 : i32, row = 0 : i32, row_num = 1 : i32}
       return
     }
   }

test/ipu-xrt/add_314_using_dma_op/aie.mlir

@@ -67,7 +67,6 @@ module {
       %c64_i64 = arith.constant 64 : i64
       aiex.ipu.dma_memcpy_nd(0, 0, %arg0[%c0_i64, %c0_i64, %c0_i64, %c0_i64][%c1_i64, %c1_i64, %c1_i64, %c64_i64][%c0_i64, %c0_i64, %c0_i64]) {id = 0 : i64, metadata = @objFifo_in0} : memref<64xi32>
       aiex.ipu.dma_memcpy_nd(0, 0, %arg2[%c0_i64, %c0_i64, %c0_i64, %c0_i64][%c1_i64, %c1_i64, %c1_i64, %c64_i64][%c0_i64, %c0_i64, %c0_i64]) {id = 1 : i64, metadata = @objFifo_out0} : memref<64xi32>
-      aiex.ipu.sync {channel = 0 : i32, column = 0 : i32, column_num = 1 : i32, direction = 0 : i32, row = 0 : i32, row_num = 1 : i32}
       return
     }
 

test/ipu-xrt/add_one_objFifo/aie.mlir

@@ -46,7 +46,6 @@ module {
       %c64 = arith.constant 64 : i64
       aiex.ipu.dma_memcpy_nd (0, 0, %out[%c0,%c0,%c0,%c0][%c1,%c1,%c1,%c64][%c0,%c0,%c0]) { metadata = @objFifo_out0, id = 1 : i64 } : memref<64xi32>
       aiex.ipu.dma_memcpy_nd (0, 0, %in[%c0,%c0,%c0,%c0][%c1,%c1,%c1,%c64][%c0,%c0,%c0]) { metadata = @objFifo_in0, id = 0 : i64 } : memref<64xi32>
-      aiex.ipu.sync { column = 0 : i32, row = 0 : i32, direction = 0 : i32, channel = 0 : i32, column_num = 1 : i32, row_num = 1 : i32 }
       return
     }
   }

test/ipu-xrt/add_one_using_dma/aie.mlir

@@ -78,7 +78,6 @@ module {
       %c64_i64 = arith.constant 64 : i64
       aiex.ipu.dma_memcpy_nd(0, 0, %arg0[%c0_i64, %c0_i64, %c0_i64, %c0_i64] [%c1_i64, %c1_i64, %c1_i64, %c64_i64] [%c0_i64, %c0_i64, %c0_i64]) {id = 0 : i64, metadata = @objFifo_in0} : memref<64xi32>
       aiex.ipu.dma_memcpy_nd(0, 0, %arg2[%c0_i64, %c0_i64, %c0_i64, %c0_i64] [%c1_i64, %c1_i64, %c1_i64, %c64_i64] [%c0_i64, %c0_i64, %c0_i64]) {id = 1 : i64, metadata = @objFifo_out0} : memref<64xi32>
-      aiex.ipu.sync {channel = 0 : i32, column = 0 : i32, column_num = 1 : i32, direction = 0 : i32, row = 0 : i32, row_num = 1 : i32}
       return
     }
 

test/ipu-xrt/cascade_flows/aie.mlir

@@ -62,7 +62,6 @@ module {
       %c64 = arith.constant 64 : i64
       aiex.ipu.dma_memcpy_nd (0, 0, %out[%c0,%c0,%c0,%c0][%c1,%c1,%c1,%c64][%c0,%c0,%c0]) { metadata = @objFifo_out0, id = 1 : i64 } : memref<64xi32>
       aiex.ipu.dma_memcpy_nd (0, 0, %in[%c0,%c0,%c0,%c0][%c1,%c1,%c1,%c64][%c0,%c0,%c0]) { metadata = @objFifo_in0, id = 0 : i64 } : memref<64xi32>
-      aiex.ipu.sync { column = 0 : i32, row = 0 : i32, direction = 0 : i32, channel = 0 : i32, column_num = 1 : i32, row_num = 1 : i32 }
       return
     }
   }

test/ipu-xrt/matrix_multiplication_using_dma/aie.mlir

@@ -116,7 +116,6 @@ module {
       aiex.ipu.dma_memcpy_nd(0, 0, %arg1[%c0_i64, %c0_i64, %c0_i64, %c0_i64] [%c2_i64, %c4_i64, %c32_i64, %c32_i64] [%c32_i64, %c2048_i64, %c64_i64]) {id = 2 : i64, metadata = @inB} : memref<8192xi32>
       aiex.ipu.dma_memcpy_nd(0, 0, %arg0[%c0_i64, %c0_i64, %c0_i64, %c4096_i64] [%c2_i64, %c4_i64, %c64_i64, %c16_i64] [%c0_i64, %c16_i64, %c64_i64]) {id = 3 : i64, metadata = @inA} : memref<8192xi32>
       aiex.ipu.dma_memcpy_nd(0, 0, %arg1[%c0_i64, %c0_i64, %c0_i64, %c0_i64] [%c2_i64, %c4_i64, %c32_i64, %c32_i64] [%c32_i64, %c2048_i64, %c64_i64]) {id = 4 : i64, metadata = @inB} : memref<8192xi32>
-      aiex.ipu.sync {channel = 0 : i32, column = 0 : i32, column_num = 1 : i32, direction = 0 : i32, row = 0 : i32, row_num = 1 : i32}
       return
     }
 

test/ipu-xrt/two_col/aie.mlir

@@ -133,7 +133,6 @@ module {
       aiex.ipu.rtp_write(1, 5, 1, 0) { buffer_sym_name = "rtp3" }
       aiex.ipu.dma_memcpy_nd (0, 0, %out[%c0,%c0,%c0,%c0][%c1,%c1,%c1,%c2048][%c0,%c0,%c0]) { metadata = @objFifo_out0, id = 1 : i64 } : memref<2048xi32>
       aiex.ipu.dma_memcpy_nd (0, 0, %in[%c0,%c0,%c0,%c0][%c1,%c1,%c1,%c2048][%c0,%c0,%c0]) { metadata = @objFifo_in0, id = 0 : i64 } : memref<2048xi32>
-      aiex.ipu.sync { column = 0 : i32, row = 0 : i32, direction = 0 : i32, channel = 0 : i32, column_num = 1 : i32, row_num = 1 : i32 }
       return
     }
   }

test/ipu-xrt/vector_scalar_using_dma/aie.mlir

@@ -68,7 +68,6 @@ module {
       %c4096_i64 = arith.constant 4096 : i64
       aiex.ipu.dma_memcpy_nd(0, 0, %arg2[%c0_i64, %c0_i64, %c0_i64, %c0_i64] [%c1_i64, %c1_i64, %c1_i64, %c4096_i64] [%c0_i64, %c0_i64, %c0_i64]) {id = 0 : i64, metadata = @out} : memref<4096xi32>
       aiex.ipu.dma_memcpy_nd(0, 0, %arg0[%c0_i64, %c0_i64, %c0_i64, %c0_i64] [%c1_i64, %c1_i64, %c1_i64, %c4096_i64] [%c0_i64, %c0_i64, %c0_i64]) {id = 1 : i64, metadata = @in} : memref<4096xi32>
-      aiex.ipu.sync {channel = 0 : i32, column = 0 : i32, column_num = 1 : i32, direction = 0 : i32, row = 0 : i32, row_num = 1 : i32}
       return
     }