LLZKTransformationPassPipelines.cpp

Go to the documentation of this file.

//===-- LLZKTransformationPassPipelines.cpp ---------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2025 Veridise Inc.
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
 
#include "llzk/Transforms/LLZKTransformationPassPipelines.h"
 
#include "llzk/Dialect/Array/Transforms/TransformationPasses.h"
#include "llzk/Dialect/POD/Transforms/TransformationPasses.h"
 
#include <mlir/Pass/PassManager.h>
#include <mlir/Pass/PassRegistry.h>
#include <mlir/Transforms/Passes.h>
 
#include <utility>
 
using namespace mlir;
 
namespace llzk {
 
//===----------------------------------------------------------------------===//
// Pipeline implementation.
//===----------------------------------------------------------------------===//
 
namespace {
 
template <typename NestedPassOptionT>
inline std::unique_ptr<Pass> createConfiguredPass(const NestedPassOptionT &options) {
  return options.getValue().createPass();
}
 
void buildFullStructInliningPipelineImpl(
    OpPassManager &pm, polymorphic::FlatteningPassOptions flattening, bool arrayToScalar,
    bool podToScalar, std::unique_ptr<Pass> inliningPass
) {
  // default to `main-as-root` if unspecified to avoid leaving parameterized templates
  // that cause the later struct inlining pass to crash
  if (flattening.cleanupMode == polymorphic::FlatteningCleanupMode::Unspecified) {
    flattening.cleanupMode = polymorphic::FlatteningCleanupMode::MainAsRoot;
  }
 
  // Circom-style loops are commonly emitted as scf.while operations whose bounds are
  // field constants converted to index values. Normalize the loops and fold those
  // conversions before flattening so the unroller can determine their trip counts.
  pm.addPass(createWhileToForPass());
  pm.addPass(mlir::createCanonicalizerPass());
  pm.addPass(polymorphic::createFlatteningPass(flattening));
 
  // Unrolling substitutes constant induction variables into array accesses. Materialize
  // those constant indices before array-to-scalar requires them to be attributes.
  pm.addPass(mlir::createCanonicalizerPass());
 
  // Run pod-to-scalar first because it is able to split `pod.type` used as array element type
  // (into parallel arrays) so it should be able to fully remove all `pod.type` usages.
  if (podToScalar) {
    pm.addPass(pod::createPodToScalarPass());
  }
  if (arrayToScalar) {
    pm.addPass(array::createArrayToScalarPass());
  }
  // Canonicalize to remove known-condition `scf.if` regions so struct inlining
  // can link "@compute" calls to struct members.
  pm.addPass(mlir::createCanonicalizerPass());
  pm.addPass(std::move(inliningPass));
 
  // Remove struct and member definitions that are no longer used after inlining.
  pm.addPass(createUnusedDeclarationEliminationPass(
      UnusedDeclarationEliminationPassOptions {.removeStructs = true}
  ));
}
 
void buildFullPolyLoweringPipelineImpl(
    OpPassManager &pm, polymorphic::FlatteningPassOptions flattening, bool arrayToScalar,
    bool podToScalar, std::unique_ptr<Pass> inliningPass, std::unique_ptr<Pass> polyLoweringPass
) {
  // 1. Struct flattening and inlining
  buildFullStructInliningPipelineImpl(
      pm, flattening, arrayToScalar, podToScalar, std::move(inliningPass)
  );
  // 2. Degree lowering
  pm.addPass(std::move(polyLoweringPass));
  // 3. Cleanup
  buildRemoveUnnecessaryOpsAndDefsPipeline(pm);
}
 
} // namespace
 
void buildRemoveUnnecessaryOpsPipeline(mlir::OpPassManager &pm) {
  pm.addPass(createRedundantReadAndWriteEliminationPass());
  pm.addPass(createRedundantOperationEliminationPass());
}
 
void buildRemoveUnnecessaryOpsAndDefsPipeline(mlir::OpPassManager &pm) {
  buildRemoveUnnecessaryOpsPipeline(pm);
  pm.addPass(createUnusedDeclarationEliminationPass());
}
 
void buildProductProgramPipeline(OpPassManager &pm) {
  pm.addPass(createComputeConstrainToProductPass());
  pm.addPass(createFuseProductControlFlowPass());
}
 
void buildFullStructInliningPipeline(OpPassManager &pm, const FullStructInliningConfig &cfg) {
  buildFullStructInliningPipelineImpl(
      pm, cfg.flattening, cfg.arrayToScalar, cfg.podToScalar,
      component::createInlineStructsPass(cfg.inlining)
  );
}
 
void buildFullInliningPipeline(OpPassManager &pm, const FullStructInliningConfig &cfg) {
  // Inline free-function bodies before struct cleanup so their struct uses
  // participate in flattening and keep the referenced definitions alive.
  pm.addPass(createInlineFreeFunctionsPass());
  buildFullStructInliningPipelineImpl(
      pm, cfg.flattening, cfg.arrayToScalar, cfg.podToScalar,
      component::createInlineStructsPass(cfg.inlining)
  );
}
 
void buildFullPolyLoweringPipeline(OpPassManager &pm, const FullPolyLoweringConfig &cfg) {
  buildFullPolyLoweringPipelineImpl(
      pm, cfg.structInlining.flattening, cfg.structInlining.arrayToScalar,
      cfg.structInlining.podToScalar,
      component::createInlineStructsPass(cfg.structInlining.inlining),
      createPolyLoweringPass(cfg.polyLowering)
  );
}
 
//===----------------------------------------------------------------------===//
// Pipeline registration.
//===----------------------------------------------------------------------===//
 
void buildConstraintFlatteningPipeline(OpPassManager &pm) {
  // Inline frontend execute_region wrappers before specialization discovery.
  pm.addPass(mlir::createCanonicalizerPass());
  pm.addPass(polymorphic::createTemplateMonomorphizationPass());
  pm.addPass(component::createExposeChildMembersPass());
  // Fold specialized constants and remove empty branches before expansion.
  pm.addPass(mlir::createCanonicalizerPass());
  pm.addPass(polymorphic::createConstraintFlatteningPass());
}
 
void registerTransformationPassPipelines() {
  PassPipelineRegistration<>(
      "llzk-constraint-flattening",
      "Monomorphize templates, expose child members, and expand main constrain",
      buildConstraintFlatteningPipeline
  );
 
  PassPipelineRegistration<>(
      "llzk-remove-unnecessary-ops",
      "Remove unnecessary operations, such as redundant reads or repeated constraints",
      buildRemoveUnnecessaryOpsPipeline
  );
 
  PassPipelineRegistration<>(
      "llzk-remove-unnecessary-ops-and-defs",
      "Remove unnecessary operations, member definitions, and struct definitions",
      buildRemoveUnnecessaryOpsAndDefsPipeline
  );
 
  PassPipelineRegistration<>(
      "llzk-product-program",
      "Convert @compute/@constrain functions to @product function and perform alignment",
      buildProductProgramPipeline
  );
 
  PassPipelineRegistration<FullStructInliningOptions>(
      "llzk-full-struct-inlining",
      "Run flattening and inlining of all struct definitions into the `main` struct. This "
      "pipeline uses the `main-as-root` cleanup mode in the flattening pass by default. It "
      "is not recommended to override this cleanup mode because other cleanup modes may "
      "leave behind parameterized templates that later cause `llzk-inline-structs` to crash.",
      [](OpPassManager &pm, const FullStructInliningOptions &opts) {
    auto flattening = opts.flattening.getValue().createOptions();
    buildFullStructInliningPipelineImpl(
        pm, flattening->createPassOptions(), opts.arrayToScalar, opts.podToScalar,
        createConfiguredPass(opts.inlining)
    );
  }
  );
 
  PassPipelineRegistration<FullStructInliningOptions>(
      "llzk-full-inlining",
      "Run free function inlining, flattening, and struct inlining. This is the "
      "recommended pipeline before any downstream pass that does not understand `function.call`.",
      [](OpPassManager &pm, const FullStructInliningOptions &opts) {
    auto flattening = opts.flattening.getValue().createOptions();
    // Inline free-function bodies before struct cleanup so their struct uses
    // participate in flattening and keep the referenced definitions alive.
    pm.addPass(createInlineFreeFunctionsPass());
    buildFullStructInliningPipelineImpl(
        pm, flattening->createPassOptions(), opts.arrayToScalar, opts.podToScalar,
        createConfiguredPass(opts.inlining)
    );
  }
  );
 
  PassPipelineRegistration<FullPolyLoweringOptions>(
      "llzk-full-poly-lowering",
      "Run flattening and inlining of all struct definitions into the `main` struct, then lower "
      "polynomial constraints to a given max degree, and finally remove unnecessary operations and "
      "definitions. This pipeline uses the `main-as-root` cleanup mode in the flattening pass by "
      "default. It is not recommended to override this cleanup mode because other cleanup modes "
      "may leave behind parameterized templates that later cause `llzk-inline-structs` to crash.",
      [](OpPassManager &pm, const FullPolyLoweringOptions &opts) {
    auto structInlining = opts.structInlining.getValue().createOptions();
    auto flattening = structInlining->flattening.getValue().createOptions();
    buildFullPolyLoweringPipelineImpl(
        pm, flattening->createPassOptions(), structInlining->arrayToScalar,
        structInlining->podToScalar, createConfiguredPass(structInlining->inlining),
        createConfiguredPass(opts.polyLowering)
    );
  }
  );
}
 
} // namespace llzk