[Mlir-commits] [mlir] [mlir][GPU] Improve handling of GPU bounds (PR #95166)

[Mehdi Amini]

================
@@ -349,27 +348,32 @@ static void populateOpPatterns(LLVMTypeConverter &converter,
 void mlir::populateGpuToROCDLConversionPatterns(
     LLVMTypeConverter &converter, RewritePatternSet &patterns,
     mlir::gpu::amd::Runtime runtime) {
+  using gpu::index_lowering::IndexKind;
+  using gpu::index_lowering::IntrType;
   using mlir::gpu::amd::Runtime;
-
   populateWithGenerated(patterns);
-  patterns
-      .add<GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, ROCDL::ThreadIdXOp,
-                                       ROCDL::ThreadIdYOp, ROCDL::ThreadIdZOp>>(
-          converter, gpu::GPUFuncOp::getKnownBlockSizeAttrName());
-  patterns.add<GPUIndexIntrinsicOpLowering<
+  patterns.add<
+      gpu::index_lowering::OpLowering<gpu::ThreadIdOp, ROCDL::ThreadIdXOp,
+                                      ROCDL::ThreadIdYOp, ROCDL::ThreadIdZOp>>(
+      converter, IndexKind::Block, IntrType::Id);
+  patterns.add<gpu::index_lowering::OpLowering<
       gpu::BlockIdOp, ROCDL::BlockIdXOp, ROCDL::BlockIdYOp, ROCDL::BlockIdZOp>>(
-      converter, gpu::GPUFuncOp::getKnownGridSizeAttrName());
-  patterns
-      .add<GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, ROCDL::BlockDimXOp,
-                                       ROCDL::BlockDimYOp, ROCDL::BlockDimZOp>,
-           GPUIndexIntrinsicOpLowering<gpu::GridDimOp, ROCDL::GridDimXOp,
-                                       ROCDL::GridDimYOp, ROCDL::GridDimZOp>,
-           GPUReturnOpLowering>(converter);
+      converter, IndexKind::Grid, IntrType::Id);
+  patterns.add<
+      gpu::index_lowering::OpLowering<gpu::BlockDimOp, ROCDL::BlockDimXOp,
+                                      ROCDL::BlockDimYOp, ROCDL::BlockDimZOp>>(
+      converter, IndexKind::Block, IntrType::Dim);
+  patterns.add<gpu::index_lowering::OpLowering<
+      gpu::GridDimOp, ROCDL::GridDimXOp, ROCDL::GridDimYOp, ROCDL::GridDimZOp>>(
+      converter, IndexKind::Grid, IntrType::Dim);
+  patterns.add<GPUReturnOpLowering>(converter);
   patterns.add<GPUFuncOpLowering>(
       converter,
       /*allocaAddrSpace=*/ROCDL::ROCDLDialect::kPrivateMemoryAddressSpace,
       /*workgroupAddrSpace=*/ROCDL::ROCDLDialect::kSharedMemoryAddressSpace,
-      ROCDL::ROCDLDialect::KernelAttrHelper(&converter.getContext()).getName());
+      ROCDL::ROCDLDialect::KernelAttrHelper(&converter.getContext()).getName(),
+      ROCDL::ROCDLDialect::ReqdWorkGroupSizeAttrHelper(&converter.getContext())
+          .getName());
----------------
joker-eph wrote:

Can you get a handle to the dialect first and then use it to get the various names?
```
      auto *rocdlDialect = converter.getContext().getLoadedDialect<ROCDLDiale```

https://github.com/llvm/llvm-project/pull/95166