This commit is contained in:
@@ -181,32 +181,32 @@ analyzeTopLevelFragmentAssemblyUses(OpResult result, RankedTensorType packedResu
|
||||
size_t elementSize = getElementTypeSizeInBytes(packedResultType.getElementType());
|
||||
|
||||
for (OpOperand& use : result.getUses()) {
|
||||
auto reconciliator = dyn_cast<spatial::SpatReconciliatorOp>(use.getOwner());
|
||||
if (!reconciliator || reconciliator->getParentOp() != reconciliator->getParentOfType<func::FuncOp>())
|
||||
auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(use.getOwner());
|
||||
if (!blueprint || blueprint->getParentOp() != blueprint->getParentOfType<func::FuncOp>())
|
||||
return failure();
|
||||
std::optional<StringRef> mode = reconciliator.getMode();
|
||||
std::optional<ArrayRef<int64_t>> operandIndicesAttr = reconciliator.getFragmentOperandIndices();
|
||||
std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = reconciliator.getFragmentSourceOffsets();
|
||||
std::optional<ArrayRef<int64_t>> stridesAttr = reconciliator.getFragmentStrides();
|
||||
std::optional<StringRef> mode = blueprint.getMode();
|
||||
std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
|
||||
std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
|
||||
std::optional<ArrayRef<int64_t>> stridesAttr = blueprint.getFragmentStrides();
|
||||
if (!mode || *mode != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr || !stridesAttr)
|
||||
return failure();
|
||||
if (!reconciliator.getOutput().hasOneUse() || !isa<func::ReturnOp>(*reconciliator.getOutput().getUsers().begin()))
|
||||
if (!blueprint.getOutput().hasOneUse() || !isa<func::ReturnOp>(*blueprint.getOutput().getUsers().begin()))
|
||||
return failure();
|
||||
|
||||
unsigned returnIndex = reconciliator.getOutput().getUses().begin()->getOperandNumber();
|
||||
auto hostResultType = dyn_cast<RankedTensorType>(reconciliator.getOutput().getType());
|
||||
unsigned returnIndex = blueprint.getOutput().getUses().begin()->getOperandNumber();
|
||||
auto hostResultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
|
||||
if (!hostResultType || !hostResultType.hasStaticShape())
|
||||
return failure();
|
||||
|
||||
ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
|
||||
ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
|
||||
ArrayRef<int64_t> flatOffsets = reconciliator.getFragmentOffsets();
|
||||
ArrayRef<int64_t> flatSizes = reconciliator.getFragmentSizes();
|
||||
ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
|
||||
ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
|
||||
ArrayRef<int64_t> flatStrides = *stridesAttr;
|
||||
int64_t rank = hostResultType.getRank();
|
||||
SmallVector<Value> fragmentOperands {reconciliator.getInput()};
|
||||
llvm::append_range(fragmentOperands, reconciliator.getFragments());
|
||||
if (failed(validateFragmentAssemblyMetadata(reconciliator,
|
||||
SmallVector<Value> fragmentOperands {blueprint.getInput()};
|
||||
llvm::append_range(fragmentOperands, blueprint.getFragments());
|
||||
if (failed(validateFragmentAssemblyMetadata(blueprint,
|
||||
rank,
|
||||
fragmentOperands.size(),
|
||||
operandIndices,
|
||||
@@ -379,34 +379,34 @@ static SmallVector<OpFoldResult, 4> buildFragmentOffsets(IRRewriter& rewriter,
|
||||
}
|
||||
|
||||
static FailureOr<Value> lowerFragmentAssemblyHostCopies(IRRewriter& rewriter,
|
||||
spatial::SpatReconciliatorOp reconciliator,
|
||||
spatial::SpatBlueprintOp blueprint,
|
||||
Value hostTarget,
|
||||
ArrayRef<OpFoldResult> baseOffsets,
|
||||
IRMapping& mapper) {
|
||||
auto hostTargetType = dyn_cast<RankedTensorType>(hostTarget.getType());
|
||||
auto resultType = dyn_cast<RankedTensorType>(reconciliator.getOutput().getType());
|
||||
auto resultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
|
||||
if (!hostTargetType || !resultType || !resultType.hasStaticShape())
|
||||
return reconciliator.emitOpError("fragment assembly lowering requires static ranked tensor results");
|
||||
return blueprint.emitOpError("fragment assembly lowering requires static ranked tensor results");
|
||||
|
||||
std::optional<ArrayRef<int64_t>> operandIndicesAttr = reconciliator.getFragmentOperandIndices();
|
||||
std::optional<ArrayRef<int64_t>> fragmentStridesAttr = reconciliator.getFragmentStrides();
|
||||
std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
|
||||
std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
|
||||
if (!operandIndicesAttr || !fragmentStridesAttr)
|
||||
return reconciliator.emitOpError(
|
||||
return blueprint.emitOpError(
|
||||
"fragment assembly lowering requires explicit operand indices and unit strides");
|
||||
|
||||
ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
|
||||
std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = reconciliator.getFragmentSourceOffsets();
|
||||
std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
|
||||
if (!sourceOffsetsAttr)
|
||||
return reconciliator.emitOpError("fragment assembly lowering requires explicit source offsets");
|
||||
return blueprint.emitOpError("fragment assembly lowering requires explicit source offsets");
|
||||
ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
|
||||
ArrayRef<int64_t> flatOffsets = reconciliator.getFragmentOffsets();
|
||||
ArrayRef<int64_t> flatSizes = reconciliator.getFragmentSizes();
|
||||
ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
|
||||
ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
|
||||
ArrayRef<int64_t> flatStrides = *fragmentStridesAttr;
|
||||
int64_t rank = resultType.getRank();
|
||||
|
||||
SmallVector<Value> fragmentOperands {reconciliator.getInput()};
|
||||
llvm::append_range(fragmentOperands, reconciliator.getFragments());
|
||||
if (failed(validateFragmentAssemblyMetadata(reconciliator,
|
||||
SmallVector<Value> fragmentOperands {blueprint.getInput()};
|
||||
llvm::append_range(fragmentOperands, blueprint.getFragments());
|
||||
if (failed(validateFragmentAssemblyMetadata(blueprint,
|
||||
rank,
|
||||
fragmentOperands.size(),
|
||||
operandIndices,
|
||||
@@ -423,14 +423,14 @@ static FailureOr<Value> lowerFragmentAssemblyHostCopies(IRRewriter& rewriter,
|
||||
for (int64_t dim = 0; dim < rank; ++dim) {
|
||||
int64_t flatIndex = fragmentIndex * rank + dim;
|
||||
if (flatStrides[flatIndex] != 1)
|
||||
return reconciliator.emitOpError("fragment assembly lowering only supports unit strides");
|
||||
return blueprint.emitOpError("fragment assembly lowering only supports unit strides");
|
||||
fragmentOffsets.push_back(flatOffsets[flatIndex]);
|
||||
}
|
||||
|
||||
Value source = mapper.lookupOrDefault(fragmentOperands[operandIndex]);
|
||||
auto sourceType = dyn_cast<ShapedType>(source.getType());
|
||||
if (!sourceType || !sourceType.hasStaticShape())
|
||||
return reconciliator.emitOpError("fragment assembly lowering requires static ranked tensor operands");
|
||||
return blueprint.emitOpError("fragment assembly lowering requires static ranked tensor operands");
|
||||
|
||||
SmallVector<int64_t, 4> fragmentShape;
|
||||
fragmentShape.reserve(rank);
|
||||
@@ -440,11 +440,11 @@ static FailureOr<Value> lowerFragmentAssemblyHostCopies(IRRewriter& rewriter,
|
||||
Value fragment = source;
|
||||
if (llvm::to_vector(sourceType.getShape()) != fragmentShape || sourceOffsets[fragmentIndex] != 0) {
|
||||
FailureOr<SmallVector<int64_t, 4>> extractOffsets = getStaticSliceOffsetsForElementOffset(
|
||||
reconciliator, sourceType, fragmentShape, sourceOffsets[fragmentIndex], "fragment assembly source slice");
|
||||
blueprint, sourceType, fragmentShape, sourceOffsets[fragmentIndex], "fragment assembly source slice");
|
||||
if (failed(extractOffsets))
|
||||
return failure();
|
||||
fragment = tensor::ExtractSliceOp::create(rewriter,
|
||||
reconciliator.getLoc(),
|
||||
blueprint.getLoc(),
|
||||
source,
|
||||
getStaticIndexAttrs(rewriter, *extractOffsets),
|
||||
getStaticIndexAttrs(rewriter, fragmentShape),
|
||||
@@ -452,11 +452,11 @@ static FailureOr<Value> lowerFragmentAssemblyHostCopies(IRRewriter& rewriter,
|
||||
}
|
||||
|
||||
hostTarget = tensor::InsertSliceOp::create(rewriter,
|
||||
reconciliator.getLoc(),
|
||||
blueprint.getLoc(),
|
||||
fragment,
|
||||
hostTarget,
|
||||
buildFragmentOffsets(rewriter,
|
||||
reconciliator.getLoc(),
|
||||
blueprint.getLoc(),
|
||||
baseOffsets,
|
||||
fragmentOffsets,
|
||||
mapper),
|
||||
@@ -585,13 +585,13 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
|
||||
if (isa<spatial::SpatYieldOp>(op))
|
||||
continue;
|
||||
|
||||
if (auto reconciliator = dyn_cast<spatial::SpatReconciliatorOp>(op)) {
|
||||
std::optional<StringRef> modeAttr = reconciliator.getMode();
|
||||
if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
|
||||
std::optional<StringRef> modeAttr = blueprint.getMode();
|
||||
if (modeAttr && *modeAttr == "fragment_assembly") {
|
||||
for (Operation* user : reconciliator.getOutput().getUsers()) {
|
||||
for (Operation* user : blueprint.getOutput().getUsers()) {
|
||||
if (!isa<tensor::ParallelInsertSliceOp>(user))
|
||||
return reconciliator.emitOpError(
|
||||
"fragment assembly reconciliator lowering expects only tensor.parallel_insert_slice users");
|
||||
return blueprint.emitOpError(
|
||||
"fragment assembly blueprint lowering expects only tensor.parallel_insert_slice users");
|
||||
}
|
||||
continue;
|
||||
}
|
||||
@@ -653,12 +653,12 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
|
||||
|
||||
Value hostTarget = getOrCreateHostOutputTensor(resultIndex, insertSlice.getLoc());
|
||||
auto hostTargetType = cast<ShapedType>(hostTarget.getType());
|
||||
if (auto reconciliator =
|
||||
insertSlice.getSource().getDefiningOp<spatial::SpatReconciliatorOp>()) {
|
||||
std::optional<StringRef> modeAttr = reconciliator.getMode();
|
||||
if (auto blueprint =
|
||||
insertSlice.getSource().getDefiningOp<spatial::SpatBlueprintOp>()) {
|
||||
std::optional<StringRef> modeAttr = blueprint.getMode();
|
||||
if (modeAttr && *modeAttr == "fragment_assembly") {
|
||||
FailureOr<Value> updatedHostTarget = lowerFragmentAssemblyHostCopies(rewriter,
|
||||
reconciliator,
|
||||
blueprint,
|
||||
hostTarget,
|
||||
insertSlice.getMixedOffsets(),
|
||||
mapper);
|
||||
|
||||
@@ -73,7 +73,7 @@ mlir::Value getBestOutputTensorFromOperandsOrAllocate(RewriterBase& rewriter, Op
|
||||
rewriter, operation->getLoc(), resultShapedType.getShape(), resultShapedType.getElementType());
|
||||
}
|
||||
|
||||
LogicalResult validateFragmentAssemblyMetadata(spatial::SpatReconciliatorOp reconciliator,
|
||||
LogicalResult validateFragmentAssemblyMetadata(spatial::SpatBlueprintOp blueprint,
|
||||
int64_t resultRank,
|
||||
size_t operandCount,
|
||||
ArrayRef<int64_t> operandIndices,
|
||||
@@ -82,19 +82,19 @@ LogicalResult validateFragmentAssemblyMetadata(spatial::SpatReconciliatorOp reco
|
||||
ArrayRef<int64_t> flatSizes,
|
||||
ArrayRef<int64_t> flatStrides) {
|
||||
if (operandIndices.size() != sourceOffsets.size())
|
||||
return reconciliator.emitOpError("fragment assembly operand index and source offset counts must match");
|
||||
return blueprint.emitOpError("fragment assembly operand index and source offset counts must match");
|
||||
if (flatOffsets.size() != flatSizes.size())
|
||||
return reconciliator.emitOpError("fragment assembly offset and size arrays must have matching lengths");
|
||||
return blueprint.emitOpError("fragment assembly offset and size arrays must have matching lengths");
|
||||
if (flatStrides.size() != flatOffsets.size())
|
||||
return reconciliator.emitOpError("fragment assembly stride and offset arrays must have matching lengths");
|
||||
return blueprint.emitOpError("fragment assembly stride and offset arrays must have matching lengths");
|
||||
if (flatOffsets.size() != operandIndices.size() * static_cast<size_t>(resultRank))
|
||||
return reconciliator.emitOpError("fragment assembly metadata must provide one rank-sized offset/size/stride tuple per fragment");
|
||||
return blueprint.emitOpError("fragment assembly metadata must provide one rank-sized offset/size/stride tuple per fragment");
|
||||
|
||||
for (auto [fragmentIndex, operandIndex] : llvm::enumerate(operandIndices)) {
|
||||
if (operandIndex < 0 || operandIndex >= static_cast<int64_t>(operandCount))
|
||||
return reconciliator.emitOpError("fragment assembly operand index is out of range");
|
||||
return blueprint.emitOpError("fragment assembly operand index is out of range");
|
||||
if (sourceOffsets[fragmentIndex] < 0)
|
||||
return reconciliator.emitOpError("fragment assembly source offsets must be nonnegative");
|
||||
return blueprint.emitOpError("fragment assembly source offsets must be nonnegative");
|
||||
}
|
||||
|
||||
return success();
|
||||
|
||||
@@ -9,7 +9,7 @@
|
||||
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
class SpatReconciliatorOp;
|
||||
class SpatBlueprintOp;
|
||||
}
|
||||
|
||||
namespace onnx_mlir {
|
||||
@@ -36,7 +36,7 @@ mlir::SmallVector<mlir::Value> getOpOperandsSortedByUses(mlir::Operation* operat
|
||||
|
||||
mlir::Value getBestOutputTensorFromOperandsOrAllocate(mlir::RewriterBase& rewriter, mlir::Operation* operation);
|
||||
|
||||
mlir::LogicalResult validateFragmentAssemblyMetadata(onnx_mlir::spatial::SpatReconciliatorOp reconciliator,
|
||||
mlir::LogicalResult validateFragmentAssemblyMetadata(onnx_mlir::spatial::SpatBlueprintOp blueprint,
|
||||
int64_t resultRank,
|
||||
size_t operandCount,
|
||||
llvm::ArrayRef<int64_t> operandIndices,
|
||||
|
||||
@@ -43,31 +43,31 @@ static Value createStaticHostTargetOffset(IRRewriter& rewriter,
|
||||
return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), byteOffset);
|
||||
}
|
||||
|
||||
static FailureOr<Value> lowerFragmentAssemblyReconciliator(IRRewriter& rewriter,
|
||||
spatial::SpatReconciliatorOp reconciliator,
|
||||
static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
|
||||
spatial::SpatBlueprintOp blueprint,
|
||||
IRMapping& mapping) {
|
||||
auto resultType = dyn_cast<ShapedType>(reconciliator.getOutput().getType());
|
||||
auto resultType = dyn_cast<ShapedType>(blueprint.getOutput().getType());
|
||||
if (!resultType || !resultType.hasStaticShape())
|
||||
return reconciliator.emitOpError("fragment assembly lowering requires a static ranked tensor result");
|
||||
return blueprint.emitOpError("fragment assembly lowering requires a static ranked tensor result");
|
||||
|
||||
std::optional<StringRef> modeAttr = reconciliator.getMode();
|
||||
std::optional<ArrayRef<int64_t>> operandIndicesAttr = reconciliator.getFragmentOperandIndices();
|
||||
std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = reconciliator.getFragmentSourceOffsets();
|
||||
std::optional<ArrayRef<int64_t>> fragmentStridesAttr = reconciliator.getFragmentStrides();
|
||||
std::optional<StringRef> modeAttr = blueprint.getMode();
|
||||
std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
|
||||
std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
|
||||
std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
|
||||
if (!modeAttr || *modeAttr != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr
|
||||
|| !fragmentStridesAttr)
|
||||
return reconciliator.emitOpError("fragment assembly lowering requires explicit fragment metadata");
|
||||
return blueprint.emitOpError("fragment assembly lowering requires explicit fragment metadata");
|
||||
|
||||
ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
|
||||
ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
|
||||
ArrayRef<int64_t> flatOffsets = reconciliator.getFragmentOffsets();
|
||||
ArrayRef<int64_t> flatSizes = reconciliator.getFragmentSizes();
|
||||
ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
|
||||
ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
|
||||
ArrayRef<int64_t> flatStrides = *fragmentStridesAttr;
|
||||
int64_t rank = resultType.getRank();
|
||||
|
||||
SmallVector<Value> fragmentOperands {reconciliator.getInput()};
|
||||
llvm::append_range(fragmentOperands, reconciliator.getFragments());
|
||||
if (failed(validateFragmentAssemblyMetadata(reconciliator,
|
||||
SmallVector<Value> fragmentOperands {blueprint.getInput()};
|
||||
llvm::append_range(fragmentOperands, blueprint.getFragments());
|
||||
if (failed(validateFragmentAssemblyMetadata(blueprint,
|
||||
rank,
|
||||
fragmentOperands.size(),
|
||||
operandIndices,
|
||||
@@ -77,7 +77,7 @@ static FailureOr<Value> lowerFragmentAssemblyReconciliator(IRRewriter& rewriter,
|
||||
flatStrides)))
|
||||
return failure();
|
||||
|
||||
Value currentOutput = createEmptyTensorFromShaped(rewriter, reconciliator.getLoc(), resultType);
|
||||
Value currentOutput = createEmptyTensorFromShaped(rewriter, blueprint.getLoc(), resultType);
|
||||
for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {
|
||||
int64_t operandIndex = operandIndices[fragmentIndex];
|
||||
|
||||
@@ -86,7 +86,7 @@ static FailureOr<Value> lowerFragmentAssemblyReconciliator(IRRewriter& rewriter,
|
||||
for (int64_t dim = 0; dim < rank; ++dim) {
|
||||
int64_t flatIndex = fragmentIndex * rank + dim;
|
||||
if (flatStrides[flatIndex] != 1)
|
||||
return reconciliator.emitOpError("fragment assembly lowering only supports unit strides");
|
||||
return blueprint.emitOpError("fragment assembly lowering only supports unit strides");
|
||||
fragmentOffsets.push_back(flatOffsets[flatIndex]);
|
||||
fragmentElements *= flatSizes[flatIndex];
|
||||
}
|
||||
@@ -94,21 +94,21 @@ static FailureOr<Value> lowerFragmentAssemblyReconciliator(IRRewriter& rewriter,
|
||||
Value source = mapping.lookupOrDefault(fragmentOperands[operandIndex]);
|
||||
auto sourceType = dyn_cast<ShapedType>(source.getType());
|
||||
if (!sourceType || !sourceType.hasStaticShape())
|
||||
return reconciliator.emitOpError("fragment assembly lowering requires static ranked tensor operands");
|
||||
return blueprint.emitOpError("fragment assembly lowering requires static ranked tensor operands");
|
||||
|
||||
int64_t fragmentBytes =
|
||||
fragmentElements * static_cast<int64_t>(getElementTypeSizeInBytes(sourceType.getElementType()));
|
||||
auto sizeAttr = pim::getCheckedI32Attr(rewriter,
|
||||
reconciliator.getOperation(),
|
||||
blueprint.getOperation(),
|
||||
fragmentBytes,
|
||||
"fragment assembly host copy size");
|
||||
if (failed(sizeAttr))
|
||||
return failure();
|
||||
|
||||
Value hostTargetOffset = createStaticHostTargetOffset(rewriter, reconciliator.getLoc(), resultType, fragmentOffsets);
|
||||
Value hostTargetOffset = createStaticHostTargetOffset(rewriter, blueprint.getLoc(), resultType, fragmentOffsets);
|
||||
auto deviceSourceOffsetBytes = pim::checkedMul(static_cast<uint64_t>(sourceOffsets[fragmentIndex]),
|
||||
static_cast<uint64_t>(getElementTypeSizeInBytes(sourceType.getElementType())),
|
||||
reconciliator,
|
||||
blueprint,
|
||||
"fragment assembly device source offset");
|
||||
if (failed(deviceSourceOffsetBytes))
|
||||
return failure();
|
||||
@@ -116,7 +116,7 @@ static FailureOr<Value> lowerFragmentAssemblyReconciliator(IRRewriter& rewriter,
|
||||
rewriter.getInsertionBlock()->getParentOp(),
|
||||
static_cast<int64_t>(*deviceSourceOffsetBytes));
|
||||
currentOutput = pim::PimMemCopyDevToHostOp::create(rewriter,
|
||||
reconciliator.getLoc(),
|
||||
blueprint.getLoc(),
|
||||
currentOutput.getType(),
|
||||
hostTargetOffset,
|
||||
deviceSourceOffset,
|
||||
@@ -230,13 +230,13 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatSchedule
|
||||
mapping.map(*weightArg, weight);
|
||||
}
|
||||
for (Operation& op : block.without_terminator()) {
|
||||
if (auto reconciliator = dyn_cast<spatial::SpatReconciliatorOp>(op)) {
|
||||
std::optional<StringRef> modeAttr = reconciliator.getMode();
|
||||
if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
|
||||
std::optional<StringRef> modeAttr = blueprint.getMode();
|
||||
if (modeAttr && *modeAttr == "fragment_assembly") {
|
||||
auto lowered = lowerFragmentAssemblyReconciliator(rewriter, reconciliator, mapping);
|
||||
auto lowered = lowerFragmentAssemblyBlueprint(rewriter, blueprint, mapping);
|
||||
if (failed(lowered))
|
||||
return false;
|
||||
mapping.map(reconciliator.getOutput(), *lowered);
|
||||
mapping.map(blueprint.getOutput(), *lowered);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -31,11 +31,11 @@ static SmallVector<OpFoldResult, 4> getUnitStrides(Builder& builder, int64_t ran
|
||||
return strides;
|
||||
}
|
||||
|
||||
struct LowerFragmentAssemblyReconciliatorPattern
|
||||
: OpConversionPattern<spatial::SpatReconciliatorOp> {
|
||||
struct LowerFragmentAssemblyBlueprintPattern
|
||||
: OpConversionPattern<spatial::SpatBlueprintOp> {
|
||||
using OpConversionPattern::OpConversionPattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatReconciliatorOp op,
|
||||
LogicalResult matchAndRewrite(spatial::SpatBlueprintOp op,
|
||||
OpAdaptor adaptor,
|
||||
ConversionPatternRewriter& rewriter) const override {
|
||||
std::optional<StringRef> modeAttr = op.getMode();
|
||||
@@ -125,7 +125,7 @@ void populateInitialPatterns(RewritePatternSet& patterns) {
|
||||
void populateCoreBodyPatterns(RewritePatternSet& patterns) {
|
||||
raptor::populateWithGenerated(patterns);
|
||||
populateTransposeLoweringPatterns(patterns);
|
||||
patterns.add<LowerFragmentAssemblyReconciliatorPattern>(patterns.getContext());
|
||||
patterns.add<LowerFragmentAssemblyBlueprintPattern>(patterns.getContext());
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir
|
||||
|
||||
@@ -149,36 +149,36 @@ static std::optional<ReturnUseInfo> analyzeReturnUse(Value value) {
|
||||
};
|
||||
}
|
||||
|
||||
static FailureOr<SmallVector<std::pair<spatial::SpatReconciliatorOp, size_t>, 4>>
|
||||
static FailureOr<SmallVector<std::pair<spatial::SpatBlueprintOp, size_t>, 4>>
|
||||
analyzeTopLevelFragmentAssemblyUses(Value value) {
|
||||
SmallVector<std::pair<spatial::SpatReconciliatorOp, size_t>, 4> uses;
|
||||
SmallVector<std::pair<spatial::SpatBlueprintOp, size_t>, 4> uses;
|
||||
for (OpOperand& use : value.getUses()) {
|
||||
auto reconciliator = dyn_cast<spatial::SpatReconciliatorOp>(use.getOwner());
|
||||
if (!reconciliator || reconciliator->getParentOp() != reconciliator->getParentOfType<func::FuncOp>())
|
||||
auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(use.getOwner());
|
||||
if (!blueprint || blueprint->getParentOp() != blueprint->getParentOfType<func::FuncOp>())
|
||||
return failure();
|
||||
std::optional<StringRef> mode = reconciliator.getMode();
|
||||
std::optional<StringRef> mode = blueprint.getMode();
|
||||
if (!mode || *mode != "fragment_assembly")
|
||||
return failure();
|
||||
if (!reconciliator.getOutput().hasOneUse() || !isa<func::ReturnOp>(*reconciliator.getOutput().getUsers().begin()))
|
||||
if (!blueprint.getOutput().hasOneUse() || !isa<func::ReturnOp>(*blueprint.getOutput().getUsers().begin()))
|
||||
return failure();
|
||||
std::optional<ArrayRef<int64_t>> operandIndicesAttr = reconciliator.getFragmentOperandIndices();
|
||||
std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = reconciliator.getFragmentSourceOffsets();
|
||||
std::optional<ArrayRef<int64_t>> stridesAttr = reconciliator.getFragmentStrides();
|
||||
auto resultType = dyn_cast<RankedTensorType>(reconciliator.getOutput().getType());
|
||||
std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
|
||||
std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
|
||||
std::optional<ArrayRef<int64_t>> stridesAttr = blueprint.getFragmentStrides();
|
||||
auto resultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
|
||||
if (!operandIndicesAttr || !sourceOffsetsAttr || !stridesAttr || !resultType || !resultType.hasStaticShape())
|
||||
return failure();
|
||||
SmallVector<Value> fragmentOperands {reconciliator.getInput()};
|
||||
llvm::append_range(fragmentOperands, reconciliator.getFragments());
|
||||
if (failed(validateFragmentAssemblyMetadata(reconciliator,
|
||||
SmallVector<Value> fragmentOperands {blueprint.getInput()};
|
||||
llvm::append_range(fragmentOperands, blueprint.getFragments());
|
||||
if (failed(validateFragmentAssemblyMetadata(blueprint,
|
||||
resultType.getRank(),
|
||||
fragmentOperands.size(),
|
||||
*operandIndicesAttr,
|
||||
*sourceOffsetsAttr,
|
||||
reconciliator.getFragmentOffsets(),
|
||||
reconciliator.getFragmentSizes(),
|
||||
blueprint.getFragmentOffsets(),
|
||||
blueprint.getFragmentSizes(),
|
||||
*stridesAttr)))
|
||||
return failure();
|
||||
uses.emplace_back(reconciliator, use.getOperandNumber());
|
||||
uses.emplace_back(blueprint, use.getOperandNumber());
|
||||
}
|
||||
return uses;
|
||||
}
|
||||
@@ -593,7 +593,7 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
|
||||
}
|
||||
}
|
||||
|
||||
FailureOr<SmallVector<std::pair<spatial::SpatReconciliatorOp, size_t>, 4>> fragmentAssemblyUses =
|
||||
FailureOr<SmallVector<std::pair<spatial::SpatBlueprintOp, size_t>, 4>> fragmentAssemblyUses =
|
||||
analyzeTopLevelFragmentAssemblyUses(producedValue);
|
||||
if (succeeded(fragmentAssemblyUses)) {
|
||||
auto sourceType = dyn_cast<RankedTensorType>(storedValue.getType());
|
||||
@@ -603,35 +603,35 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
|
||||
}
|
||||
|
||||
size_t elementSize = getElementTypeSizeInBytes(sourceType.getElementType());
|
||||
for (auto [reconciliator, operandNumber] : *fragmentAssemblyUses) {
|
||||
for (auto [blueprint, operandNumber] : *fragmentAssemblyUses) {
|
||||
rewriter.setInsertionPointAfterValue(storedValue);
|
||||
std::optional<ArrayRef<int64_t>> operandIndicesAttr = reconciliator.getFragmentOperandIndices();
|
||||
std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = reconciliator.getFragmentSourceOffsets();
|
||||
std::optional<ArrayRef<int64_t>> stridesAttr = reconciliator.getFragmentStrides();
|
||||
std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
|
||||
std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
|
||||
std::optional<ArrayRef<int64_t>> stridesAttr = blueprint.getFragmentStrides();
|
||||
if (!operandIndicesAttr || !sourceOffsetsAttr || !stridesAttr) {
|
||||
reconciliator.emitOpError(
|
||||
blueprint.emitOpError(
|
||||
"fragment assembly lowering requires explicit operand, source-offset, and stride metadata");
|
||||
return ReturnPathLoweringResult::Failure;
|
||||
}
|
||||
|
||||
size_t returnIndex = reconciliator.getOutput().getUses().begin()->getOperandNumber();
|
||||
size_t returnIndex = blueprint.getOutput().getUses().begin()->getOperandNumber();
|
||||
Value outputTensor = outputTensors[returnIndex](rewriter, loc);
|
||||
auto outputType = dyn_cast<RankedTensorType>(outputTensor.getType());
|
||||
auto resultType = dyn_cast<RankedTensorType>(reconciliator.getOutput().getType());
|
||||
auto resultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
|
||||
if (!outputType || !resultType || !resultType.hasStaticShape()) {
|
||||
reconciliator.emitOpError("fragment assembly lowering requires static ranked host outputs");
|
||||
blueprint.emitOpError("fragment assembly lowering requires static ranked host outputs");
|
||||
return ReturnPathLoweringResult::Failure;
|
||||
}
|
||||
|
||||
ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
|
||||
ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
|
||||
ArrayRef<int64_t> flatOffsets = reconciliator.getFragmentOffsets();
|
||||
ArrayRef<int64_t> flatSizes = reconciliator.getFragmentSizes();
|
||||
ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
|
||||
ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
|
||||
ArrayRef<int64_t> flatStrides = *stridesAttr;
|
||||
int64_t rank = resultType.getRank();
|
||||
if (failed(validateFragmentAssemblyMetadata(reconciliator,
|
||||
if (failed(validateFragmentAssemblyMetadata(blueprint,
|
||||
rank,
|
||||
1 + reconciliator.getFragments().size(),
|
||||
1 + blueprint.getFragments().size(),
|
||||
operandIndices,
|
||||
sourceOffsets,
|
||||
flatOffsets,
|
||||
@@ -647,7 +647,7 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
|
||||
for (int64_t dim = 0; dim < rank; ++dim) {
|
||||
int64_t flatIndex = fragmentIndex * rank + dim;
|
||||
if (flatStrides[flatIndex] != 1) {
|
||||
reconciliator.emitOpError("fragment assembly lowering only supports unit strides");
|
||||
blueprint.emitOpError("fragment assembly lowering only supports unit strides");
|
||||
return ReturnPathLoweringResult::Failure;
|
||||
}
|
||||
fragmentOffsets.push_back(flatOffsets[flatIndex]);
|
||||
@@ -684,7 +684,7 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
|
||||
|
||||
outputTensor =
|
||||
pim::PimMemCopyDevToHostOp::create(rewriter,
|
||||
reconciliator.getLoc(),
|
||||
blueprint.getLoc(),
|
||||
outputTensor.getType(),
|
||||
getOrCreateIndexConstant(rewriter, producerOp, *hostOffset),
|
||||
getOrCreateIndexConstant(rewriter, producerOp, *sourceOffset),
|
||||
@@ -698,7 +698,7 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
|
||||
if (failedChunk)
|
||||
return ReturnPathLoweringResult::Failure;
|
||||
}
|
||||
markOpToRemove(reconciliator.getOperation());
|
||||
markOpToRemove(blueprint.getOperation());
|
||||
}
|
||||
return ReturnPathLoweringResult::Handled;
|
||||
}
|
||||
@@ -813,11 +813,11 @@ void raptor::SpatialToPimPass::replaceReturnWithOutputBuffers(func::ReturnOp ret
|
||||
return;
|
||||
}
|
||||
|
||||
if (auto reconciliator = dyn_cast<spatial::SpatReconciliatorOp>(op)) {
|
||||
std::optional<StringRef> mode = reconciliator.getMode();
|
||||
if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
|
||||
std::optional<StringRef> mode = blueprint.getMode();
|
||||
if (mode && *mode == "fragment_assembly") {
|
||||
markOpToRemove(reconciliator.getOperation());
|
||||
for (Value operand : reconciliator->getOperands())
|
||||
markOpToRemove(blueprint.getOperation());
|
||||
for (Value operand : blueprint->getOperands())
|
||||
markOwnedReturnChain(operand.getDefiningOp(), markOwnedReturnChain);
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -203,7 +203,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
|
||||
func::FuncOp funcOp = *entryFunc;
|
||||
if (failed(verifyScheduledSpatialInvariants(funcOp))) {
|
||||
funcOp.emitOpError(
|
||||
"RAPTOR_PHASE_CHECK scheduled Spatial verification failed at the start of SpatialToPim");
|
||||
"scheduled Spatial verification failed at the start of SpatialToPim");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user