Prune unneeded code when preparing IR to run in interpreter

zirgen/dsl/driver.cpp

@@ -293,6 +293,8 @@ int runTests(mlir::ModuleOp& module) {
     return 1;
   }
   pm.enableVerifier(true);
+  pm.addPass(zirgen::dsl::createEraseUnusedAspectsPass(/*forTests=*/true));
+  pm.addPass(mlir::createSymbolDCEPass());
   pm.addPass(mlir::createInlinerPass());
   mlir::OpPassManager& opm = pm.nest<zirgen::Zhlt::StepFuncOp>();
   opm.addPass(zirgen::ZStruct::createUnrollPass());

zirgen/dsl/passes/BUILD.bazel

@@ -34,6 +34,7 @@ cc_library(
     srcs = [
         "CommonRewrites.cpp",
         "ElideTrivialStructs.cpp",
+        "EraseUnusedAspects.cpp",
         "GenerateAccum.cpp",
         "GenerateCheckLayout.cpp",
         "GenerateGlobals.cpp",

zirgen/dsl/passes/EraseUnusedAspects.cpp

@@ -0,0 +1,54 @@
+// Copyright 2024 RISC Zero, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "zirgen/dsl/passes/PassDetail.h"
+
+using namespace mlir;
+
+namespace zirgen::dsl {
+
+namespace {
+
+struct EraseUnusedAspectsPass : public EraseUnusedAspectsBase<EraseUnusedAspectsPass> {
+  EraseUnusedAspectsPass(bool forTests) : forTests(forTests) {}
+
+  void runOnOperation() override {
+    if (forTests) {
+      getOperation().walk([&](Operation* op) {
+        // keep steps and execs, except for Top
+        if (isa<Zhlt::ComponentOp>(op) || isa<Zhlt::CheckLayoutFuncOp>(op) ||
+            isa<Zhlt::CheckFuncOp>(op) || isa<Zhlt::ValidityRegsFuncOp>(op) ||
+            isa<Zhlt::ValidityTapsFuncOp>(op)) {
+          op->erase();
+          return;
+        }
+        auto step = dyn_cast<Zhlt::StepFuncOp>(op);
+        if (step && step.getName() == "step$Top") {
+          step->erase();
+          return;
+        }
+      });
+    }
+  }
+
+  bool forTests;
+};
+
+} // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> createEraseUnusedAspectsPass(bool forTests) {
+  return std::make_unique<EraseUnusedAspectsPass>(forTests);
+}
+
+} // namespace zirgen::dsl

zirgen/dsl/passes/Passes.h

@@ -22,6 +22,7 @@ namespace zirgen {
 namespace dsl {
 
 // Pass constructors
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> createEraseUnusedAspectsPass(bool forTests = false);
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> createGenerateBackPass();
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> createGenerateExecPass();
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> createGenerateLayoutPass();

zirgen/dsl/passes/Passes.td

@@ -18,6 +18,15 @@
 include "mlir/Pass/PassBase.td"
 include "mlir/Rewrite/PassUtil.td"
 
+def EraseUnusedAspects : Pass<"erase-unused-aspects", "mlir::ModuleOp"> {
+  let summary = "Removes aspects that aren't necessary for tests or codegen";
+  let constructor = "zirgen::dsl::createEraseUnusedAspectsPass()";
+  let dependentDialects = ["zirgen::Zhlt::ZhltDialect"];
+  let options = [
+    Option<"forTests", "for-tests", "bool", /*default=*/"false", "leave only aspects for tests (rather than codegen)">
+  ];
+}
+
 def GenerateExec : Pass<"generate-exec", "mlir::ModuleOp"> {
   let summary = "Lowers the ZHLT dialect to produce circuit execution functions";
   let description = [{