Raptor/src/PIM/Pass/PimCodegen/MaterializeHostConstantsPass.cpp

#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"

#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/MathExtras.h"

#include <type_traits>

#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

using namespace mlir;

namespace onnx_mlir {

namespace {

static bool isExplicitHostOperand(Operation* op, unsigned operandIndex) {
  if (isa<pim::PimMemCopyHostToDevOp>(op))
    return operandIndex == 1;
  if (isa<pim::PimMemCopyHostToDevBatchOp>(op))
    return operandIndex == 1;
  if (isa<pim::PimMemCopyDevToHostOp>(op))
    return operandIndex == 0;
  return false;
}

static int64_t getValueSizeInBytes(Value value) {
  auto type = dyn_cast<ShapedType>(value.getType());
  if (!type || !type.hasStaticShape())
    return -1;
  return type.getNumElements() * type.getElementTypeBitWidth() / 8;
}

template <typename CoreOpTy>
static void materializeHostConstantsInCore(CoreOpTy coreOp, IRRewriter& rewriter, bool& hasFailure) {
  DenseMap<Value, DenseMap<int64_t, DenseMap<Type, Value>>> materializedValues;
  SmallVector<Operation*> ops;
  coreOp.getBody().front().walk([&](Operation* op) {
    if (!isa<pim::PimHaltOp, scf::YieldOp>(op))
      ops.push_back(op);
  });

  for (Operation* op : ops) {
    for (OpOperand& operand : op->getOpOperands()) {
      Value originalValue = operand.get();
      if (!isa<BaseMemRefType>(originalValue.getType()) || isExplicitHostOperand(op, operand.getOperandNumber()))
        continue;

      auto resolvedAddress = resolveContiguousAddress(originalValue);
      if (failed(resolvedAddress))
        continue;

      auto getGlobalOp = dyn_cast_or_null<memref::GetGlobalOp>(resolvedAddress->base.getDefiningOp());
      if (!getGlobalOp)
        continue;

      auto originalType = dyn_cast<MemRefType>(originalValue.getType());
      if (!originalType || !originalType.hasStaticShape()) {
        op->emitOpError("host constant materialization requires a static memref operand");
        hasFailure = true;
        continue;
      }

      auto& cachedByOffset = materializedValues[resolvedAddress->base];
      auto& cachedByType = cachedByOffset[resolvedAddress->byteOffset];
      auto cachedValue = cachedByType.find(originalType);
      if (cachedValue != cachedByType.end()) {
        operand.set(cachedValue->second);
        continue;
      }

      int64_t totalBytes = getValueSizeInBytes(originalValue);
      if (totalBytes < 0 || !llvm::isInt<32>(totalBytes) || !llvm::isInt<32>(resolvedAddress->byteOffset)) {
        op->emitOpError("host constant materialization requires 32-bit copy sizes and offsets");
        hasFailure = true;
        continue;
      }

      auto contiguousType = MemRefType::get(originalType.getShape(), originalType.getElementType());

      rewriter.setInsertionPoint(op);
      Value localAlloc = memref::AllocOp::create(rewriter, op->getLoc(), contiguousType);
      Value deviceDst = localAlloc;
      if (contiguousType != originalType)
        deviceDst = memref::CastOp::create(rewriter, op->getLoc(), originalType, localAlloc);

      Value copiedValue;
      if constexpr (std::is_same_v<CoreOpTy, pim::PimCoreBatchOp>) {
        copiedValue = pim::PimMemCopyHostToDevBatchOp::create(
                        rewriter,
                        op->getLoc(),
                        originalType,
                        deviceDst,
                        getGlobalOp.getResult(),
                        rewriter.getI32IntegerAttr(0),
                        rewriter.getI32IntegerAttr(static_cast<int32_t>(resolvedAddress->byteOffset)),
                        rewriter.getI32IntegerAttr(static_cast<int32_t>(totalBytes)))
                        .getOutput();
      }
      else {
        copiedValue = pim::PimMemCopyHostToDevOp::create(
                        rewriter,
                        op->getLoc(),
                        originalType,
                        deviceDst,
                        getGlobalOp.getResult(),
                        rewriter.getI32IntegerAttr(0),
                        rewriter.getI32IntegerAttr(static_cast<int32_t>(resolvedAddress->byteOffset)),
                        rewriter.getI32IntegerAttr(static_cast<int32_t>(totalBytes)))
                        .getOutput();
      }

      cachedByType[originalType] = copiedValue;
      operand.set(copiedValue);
    }
  }
}

struct MaterializeHostConstantsPass : PassWrapper<MaterializeHostConstantsPass, OperationPass<ModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MaterializeHostConstantsPass)

  StringRef getArgument() const override { return "materialize-pim-host-constants"; }
  StringRef getDescription() const override {
    return "Materialize explicit host-to-device copies for constant globals used by PIM runtime ops";
  }

  void runOnOperation() override {
    ModuleOp moduleOp = getOperation();
    IRRewriter rewriter(moduleOp.getContext());
    bool hasFailure = false;

    for (func::FuncOp funcOp : moduleOp.getOps<func::FuncOp>()) {
      if (funcOp.isExternal())
        continue;

      for (pim::PimCoreOp coreOp : funcOp.getOps<pim::PimCoreOp>())
        materializeHostConstantsInCore(coreOp, rewriter, hasFailure);

      for (pim::PimCoreBatchOp coreBatchOp : funcOp.getOps<pim::PimCoreBatchOp>())
        materializeHostConstantsInCore(coreBatchOp, rewriter, hasFailure);

      SmallVector<Operation*> hostCompactOps;
      for (Operation& op : funcOp.getBody().front())
        if (isa<pim::PimConcatOp>(op))
          hostCompactOps.push_back(&op);

      for (Operation* op : hostCompactOps) {
        rewriter.setInsertionPoint(op);
        auto concatOp = cast<pim::PimConcatOp>(op);
        concatOp.emitOpError("host-side concat must be folded away or lowered into pim.core before materialization");
        hasFailure = true;
      }
    }

    if (hasFailure) {
      signalPassFailure();
      return;
    }

    dumpModule(moduleOp, "pim4_materialized");
  }
};

} // namespace

std::unique_ptr<Pass> createPimMaterializeHostConstantsPass() {
  return std::make_unique<MaterializeHostConstantsPass>();
}

} // namespace onnx_mlir