[Mlir-commits] [mlir] [mlir] Let GPU ID bounds work on any FunctionOpInterfaces (PR #95166)

[Mehdi Amini]

================
@@ -74,35 +74,41 @@ GPUFuncOpLowering::matchAndRewrite(gpu::GPUFuncOp gpuFuncOp, OpAdaptor adaptor,
         attr.getName() ==
             gpu::GPUFuncOp::getNumWorkgroupAttributionsAttrName() ||
         attr.getName() == gpuFuncOp.getWorkgroupAttribAttrsAttrName() ||
-        attr.getName() == gpuFuncOp.getPrivateAttribAttrsAttrName())
+        attr.getName() == gpuFuncOp.getPrivateAttribAttrsAttrName() ||
+        attr.getName() == gpuFuncOp.getKnownBlockSizeAttrName() ||
+        attr.getName() == gpuFuncOp.getKnownGridSizeAttrName())
       continue;
     if (attr.getName() == gpuFuncOp.getArgAttrsAttrName()) {
       argAttrs = gpuFuncOp.getArgAttrsAttr();
       continue;
     }
     attributes.push_back(attr);
   }
+
+  DenseI32ArrayAttr knownBlockSize = gpuFuncOp.getKnownBlockSizeAttr();
+  DenseI32ArrayAttr knownGridSize = gpuFuncOp.getKnownGridSizeAttr();
+  // Ensure we don't lose information if the function is lowered before its
+  // surrounding context.
+  if (knownBlockSize)
+    attributes.emplace_back(
+        rewriter.getStringAttr(
+            gpu::GPUDialect::KnownBlockSizeAttrHelper::getNameStr()),
+        knownBlockSize);
+  if (knownGridSize)
+    attributes.emplace_back(
+        rewriter.getStringAttr(
+            gpu::GPUDialect::KnownGridSizeAttrHelper::getNameStr()),
+        knownGridSize);
----------------
joker-eph wrote:

```suggestion
  auto *gpuDialect = cast<GPUDialect>(gpuFuncOp.getDialect());
  if (knownBlockSize)
    attributes.emplace_back(gpuDialect.getKnownBlockSizeAttrHelper().getName()),
        knownBlockSize);
  if (knownGridSize)
    attributes.emplace_back(gpuDialect.getKnownGridSizeAttrHelper().getName()),
        knownGridSize);
```

The point of the ODS discardableAttrs support is (also) about avoiding to recreate StringAttr constantly :)


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