diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -532,23 +532,19 @@
                  parser.parseOptionalAttrDict(result.attributes));
 }
 
-/// Simplify the gpu.launch when the range of the thread and block IDs is
+/// Simplify the gpu.launch when the range of a thread or block ID is
 /// trivially known to be one.
 struct FoldLaunchArguments : public OpRewritePattern<LaunchOp> {
   using OpRewritePattern<LaunchOp>::OpRewritePattern;
   LogicalResult matchAndRewrite(LaunchOp op,
                                 PatternRewriter &rewriter) const override {
-    auto isTriviallyOne = [](Value size) {
-      IntegerAttr cst;
-      return matchPattern(size, m_Constant(&cst)) && cst.getInt() == 1;
-    };
-
     // If the range implies a single value for `id`, replace `id`'s uses by
     // zero.
     Value zero;
     bool simplified = false;
     auto constPropIdUses = [&](Value id, Value size) {
-      if (!isTriviallyOne(size))
+      // Check if size is trivially one.
+      if (!matchPattern(size, m_One()))
         return;
       if (!simplified) {
         // Create a zero value the first time.