[CIR][NFC] Generalize IdiomRecognizer (#1484)

The comments suggested that we should use TableGen to generate the
recognizing functions. However, I think templates might be more suitable
for generating them -- and I can't find any existing TableGen backends
that let us generate arbitrary functions.

My choice of design is to offer a template to match standard library
functions:
```cpp
// matches std::find with 3 arguments, and raise it into StdFindOp
StdRecognizer<3, StdFindOp, StdFuncsID::Find>
```
I have to use a TableGen'd enum to map names to IDs, as we can't pass
string literals to template arguments easily in C++17.

This also constraints design of future `StdXXXOp`s: they must take
operands the same way of StdFindOp, where the first one is the original
function, and the rest are function arguments.

I'm not sure if this approach is the best way. Please tell me if you
have concerns or any alternative ways.

Author: AdUhTkJm

clang/include/clang/CIR/Dialect/IR/CIROps.td

@@ -4700,72 +4700,6 @@ def FrameAddrOp : FuncAddrBuiltinOp<"frame_address"> {
   }];
 }
 
-//===----------------------------------------------------------------------===//
-// StdFindOp
-//===----------------------------------------------------------------------===//
-
-def StdFindOp : CIR_Op<"std.find", [SameFirstSecondOperandAndResultType]> {
-  let arguments = (ins FlatSymbolRefAttr:$original_fn,
-                       CIR_AnyType:$first,
-                       CIR_AnyType:$last,
-                       CIR_AnyType:$pattern);
-  let summary = "std:find()";
-  let results = (outs CIR_AnyType:$result);
-
-  let description = [{
-    Search for `pattern` in data range from `first` to `last`. This currently
-    maps to only one form of `std::find`. The `original_fn` operand tracks the
-    mangled named that can be used when lowering to a `cir.call`.
-
-    Example:
-
-    ```mlir
-    ...
-    %result = cir.std.find(@original_fn,
-                           %first : !T, %last : !T, %pattern : !P) -> !T
-    ```
-  }];
-
-  let assemblyFormat = [{
-    `(`
-      $original_fn
-      `,` $first `:` type($first)
-      `,` $last `:` type($last)
-      `,` $pattern `:` type($pattern)
-    `)` `->` type($result) attr-dict
-  }];
-  let hasVerifier = 0;
-}
-
-//===----------------------------------------------------------------------===//
-// IterBegin/End
-//===----------------------------------------------------------------------===//
-
-def IterBeginOp : CIR_Op<"iterator_begin"> {
-  let arguments = (ins FlatSymbolRefAttr:$original_fn, CIR_AnyType:$container);
-  let summary = "Returns an iterator to the first element of a container";
-  let results = (outs CIR_AnyType:$result);
-  let assemblyFormat = [{
-    `(`
-      $original_fn `,` $container `:` type($container)
-    `)` `->` type($result) attr-dict
-  }];
-  let hasVerifier = 0;
-}
-
-def IterEndOp : CIR_Op<"iterator_end"> {
-  let arguments = (ins FlatSymbolRefAttr:$original_fn, CIR_AnyType:$container);
-  let summary = "Returns an iterator to the element following the last element"
-                " of a container";
-  let results = (outs CIR_AnyType:$result);
-  let assemblyFormat = [{
-    `(`
-      $original_fn `,` $container `:` type($container)
-    `)` `->` type($result) attr-dict
-  }];
-  let hasVerifier = 0;
-}
-
 //===----------------------------------------------------------------------===//
 // Floating Point Ops
 //===----------------------------------------------------------------------===//
@@ -5755,4 +5689,10 @@ def SignBitOp : CIR_Op<"signbit", [Pure]> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Standard library function calls
+//===----------------------------------------------------------------------===//
+
+include "clang/CIR/Dialect/IR/CIRStdOps.td"
+
 #endif // LLVM_CLANG_CIR_DIALECT_IR_CIROPS

clang/include/clang/CIR/Dialect/IR/CIRStdOps.td

@@ -0,0 +1,66 @@
+//===-- CIRStdOps.td - CIR standard library ops ------------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// Defines ops representing standard library calls
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_CIR_DIALECT_IR_CIRSTDOPS
+#define LLVM_CLANG_CIR_DIALECT_IR_CIRSTDOPS
+
+class CIRStdOp<string functionName, dag args, dag res, list<Trait> traits = []>:
+    CIR_Op<"std." # functionName, traits> {
+  string funcName = functionName;
+  
+  let arguments = !con((ins FlatSymbolRefAttr:$original_fn), args);
+
+  let summary = "std::" # functionName # "()";
+  let results = res;
+
+  let extraClassDeclaration = [{
+    static constexpr unsigned getNumArgs() {
+      return }] # !size(args) # [{;
+    }
+    static llvm::StringRef getFunctionName() {
+      return "}] # functionName # [{";
+    }
+  }];
+
+  string argsAssemblyFormat = !interleave(
+    !foreach(
+      name,
+      !foreach(i, !range(!size(args)), !getdagname(args, i)),
+      !strconcat("$", name, " `:` type($", name, ")")
+    ), " `,` "
+  );
+
+  string resultAssemblyFormat = !if(
+    !empty(res),
+    "",
+    " `->` type($" # !getdagname(res, 0) # ")"
+  );
+
+  let assemblyFormat = !strconcat("`(` ", argsAssemblyFormat,
+                                  " `,` $original_fn `)`", resultAssemblyFormat,
+                                  " attr-dict");
+
+  let hasVerifier = 0;
+}
+
+def StdFindOp : CIRStdOp<"find",
+  (ins CIR_AnyType:$first, CIR_AnyType:$last, CIR_AnyType:$pattern),
+  (outs CIR_AnyType:$result),
+  [SameFirstSecondOperandAndResultType]>;
+def IterBeginOp: CIRStdOp<"begin",
+  (ins CIR_AnyType:$container),
+  (outs CIR_AnyType:$result)>;
+def IterEndOp: CIRStdOp<"end", 
+  (ins CIR_AnyType:$container),
+  (outs CIR_AnyType:$result)>;
+
+#endif

clang/lib/CIR/Dialect/Transforms/IdiomRecognizer.cpp

@@ -32,11 +32,52 @@ using namespace cir;
 
 namespace {
 
+// Recognizes a cir.call that calls a standard library function represented
+// by `TargetOp`, and raise it to that operation.
+template <typename TargetOp> class StdRecognizer {
+private:
+  // Reserved for template specialization.
+  static bool checkArguments(mlir::ValueRange) { return true; }
+
+  template <size_t... Indices>
+  static TargetOp buildCall(CIRBaseBuilderTy &builder, CallOp call,
+                            std::index_sequence<Indices...>) {
+    return builder.create<TargetOp>(call.getLoc(), call.getResult().getType(),
+                                    call.getCalleeAttr(),
+                                    call.getOperand(Indices)...);
+  }
+
+public:
+  static bool raise(CallOp call, mlir::MLIRContext &context, bool remark) {
+    constexpr int numArgs = TargetOp::getNumArgs();
+    if (call.getNumOperands() != numArgs)
+      return false;
+
+    auto callExprAttr = call.getAstAttr();
+    llvm::StringRef stdFuncName = TargetOp::getFunctionName();
+    if (!callExprAttr || !callExprAttr.isStdFunctionCall(stdFuncName))
+      return false;
+
+    if (!checkArguments(call.getArgOperands()))
+      return false;
+
+    if (remark)
+      mlir::emitRemark(call.getLoc())
+          << "found call to std::" << stdFuncName << "()";
+
+    CIRBaseBuilderTy builder(context);
+    builder.setInsertionPointAfter(call.getOperation());
+    TargetOp op = buildCall(builder, call, std::make_index_sequence<numArgs>());
+    call.replaceAllUsesWith(op);
+    call.erase();
+    return true;
+  }
+};
+
 struct IdiomRecognizerPass : public IdiomRecognizerBase<IdiomRecognizerPass> {
   IdiomRecognizerPass() = default;
   void runOnOperation() override;
   void recognizeCall(CallOp call);
-  bool raiseStdFind(CallOp call);
   bool raiseIteratorBeginEnd(CallOp call);
 
   // Handle pass options
@@ -88,30 +129,6 @@ struct IdiomRecognizerPass : public IdiomRecognizerBase<IdiomRecognizerPass> {
 };
 } // namespace
 
-bool IdiomRecognizerPass::raiseStdFind(CallOp call) {
-  // FIXME: tablegen all of this function.
-  if (call.getNumOperands() != 3)
-    return false;
-
-  auto callExprAttr = call.getAstAttr();
-  if (!callExprAttr || !callExprAttr.isStdFunctionCall("find")) {
-    return false;
-  }
-
-  if (opts.emitRemarkFoundCalls())
-    emitRemark(call.getLoc()) << "found call to std::find()";
-
-  CIRBaseBuilderTy builder(getContext());
-  builder.setInsertionPointAfter(call.getOperation());
-  auto findOp = builder.create<cir::StdFindOp>(
-      call.getLoc(), call.getResult().getType(), call.getCalleeAttr(),
-      call.getOperand(0), call.getOperand(1), call.getOperand(2));
-
-  call.replaceAllUsesWith(findOp);
-  call.erase();
-  return true;
-}
-
 static bool isIteratorLikeType(mlir::Type t) {
   // TODO: some iterators are going to be represented with structs,
   // in which case we could look at ASTRecordDeclInterface for more
@@ -175,8 +192,16 @@ void IdiomRecognizerPass::recognizeCall(CallOp call) {
   if (raiseIteratorBeginEnd(call))
     return;
 
-  if (raiseStdFind(call))
-    return;
+  bool remark = opts.emitRemarkFoundCalls();
+
+  using StdFunctionsRecognizer = std::tuple<StdRecognizer<StdFindOp>>;
+
+  // MSVC requires explicitly capturing these variables.
+  std::apply(
+      [&, call, remark, this](auto... recognizers) {
+        (decltype(recognizers)::raise(call, this->getContext(), remark) || ...);
+      },
+      StdFunctionsRecognizer());
 }
 
 void IdiomRecognizerPass::runOnOperation() {

clang/lib/CIR/Dialect/Transforms/LoweringPrepare.cpp

@@ -1481,8 +1481,7 @@ void LoweringPreparePass::lowerIterBeginOp(IterBeginOp op) {
   CIRBaseBuilderTy builder(getContext());
   builder.setInsertionPointAfter(op.getOperation());
   auto call = builder.createCallOp(op.getLoc(), op.getOriginalFnAttr(),
-                                   op.getResult().getType(),
-                                   mlir::ValueRange{op.getOperand()});
+                                   op.getResult().getType(), op.getOperand());
 
   op.replaceAllUsesWith(call);
   op.erase();
@@ -1492,8 +1491,7 @@ void LoweringPreparePass::lowerIterEndOp(IterEndOp op) {
   CIRBaseBuilderTy builder(getContext());
   builder.setInsertionPointAfter(op.getOperation());
   auto call = builder.createCallOp(op.getLoc(), op.getOriginalFnAttr(),
-                                   op.getResult().getType(),
-                                   mlir::ValueRange{op.getOperand()});
+                                   op.getResult().getType(), op.getOperand());
 
   op.replaceAllUsesWith(call);
   op.erase();

clang/test/CIR/Transforms/idiom-recognizer.cpp

@@ -18,15 +18,15 @@ int test_find(unsigned char n = 3)
                                                // expected-remark@-2 {{found call to end() iterator}}
 
     // BEFORE-IDIOM: {{.*}} cir.call @_ZNSt5arrayIhLj9EE5beginEv(
-    // AFTER-IDIOM: {{.*}} cir.iterator_begin(@_ZNSt5arrayIhLj9EE5beginEv,
+    // AFTER-IDIOM: {{.*}} cir.std.begin({{.*}}, @_ZNSt5arrayIhLj9EE5beginEv
     // AFTER-LOWERING-PREPARE: {{.*}} cir.call @_ZNSt5arrayIhLj9EE5beginEv(
 
     // BEFORE-IDIOM: {{.*}} cir.call @_ZNSt5arrayIhLj9EE3endEv(
-    // AFTER-IDIOM: {{.*}} cir.iterator_end(@_ZNSt5arrayIhLj9EE3endEv,
+    // AFTER-IDIOM: {{.*}} cir.std.end({{.*}}, @_ZNSt5arrayIhLj9EE3endEv
     // AFTER-LOWERING-PREPARE: {{.*}} cir.call @_ZNSt5arrayIhLj9EE3endEv(
 
     // BEFORE-IDIOM: {{.*}} cir.call @_ZSt4findIPhhET_S1_S1_RKT0_(
-    // AFTER-IDIOM: {{.*}} cir.std.find(@_ZSt4findIPhhET_S1_S1_RKT0_,
+    // AFTER-IDIOM: {{.*}} cir.std.find({{.*}}, @_ZSt4findIPhhET_S1_S1_RKT0_
     // AFTER-LOWERING-PREPARE: {{.*}} cir.call @_ZSt4findIPhhET_S1_S1_RKT0_(
 
     if (f != v.end()) // expected-remark {{found call to end() iterator}}
@@ -43,8 +43,7 @@ template<typename T, unsigned N> struct array {
 };
 }
 
-int iter_test()
-{
+void iter_test() {
   yolo::array<unsigned char, 3> v = {1, 2, 3};
   (void)v.begin(); // no remark should be produced.
-}
+}