ConstGlobalPropagationPass.cpp

Go to the documentation of this file.

//===-- ConstGlobalPropagationPass.cpp --------------------------*- 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
//
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
 
#include "llzk/Analysis/SymbolUseGraph.h"
#include "llzk/Dialect/Felt/IR/Ops.h"
#include "llzk/Dialect/Global/IR/Ops.h"
#include "llzk/Dialect/Global/Transforms/TransformationPasses.h"
#include "llzk/Dialect/Polymorphic/IR/Ops.h"
#include "llzk/Dialect/String/IR/Ops.h"
#include "llzk/Util/Debug.h"
#include "llzk/Util/SymbolTableLLZK.h"
#include "llzk/Util/TypeHelper.h"
#include "llzk/Util/Walk.h"
 
#include <mlir/Dialect/Arith/IR/Arith.h>
 
#include <llvm/ADT/DenseMap.h>
#include <llvm/ADT/SmallVector.h>
 
#include <optional>
 
// Include the generated base pass class definitions.
namespace llzk::global {
#define GEN_PASS_DEF_CONSTGLOBALPROPAGATIONPASS
#include "llzk/Dialect/Global/Transforms/TransformationPasses.h.inc"
} // namespace llzk::global
 
#define DEBUG_TYPE "llzk-const-global-propagation"
 
using namespace mlir;
using namespace llzk;
using namespace llzk::global;
 
namespace {
 
static Value convertToConstantValue(GlobalReadOp readOp, Attribute attr) {
  OpBuilder bldr(readOp);
  Location loc = readOp.getLoc();
  if (auto intAttr = llvm::dyn_cast<IntegerAttr>(attr)) {
    return arith::ConstantOp::create(bldr, loc, intAttr).getResult();
  } else if (auto feltAttr = llvm::dyn_cast<felt::FeltConstAttr>(attr)) {
    Value constant = felt::FeltConstantOp::create(bldr, loc, feltAttr).getResult();
    if (constant.getType() == readOp.getType()) {
      return constant;
    }
    return polymorphic::UnifiableCastOp::create(bldr, loc, readOp.getType(), constant).getResult();
  } else if (auto strAttr = llvm::dyn_cast<StringAttr>(attr)) {
    return string::LitStringOp::create(bldr, loc, readOp.getType(), strAttr).getResult();
  } else {
    llvm::outs() << "Encountered: " << attr.getAbstractAttribute().getName() << '\n';
    llvm_unreachable("Unsupported constant attribute type");
  }
}
 
class PassImpl : public llzk::global::impl::ConstGlobalPropagationPassBase<PassImpl> {
  using Base = ConstGlobalPropagationPassBase<PassImpl>;
 
public:
  using Base::Base;
 
  void runOnOperation() override {
    ModuleOp root = getOperation();
    auto constGlobals = walkCollect<GlobalDefOp>(root, [](auto g) { return g.isConstant(); });
    if (constGlobals.empty()) {
      return;
    }
 
    // Keep track of erased operations to skip them during the inner loop iteration and avoid
    // use-after-free errors. These MUST NOT be dereferenced because the ops have been erased.
    llvm::DenseSet<void *> erased;
    SymbolUseGraph &useGraph = getAnalysis<SymbolUseGraph>();
    llvm::DenseMap<Operation *, SmallVector<const SymbolUseGraphNode *>> useNodesByDefinition;
    SymbolTableCollection tables;
    for (const SymbolUseGraphNode *candidate : useGraph.nodesIter()) {
      auto target = candidate->lookupSymbol(tables, /*reportMissing=*/false);
      if (succeeded(target)) {
        useNodesByDefinition[target->get()].push_back(candidate);
      }
    }
 
    for (GlobalDefOp globalDef : constGlobals) {
      Attribute constValue = globalDef.getInitialValueAttr();
 
      // Array constants cannot yet be materialized as operations. Leave both
      // the global and all of its uses unchanged until they are supported.
      if (llvm::isa<ArrayAttr>(constValue)) {
        continue;
      }
 
      // A global definition and its uses can have different closest LLZK root
      // modules after include inlining. Find every use-graph node that resolves
      // to this definition instead of looking it up by a single symbol path.
      for (const SymbolUseGraphNode *node : useNodesByDefinition[globalDef.getOperation()]) {
        SymbolRefAttr symbolAttr = node->getSymbolPath();
        for (Operation *userOp : node->getUserOps()) {
          // Skip erased ops to avoid use-after-free errors
          if (erased.contains(userOp)) {
            continue;
          }
          LLVM_DEBUG(
              llvm::dbgs() << "Replacing '" << symbolAttr << "' with '" << constValue << "' in "
                           << *userOp << '\n'
          );
 
          // Special handling to fully replace GlobalReadOp with the appropriate constant op.
          if (auto readOp = llvm::dyn_cast<GlobalReadOp>(userOp);
              readOp && readOp.getNameRef() == symbolAttr) {
            Value constantResult = convertToConstantValue(readOp, constValue);
            readOp.getResult().replaceAllUsesWith(constantResult);
            readOp.erase();
            erased.insert(readOp);
            continue;
          }
 
          // Standard case that just replaces all uses of the symbol with the constant value.
          AttrTypeReplacer replacer;
          replacer.addReplacement([symbolAttr, constValue](SymbolRefAttr a) {
            return (a == symbolAttr) ? std::make_optional(constValue) : std::nullopt;
          });
          replacer.replaceElementsIn(
              userOp,
              /*replaceAttrs=*/true,
              /*replaceLocs=*/false,
              /*replaceTypes=*/true
          );
        }
      }
      // Delete the const global since there are no (remaining) uses.
      globalDef.erase();
      erased.insert(globalDef);
    }
  }
};
 
} // namespace