[Mlir-commits] [mlir] [mlir][GPU] Refactor, improve constant size information handling (PR #186907)

[Fabian Mora]

================
@@ -54,20 +55,53 @@ struct GpuSubgroupIdRewriter final : OpRewritePattern<gpu::SubgroupIdOp> {
     Location loc = op->getLoc();
     Type indexType = rewriter.getIndexType();
 
-    Value dimX = gpu::BlockDimOp::create(rewriter, loc, gpu::Dimension::x);
-    Value dimY = gpu::BlockDimOp::create(rewriter, loc, gpu::Dimension::y);
-    Value tidX = gpu::ThreadIdOp::create(rewriter, loc, gpu::Dimension::x);
-    Value tidY = gpu::ThreadIdOp::create(rewriter, loc, gpu::Dimension::y);
-    Value tidZ = gpu::ThreadIdOp::create(rewriter, loc, gpu::Dimension::z);
+    auto asMaybeIndexAttr = [&](std::optional<uint32_t> bound) -> IntegerAttr {
+      if (!bound)
+        return IntegerAttr();
+      return IntegerAttr::get(
+          indexType, static_cast<int64_t>(static_cast<uint64_t>(*bound)));
+    };
 
+    IntegerAttr maybeKnownDimX =
+        asMaybeIndexAttr(gpu::getKnownDimensionSizeAround(
+            op, gpu::DimensionKind::Block, gpu::Dimension::x));
+    IntegerAttr maybeKnownDimY =
+        asMaybeIndexAttr(gpu::getKnownDimensionSizeAround(
+            op, gpu::DimensionKind::Block, gpu::Dimension::y));
+    IntegerAttr maybeKnownDimZ =
+        asMaybeIndexAttr(gpu::getKnownDimensionSizeAround(
+            op, gpu::DimensionKind::Block, gpu::Dimension::z));
+
+    Value dimX, dimY;
+    if (maybeKnownDimX)
+      dimX = arith::ConstantOp::create(rewriter, loc, maybeKnownDimX);
+    else
+      dimX = gpu::BlockDimOp::create(rewriter, loc, gpu::Dimension::x);
+    if (maybeKnownDimY)
+      dimY = arith::ConstantOp::create(rewriter, loc, maybeKnownDimY);
+    else
+      dimY = gpu::BlockDimOp::create(rewriter, loc, gpu::Dimension::y);
+
+    Value tidX = gpu::ThreadIdOp::create(rewriter, loc, gpu::Dimension::x,
+                                         maybeKnownDimX);
+    Value tidY = gpu::ThreadIdOp::create(rewriter, loc, gpu::Dimension::y,
+                                         maybeKnownDimY);
+    Value tidZ = gpu::ThreadIdOp::create(rewriter, loc, gpu::Dimension::z,
+                                         maybeKnownDimZ);
+
+    // Block dimensions don't exceed a signed int32_t maximum, and neither does
+    // their product, on any realistic hardware, nor would any targets compile
+    // with index < 32 bits, so we can assert no overflow.
+    auto flags =
+        arith::IntegerOverflowFlags::nsw | arith::IntegerOverflowFlags::nuw;
     Value dimYxIdZ =
-        arith::MulIOp::create(rewriter, loc, indexType, dimY, tidZ);
+        arith::MulIOp::create(rewriter, loc, indexType, dimY, tidZ, flags);
     Value dimYxIdZPlusIdY =
-        arith::AddIOp::create(rewriter, loc, indexType, dimYxIdZ, tidY);
-    Value dimYxIdZPlusIdYTimesDimX =
-        arith::MulIOp::create(rewriter, loc, indexType, dimX, dimYxIdZPlusIdY);
+        arith::AddIOp::create(rewriter, loc, indexType, dimYxIdZ, tidY, flags);
+    Value dimYxIdZPlusIdYTimesDimX = arith::MulIOp::create(
+        rewriter, loc, indexType, dimX, dimYxIdZPlusIdY, flags);
     Value idXPlusDimYxIdZPlusIdYTimesDimX = arith::AddIOp::create(
-        rewriter, loc, indexType, tidX, dimYxIdZPlusIdYTimesDimX);
+        rewriter, loc, indexType, tidX, dimYxIdZPlusIdYTimesDimX, flags);
----------------
fabianmcg wrote:

Can we split adding `nsw`, `nuw` into a different change? Seems unrelated to the description.

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