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

[Krzysztof Drewniak]

================
@@ -73,12 +85,16 @@ static std::optional<uint64_t> getKnownLaunchDim(Op op, LaunchDims type) {
       return value.getZExtValue();
   }
 
-  if (auto func = op->template getParentOfType<GPUFuncOp>()) {
+  if (auto func = op->template getParentOfType<FunctionOpInterface>()) {
     switch (type) {
     case LaunchDims::Block:
-      return llvm::transformOptional(func.getKnownBlockSize(dim), zext);
+      return llvm::transformOptional(
+          getKnownLaunchAttr(func, GPUFuncOp::getKnownBlockSizeAttrName(), dim),
+          zext);
     case LaunchDims::Grid:
-      return llvm::transformOptional(func.getKnownGridSize(dim), zext);
+      return llvm::transformOptional(
+          getKnownLaunchAttr(func, GPUFuncOp::getKnownGridSizeAttrName(), dim),
+          zext);
----------------
krzysz00 wrote:

Though I think we just drifted from functions to modules.

As far as I'm concerned, the only difference between

```
gpu.func @f(...) known_block_sizes = [64, 2, 1] {
  ...
}
```

and
```
func.func @f(...) attributes {gpu.known_block_sizes = [64, 2, 1]} {
  ...
}
```
(or even `my_custom.func` in the above) is the somewhat higher risk of discardability on the optimization hint.

This PR's trying to generalize some already general handling for this particular hint so that we're not locking downstream into needing a very particular infrastructure (that might not compose with other compilation infrastructures because there's an extra level of module involved) just to get `ThreadIdOp::inferIntRanges()` working.

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