From 14b091263faeac7e66aef97e1e2d8f7229ddbd71 Mon Sep 17 00:00:00 2001
From: Jeremy Kun <j2kun@users.noreply.github.com>
Date: Fri, 10 Nov 2023 17:01:23 -0800
Subject: [PATCH] add noise verification analysis check

---
 lib/Transform/Noisy/BUILD                    |  2 +
 lib/Transform/Noisy/ReduceNoiseOptimizer.cpp | 48 +++++++++++++++++++-
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/lib/Transform/Noisy/BUILD b/lib/Transform/Noisy/BUILD
index 6b4dcdd..369c4f0 100644
--- a/lib/Transform/Noisy/BUILD
+++ b/lib/Transform/Noisy/BUILD
@@ -40,6 +40,8 @@ cc_library(
         ":pass_inc_gen",
         "//lib/Analysis/ReduceNoiseAnalysis",
         "//lib/Dialect/Noisy",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Transforms",
     ],
diff --git a/lib/Transform/Noisy/ReduceNoiseOptimizer.cpp b/lib/Transform/Noisy/ReduceNoiseOptimizer.cpp
index 02450fd..6939617 100644
--- a/lib/Transform/Noisy/ReduceNoiseOptimizer.cpp
+++ b/lib/Transform/Noisy/ReduceNoiseOptimizer.cpp
@@ -3,6 +3,10 @@
 #include "lib/Analysis/ReduceNoiseAnalysis/ReduceNoiseAnalysis.h"
 #include "lib/Dialect/Noisy/NoisyOps.h"
 #include "lib/Dialect/Noisy/NoisyTypes.h"
+#include "mlir/include/mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
+#include "mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h"
+#include "mlir/include/mlir/Analysis/DataFlowFramework.h"
+#include "mlir/include/mlir/IR/Visitors.h"
 #include "mlir/include/mlir/Pass/Pass.h"
 
 namespace mlir {
@@ -21,7 +25,9 @@ struct ReduceNoiseOptimizer
     ReduceNoiseAnalysis analysis(getOperation());
     OpBuilder b(&getContext());
 
-    getOperation()->walk([&](Operation *op) {
+    Operation *module = getOperation();
+
+    module->walk([&](Operation *op) {
       if (!analysis.shouldInsertReduceNoise(op))
         return;
 
@@ -29,6 +35,46 @@ struct ReduceNoiseOptimizer
       auto reduceOp = b.create<ReduceNoiseOp>(op->getLoc(), op->getResult(0));
       op->getResult(0).replaceAllUsesExcept(reduceOp.getResult(), {reduceOp});
     });
+
+    // Afterwards, use the int range analysis to confirm the noise is always
+    // below the maximum.
+    DataFlowSolver solver;
+    // The IntegerRangeAnalysis depends on DeadCodeAnalysis, but this
+    // dependence is not automatic and fails silently.
+    solver.load<dataflow::DeadCodeAnalysis>();
+    solver.load<dataflow::IntegerRangeAnalysis>();
+    if (failed(solver.initializeAndRun(module)))
+      signalPassFailure();
+
+    auto result = module->walk([&](Operation *op) {
+      if (!llvm::isa<noisy::AddOp, noisy::SubOp, noisy::MulOp,
+                     noisy::ReduceNoiseOp>(*op)) {
+        return WalkResult::advance();
+      }
+      const dataflow::IntegerValueRangeLattice *opRange =
+          solver.lookupState<dataflow::IntegerValueRangeLattice>(
+              op->getResult(0));
+      if (!opRange || opRange->getValue().isUninitialized()) {
+        op->emitOpError()
+            << "Found op without a set integer range; did the analysis fail?";
+        return WalkResult::interrupt();
+      }
+
+      ConstantIntRanges range = opRange->getValue().getValue();
+      if (range.umax().getZExtValue() > MAX_NOISE) {
+        op->emitOpError() << "Found op after which the noise exceeds the "
+                             "allowable maximum of "
+                          << MAX_NOISE
+                          << "; it was: " << range.umax().getZExtValue()
+                          << "\n";
+        return WalkResult::interrupt();
+      }
+
+      return WalkResult::advance();
+    });
+
+    if (result.wasInterrupted())
+      signalPassFailure();
   }
 };