[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:

I've found at least one downstream context (IREE, where I want to use `-int-range-optimizations`) that has decided the answer is "no" (and even in rocMLIR, we keep our GPU-specific code in a `func.func` until relatively late in our pipeline).

To try and cook up an argument for no, I'd say that GPU-specific code can arise from some notional `-specialize-to-target` pass that takes a `func.func` (or some other `*.func`) and makes clones of it for different targets. Forcing such a pass to change the function operation (and, worse, the nesting structure, because a gpu.func must be in a gpu.module, which is isolated from above, so you'd have to copy any referenced globals in etc) just to get access to `gpu.thread_id` seems to be to be a rather strange decision.

That is, the state of the world is "upstream MLIR provides a GPU compilation infrastructure that's intended to be used as a whole, including gpu.module/gpu.func. However, the wrappers around GPU primitives (like gpu.thread_id or the matrix ops) can be used independently of that infrastructure"

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