Another update

Author: lialan

mlir/include/mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h

@@ -20,6 +20,10 @@ class RewritePatternSet;
 template <typename OpT>
 class OperationPass;
 
+namespace amdgpu {
+struct Chipset;
+} // namespace amdgpu
+
 namespace gpu {
 class GPUModuleOp;
 } // namespace gpu
@@ -32,7 +36,8 @@ class GPUModuleOp;
 /// The resulting pattern set should be run over a gpu.module op
 void populateGpuToROCDLConversionPatterns(const LLVMTypeConverter &converter,
                                           RewritePatternSet &patterns,
-                                          gpu::amd::Runtime runtime);
+                                          gpu::amd::Runtime runtime,
+                                          amdgpu::Chipset chipset);
 
 /// Configure target to convert from the GPU dialect to ROCDL.
 void configureGpuToROCDLConversionLegality(ConversionTarget &target);

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

@@ -134,14 +134,21 @@ struct GPULaneIdOpToROCDL : ConvertOpToLLVMPattern<gpu::LaneIdOp> {
 
 struct GPUSubgroupSizeOpToROCDL : ConvertOpToLLVMPattern<gpu::SubgroupSizeOp> {
   using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
+
+  GPUSubgroupSizeOpToROCDL(const LLVMTypeConverter &converter,
+                           amdgpu::Chipset chipset)
+      : ConvertOpToLLVMPattern<gpu::SubgroupSizeOp>(converter),
+        chipset(chipset) {}
+
   LogicalResult
   matchAndRewrite(gpu::SubgroupSizeOp op, gpu::SubgroupSizeOp::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     LLVM::ConstantRangeAttr bounds = nullptr;
+    bool isBeforeGfx10 = chipset.majorVersion < 10;
     if (auto upperBoundAttr = op.getUpperBoundAttr()) {
       bounds = rewriter.getAttr<LLVM::ConstantRangeAttr>(
-          /*bitWidth=*/32, /*lower=*/32,
-          /*upper=*/op.getUpperBoundAttr().getInt());
+          /*bitWidth=*/32, /*lower=*/isBeforeGfx10 ? 64 : 32,
+          /*upper=*/op.getUpperBoundAttr().getInt() + 1);
     }
     Value wavefrontOp = rewriter.create<ROCDL::WavefrontSizeOp>(
         op.getLoc(), rewriter.getI32Type(), bounds);
@@ -150,6 +157,8 @@ struct GPUSubgroupSizeOpToROCDL : ConvertOpToLLVMPattern<gpu::SubgroupSizeOp> {
     rewriter.replaceOp(op, {wavefrontOp});
     return success();
   }
+
+  const amdgpu::Chipset chipset;
 };
 
 struct GPUShuffleOpLowering : public ConvertOpToLLVMPattern<gpu::ShuffleOp> {
@@ -358,7 +367,8 @@ struct LowerGpuOpsToROCDLOpsPass final
 
     populateAMDGPUToROCDLConversionPatterns(converter, llvmPatterns,
                                             *maybeChipset);
-    populateGpuToROCDLConversionPatterns(converter, llvmPatterns, runtime);
+    populateGpuToROCDLConversionPatterns(converter, llvmPatterns, runtime,
+                                         *maybeChipset);
     configureGpuToROCDLConversionLegality(target);
     if (failed(applyPartialConversion(m, target, std::move(llvmPatterns))))
       signalPassFailure();
@@ -406,7 +416,7 @@ void mlir::configureGpuToROCDLConversionLegality(ConversionTarget &target) {
 
 void mlir::populateGpuToROCDLConversionPatterns(
     const LLVMTypeConverter &converter, RewritePatternSet &patterns,
-    mlir::gpu::amd::Runtime runtime) {
+    mlir::gpu::amd::Runtime runtime, amdgpu::Chipset chipset) {
   using gpu::index_lowering::IndexKind;
   using gpu::index_lowering::IntrType;
   using mlir::gpu::amd::Runtime;
@@ -447,6 +457,7 @@ void mlir::populateGpuToROCDLConversionPatterns(
   patterns
       .add<GPUShuffleOpLowering, GPULaneIdOpToROCDL, GPUSubgroupSizeOpToROCDL>(
           converter);
+  patterns.add<GPUSubgroupSizeOpToROCDL>(converter, chipset);
 
   populateMathToROCDLConversionPatterns(converter, patterns);
 }

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

@@ -63,7 +63,7 @@ gpu.module @test_module {
     // CHECK: = llvm.sext %{{.*}} : i32 to i64
     %subgroupSize = gpu.subgroup_size : index
 
-    // CHECK: = rocdl.wavefrontsize range <i32, 32, 64> : i32
+    // CHECK: = rocdl.wavefrontsize range <i32, 64, 65> : i32
     // CHECK: = llvm.sext %{{.*}} : i32 to i64
     %subgroupSize2 = gpu.subgroup_size upper_bound 64 : index
 

mlir/test/Target/LLVMIR/rocdl.mlir

@@ -36,8 +36,8 @@ llvm.func @rocdl_special_regs() -> i32 {
   // CHECK: call i32 @llvm.amdgcn.wavefrontsize()
   %15 = rocdl.wavefrontsize : i32
 
-  // CHECK: call range(i32 32, 64) i32 @llvm.amdgcn.wavefrontsize()
-  %16 = rocdl.wavefrontsize range <i32, 32, 64> : i32
+  // CHECK: call range(i32 32, 65) i32 @llvm.amdgcn.wavefrontsize()
+  %16 = rocdl.wavefrontsize range <i32, 32, 65> : i32
 
   llvm.return %1 : i32
 }