SpecializedMemoryPasses.h

Go to the documentation of this file.

//===-- SpecializedMemoryPasses.h - Targeted SROA / mem2reg -----*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2026 Project LLZK
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include "llzk/Dialect/LLZK/IR/Dialect.h"
#include "llzk/Dialect/LLZK/IR/Ops.h"
#include "llzk/Util/Walk.h"
 
#include <mlir/Analysis/DataLayoutAnalysis.h>
#include <mlir/Dialect/SCF/IR/SCF.h>
#include <mlir/Dialect/UB/IR/UBOps.h>
#include <mlir/IR/Dominance.h>
#include <mlir/Interfaces/MemorySlotInterfaces.h>
#include <mlir/Pass/Pass.h>
#include <mlir/Pass/PassManager.h>
#include <mlir/Transforms/Mem2Reg.h>
#include <mlir/Transforms/Passes.h>
#include <mlir/Transforms/SROA.h>
 
#include <llvm/ADT/StringRef.h>
 
#include <memory>
 
namespace llzk {

template <typename AllocOpTy>
struct SpecializedSROA : mlir::PassWrapper<SpecializedSROA<AllocOpTy>, mlir::OperationPass<>> {
 
  mlir::StringRef getArgument() const override { return "llzk-specialized-sroa"; }
 
  mlir::StringRef getDescription() const override {
    return "Scalar replacement of aggregates for a specific allocator op type";
  }
 
  void runOnOperation() override {
    mlir::Operation *scopeOp = this->getOperation();
 
    auto &dataLayoutAnalysis = this->template getAnalysis<mlir::DataLayoutAnalysis>();
    const mlir::DataLayout &dataLayout = dataLayoutAnalysis.getAtOrAbove(scopeOp);
 
    bool changed = false;
 
    for (mlir::Region &region : scopeOp->getRegions()) {
      if (region.getBlocks().empty()) {
        continue;
      }
 
      mlir::OpBuilder builder(&region.front(), region.front().begin());
 
      auto allocators = walkCollectMapped<AllocOpTy>(region, [](auto allocator) {
        return mlir::DestructurableAllocationOpInterface(allocator);
      });
 
      if (mlir::succeeded(mlir::tryToDestructureMemorySlots(allocators, builder, dataLayout))) {
        changed = true;
      }
    }
 
    if (!changed) {
      this->markAllAnalysesPreserved();
    }
  }
};
 
// Pass factory for `SpecializedSROA`.
template <typename AllocOpTy>
std::unique_ptr<SpecializedSROA<AllocOpTy>> createSpecializedSROAPass() {
  return std::make_unique<SpecializedSROA<AllocOpTy>>();
}

template <typename AllocOpTy>
struct SpecializedMem2Reg
    : mlir::PassWrapper<SpecializedMem2Reg<AllocOpTy>, mlir::OperationPass<>> {
 
  mlir::StringRef getArgument() const override { return "llzk-specialized-mem2reg"; }
 
  mlir::StringRef getDescription() const override {
    return "Promotes memory slots of a specific allocator op type into values";
  }
 
  void runOnOperation() override {
    mlir::Operation *scopeOp = this->getOperation();
 
    auto &dataLayoutAnalysis = this->template getAnalysis<mlir::DataLayoutAnalysis>();
    const mlir::DataLayout &dataLayout = dataLayoutAnalysis.getAtOrAbove(scopeOp);
    auto &dominance = this->template getAnalysis<mlir::DominanceInfo>();
 
    bool changed = false;
 
    for (mlir::Region &region : scopeOp->getRegions()) {
      if (region.getBlocks().empty()) {
        continue;
      }
 
      mlir::OpBuilder builder(&region.front(), region.front().begin());
 
      auto allocators = walkCollectMapped<AllocOpTy>(region, [](auto allocator) {
        return mlir::PromotableAllocationOpInterface(allocator);
      });
 
      auto promoteRes = mlir::tryToPromoteMemorySlots(allocators, builder, dataLayout, dominance);
      if (mlir::succeeded(promoteRes)) {
        changed = true;
      }
    }
 
    if (!changed) {
      this->markAllAnalysesPreserved();
    }
  }
};
 
// Pass factory for `SpecializedMem2Reg`.
template <typename AllocOpTy>
std::unique_ptr<SpecializedMem2Reg<AllocOpTy>> createSpecializedMem2RegPass() {
  return std::make_unique<SpecializedMem2Reg<AllocOpTy>>();
}
 
namespace detail {

class RemoveDeadValuesWorkaroundPass
    : public mlir::PassWrapper<RemoveDeadValuesWorkaroundPass, mlir::OperationPass<>> {
 
  static inline std::unique_ptr<mlir::Pass> createRemoveDeadValuesPass() {
    return mlir::createRemoveDeadValuesPass(
        mlir::RemoveDeadValuesPassOptions {.canonicalize = true}
    );
  }
 
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RemoveDeadValuesWorkaroundPass)
 
  llvm::StringRef getArgument() const override { return "remove-dead-values"; }
  llvm::StringRef getDescription() const override { return "Remove dead values"; }
 
  void getDependentDialects(mlir::DialectRegistry &registry) const override {
    createRemoveDeadValuesPass()->getDependentDialects(registry);
    registry.insert<LLZKDialect>();
  }
 
  void runOnOperation() final {
    mlir::Operation *scopeOp = this->getOperation();
 
    mlir::OpPassManager pm(scopeOp->getName().getStringRef());
    pm.addPass(createRemoveDeadValuesPass());
    if (mlir::failed(runPipeline(pm, scopeOp))) {
      signalPassFailure();
    }
 
    // Upstream RDV only tracks poison values created during its current invocation. This wrapper
    // may run repeatedly, so values created by an earlier invocation can become unused later.
    // Additionally, replace any remaining uses of `ub.poison` with LLZK native `llzk.nondet` op.
    for (mlir::ub::PoisonOp poisonOp : walkCollect<mlir::ub::PoisonOp>(*scopeOp)) {
      auto poisonResult = poisonOp.getResult();
      if (!poisonResult.use_empty()) {
        mlir::OpBuilder builder(poisonOp);
        auto nondetOp = NonDetOp::create(builder, poisonOp.getLoc(), poisonOp.getType());
        poisonResult.replaceAllUsesWith(nondetOp.getResult());
      }
      poisonOp.erase();
    }
 
    // Finally, those `ub.poison` removals can leave behind trivial `else` blocks that would have
    // otherwise been removed by RDV so go ahead and remove them now.
    scopeOp->walk([](mlir::scf::IfOp ifOp) {
      if (ifOp.getResults().empty()) {
        mlir::Region &elseRegion = ifOp.getElseRegion();
        if (!llvm::hasSingleElement(elseRegion)) {
          return;
        }
        mlir::Block &elseBlock = elseRegion.front();
        if (!llvm::hasSingleElement(elseBlock)) {
          return;
        }
        if (!llvm::isa<mlir::scf::YieldOp>(elseBlock.front())) {
          return;
        }
        elseRegion.dropAllReferences();
        elseBlock.clear();
        elseRegion.getBlocks().clear();
      }
    });
  }
};
 
} // namespace detail
 
inline std::unique_ptr<mlir::Pass> createRemoveDeadValuesWorkaroundPass() {
  return std::make_unique<detail::RemoveDeadValuesWorkaroundPass>();
}
 
} // namespace llzk