Add lowering of aie.event for aie2 and aie2p (#2652)

Co-authored-by: Sylvain Noiry <[email protected]>
Co-authored-by: Joseph Melber <[email protected]>

Author: 3 people

lib/Dialect/AIE/Transforms/AIECoreToStandard.cpp

@@ -92,6 +92,7 @@ static auto getAIE2Intrinsics(OpBuilder &builder) {
   Type accType = VectorType::get({16}, int32Type);
   IntrinsicDecls functions = {
       {"debug_i32", {int32Type}, {}},
+      {"llvm.aie2.event", {int32Type}, {}},
       {"llvm.aie2.put.ms", {int32Type, int32Type}, {}}, //(%value, %tlast) -> ()
       {"llvm.aie2.get.ss", {}, {int32Type, int32Type}}, //() -> (%value, %tlast)
       {"llvm.aie2.mcd.write.vec",
@@ -115,6 +116,7 @@ static auto getAIE2pIntrinsics(OpBuilder &builder) {
   Type accType = VectorType::get({16}, int32Type);
   IntrinsicDecls functions = {
       {"debug_i32", {int32Type}, {}},
+      {"llvm.aie2p.event", {int32Type}, {}},
       {"llvm.aie2p.put.ms",
        {int32Type, int32Type},
        {}}, //(%value, %tlast) -> ()
@@ -577,13 +579,33 @@ struct AIEEventOpToStdLowering : OpConversionPattern<EventOp> {
   LogicalResult
   matchAndRewrite(EventOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    std::string funcName = "llvm.aie.event" + std::to_string(op.getVal());
+    auto device = op->getParentOfType<DeviceOp>();
+    std::string funcName;
+    SmallVector<Value, 1> args;
+    switch (device.getTargetModel().getTargetArch()) {
+    case AIEArch::AIE1:
+      funcName = "llvm.aie.event" + std::to_string(op.getVal());
+      break;
+    case AIEArch::AIE2:
+      funcName = "llvm.aie2.event";
+      args.push_back(rewriter.create<arith::ConstantOp>(
+          op.getLoc(), rewriter.getI32Type(),
+          rewriter.getI32IntegerAttr(op.getVal())));
+      break;
+    case AIEArch::AIE2p:
+      funcName = "llvm.aie2p.event";
+      args.push_back(rewriter.create<arith::ConstantOp>(
+          op.getLoc(), rewriter.getI32Type(),
+          rewriter.getI32IntegerAttr(op.getVal())));
+      break;
+    default:
+      return op->emitOpError("Unsupported AIEArch for EventOp lowering");
+    }
     auto eventFunc = module.lookupSymbol<func::FuncOp>(funcName);
     if (!eventFunc)
       return op.emitOpError("Could not find the intrinsic function ")
              << funcName;
-    rewriter.create<func::CallOp>(rewriter.getUnknownLoc(), eventFunc,
-                                  ValueRange({}));
+    rewriter.create<func::CallOp>(rewriter.getUnknownLoc(), eventFunc, args);
     rewriter.eraseOp(op);
     return success();
   }

test/lower-to-standard/lower_event.mlir

@@ -8,11 +8,12 @@
 //
 //===----------------------------------------------------------------------===//
 
-// RUN: aie-opt --aie-standard-lowering %s | FileCheck %s
+// RUN: aie-opt --aie-standard-lowering --split-input-file %s | FileCheck %s
 
+// CHECK-LABEL: test_aie1
 // CHECK: call @llvm.aie.event0()
 // CHECK: call @llvm.aie.event1()
-module @test {
+module @test_aie1 {
  aie.device(xcvc1902) {
   %tile11 = aie.tile(1, 1)
   %core11 = aie.core(%tile11) {
@@ -22,3 +23,35 @@ module @test {
   }
  }
 }
+
+// -----
+
+// CHECK-LABEL: test_aie2
+// CHECK: call @llvm.aie2.event(%{{.*}})
+// CHECK: call @llvm.aie2.event(%{{.*}})
+module @test_aie2 {
+ aie.device(npu1) {
+  %tile02 = aie.tile(0, 2)
+  %core02 = aie.core(%tile02) {
+    aie.event(0)
+    aie.event(1)
+    aie.end
+  }
+ }
+}
+
+// -----
+
+// CHECK-LABEL: test_aie2p
+// CHECK: call @llvm.aie2p.event(%{{.*}})
+// CHECK: call @llvm.aie2p.event(%{{.*}})
+module @test_aie2p {
+ aie.device(npu2) {
+  %tile02 = aie.tile(0, 2)
+  %core02 = aie.core(%tile02) {
+    aie.event(0)
+    aie.event(1)
+    aie.end
+  }
+ }
+}