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

[Krzysztof Drewniak]

================
@@ -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);
----------------
krzysz00 wrote:

... Huh, I knew I was missing some obvious way to do that, thanks!

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