[MLIR] NFC. gpu.launch op argument const folder cleanup

[bondhugula]

Summary

NFC updates to gpu.launch op argument const folder.

[bondhugula]

Addressed NFC comments from @rriddle and @herhut on D109994. Self-accepting.

[herhut]

Thanks!

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

static ParseResult parseLaunchOp(OpAsmParser &parser, OperationState &result) {
   Type index = parser.getBuilder().getIndexType();
   SmallVector<Type, LaunchOp::kNumConfigRegionAttributes> dataTypes(
       LaunchOp::kNumConfigRegionAttributes, index);
   Region *body = result.addRegion();
   return failure(parser.parseRegion(*body, regionArgs, dataTypes) ||
                  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.
         OpBuilder::InsertionGuard guard(rewriter);
         rewriter.setInsertionPointToStart(&op.body().front());
         zero = rewriter.create<ConstantIndexOp>(op.getLoc(), /*value=*/0);
       }
       id.replaceAllUsesWith(zero);