DenseAnalysis.h

Go to the documentation of this file.

//===-- DenseAnalysis.h - Dense data-flow analysis --------------*- 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
//
// Adapted from mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h.
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include <mlir/Analysis/DataFlow/DenseAnalysis.h>
#include <mlir/Analysis/DataFlowFramework.h>
#include <mlir/IR/SymbolTable.h>
#include <mlir/Interfaces/CallInterfaces.h>
#include <mlir/Interfaces/ControlFlowInterfaces.h>
 
namespace llzk {
namespace dataflow {
 
//===----------------------------------------------------------------------===//
// Utilities
//===----------------------------------------------------------------------===//

void markAllOpsAsLive(mlir::DataFlowSolver &solver, mlir::Operation *top);
 
//===----------------------------------------------------------------------===//
// AbstractDenseForwardDataFlowAnalysis
//===----------------------------------------------------------------------===//
 
using AbstractDenseLattice = mlir::dataflow::AbstractDenseLattice;
using CallControlFlowAction = mlir::dataflow::CallControlFlowAction;

class AbstractDenseForwardDataFlowAnalysis : public mlir::DataFlowAnalysis {
public:
  using mlir::DataFlowAnalysis::DataFlowAnalysis;

  mlir::LogicalResult initialize(mlir::Operation *top) override;

  mlir::LogicalResult visit(mlir::ProgramPoint point) override;
 
protected:
  virtual void visitOperationImpl(
      mlir::Operation *op, const AbstractDenseLattice &before, AbstractDenseLattice *after
  ) = 0;

  virtual AbstractDenseLattice *getLattice(mlir::ProgramPoint point) = 0;

  const AbstractDenseLattice *getLatticeFor(mlir::ProgramPoint dependent, mlir::ProgramPoint point);

  virtual void setToEntryState(AbstractDenseLattice *lattice) = 0;

  void join(AbstractDenseLattice *lhs, const AbstractDenseLattice &rhs) {
    propagateIfChanged(lhs, lhs->join(rhs));
  }

  virtual void processOperation(mlir::Operation *op);

  virtual void visitRegionBranchControlFlowTransfer(
      mlir::RegionBranchOpInterface branch, std::optional<unsigned> regionFrom,
      std::optional<unsigned> regionTo, const AbstractDenseLattice &before,
      AbstractDenseLattice *after
  ) {
    join(after, before);
  }

  virtual void visitCallControlFlowTransfer(
      mlir::CallOpInterface call, CallControlFlowAction action, const AbstractDenseLattice &before,
      AbstractDenseLattice *after
  ) {
    join(after, before);
    // Note that `setToEntryState` may be a "partial fixpoint" for some
    // lattices, e.g., lattices that are lists of maps of other lattices will
    // only set fixpoint for "known" lattices.
    if (action == CallControlFlowAction::ExternalCallee) {
      setToEntryState(after);
    }
  }

  void visitRegionBranchOperation(
      mlir::ProgramPoint point, mlir::RegionBranchOpInterface branch, AbstractDenseLattice *after
  );

  mlir::SymbolTableCollection tables;
 
private:
  void visitBlock(mlir::Block *block);

  void visitCallOperation(
      mlir::CallOpInterface call, const AbstractDenseLattice &before, AbstractDenseLattice *after
  );
};
 
//===----------------------------------------------------------------------===//
// DenseForwardDataFlowAnalysis
//===----------------------------------------------------------------------===//

template <typename LatticeT>
class DenseForwardDataFlowAnalysis : public AbstractDenseForwardDataFlowAnalysis {
  static_assert(
      std::is_base_of<AbstractDenseLattice, LatticeT>::value,
      "analysis state class expected to subclass AbstractDenseLattice"
  );
 
public:
  using AbstractDenseForwardDataFlowAnalysis::AbstractDenseForwardDataFlowAnalysis;

  virtual void visitOperation(mlir::Operation *op, const LatticeT &before, LatticeT *after) = 0;

  virtual void visitCallControlFlowTransfer(
      mlir::CallOpInterface call, CallControlFlowAction action, const LatticeT &before,
      LatticeT *after
  ) {
    AbstractDenseForwardDataFlowAnalysis::visitCallControlFlowTransfer(call, action, before, after);
  }

  virtual void visitRegionBranchControlFlowTransfer(
      mlir::RegionBranchOpInterface branch, std::optional<unsigned> regionFrom,
      std::optional<unsigned> regionTo, const LatticeT &before, LatticeT *after
  ) {
    AbstractDenseForwardDataFlowAnalysis::visitRegionBranchControlFlowTransfer(
        branch, regionFrom, regionTo, before, after
    );
  }
 
protected:
  LatticeT *getLattice(mlir::ProgramPoint point) override { return getOrCreate<LatticeT>(point); }

  virtual void setToEntryState(LatticeT *lattice) = 0;
  void setToEntryState(AbstractDenseLattice *lattice) override {
    setToEntryState(static_cast<LatticeT *>(lattice));
  }

  void visitOperationImpl(
      mlir::Operation *op, const AbstractDenseLattice &before, AbstractDenseLattice *after
  ) final {
    visitOperation(op, static_cast<const LatticeT &>(before), static_cast<LatticeT *>(after));
  }
  void visitCallControlFlowTransfer(
      mlir::CallOpInterface call, CallControlFlowAction action, const AbstractDenseLattice &before,
      AbstractDenseLattice *after
  ) final {
    visitCallControlFlowTransfer(
        call, action, static_cast<const LatticeT &>(before), static_cast<LatticeT *>(after)
    );
  }
  void visitRegionBranchControlFlowTransfer(
      mlir::RegionBranchOpInterface branch, std::optional<unsigned> regionFrom,
      std::optional<unsigned> regionTo, const AbstractDenseLattice &before,
      AbstractDenseLattice *after
  ) final {
    visitRegionBranchControlFlowTransfer(
        branch, regionFrom, regionTo, static_cast<const LatticeT &>(before),
        static_cast<LatticeT *>(after)
    );
  }
};
 
} // namespace dataflow
} // namespace llzk