SparseAnalysis.h

Go to the documentation of this file.

//===- SparseAnalysis.h - LLZK sparse data-flow adapter -------------------===//
//
// 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
//
// Portions adapted from mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.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/DeadCodeAnalysis.h>
#include <mlir/Analysis/DataFlow/SparseAnalysis.h>
#include <mlir/Analysis/DataFlowFramework.h>
#include <mlir/IR/Block.h>
#include <mlir/IR/Operation.h>
#include <mlir/IR/Region.h>
#include <mlir/IR/SymbolTable.h>
#include <mlir/IR/Value.h>
#include <mlir/Interfaces/CallInterfaces.h>
#include <mlir/Interfaces/ControlFlowInterfaces.h>
#include <mlir/Support/LLVM.h>
 
#include <llvm/ADT/SmallVector.h>
#include <llvm/Support/Casting.h>
 
#include <type_traits>
 
namespace llzk::dataflow {
 
using AbstractSparseLattice = mlir::dataflow::AbstractSparseLattice;
 
//===----------------------------------------------------------------------===//
// AbstractSparseForwardDataFlowAnalysis
//===----------------------------------------------------------------------===//

class AbstractSparseForwardDataFlowAnalysis
    : public mlir::dataflow::AbstractSparseForwardDataFlowAnalysis {
  using Base = mlir::dataflow::AbstractSparseForwardDataFlowAnalysis;
 
public:
  mlir::LogicalResult initialize(mlir::Operation *top) override;

  mlir::LogicalResult visit(mlir::ProgramPoint *point) override;
 
protected:
  explicit AbstractSparseForwardDataFlowAnalysis(mlir::DataFlowSolver &s);

  mlir::SymbolTableCollection tables;
 
private:
  mlir::LogicalResult initializeRecursivelyInProgramOrder(mlir::Operation *op);

  mlir::LogicalResult visitOperationDuringInitialization(mlir::Operation *op);

  bool isOperationLive(mlir::Operation *op);

  llvm::SmallVector<const AbstractSparseLattice *, 4>
  collectOperandLatticesAndSubscribe(mlir::Operation *op);

  mlir::LogicalResult visitZeroResultCallOperation(
      mlir::CallOpInterface call, mlir::ArrayRef<const AbstractSparseLattice *> operandLattices
  );

  mlir::LogicalResult visitZeroResultOperation(mlir::Operation *op);
};
 
//===----------------------------------------------------------------------===//
// SparseForwardDataFlowAnalysis
//===----------------------------------------------------------------------===//

template <typename StateT>
class SparseForwardDataFlowAnalysis : public AbstractSparseForwardDataFlowAnalysis {
  static_assert(
      std::is_base_of<AbstractSparseLattice, StateT>::value,
      "analysis state class expected to subclass AbstractSparseLattice"
  );
 
public:
  explicit SparseForwardDataFlowAnalysis(mlir::DataFlowSolver &s)
      : AbstractSparseForwardDataFlowAnalysis(s) {}

  virtual mlir::LogicalResult visitOperation(
      mlir::Operation *op, mlir::ArrayRef<const StateT *> operands, mlir::ArrayRef<StateT *> results
  ) = 0;

  virtual void visitExternalCall(
      mlir::CallOpInterface /*call*/, mlir::ArrayRef<const StateT *> /*argumentLattices*/,
      mlir::ArrayRef<StateT *> resultLattices
  ) {
    setAllToEntryStates(resultLattices);
  }

  virtual void visitNonControlFlowArguments(
      mlir::Operation * /*op*/, const mlir::RegionSuccessor & /*successor*/,
      mlir::ValueRange nonSuccessorInputs, mlir::ArrayRef<StateT *> nonSuccessorInputLattices
  ) {
    assert(nonSuccessorInputs.size() == nonSuccessorInputLattices.size() && "size mismatch");
    setAllToEntryStates(nonSuccessorInputLattices);
  }
 
protected:
  StateT *getLatticeElement(mlir::Value value) override { return getOrCreate<StateT>(value); }

  // NOLINTNEXTLINE(bugprone-derived-method-shadowing-base-method)
  const StateT *getLatticeElementFor(mlir::ProgramPoint *point, mlir::Value value) {
    return static_cast<const StateT *>(
        mlir::dataflow::AbstractSparseForwardDataFlowAnalysis::getLatticeElementFor(point, value)
    );
  }

  virtual void setToEntryState(StateT *lattice) = 0;
  void setAllToEntryStates(mlir::ArrayRef<StateT *> lattices) {
    mlir::dataflow::AbstractSparseForwardDataFlowAnalysis::setAllToEntryStates(
        {reinterpret_cast<AbstractSparseLattice *const *>(lattices.begin()), lattices.size()}
    );
  }
 
private:
  mlir::LogicalResult visitOperationImpl(
      mlir::Operation *op, mlir::ArrayRef<const AbstractSparseLattice *> operandLattices,
      mlir::ArrayRef<AbstractSparseLattice *> resultLattices
  ) override {
    return visitOperation(
        op,
        {reinterpret_cast<const StateT *const *>(operandLattices.begin()), operandLattices.size()},
        {reinterpret_cast<StateT *const *>(resultLattices.begin()), resultLattices.size()}
    );
  }
 
  void visitExternalCallImpl(
      mlir::CallOpInterface call, mlir::ArrayRef<const AbstractSparseLattice *> argumentLattices,
      mlir::ArrayRef<AbstractSparseLattice *> resultLattices
  ) override {
    visitExternalCall(
        call,
        {reinterpret_cast<const StateT *const *>(argumentLattices.begin()),
         argumentLattices.size()},
        {reinterpret_cast<StateT *const *>(resultLattices.begin()), resultLattices.size()}
    );
  }
 
  void visitNonControlFlowArgumentsImpl(
      mlir::Operation *op, const mlir::RegionSuccessor &successor,
      mlir::ValueRange nonSuccessorInputs,
      mlir::ArrayRef<AbstractSparseLattice *> nonSuccessorInputLattices
  ) override {
    visitNonControlFlowArguments(
        op, successor, nonSuccessorInputs,
        {reinterpret_cast<StateT *const *>(nonSuccessorInputLattices.begin()),
         nonSuccessorInputLattices.size()}
    );
  }
 
  void setToEntryState(AbstractSparseLattice *lattice) override {
    return setToEntryState(reinterpret_cast<StateT *>(lattice));
  }
};
 
} // namespace llzk::dataflow