LLZKComputeConstrainToProductPass.cpp

Go to the documentation of this file.

//===-- LLZKComputeConstrainToProductPass.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/LLZKComputeConstrainToProductPass.h"
 
#include "llzk/Analysis/LightweightSignalEquivalenceAnalysis.h"
#include "llzk/Dialect/Function/IR/Ops.h"
#include "llzk/Dialect/Struct/IR/Ops.h"
#include "llzk/Dialect/Struct/Transforms/InlineStructsPass.h"
#include "llzk/Transforms/LLZKTransformationPasses.h"
#include "llzk/Util/Constants.h"
#include "llzk/Util/SymbolHelper.h"
 
#include <mlir/IR/Builders.h>
#include <mlir/IR/SymbolTable.h>
#include <mlir/Transforms/Inliner.h>
#include <mlir/Transforms/InliningUtils.h>
 
#include <llvm/Support/Debug.h>
 
#include <memory>
#include <ranges>
 
namespace llzk {
#define GEN_PASS_DEF_COMPUTECONSTRAINTOPRODUCTPASS
#include "llzk/Transforms/LLZKTransformationPasses.h.inc"
} // namespace llzk
 
#define DEBUG_TYPE "llzk-compute-constrain-to-product-pass"
 
using namespace mlir;
using namespace llzk;
using namespace llzk::component;
using namespace llzk::function;
 
FuncDefOp ProductAligner::alignFuncs(StructDefOp root, FuncDefOp compute, FuncDefOp constrain) {
 
  if (auto prod = root.getProductFuncOp()) {
    return prod;
  }
 
  OpBuilder funcBuilder(compute);
 
  // Add compute/constrain attributes
  compute.walk([&funcBuilder](Operation *op) {
    op->setAttr(PRODUCT_SOURCE, funcBuilder.getStringAttr(FUNC_NAME_COMPUTE));
  });
 
  constrain.walk([&funcBuilder](Operation *op) {
    op->setAttr(PRODUCT_SOURCE, funcBuilder.getStringAttr(FUNC_NAME_CONSTRAIN));
  });
 
  // Create an empty @product func...
  FuncDefOp productFunc = FuncDefOp::create(
      funcBuilder, funcBuilder.getFusedLoc({compute.getLoc(), constrain.getLoc()}),
      FUNC_NAME_PRODUCT, compute.getFunctionType()
  );
  productFunc->setAttr(DERIVED_ATTR_NAME, UnitAttr::get(funcBuilder.getContext()));
  Block *entryBlock = productFunc.addEntryBlock();
  funcBuilder.setInsertionPointToStart(entryBlock);
 
  productFunc.setAllowNonNativeFieldOpsAttr(
      compute.hasAllowNonNativeFieldOpsAttr() || constrain.hasAllowNonNativeFieldOpsAttr()
  );
 
  // ...with the right arguments
  llvm::SmallVector<Value> args {productFunc.getArguments()};
 
  // Add calls to @compute and @constrain...
  CallOp computeCall = CallOp::create(funcBuilder, funcBuilder.getUnknownLoc(), compute, args);
  args.insert(args.begin(), computeCall->getResult(0));
  CallOp constrainCall = CallOp::create(funcBuilder, funcBuilder.getUnknownLoc(), constrain, args);
  ReturnOp::create(funcBuilder, funcBuilder.getUnknownLoc(), computeCall->getResult(0));
 
  // ..and inline them
  InlinerInterface inliner(productFunc.getContext());
  InlinerConfig inlinerConfig;
  if (failed(inlineCall(
          inliner, inlinerConfig.getCloneCallback(), computeCall, compute, &compute.getBody(), true
      ))) {
    root->emitError().append("failed to inline ", FUNC_NAME_COMPUTE).report();
    return nullptr;
  }
  if (failed(inlineCall(
          inliner, inlinerConfig.getCloneCallback(), constrainCall, constrain, &constrain.getBody(),
          true
      ))) {
    root->emitError().append("failed to inline ", FUNC_NAME_CONSTRAIN).report();
    return nullptr;
  }
  computeCall->erase();
  constrainCall->erase();
 
  // Mark the compute/constrain for deletion
  alignedStructs.push_back(root);
 
  // Make sure we can align sub-calls to @compute and @constrain
  if (failed(alignCalls(productFunc))) {
    return nullptr;
  }
  return productFunc;
}
 
LogicalResult ProductAligner::alignCalls(FuncDefOp product) {
  // Gather up all the remaining calls to @compute and @constrain
  llvm::SetVector<CallOp> computeCalls, constrainCalls;
  product.walk([&](CallOp callOp) {
    if (callOp.calleeIsStructCompute()) {
      computeCalls.insert(callOp);
    } else if (callOp.calleeIsStructConstrain()) {
      constrainCalls.insert(callOp);
    }
  });
 
  llvm::SetVector<std::pair<CallOp, CallOp>> alignedCalls;
 
  // A @compute matches a @constrain if they belong to the same struct and all their input signals
  // are pairwise equivalent
  auto doCallsMatch = [&](CallOp compute, CallOp constrain) -> bool {
    LLVM_DEBUG({
      llvm::outs() << "Asking for equivalence between calls\n"
                   << compute << "\nand\n"
                   << constrain << "\n\n";
      llvm::outs() << "In block:\n\n" << *compute->getBlock() << '\n';
    });
 
    auto computeStruct = getPrefixAsSymbolRefAttr(compute.getCallee());
    auto constrainStruct = getPrefixAsSymbolRefAttr(constrain.getCallee());
    if (computeStruct != constrainStruct) {
      return false;
    }
    if (compute.getNumOperands() == 0) {
      return true;
    }
    for (unsigned i = 0, e = compute->getNumOperands() - 1; i < e; i++) {
      if (!equivalence.areSignalsEquivalent(compute->getOperand(i), constrain->getOperand(i + 1))) {
        return false;
      }
    }
 
    return true;
  };
 
  for (auto compute : computeCalls) {
    // If there is exactly one @compute that matches a given @constrain, we can align them
    auto matches = llvm::filter_to_vector(constrainCalls, [&](CallOp constrain) {
      return doCallsMatch(compute, constrain);
    });
 
    if (matches.size() == 1) {
      alignedCalls.insert({compute, matches[0]});
      computeCalls.remove(compute);
      constrainCalls.remove(matches[0]);
    }
  }
 
  if (!computeCalls.empty() && constrainCalls.empty()) {
    product.emitWarning()
        .append("failed to align some @", FUNC_NAME_COMPUTE, " and @", FUNC_NAME_CONSTRAIN)
        .report();
  }
 
  for (auto [compute, constrain] : alignedCalls) {
    // If @A::@compute matches @A::@constrain, recursively align the functions in @A...
    auto calleeTgt = compute.getCalleeTarget(tables);
    if (failed(calleeTgt)) {
      return failure();
    }
    auto newRoot = calleeTgt->get()->getParentOfType<StructDefOp>();
    assert(newRoot);
    FuncDefOp newProduct =
        alignFuncs(newRoot, newRoot.getComputeFuncOp(), newRoot.getConstrainFuncOp());
    if (!newProduct) {
      return failure();
    }
 
    // ...and replace the two calls with a single call to @A::@product
    OpBuilder callBuilder(compute);
    CallOp newCall = CallOp::create(
        callBuilder, callBuilder.getFusedLoc({compute.getLoc(), constrain.getLoc()}), newProduct,
        compute.getOperands()
    );
    compute->replaceAllUsesWith(newCall.getResults());
  }
 
  return success();
}
 
namespace {
 
bool isValidRoot(StructDefOp root) {
  FuncDefOp computeFunc = root.getComputeFuncOp();
  FuncDefOp constrainFunc = root.getConstrainFuncOp();
 
  if (!computeFunc || !constrainFunc) {
    root->emitError()
        .append("no ", FUNC_NAME_COMPUTE, "/", FUNC_NAME_CONSTRAIN, " to align")
        .report();
    return false;
  }

 
  return true;
}
 
LogicalResult alignStartingAt(
    component::StructDefOp root, SymbolTableCollection &tables,
    LightweightSignalEquivalenceAnalysis &equivalence
) {
  if (!isValidRoot(root)) {
    return failure();
  }
 
  ProductAligner aligner {tables, equivalence};
  if (!aligner.alignFuncs(root, root.getComputeFuncOp(), root.getConstrainFuncOp())) {
    return failure();
  }
 
  return success();
}
 
class PassImpl : public llzk::impl::ComputeConstrainToProductPassBase<PassImpl> {
  using Base = ComputeConstrainToProductPassBase<PassImpl>;
  using Base::Base;
 
  void runOnOperation() override {
    ModuleOp mod = getOperation();
    StructDefOp root;
 
    SymbolTableCollection tables;
    LightweightSignalEquivalenceAnalysis equivalence {
        getAnalysis<LightweightSignalEquivalenceAnalysis>()
    };
 
    // Find the indicated root struct and make sure its a valid place to start aligning
    mod.walk([&root, this](StructDefOp structDef) {
      if (structDef.getSymName() == rootStruct) {
        root = structDef;
      }
    });
 
    if (!root) {
      mod.emitError().append("could not find root struct \"", rootStruct, "\"").report();
      signalPassFailure();
      return;
    }
 
    if (failed(alignStartingAt(root, tables, equivalence))) {
      signalPassFailure();
    }
  }
};
 
} // namespace