40 lines
1.2 KiB
C++
40 lines
1.2 KiB
C++
#include "mlir/Transforms/DialectConversion.h"
|
|
|
|
#include "llvm/ADT/SmallVector.h"
|
|
|
|
#include "ComputeRegionBuilder.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace onnx_mlir {
|
|
|
|
Value sumTensors(ArrayRef<Value> tensors, ConversionPatternRewriter& rewriter) {
|
|
if (tensors.size() == 1)
|
|
return tensors[0];
|
|
|
|
SmallVector<Value> tensors1 = {tensors.begin(), tensors.end()};
|
|
SmallVector<Value> tensors2;
|
|
tensors2.reserve(tensors.size() / 2);
|
|
|
|
auto* currTensors = &tensors1;
|
|
auto* nextTensors = &tensors2;
|
|
while (currTensors->size() > 1) {
|
|
for (size_t i = 0; i < currTensors->size() - 1; i += 2) {
|
|
Value a = (*currTensors)[i];
|
|
Value b = (*currTensors)[i + 1];
|
|
rewriter.setInsertionPointAfterValue(b);
|
|
auto addedValue = spatial::SpatVAddOp::create(rewriter, a.getLoc(), a.getType(), a, b);
|
|
nextTensors->push_back(addedValue);
|
|
}
|
|
if (currTensors->size() % 2 == 1)
|
|
nextTensors->push_back(currTensors->back());
|
|
std::swap(currTensors, nextTensors);
|
|
nextTensors->clear();
|
|
}
|
|
assert(currTensors->size() == 1 && "Expected a single input at this point.");
|
|
return (*currTensors)[0];
|
|
}
|
|
|
|
} // namespace onnx_mlir
|