NiccoloN
58 lines
2.1 KiB
C++
58 lines
2.1 KiB
C++
#include <cassert>
|
|
|
|
#include "src/Accelerators/PIM/Conversion/SpatialToPim/ComputeLikeRegionUtils.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace onnx_mlir {
|
|
|
|
std::optional<unsigned> getDirectComputeLikeInputIndex(Operation* owner, unsigned operandNumber) {
|
|
auto getInputIndex = [operandNumber](Operation* op, unsigned inputCount) -> std::optional<unsigned> {
|
|
if (inputCount == 0)
|
|
return std::nullopt;
|
|
unsigned inputBegin = op->getNumOperands() - inputCount;
|
|
if (operandNumber < inputBegin)
|
|
return std::nullopt;
|
|
return operandNumber - inputBegin;
|
|
};
|
|
|
|
if (auto compute = dyn_cast<spatial::SpatCompute>(owner))
|
|
return getInputIndex(owner, compute.getInputs().size());
|
|
|
|
if (auto computeBatch = dyn_cast<spatial::SpatComputeBatch>(owner))
|
|
return getInputIndex(owner, computeBatch.getInputs().size());
|
|
|
|
return std::nullopt;
|
|
}
|
|
|
|
void replaceAndEraseDirectComputeLikeInput(PatternRewriter& rewriter,
|
|
Operation* owner,
|
|
unsigned inputIndex,
|
|
Value replacement) {
|
|
Block& body = owner->getRegion(0).front();
|
|
BlockArgument bodyArgument;
|
|
if (auto compute = dyn_cast<spatial::SpatCompute>(owner)) {
|
|
auto computeArg = compute.getInputArgument(inputIndex);
|
|
assert(computeArg && "expected compute input block argument");
|
|
bodyArgument = *computeArg;
|
|
}
|
|
else {
|
|
auto batchArg = cast<spatial::SpatComputeBatch>(owner).getInputArgument(inputIndex);
|
|
assert(batchArg && "expected compute_batch input block argument");
|
|
bodyArgument = *batchArg;
|
|
}
|
|
unsigned bodyArgIndex = bodyArgument.getArgNumber();
|
|
|
|
rewriter.startOpModification(owner);
|
|
bodyArgument.replaceAllUsesWith(replacement);
|
|
if (auto compute = dyn_cast<spatial::SpatCompute>(owner))
|
|
compute.getInputsMutable().erase(inputIndex);
|
|
else
|
|
cast<spatial::SpatComputeBatch>(owner).getInputsMutable().erase(inputIndex);
|
|
body.eraseArgument(bodyArgIndex);
|
|
rewriter.finalizeOpModification(owner);
|
|
}
|
|
|
|
} // namespace onnx_mlir
|