[MLIR][NVVM] Update PMEvent lowering to intrinsics

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

@@ -543,7 +543,7 @@ def NVVM_NanosleepOp : NVVM_Op<"nanosleep">,
 // NVVM Performance Monitor events
 //===----------------------------------------------------------------------===//
 
-def NVVM_PMEventOp : NVVM_PTXBuilder_Op<"pmevent">,
+def NVVM_PMEventOp : NVVM_Op<"pmevent">,
   Arguments<(ins OptionalAttr<I16Attr>:$maskedEventId, 
                  OptionalAttr<I32Attr>:$eventId)> {
   let summary = "Trigger one or more Performance Monitor events.";
@@ -561,20 +561,20 @@ def NVVM_PMEventOp : NVVM_PTXBuilder_Op<"pmevent">,
     [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#miscellaneous-instructions-pmevent)
   }];
 
-  string llvmBuilder = [{
-      llvm::Value *mId = builder.getInt16(* $maskedEventId);
-      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_pm_event_mask, {mId});
-  }];
-
   let assemblyFormat = "attr-dict (`id` `=` $eventId^)? (`mask` `=` $maskedEventId^)?";
+  let hasVerifier = 1;
 
   let extraClassDeclaration = [{
-    bool hasIntrinsic() { return !getEventId(); }
+    static mlir::NVVM::IDArgPair
+    getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
+                          llvm::IRBuilderBase& builder);
   }];
-  let extraClassDefinition = [{
-    std::string $cppClass::getPtx() { return std::string("pmevent %0;"); }
+
+  string llvmBuilder = [{
+    auto [id, args] = NVVM::PMEventOp::getIntrinsicIDAndArgs(
+                      *op, moduleTranslation, builder);
+    createIntrinsicCall(builder, id, args);
   }];
-  let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp

@@ -2476,6 +2476,25 @@ mlir::NVVM::IDArgPair NVVM::BarrierOp::getIntrinsicIDAndArgs(
   return {id, std::move(args)};
 }
 
+mlir::NVVM::IDArgPair
+PMEventOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
+                                 llvm::IRBuilderBase &builder) {
+  auto thisOp = cast<NVVM::PMEventOp>(op);
+  llvm::Type *i16Ty = llvm::Type::getInt16Ty(mt.getLLVMContext());
+
+  // With event-id, mask is generated as (1 << event-id)
+  llvm::Value *maskVal;
+  if (auto eventAttr = thisOp.getEventIdAttr()) {
+    uint16_t mask = static_cast<uint16_t>(1u << eventAttr.getInt());
+    maskVal = llvm::ConstantInt::get(i16Ty, mask);
+  } else {
+    maskVal =
+        llvm::ConstantInt::get(i16Ty, thisOp.getMaskedEventIdAttr().getValue());
+  }
+
+  return {llvm::Intrinsic::nvvm_pm_event_mask, {maskVal}};
+}
+
 mlir::NVVM::IDArgPair MBarrierInitOp::getIntrinsicIDAndArgs(
     Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
   auto thisOp = cast<NVVM::MBarrierInitOp>(op);

mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir

@@ -678,21 +678,6 @@ llvm.func @inline_ptx_multi_rw_r(%a : i32, %b : i32,  %rw_c : f32, %rw_d : f32)
    llvm.return %r5 : f32
 }
 
-
-// -----
-
-// CHECK-LABEL: @nvvm_pmevent
-llvm.func @nvvm_pmevent() {
-  // CHECK: %[[S0:.+]] = llvm.mlir.constant(10 : i32) : i32
-  // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "pmevent $0;", "n" %[[S0]] : (i32) -> ()
-
-  nvvm.pmevent id = 10
-  // CHECK: %[[S1:.+]] = llvm.mlir.constant(4 : i32) : i32
-  // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "pmevent $0;", "n" %[[S1]] : (i32) -> ()
-  nvvm.pmevent id = 4
-  llvm.return
-}
-
 // -----
 
 llvm.func @inline_ptx_pack_4i8(%src : vector<4xi8>,  %mask : i32, %zero: i32)  {

mlir/test/Target/LLVMIR/nvvm/pm_event.mlir

@@ -0,0 +1,23 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+llvm.func @nvvm_pmevent_mask() {
+  // CHECK-LABEL: define void @nvvm_pmevent_mask() {
+  // CHECK-NEXT: call void @llvm.nvvm.pm.event.mask(i16 15000)
+  // CHECK-NEXT: call void @llvm.nvvm.pm.event.mask(i16 4)
+  // CHECK-NEXT: ret void
+  // CHECK-NEXT: }
+  nvvm.pmevent mask = 15000
+  nvvm.pmevent mask = 4
+  llvm.return
+}
+
+llvm.func @nvvm_pmevent_id() {
+  // CHECK-LABEL: define void @nvvm_pmevent_id() {
+  // CHECK-NEXT: call void @llvm.nvvm.pm.event.mask(i16 1024)
+  // CHECK-NEXT: call void @llvm.nvvm.pm.event.mask(i16 16)
+  // CHECK-NEXT: ret void
+  // CHECK-NEXT: }
+  nvvm.pmevent id = 10
+  nvvm.pmevent id = 4
+  llvm.return
+}

mlir/test/Target/LLVMIR/nvvm/pm_event_invalid.mlir

@@ -0,0 +1,21 @@
+// RUN: mlir-translate -verify-diagnostics -split-input-file -mlir-to-llvmir %s
+
+llvm.func @pmevent_no_id() {
+  // expected-error @below {{either `id` or `mask` must be set}}
+  nvvm.pmevent 
+}
+
+// -----
+
+llvm.func @pmevent_bigger15() {
+  // expected-error @below {{`id` must be between 0 and 15}}
+  nvvm.pmevent id  = 16
+}
+
+// -----
+
+llvm.func @pmevent_many_ids() {
+  // expected-error @below {{`id` and `mask` cannot be set at the same time}}
+  nvvm.pmevent id = 1 mask = 1
+}
+

mlir/test/Target/LLVMIR/nvvmir-invalid.mlir

@@ -1,26 +1,5 @@
 // RUN: mlir-translate -verify-diagnostics -split-input-file -mlir-to-llvmir %s
 
-llvm.func @pmevent_no_id() {
-  // expected-error @below {{either `id` or `mask` must be set}}
-  nvvm.pmevent 
-}
-
-// -----
-
-llvm.func @pmevent_bigger15() {
-  // expected-error @below {{`id` must be between 0 and 15}}
-  nvvm.pmevent id  = 141
-}
-
-// -----
-
-llvm.func @pmevent_many_ids() {
-  // expected-error @below {{`id` and `mask` cannot be set at the same time}}
-  nvvm.pmevent id = 1 mask = 1
-}
-
-// -----
-
 llvm.func @kernel_func(%numberOfThreads : i32) {
   // expected-error @below {{'nvvm.barrier' op barrier id is missing, it should be set between 0 to 15}}
   nvvm.barrier number_of_threads = %numberOfThreads

mlir/test/Target/LLVMIR/nvvmir.mlir

@@ -903,17 +903,6 @@ llvm.func @nvvm_dot_accumulate_2way(%a: vector<2xi16>, %b: vector<4xi8>, %c: i32
 
 // -----
 
-// CHECK-LABEL: @nvvm_pmevent
-llvm.func @nvvm_pmevent() {
-  // CHECK: call void @llvm.nvvm.pm.event.mask(i16 15000)
-  nvvm.pmevent mask = 15000
-  // CHECK: call void @llvm.nvvm.pm.event.mask(i16 4)
-  nvvm.pmevent mask = 4
-  llvm.return
-}
-
-// -----
-
 // CHECK-LABEL: @nanosleep
 llvm.func @nanosleep(%duration: i32) {
   // CHECK: call void @llvm.nvvm.nanosleep(i32 %{{.*}})