updates

Author: lialan

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

@@ -216,7 +216,10 @@ def ROCDL_BlockIdXOp : ROCDL_SpecialIdRegisterOp<"workgroup.id.x">;
 def ROCDL_BlockIdYOp : ROCDL_SpecialIdRegisterOp<"workgroup.id.y">;
 def ROCDL_BlockIdZOp : ROCDL_SpecialIdRegisterOp<"workgroup.id.z">;
 
-def ROCDL_WavefrontSizeOp : ROCDL_SpecialIdRegisterOp<"wavefrontsize">;
+def ROCDL_WavefrontSizeOp : ROCDL_IntrPure1Op<"wavefrontsize">,
+    Arguments<(ins OptionalAttr<IndexAttr>:$upper_bound)> {
+  let assemblyFormat = "(`upper_bound` $upper_bound^)? attr-dict `:` type($res)";
+}
 
 //===----------------------------------------------------------------------===//
 // Thread range and Block range

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

@@ -138,7 +138,7 @@ struct GPUSubgroupSizeOpToROCDL : ConvertOpToLLVMPattern<gpu::SubgroupSizeOp> {
   matchAndRewrite(gpu::SubgroupSizeOp op, gpu::SubgroupSizeOp::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Value wavefrontOp = rewriter.create<ROCDL::WavefrontSizeOp>(
-        op.getLoc(), IntegerType::get(rewriter.getContext(), 32));
+        op.getLoc(), rewriter.getI32Type(), op.getUpperBoundAttr());
     wavefrontOp = truncOrExtToLLVMType(rewriter, op.getLoc(), wavefrontOp,
                                        *getTypeConverter());
     rewriter.replaceOp(op, {wavefrontOp});

mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir

@@ -11,7 +11,7 @@ gpu.module @test_module {
   func.func @gpu_index_ops()
       -> (index, index, index, index, index, index,
           index, index, index, index, index, index,
-          index, index) {
+          index, index, index) {
     // CHECK32-NOT: = llvm.sext %{{.*}} : i32 to i64
 
     // CHECK: rocdl.workitem.id.x : i32
@@ -63,12 +63,16 @@ gpu.module @test_module {
     // CHECK: = llvm.sext %{{.*}} : i32 to i64
     %subgroupSize = gpu.subgroup_size : index
 
+    // CHECK: = rocdl.wavefrontsize upper_bound 64 : i32
+    // CHECK: = llvm.sext %{{.*}} : i32 to i64
+    %subgroupSize2 = gpu.subgroup_size upper_bound 64 : index
+
     func.return %tIdX, %tIdY, %tIdZ, %bDimX, %bDimY, %bDimZ,
                %bIdX, %bIdY, %bIdZ, %gDimX, %gDimY, %gDimZ,
-               %laneId, %subgroupSize
+               %laneId, %subgroupSize, %subgroupSize2
         : index, index, index, index, index, index,
           index, index, index, index, index, index,
-          index, index
+          index, index, index
   }
 }
 

mlir/test/Target/LLVMIR/rocdl.mlir

@@ -33,9 +33,12 @@ llvm.func @rocdl_special_regs() -> i32 {
   // CHECK: call range(i64 1, 65) i64 @__ockl_get_local_size(i32 0)
   %14 = rocdl.workgroup.dim.x range <i32, 1, 65> : i64
 
-  // CHECK: call i32 $llvm.amdgcn.wavefrontsize()
+  // CHECK: call i32 @llvm.amdgcn.wavefrontsize()
   %15 = rocdl.wavefrontsize : i32
 
+  // CHECK: call i32 @llvm.amdgcn.wavefrontsize()
+  %16 = rocdl.wavefrontsize upper_bound 32 : i32
+
   llvm.return %1 : i32
 }