[mlir] Add filtering callback to GenerateRuntimeVerification pass

mlir/include/mlir/Transforms/Passes.h

@@ -26,6 +26,7 @@
 namespace mlir {
 
 class GreedyRewriteConfig;
+class RuntimeVerifiableOpInterface;
 
 //===----------------------------------------------------------------------===//
 // Passes
@@ -77,6 +78,13 @@ std::unique_ptr<Pass> createPrintIRPass(const PrintIRPassOptions & = {});
 /// Creates a pass that generates IR to verify ops at runtime.
 std::unique_ptr<Pass> createGenerateRuntimeVerificationPass();
 
+/// Create an instance of the generate runtime verification pass, and
+/// use the provided filter function to skip certain verifiable ops.
+/// The default implementation does not filter any ops.
+std::unique_ptr<Pass> createGenerateRuntimeVerificationPass(
+    std::function<bool(RuntimeVerifiableOpInterface)>
+        shouldHandleVerifiableOpFn);
+
 /// Creates a loop invariant code motion pass that hoists loop invariant
 /// instructions out of the loop.
 std::unique_ptr<Pass> createLoopInvariantCodeMotionPass();

mlir/lib/Transforms/GenerateRuntimeVerification.cpp

@@ -17,16 +17,46 @@ namespace mlir {
 #include "mlir/Transforms/Passes.h.inc"
 } // namespace mlir
 
+#define DEBUG_TYPE "generate-runtime-verification"
+#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
+#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
+
 using namespace mlir;
 
+static bool defaultShouldHandleVerifiableOpFn(RuntimeVerifiableOpInterface op) {
+  // By default, all verifiable ops are considered
+  return true;
+}
+
 namespace {
 struct GenerateRuntimeVerificationPass
     : public impl::GenerateRuntimeVerificationBase<
           GenerateRuntimeVerificationPass> {
+
+  GenerateRuntimeVerificationPass();
+  GenerateRuntimeVerificationPass(const GenerateRuntimeVerificationPass &) =
+      default;
+  GenerateRuntimeVerificationPass(
+      std::function<bool(RuntimeVerifiableOpInterface)>
+          shouldHandleVerifiableOpFn);
+
   void runOnOperation() override;
+
+private:
+  // A filter function to select verifiable ops to generate verification for.
+  // If empty, all verifiable ops are considered.
+  std::function<bool(RuntimeVerifiableOpInterface)> shouldHandleVerifiableOpFn;
 };
 } // namespace
 
+GenerateRuntimeVerificationPass::GenerateRuntimeVerificationPass()
+    : shouldHandleVerifiableOpFn(defaultShouldHandleVerifiableOpFn) {}
+
+GenerateRuntimeVerificationPass::GenerateRuntimeVerificationPass(
+    std::function<bool(RuntimeVerifiableOpInterface)>
+        shouldHandleVerifiableOpFn)
+    : shouldHandleVerifiableOpFn(std::move(shouldHandleVerifiableOpFn)) {}
+
 void GenerateRuntimeVerificationPass::runOnOperation() {
   // The implementation of the RuntimeVerifiableOpInterface may create ops that
   // can be verified. We don't want to generate verification for IR that
@@ -38,11 +68,22 @@ void GenerateRuntimeVerificationPass::runOnOperation() {
 
   OpBuilder builder(getOperation()->getContext());
   for (RuntimeVerifiableOpInterface verifiableOp : ops) {
-    builder.setInsertionPoint(verifiableOp);
-    verifiableOp.generateRuntimeVerification(builder, verifiableOp.getLoc());
-  };
+    if (shouldHandleVerifiableOpFn(verifiableOp)) {
+      builder.setInsertionPoint(verifiableOp);
+      verifiableOp.generateRuntimeVerification(builder, verifiableOp.getLoc());
+    } else {
+      LDBG("Skipping operation: " << verifiableOp.getOperation());
+    }
+  }
 }
 
 std::unique_ptr<Pass> mlir::createGenerateRuntimeVerificationPass() {
   return std::make_unique<GenerateRuntimeVerificationPass>();
 }
+
+std::unique_ptr<Pass> mlir::createGenerateRuntimeVerificationPass(
+    std::function<bool(RuntimeVerifiableOpInterface)>
+        shouldHandleVerifiableOpFn) {
+  return std::make_unique<GenerateRuntimeVerificationPass>(
+      std::move(shouldHandleVerifiableOpFn));
+}