This commit is contained in:
@@ -10,15 +10,11 @@ namespace onnx_mlir {
|
||||
|
||||
mlir::Block* getConstantInsertionBlock(mlir::Operation* anchorOp);
|
||||
|
||||
mlir::Value getOrCreateConstant(mlir::OperationFolder& folder,
|
||||
mlir::Operation* anchorOp,
|
||||
mlir::Attribute value,
|
||||
mlir::Type type);
|
||||
mlir::Value
|
||||
getOrCreateConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);
|
||||
|
||||
mlir::Value getOrCreateConstant(mlir::RewriterBase& rewriter,
|
||||
mlir::Operation* anchorOp,
|
||||
mlir::Attribute value,
|
||||
mlir::Type type);
|
||||
mlir::Value
|
||||
getOrCreateConstant(mlir::RewriterBase& rewriter, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);
|
||||
|
||||
mlir::Value getOrCreateConstantLike(mlir::OperationFolder& folder, mlir::arith::ConstantOp constantOp);
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
#include "mlir/Dialect/Linalg/IR/Linalg.h"
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
#include "mlir/IR/BuiltinAttributes.h"
|
||||
#include "mlir/IR/BuiltinTypes.h"
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@ void dumpModule(mlir::ModuleOp moduleOp, const std::string& name) {
|
||||
std::fstream file(dialectsDir + "/" + name + ".mlir", std::ios::out);
|
||||
llvm::raw_os_ostream os(file);
|
||||
mlir::OpPrintingFlags flags;
|
||||
flags.elideLargeElementsAttrs().enableDebugInfo(true,false);
|
||||
flags.elideLargeElementsAttrs().enableDebugInfo(true, false);
|
||||
moduleOp.print(os, flags);
|
||||
os.flush();
|
||||
file.close();
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
#include "AttributeUtils.hpp"
|
||||
|
||||
#include "mlir/IR/BuiltinAttributes.h"
|
||||
|
||||
#include "AttributeUtils.hpp"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
#pragma once
|
||||
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
#include "mlir/IR/Block.h"
|
||||
#include "mlir/IR/BuiltinTypes.h"
|
||||
#include "mlir/IR/ValueRange.h"
|
||||
@@ -11,8 +12,6 @@
|
||||
#include <type_traits>
|
||||
#include <utility>
|
||||
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
|
||||
@@ -209,8 +208,7 @@ auto createSpatComputeBatch(RewriterT& rewriter,
|
||||
block->getArgument(0),
|
||||
mlir::ValueRange(block->getArguments()).slice(1, weights.size()),
|
||||
mlir::ValueRange(block->getArguments()).slice(1 + weights.size(), inputs.size()),
|
||||
mlir::ValueRange(block->getArguments()).drop_front(1 + weights.size() + inputs.size())
|
||||
};
|
||||
mlir::ValueRange(block->getArguments()).drop_front(1 + weights.size() + inputs.size())};
|
||||
|
||||
using BodyResult = std::invoke_result_t<BodyFn, detail::SpatComputeBatchBodyArgs>;
|
||||
if constexpr (std::is_same_v<BodyResult, void>) {
|
||||
@@ -252,8 +250,8 @@ mlir::Value materializeOrComputeUnary(mlir::Value input,
|
||||
if (isCompileTimeComputable(input))
|
||||
return buildFn(input);
|
||||
|
||||
auto computeOp =
|
||||
createSpatCompute<1>(rewriter, loc, mlir::TypeRange {resultType}, {}, mlir::ValueRange {input}, [&](mlir::Value computeInput) {
|
||||
auto computeOp = createSpatCompute<1>(
|
||||
rewriter, loc, mlir::TypeRange {resultType}, {}, mlir::ValueRange {input}, [&](mlir::Value computeInput) {
|
||||
mlir::Value result = buildFn(computeInput);
|
||||
spatial::SpatYieldOp::create(rewriter, loc, result);
|
||||
});
|
||||
|
||||
@@ -1,11 +1,10 @@
|
||||
#include "IndexingUtils.hpp"
|
||||
|
||||
#include "mlir/Dialect/Arith/IR/Arith.h"
|
||||
|
||||
#include "llvm/ADT/APInt.h"
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
#include "IndexingUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
|
||||
|
||||
using namespace mlir;
|
||||
@@ -85,7 +84,8 @@ Value floorDivIndexByConstant(PatternRewriter& rewriter, Location loc, Value val
|
||||
|
||||
Value getOrMaterializeIndexValue(PatternRewriter& rewriter, OpFoldResult value) {
|
||||
if (auto attr = dyn_cast<Attribute>(value))
|
||||
return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), cast<IntegerAttr>(attr).getInt());
|
||||
return getOrCreateIndexConstant(
|
||||
rewriter, rewriter.getInsertionBlock()->getParentOp(), cast<IntegerAttr>(attr).getInt());
|
||||
return cast<Value>(value);
|
||||
}
|
||||
|
||||
|
||||
@@ -26,8 +26,10 @@ mlir::Value createAffineApplyOrFoldedConstant(mlir::PatternRewriter& rewriter,
|
||||
mlir::AffineExpr expr,
|
||||
mlir::ValueRange operands);
|
||||
|
||||
mlir::Value
|
||||
multiplyIndexByConstant(mlir::PatternRewriter& rewriter, mlir::Operation* anchorOp, mlir::Value value, int64_t multiplier);
|
||||
mlir::Value multiplyIndexByConstant(mlir::PatternRewriter& rewriter,
|
||||
mlir::Operation* anchorOp,
|
||||
mlir::Value value,
|
||||
int64_t multiplier);
|
||||
|
||||
mlir::Value modIndexByConstant(mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value value, int64_t divisor);
|
||||
|
||||
|
||||
@@ -8,8 +8,8 @@
|
||||
#include <algorithm>
|
||||
#include <functional>
|
||||
|
||||
#include "ShapeTilingUtils.hpp"
|
||||
#include "IndexingUtils.hpp"
|
||||
#include "ShapeTilingUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
|
||||
@@ -53,7 +53,9 @@ bool hasStaticPositiveShape(ArrayRef<int64_t> shape) {
|
||||
return llvm::all_of(shape, [](int64_t dim) { return dim > 0; });
|
||||
}
|
||||
|
||||
bool hasStaticPositiveShape(RankedTensorType type) { return type.hasStaticShape() && hasStaticPositiveShape(type.getShape()); }
|
||||
bool hasStaticPositiveShape(RankedTensorType type) {
|
||||
return type.hasStaticShape() && hasStaticPositiveShape(type.getShape());
|
||||
}
|
||||
|
||||
int64_t getStaticShapeElementCount(ArrayRef<int64_t> shape) {
|
||||
return std::accumulate(shape.begin(), shape.end(), int64_t {1}, std::multiplies<int64_t> {});
|
||||
@@ -98,11 +100,8 @@ FailureOr<SmallVector<int64_t>> getTransposePermutationChecked(std::optional<Arr
|
||||
return permutation;
|
||||
}
|
||||
|
||||
Value transposeMaybeInCompute(Value value,
|
||||
RankedTensorType resultType,
|
||||
ArrayRef<int64_t> permutation,
|
||||
PatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
Value transposeMaybeInCompute(
|
||||
Value value, RankedTensorType resultType, ArrayRef<int64_t> permutation, PatternRewriter& rewriter, Location loc) {
|
||||
auto buildTranspose = [&](Value input) -> Value {
|
||||
return ONNXTransposeOp::create(rewriter, loc, resultType, input, rewriter.getI64ArrayAttr(permutation)).getResult();
|
||||
};
|
||||
@@ -127,7 +126,8 @@ SmallVector<OpFoldResult> getStaticSizes(PatternRewriter& rewriter, ArrayRef<int
|
||||
|
||||
static bool isContiguousTensorSlice(Value source, RankedTensorType resultType, ArrayRef<OpFoldResult> strides) {
|
||||
auto sourceType = dyn_cast<RankedTensorType>(source.getType());
|
||||
if (!sourceType || !sourceType.hasStaticShape() || !resultType.hasStaticShape() || sourceType.getRank() != resultType.getRank())
|
||||
if (!sourceType || !sourceType.hasStaticShape() || !resultType.hasStaticShape()
|
||||
|| sourceType.getRank() != resultType.getRank())
|
||||
return false;
|
||||
|
||||
for (OpFoldResult stride : strides) {
|
||||
@@ -290,7 +290,8 @@ Value materializeContiguousTensorSlice(Value source,
|
||||
}
|
||||
|
||||
Value lower = zeroIndices[dim];
|
||||
Value upper = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), resultType.getDimSize(dim));
|
||||
Value upper =
|
||||
getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), resultType.getDimSize(dim));
|
||||
Value step = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
|
||||
auto loop = scf::ForOp::create(rewriter, loc, lower, upper, step, ValueRange {accumulator});
|
||||
rewriter.setInsertionPointToStart(loop.getBody());
|
||||
@@ -316,7 +317,8 @@ Value extractAxisSlice(
|
||||
SmallVector<OpFoldResult> sizes = getStaticSizes(rewriter, sourceType.getShape());
|
||||
offsets[axis] = rewriter.getIndexAttr(offset);
|
||||
sizes[axis] = rewriter.getIndexAttr(size);
|
||||
return tensor::ExtractSliceOp::create(rewriter, loc, resultType, source, offsets, sizes, getUnitStrides(rewriter, sourceType.getRank()))
|
||||
return tensor::ExtractSliceOp::create(
|
||||
rewriter, loc, resultType, source, offsets, sizes, getUnitStrides(rewriter, sourceType.getRank()))
|
||||
.getResult();
|
||||
}
|
||||
|
||||
|
||||
@@ -115,12 +115,8 @@ mlir::Value materializeContiguousTensorSlice(mlir::Value source,
|
||||
mlir::ConversionPatternRewriter& rewriter,
|
||||
mlir::Location loc);
|
||||
|
||||
mlir::Value extractAxisSlice(mlir::PatternRewriter& rewriter,
|
||||
mlir::Location loc,
|
||||
mlir::Value source,
|
||||
int64_t axis,
|
||||
int64_t offset,
|
||||
int64_t size);
|
||||
mlir::Value extractAxisSlice(
|
||||
mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value source, int64_t axis, int64_t offset, int64_t size);
|
||||
|
||||
mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter,
|
||||
mlir::Location loc,
|
||||
|
||||
@@ -15,8 +15,8 @@
|
||||
#include "Common/Common.hpp"
|
||||
#include "Common/PimCommon.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
#include "src/Dialect/ONNX/ONNXOps.hpp"
|
||||
|
||||
|
||||
@@ -69,9 +69,9 @@ LogicalResult verifyComputeLikeInputs(Operation* computeLikeOp,
|
||||
continue;
|
||||
|
||||
diagnostics.report(computeLikeOp, [&](Operation* illegalOp) {
|
||||
InFlightDiagnostic diagnostic = illegalOp->emitOpError() << kind << " input #" << currentInputIndex
|
||||
<< (allowChannelReceiveInputs
|
||||
? " must come from the host or an explicit "
|
||||
InFlightDiagnostic diagnostic = illegalOp->emitOpError()
|
||||
<< kind << " input #" << currentInputIndex
|
||||
<< (allowChannelReceiveInputs ? " must come from the host or an explicit "
|
||||
"spat.channel_receive"
|
||||
: " must come from the host");
|
||||
if (definingOp)
|
||||
@@ -135,13 +135,13 @@ LogicalResult verifySpatialCommunicationInvariants(func::FuncOp funcOp) {
|
||||
pim::CappedDiagnosticReporter diagnostics;
|
||||
|
||||
for (auto computeOp : funcOp.getOps<spatial::SpatCompute>()) {
|
||||
(void)verifyComputeLikeInputs(
|
||||
(void) verifyComputeLikeInputs(
|
||||
computeOp.getOperation(), computeOp.getInputs(), /*allowChannelReceiveInputs=*/true, "spat.compute", diagnostics);
|
||||
verifyNoExternalTensorCaptures(computeOp.getOperation(), computeOp.getBody(), "spat.compute", diagnostics);
|
||||
}
|
||||
|
||||
for (auto computeBatchOp : funcOp.getOps<spatial::SpatComputeBatch>()) {
|
||||
(void)verifyComputeLikeInputs(computeBatchOp.getOperation(),
|
||||
(void) verifyComputeLikeInputs(computeBatchOp.getOperation(),
|
||||
computeBatchOp.getInputs(),
|
||||
/*allowChannelReceiveInputs=*/false,
|
||||
"spat.compute_batch",
|
||||
|
||||
@@ -5,9 +5,7 @@ using namespace mlir;
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
void populatePrePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
|
||||
populateGeneratedPrePatterns(patterns, ctx);
|
||||
}
|
||||
void populatePrePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { populateGeneratedPrePatterns(patterns, ctx); }
|
||||
|
||||
void populateConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
|
||||
populateGeneratedConversionPatterns(patterns, ctx);
|
||||
|
||||
@@ -51,8 +51,8 @@ static Value createPaddedRows(Value tensorValue,
|
||||
if (tensorType.getDimSize(0) == paddedRows)
|
||||
return tensorValue;
|
||||
|
||||
auto paddedType =
|
||||
RankedTensorType::get({paddedRows, tensorType.getDimSize(1)}, tensorType.getElementType(), tensorType.getEncoding());
|
||||
auto paddedType = RankedTensorType::get(
|
||||
{paddedRows, tensorType.getDimSize(1)}, tensorType.getElementType(), tensorType.getEncoding());
|
||||
SmallVector<OpFoldResult> lowPads = {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
|
||||
SmallVector<OpFoldResult> highPads = {rewriter.getIndexAttr(paddedRows - tensorType.getDimSize(0)),
|
||||
rewriter.getIndexAttr(0)};
|
||||
@@ -62,20 +62,15 @@ static Value createPaddedRows(Value tensorValue,
|
||||
padBlock->addArgument(rewriter.getIndexType(), loc);
|
||||
padOp.getRegion().push_back(padBlock);
|
||||
rewriter.setInsertionPointToStart(padBlock);
|
||||
auto zero = getOrCreateConstant(rewriter,
|
||||
padOp.getOperation(),
|
||||
rewriter.getZeroAttr(tensorType.getElementType()),
|
||||
tensorType.getElementType());
|
||||
auto zero = getOrCreateConstant(
|
||||
rewriter, padOp.getOperation(), rewriter.getZeroAttr(tensorType.getElementType()), tensorType.getElementType());
|
||||
tensor::YieldOp::create(rewriter, loc, zero);
|
||||
rewriter.setInsertionPointAfter(padOp);
|
||||
return padOp.getResult();
|
||||
}
|
||||
|
||||
static Value packRowsForParallelGemm(Value rows,
|
||||
RankedTensorType rowsType,
|
||||
int64_t packFactor,
|
||||
ConversionPatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
static Value packRowsForParallelGemm(
|
||||
Value rows, RankedTensorType rowsType, int64_t packFactor, ConversionPatternRewriter& rewriter, Location loc) {
|
||||
if (packFactor == 1)
|
||||
return rows;
|
||||
|
||||
@@ -118,10 +113,8 @@ static Value unpackRowsFromParallelGemm(Value packedRows,
|
||||
|
||||
const int64_t packedNumRows = packedRowsType.getDimSize(0);
|
||||
const int64_t paddedNumRows = packedNumRows * packFactor;
|
||||
auto expandedType =
|
||||
RankedTensorType::get({packedNumRows, packFactor, rowWidth},
|
||||
packedRowsType.getElementType(),
|
||||
packedRowsType.getEncoding());
|
||||
auto expandedType = RankedTensorType::get(
|
||||
{packedNumRows, packFactor, rowWidth}, packedRowsType.getElementType(), packedRowsType.getEncoding());
|
||||
auto paddedType =
|
||||
RankedTensorType::get({paddedNumRows, rowWidth}, packedRowsType.getElementType(), packedRowsType.getEncoding());
|
||||
auto unpackedType =
|
||||
@@ -193,11 +186,8 @@ static Value buildPackedWeight(DenseElementsAttr wDenseAttr,
|
||||
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), packedAttr, packedWeightType);
|
||||
}
|
||||
|
||||
static Value createConvWeightMatrix(Value w,
|
||||
RankedTensorType wFlatType,
|
||||
RankedTensorType wTransType,
|
||||
ConversionPatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
static Value createConvWeightMatrix(
|
||||
Value w, RankedTensorType wFlatType, RankedTensorType wTransType, ConversionPatternRewriter& rewriter, Location loc) {
|
||||
auto buildWeightMatrix = [&](Value weight) -> Value {
|
||||
Value wFlat = tensor::CollapseShapeOp::create(rewriter,
|
||||
loc,
|
||||
@@ -360,9 +350,8 @@ static Value createIm2colRowComputes(Value x,
|
||||
Value im2col = im2colLoop.getResult(0);
|
||||
|
||||
Value gemmInputRows = im2col;
|
||||
if (packFactor != 1) {
|
||||
if (packFactor != 1)
|
||||
gemmInputRows = packRowsForParallelGemm(im2col, im2colType, packFactor, rewriter, loc);
|
||||
}
|
||||
|
||||
spatial::SpatYieldOp::create(rewriter, loc, gemmInputRows);
|
||||
});
|
||||
@@ -387,8 +376,13 @@ static Value createCollectedConvOutput(ValueRange gemmRows,
|
||||
}
|
||||
else {
|
||||
Value packedOutput = createSpatConcat(rewriter, loc, /*axis=*/0, gemmRowArgs);
|
||||
gemmOut = unpackRowsFromParallelGemm(
|
||||
packedOutput, cast<RankedTensorType>(packedOutput.getType()), numPatches, numChannelsOut, packFactor, rewriter, loc);
|
||||
gemmOut = unpackRowsFromParallelGemm(packedOutput,
|
||||
cast<RankedTensorType>(packedOutput.getType()),
|
||||
numPatches,
|
||||
numChannelsOut,
|
||||
packFactor,
|
||||
rewriter,
|
||||
loc);
|
||||
}
|
||||
|
||||
// Restore to NCHW layout:
|
||||
|
||||
@@ -252,7 +252,13 @@ static spatial::SpatComputeBatch createVmmBatch(Value a,
|
||||
Location loc) {
|
||||
const int64_t laneCount = partialPiecesType.getDimSize(0);
|
||||
auto batchOp = createSpatComputeBatch(
|
||||
rewriter, loc, TypeRange {partialPiecesType}, laneCount, ValueRange {b}, ValueRange {a}, [&](detail::SpatComputeBatchBodyArgs args) {
|
||||
rewriter,
|
||||
loc,
|
||||
TypeRange {partialPiecesType},
|
||||
laneCount,
|
||||
ValueRange {b},
|
||||
ValueRange {a},
|
||||
[&](detail::SpatComputeBatchBodyArgs args) {
|
||||
Value row = onnx_mlir::modIndexByConstant(rewriter, loc, args.lane, numOutRows);
|
||||
Value kOffset = createGemmBatchKOffset(args.lane, numOutRows, numKSlices, rewriter, loc);
|
||||
Value hOffset = createGemmBatchHOffset(args.lane, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
|
||||
@@ -284,8 +290,8 @@ static spatial::SpatComputeBatch createVmmBatch(Value a,
|
||||
return *batchOp;
|
||||
}
|
||||
|
||||
static Value createDynamicGemmBatchRow(
|
||||
Value lane, int64_t numOutCols, ConversionPatternRewriter& rewriter, Location loc) {
|
||||
static Value
|
||||
createDynamicGemmBatchRow(Value lane, int64_t numOutCols, ConversionPatternRewriter& rewriter, Location loc) {
|
||||
if (numOutCols == 1)
|
||||
return lane;
|
||||
|
||||
@@ -294,17 +300,21 @@ static Value createDynamicGemmBatchRow(
|
||||
return createAffineApplyOrFoldedConstant(rewriter, loc, d0.floorDiv(numOutCols), ValueRange {lane});
|
||||
}
|
||||
|
||||
static Value
|
||||
extractDynamicGemmBColumn(Value matrix, Value column, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) {
|
||||
static Value extractDynamicGemmBColumn(
|
||||
Value matrix, Value column, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) {
|
||||
SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), column};
|
||||
SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
|
||||
auto columnSliceType = RankedTensorType::get({vectorType.getDimSize(1), 1}, vectorType.getElementType());
|
||||
Value columnSlice = materializeContiguousTensorSlice(matrix, columnSliceType, offsets, strides, rewriter, loc);
|
||||
SmallVector<ReassociationIndices> collapseReassociation {ReassociationIndices {0, 1}};
|
||||
SmallVector<ReassociationIndices> collapseReassociation {
|
||||
ReassociationIndices {0, 1}
|
||||
};
|
||||
auto collapsedType = RankedTensorType::get({vectorType.getDimSize(1)}, vectorType.getElementType());
|
||||
Value collapsed =
|
||||
tensor::CollapseShapeOp::create(rewriter, loc, collapsedType, columnSlice, collapseReassociation).getResult();
|
||||
SmallVector<ReassociationIndices> expandReassociation {ReassociationIndices {0, 1}};
|
||||
SmallVector<ReassociationIndices> expandReassociation {
|
||||
ReassociationIndices {0, 1}
|
||||
};
|
||||
return tensor::ExpandShapeOp::create(rewriter, loc, vectorType, collapsed, expandReassociation).getResult();
|
||||
}
|
||||
|
||||
@@ -371,13 +381,15 @@ static Value createBroadcastedBiasScalar(Value bias,
|
||||
Location loc) {
|
||||
SmallVector<OpFoldResult> unitStrides(biasType.getRank(), rewriter.getIndexAttr(1));
|
||||
if (biasType.getRank() == 1) {
|
||||
SmallVector<OpFoldResult> offsets {
|
||||
biasType.getDimSize(0) == 1 ? OpFoldResult(rewriter.getIndexAttr(0)) : OpFoldResult(column)};
|
||||
SmallVector<OpFoldResult> offsets {biasType.getDimSize(0) == 1 ? OpFoldResult(rewriter.getIndexAttr(0))
|
||||
: OpFoldResult(column)};
|
||||
SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1)};
|
||||
auto vectorType = RankedTensorType::get({1}, scalarType.getElementType());
|
||||
Value vector = tensor::ExtractSliceOp::create(rewriter, loc, vectorType, bias, offsets, sizes, unitStrides)
|
||||
.getResult();
|
||||
SmallVector<ReassociationIndices> reassociation {ReassociationIndices {0, 1}};
|
||||
Value vector =
|
||||
tensor::ExtractSliceOp::create(rewriter, loc, vectorType, bias, offsets, sizes, unitStrides).getResult();
|
||||
SmallVector<ReassociationIndices> reassociation {
|
||||
ReassociationIndices {0, 1}
|
||||
};
|
||||
return tensor::ExpandShapeOp::create(rewriter, loc, scalarType, vector, reassociation).getResult();
|
||||
}
|
||||
|
||||
@@ -407,15 +419,20 @@ static spatial::SpatComputeBatch createVvdmulBatch(Value a,
|
||||
const int64_t reductionSize = aType.getDimSize(1);
|
||||
const int64_t laneCount = numOutRows * numOutCols;
|
||||
auto batchOp = createSpatComputeBatch(
|
||||
rewriter, loc, TypeRange {scalarPiecesType}, laneCount, ValueRange {}, ValueRange {a, b}, [&](detail::SpatComputeBatchBodyArgs args) {
|
||||
rewriter,
|
||||
loc,
|
||||
TypeRange {scalarPiecesType},
|
||||
laneCount,
|
||||
ValueRange {},
|
||||
ValueRange {a, b},
|
||||
[&](detail::SpatComputeBatchBodyArgs args) {
|
||||
Value row = createDynamicGemmBatchRow(args.lane, numOutCols, rewriter, loc);
|
||||
Value column = onnx_mlir::modIndexByConstant(rewriter, loc, args.lane, numOutCols);
|
||||
|
||||
auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
|
||||
auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
|
||||
Value aVector = extractDynamicGemmRowVector(args.inputs[0], row, vectorType, rewriter, loc);
|
||||
Value bVector = bAlreadyTransposed
|
||||
? extractTransposedBRow(args.inputs[1], column, vectorType, rewriter, loc)
|
||||
Value bVector = bAlreadyTransposed ? extractTransposedBRow(args.inputs[1], column, vectorType, rewriter, loc)
|
||||
: extractDynamicGemmBColumn(args.inputs[1], column, vectorType, rewriter, loc);
|
||||
Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
|
||||
|
||||
@@ -578,9 +595,8 @@ static spatial::SpatCompute createReductionCompute(Value partialPieces,
|
||||
auto buildOutputSlice = [&](Value outputAcc, Value hSlice) -> Value {
|
||||
Value reduced =
|
||||
reducePartialPiecesForHSlice(partialPiecesArg, hSlice, pieceType, numKSlices, numOutRows, rewriter, loc);
|
||||
Value hOffset =
|
||||
onnx_mlir::multiplyIndexByConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), hSlice,
|
||||
crossbarSize.getValue());
|
||||
Value hOffset = onnx_mlir::multiplyIndexByConstant(
|
||||
rewriter, rewriter.getInsertionBlock()->getParentOp(), hSlice, crossbarSize.getValue());
|
||||
if (biasArg) {
|
||||
SmallVector<OpFoldResult> biasOffsets {rewriter.getIndexAttr(0), hOffset};
|
||||
Value biasSlice =
|
||||
@@ -721,8 +737,8 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
|
||||
}
|
||||
|
||||
auto scalarPiecesType = RankedTensorType::get({laneCount64, 1}, outType.getElementType());
|
||||
auto batchOp = createVvdmulBatch(
|
||||
a, b, aType, bType, scalarPiecesType, outType, gemmOpAdaptor.getTransB(), rewriter, loc);
|
||||
auto batchOp =
|
||||
createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, gemmOpAdaptor.getTransB(), rewriter, loc);
|
||||
auto outputCompute = createDynamicGemmOutputCompute(
|
||||
batchOp.getResult(0), hasC ? c : Value(), scalarPiecesType, biasType, outType, alpha, beta, rewriter, loc);
|
||||
rewriter.replaceOp(gemmOp, outputCompute.getResults());
|
||||
|
||||
@@ -70,11 +70,8 @@ static SmallVector<int64_t> getKeptAxes(ArrayRef<bool> reducedAxes) {
|
||||
return keptAxes;
|
||||
}
|
||||
|
||||
static Value computeLaneIndex(Value lane,
|
||||
int64_t stride,
|
||||
int64_t dimSize,
|
||||
ConversionPatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
static Value
|
||||
computeLaneIndex(Value lane, int64_t stride, int64_t dimSize, ConversionPatternRewriter& rewriter, Location loc) {
|
||||
if (dimSize == 1)
|
||||
return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
|
||||
|
||||
@@ -119,8 +116,14 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
|
||||
sliceSizes.reserve(inputType.getRank());
|
||||
insertOffsets.reserve(inputType.getRank());
|
||||
|
||||
auto batchOp = createSpatComputeBatch(
|
||||
rewriter, loc, TypeRange {batchType}, laneCount, {}, ValueRange {input}, [&](detail::SpatComputeBatchBodyArgs args) {
|
||||
auto batchOp =
|
||||
createSpatComputeBatch(rewriter,
|
||||
loc,
|
||||
TypeRange {batchType},
|
||||
laneCount,
|
||||
{},
|
||||
ValueRange {input},
|
||||
[&](detail::SpatComputeBatchBodyArgs args) {
|
||||
size_t keptAxisIndex = 0;
|
||||
sliceOffsets.clear();
|
||||
sliceSizes.clear();
|
||||
@@ -132,8 +135,8 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
|
||||
continue;
|
||||
}
|
||||
|
||||
Value axisIndex =
|
||||
computeLaneIndex(args.lane, keptAxisStrides[keptAxisIndex], inputType.getDimSize(axis), rewriter, loc);
|
||||
Value axisIndex = computeLaneIndex(
|
||||
args.lane, keptAxisStrides[keptAxisIndex], inputType.getDimSize(axis), rewriter, loc);
|
||||
++keptAxisIndex;
|
||||
sliceOffsets.push_back(axisIndex);
|
||||
sliceSizes.push_back(rewriter.getIndexAttr(1));
|
||||
@@ -142,8 +145,8 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
|
||||
insertOffsets.push_back(args.lane);
|
||||
insertOffsets.append(inputType.getRank() - 1, rewriter.getIndexAttr(0));
|
||||
|
||||
Value slice =
|
||||
tensor::ExtractSliceOp::create(rewriter, loc, sliceType, args.inputs.front(), sliceOffsets, sliceSizes, unitStrides);
|
||||
Value slice = tensor::ExtractSliceOp::create(
|
||||
rewriter, loc, sliceType, args.inputs.front(), sliceOffsets, sliceSizes, unitStrides);
|
||||
Value reduced = spatial::SpatVAvgOp::create(rewriter, loc, leafType, slice).getResult();
|
||||
createParallelInsertSliceIntoBatchOutput(
|
||||
rewriter, loc, reduced, args.outputs.front(), insertOffsets, insertSizes, unitStrides);
|
||||
@@ -193,15 +196,15 @@ static Value buildKeepdimsFromLanePackedBatch(Value batchValue,
|
||||
|
||||
auto reshapeCompute =
|
||||
createSpatCompute<1>(rewriter, loc, TypeRange {keepdimsType}, {}, ValueRange {batchValue}, [&](Value input) {
|
||||
auto flatType = RankedTensorType::get({batchType.getDimSize(0)}, batchType.getElementType(), batchType.getEncoding());
|
||||
auto flatType =
|
||||
RankedTensorType::get({batchType.getDimSize(0)}, batchType.getElementType(), batchType.getEncoding());
|
||||
Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, input, collapseToFlat);
|
||||
Value compact = flat;
|
||||
if (compactKeptType != flatType)
|
||||
compact = tensor::ExpandShapeOp::create(rewriter, loc, compactKeptType, flat, expandFlatToCompact);
|
||||
Value keepdims = compact;
|
||||
if (keepdimsType != compactKeptType)
|
||||
keepdims =
|
||||
tensor::ExpandShapeOp::create(rewriter, loc, keepdimsType, compact, expandCompactToKeepdims);
|
||||
keepdims = tensor::ExpandShapeOp::create(rewriter, loc, keepdimsType, compact, expandCompactToKeepdims);
|
||||
spatial::SpatYieldOp::create(rewriter, loc, keepdims);
|
||||
});
|
||||
return reshapeCompute.getResult(0);
|
||||
|
||||
@@ -121,11 +121,9 @@ struct SoftmaxToSpatialCompute : OpConversionPattern<ONNXSoftmaxOp> {
|
||||
|
||||
auto transposedType = RankedTensorType::get(
|
||||
permuteShape(inputType.getShape(), permutation), inputType.getElementType(), inputType.getEncoding());
|
||||
Value transposedInput =
|
||||
transposeMaybeInCompute(input, transposedType, permutation, rewriter, softmaxOp.getLoc());
|
||||
Value transposedInput = transposeMaybeInCompute(input, transposedType, permutation, rewriter, softmaxOp.getLoc());
|
||||
Value transposedResult = createLoopSoftmaxCompute(transposedInput, rewriter, softmaxOp.getLoc());
|
||||
result = transposeMaybeInCompute(
|
||||
transposedResult, inputType, inversePermutation, rewriter, softmaxOp.getLoc());
|
||||
result = transposeMaybeInCompute(transposedResult, inputType, inversePermutation, rewriter, softmaxOp.getLoc());
|
||||
}
|
||||
|
||||
rewriter.replaceOp(softmaxOp, result);
|
||||
|
||||
@@ -77,7 +77,7 @@ static FailureOr<PromotedOperands> computePromotedOperands(ComputeOpTy compute)
|
||||
needsRewrite = true;
|
||||
continue;
|
||||
|
||||
keep_input:
|
||||
keep_input:
|
||||
promoted.newInputs.push_back(input);
|
||||
promoted.newInputTypes.push_back(input.getType());
|
||||
promoted.newInputLocs.push_back(input.getLoc());
|
||||
@@ -127,8 +127,8 @@ struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatCom
|
||||
Block& oldBlock = compute.getBody().front();
|
||||
|
||||
rewriter.setInsertionPointAfter(compute);
|
||||
auto newCompute =
|
||||
spatial::SpatCompute::create(rewriter, compute.getLoc(), compute.getResultTypes(), promoted->newWeights, promoted->newInputs);
|
||||
auto newCompute = spatial::SpatCompute::create(
|
||||
rewriter, compute.getLoc(), compute.getResultTypes(), promoted->newWeights, promoted->newInputs);
|
||||
SmallVector<Type> newBlockArgTypes;
|
||||
SmallVector<Location> newBlockArgLocs;
|
||||
for (Value weight : promoted->newWeights) {
|
||||
@@ -155,7 +155,12 @@ struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatCom
|
||||
mapper.map(*oldWeightArg, *newWeightArg);
|
||||
}
|
||||
if (failed(mapPromotedInputArguments(
|
||||
compute, *promoted, bodyRewriter, mapper, [&](size_t index) { return newCompute.getInputArgument(index); }, rewriter)))
|
||||
compute,
|
||||
*promoted,
|
||||
bodyRewriter,
|
||||
mapper,
|
||||
[&](size_t index) { return newCompute.getInputArgument(index); },
|
||||
rewriter)))
|
||||
return failure();
|
||||
|
||||
for (Operation& op : oldBlock.without_terminator())
|
||||
@@ -199,7 +204,8 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
|
||||
return rewriter.notifyMatchFailure(compute, "missing compute_batch lane block argument");
|
||||
SmallVector<Type> newBlockArgTypes;
|
||||
SmallVector<Location> newBlockArgLocs;
|
||||
newBlockArgTypes.reserve(1 + promoted->newWeights.size() + promoted->newInputTypes.size() + compute.getNumResults());
|
||||
newBlockArgTypes.reserve(1 + promoted->newWeights.size() + promoted->newInputTypes.size()
|
||||
+ compute.getNumResults());
|
||||
newBlockArgLocs.reserve(1 + promoted->newWeights.size() + promoted->newInputLocs.size() + compute.getNumResults());
|
||||
newBlockArgTypes.push_back(laneArg->getType());
|
||||
newBlockArgLocs.push_back(laneArg->getLoc());
|
||||
@@ -239,7 +245,12 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
|
||||
mapper.map(*oldWeightArg, *newWeightArg);
|
||||
}
|
||||
if (failed(mapPromotedInputArguments(
|
||||
compute, *promoted, bodyRewriter, mapper, [&](size_t index) { return newCompute.getInputArgument(index); }, rewriter)))
|
||||
compute,
|
||||
*promoted,
|
||||
bodyRewriter,
|
||||
mapper,
|
||||
[&](size_t index) { return newCompute.getInputArgument(index); },
|
||||
rewriter)))
|
||||
return failure();
|
||||
for (auto resultIndex : llvm::seq<size_t>(0, compute.getNumResults())) {
|
||||
auto outputArg = compute.getOutputArgument(resultIndex);
|
||||
|
||||
@@ -111,7 +111,8 @@ struct Reshape : OpConversionPattern<ONNXReshapeOp> {
|
||||
}
|
||||
|
||||
auto replaceWithReshape = [&](auto buildReshape) -> LogicalResult {
|
||||
Value reshaped = materializeOrComputeUnary(adaptor.getData(), resultType, rewriter, reshapeOp.getLoc(), buildReshape);
|
||||
Value reshaped =
|
||||
materializeOrComputeUnary(adaptor.getData(), resultType, rewriter, reshapeOp.getLoc(), buildReshape);
|
||||
rewriter.replaceOp(reshapeOp, reshaped);
|
||||
return success();
|
||||
};
|
||||
|
||||
@@ -44,8 +44,7 @@ struct Split : OpConversionPattern<ONNXSplitOp> {
|
||||
|
||||
if (isCompileTimeComputable(adaptor.getInput())) {
|
||||
for (int64_t sliceSize : sliceSizes) {
|
||||
outputs.push_back(
|
||||
extractAxisSlice(rewriter, splitOp.getLoc(), adaptor.getInput(), *axis, offset, sliceSize));
|
||||
outputs.push_back(extractAxisSlice(rewriter, splitOp.getLoc(), adaptor.getInput(), *axis, offset, sliceSize));
|
||||
offset += sliceSize;
|
||||
}
|
||||
rewriter.replaceOp(splitOp, outputs);
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
#include "mlir/Dialect/Linalg/IR/Linalg.h"
|
||||
#include "mlir/Dialect/Arith/IR/Arith.h"
|
||||
#include "mlir/Dialect/Linalg/IR/Linalg.h"
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
#include "mlir/Transforms/DialectConversion.h"
|
||||
|
||||
@@ -104,8 +104,7 @@ struct TransposeToLinalgTranspose : OpConversionPattern<ONNXTransposeOp> {
|
||||
}
|
||||
Value init = createTransposeInit(adaptor.getData(), resultType, *permutation, rewriter, transposeOp.getLoc());
|
||||
Value transposed =
|
||||
linalg::TransposeOp::create(rewriter, transposeOp.getLoc(), adaptor.getData(), init, *permutation)
|
||||
.getResult()[0];
|
||||
linalg::TransposeOp::create(rewriter, transposeOp.getLoc(), adaptor.getData(), init, *permutation).getResult()[0];
|
||||
rewriter.replaceOp(transposeOp, transposed);
|
||||
return success();
|
||||
}
|
||||
|
||||
@@ -7,8 +7,8 @@
|
||||
#include "mlir/IR/Matchers.h"
|
||||
|
||||
#include "Conversion/ONNXToSpatial/Common/Common.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
|
||||
#include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
|
||||
#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
|
||||
|
||||
@@ -334,8 +334,8 @@ LogicalResult raptor::SpatialToPimPass::allocateAndInitializeCoreLocalVariables(
|
||||
loc,
|
||||
tensorType,
|
||||
getOrCreateIndexConstant(constantFolder, deviceTensor.getOperation(), 0),
|
||||
getOrCreateIndexConstant(constantFolder,
|
||||
deviceTensor.getOperation(), static_cast<int64_t>(elementsOffset * elementByteSize) ),
|
||||
getOrCreateIndexConstant(
|
||||
constantFolder, deviceTensor.getOperation(), static_cast<int64_t>(elementsOffset * elementByteSize)),
|
||||
deviceTensor,
|
||||
inputTensor,
|
||||
rewriter.getI32IntegerAttr(static_cast<int32_t>(tensorType.getNumElements() * elementByteSize)));
|
||||
|
||||
@@ -1482,7 +1482,8 @@ void appendScalarSendLoop(MaterializerState& state,
|
||||
state.rewriter.setInsertionPoint(sourceClass.body->getTerminator());
|
||||
|
||||
Value lowerBound = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 0);
|
||||
Value upperBound = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, static_cast<int64_t>(channelIds.size()));
|
||||
Value upperBound =
|
||||
getOrCreateIndexConstant(state.constantFolder, sourceClass.op, static_cast<int64_t>(channelIds.size()));
|
||||
Value step = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 1);
|
||||
|
||||
auto loop = scf::ForOp::create(state.rewriter, loc, lowerBound, upperBound, step, ValueRange {});
|
||||
@@ -1577,7 +1578,8 @@ void appendProjectedScalarSendLoop(MaterializerState& state,
|
||||
}
|
||||
|
||||
Value lowerBound = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 0);
|
||||
Value upperBound = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, static_cast<int64_t>(channelIds.size()));
|
||||
Value upperBound =
|
||||
getOrCreateIndexConstant(state.constantFolder, sourceClass.op, static_cast<int64_t>(channelIds.size()));
|
||||
Value step = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 1);
|
||||
|
||||
auto loop = scf::ForOp::create(state.rewriter, loc, lowerBound, upperBound, step, ValueRange {});
|
||||
@@ -2342,7 +2344,8 @@ FailureOr<Value> insertPackedScalarRunIntoWholeBatch(MaterializerState& state,
|
||||
}
|
||||
|
||||
Value lowerBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 0);
|
||||
Value upperBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, static_cast<int64_t>(run.slots.size()));
|
||||
Value upperBound =
|
||||
getOrCreateIndexConstant(state.constantFolder, targetClass.op, static_cast<int64_t>(run.slots.size()));
|
||||
Value step = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 1);
|
||||
|
||||
state.rewriter.setInsertionPoint(targetClass.body->getTerminator());
|
||||
|
||||
@@ -99,13 +99,12 @@ static void materializeHostConstantsInCore(CoreOpTy coreOp,
|
||||
.getOutput();
|
||||
}
|
||||
else {
|
||||
copiedValue =
|
||||
pim::PimMemCopyHostToDevOp::create(
|
||||
copiedValue = pim::PimMemCopyHostToDevOp::create(
|
||||
rewriter,
|
||||
op->getLoc(),
|
||||
originalType,
|
||||
getOrCreateIndexConstant(constantFolder, op, 0),
|
||||
getOrCreateIndexConstant(constantFolder, op, static_cast<int64_t>(resolvedAddress->byteOffset) ),
|
||||
getOrCreateIndexConstant(constantFolder, op, static_cast<int64_t>(resolvedAddress->byteOffset)),
|
||||
deviceDst,
|
||||
getGlobalOp.getResult(),
|
||||
rewriter.getI32IntegerAttr(static_cast<int32_t>(totalBytes)))
|
||||
|
||||
Reference in New Issue
Block a user