Change waitcnt to take an attribute because it's an immarg in LLVM

Author: krzysz00

mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td

@@ -196,10 +196,10 @@ def ROCDL_GridDimZOp : ROCDL_DeviceFunctionOp<"grid.dim.z",
 
 // Emits the waintcnt instruction. The bitfield's semantics depend
 // on the target chipset
-def ROCDL_WaitcntOp : ROCDL_Op<"waitcnt">, Arguments<(ins I32:$bitfield)> {
+def ROCDL_WaitcntOp : ROCDL_Op<"waitcnt">, Arguments<(ins I32Attr:$bitfield)> {
   string llvmBuilder = [{
     createIntrinsicCall(builder, llvm::Intrinsic::amdgcn_s_waitcnt,
-      {$bitfield});
+      {builder.getInt32($bitfield)});
   }];
   let assemblyFormat = "attr-dict $bitfield";
 }

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

@@ -316,8 +316,7 @@ struct LDSBarrierOpLowering : public ConvertOpToLLVMPattern<LDSBarrierOp> {
              << chipset.majorVersion;
 
     Location loc = op->getLoc();
-    Value constant = createI32Constant(rewriter, loc, ldsOnlyBits);
-    rewriter.create<ROCDL::WaitcntOp>(loc, constant);
+    rewriter.create<ROCDL::WaitcntOp>(loc, ldsOnlyBits);
     rewriter.replaceOpWithNewOp<ROCDL::SBarrierOp>(op);
     return success();
   }

mlir/test/Conversion/AMDGPUToROCDL/amdgpu-to-rocdl.mlir

@@ -217,11 +217,9 @@ func.func @amdgpu_raw_buffer_atomic_cmpswap_i64(%src : i64, %cmp : i64, %buf : m
 func.func @lds_barrier() {
   // GFX908: llvm.inline_asm has_side_effects asm_dialect = att
   // GFX908-SAME: ";;;WARNING: BREAKS DEBUG WATCHES\0As_waitcnt lgkmcnt(0)\0As_barrier"
-  // GFX90A:  %[[cst:.*]] = llvm.mlir.constant(-7937 : i32) : i32
-  // GFX90A: rocdl.waitcnt %[[cst]]
+  // GFX90A: rocdl.waitcnt -7937
   // GFX90A-NEXT: rocdl.s.barrier
-  // GFX10:  %[[cst:.*]] = llvm.mlir.constant(-16129 : i32) : i32
-  // GFX10:  rocdl.waitcnt %[[cst]]
+  // GFX10:  rocdl.waitcnt -16129
   // GFX10-NEXT: rocdl.s.barrier
   // GFX11:  llvm.inline_asm has_side_effects asm_dialect = att
   // GFX11-SAME: ";;;WARNING: BREAKS DEBUG WATCHES\0As_waitcnt lgkmcnt(0)\0As_barrier"

mlir/test/Dialect/LLVMIR/rocdl.mlir

@@ -363,10 +363,10 @@ llvm.func @rocdl_8bit_floats(%source: i32, %stoch: i32) -> i32 {
   llvm.return %source5 : i32
 }
 
-llvm.func @rocdl.waitcnt(%arg0 : i32) {
+llvm.func @rocdl.waitcnt() {
   // CHECK-LABEL: rocdl.waitcnt
-  // CHECK: rocdl.waitcnt
-  rocdl.waitcnt %arg0
+  // CHECK: rocdl.waitcnt 0
+  rocdl.waitcnt 0
   llvm.return
 }
 

mlir/test/Target/LLVMIR/rocdl.mlir

@@ -91,8 +91,7 @@ llvm.func @rocdl.bpermute(%src : i32) -> i32 {
 llvm.func @rocdl.waitcnt() {
   // CHECK-LABEL: rocdl.waitcnt
   // CHECK-NEXT: call void @llvm.amdgcn.s.waitcnt(i32 0)
-  %0 = llvm.mlir.constant(0 : i32) : i32
-  rocdl.waitcnt %0
+  rocdl.waitcnt 0
   llvm.return
 }