Unexpected invariant now it's clear (batched in the first tensor rank)
Validate Operations / validate-operations (push) Has been cancelled
Validate Operations / validate-operations (push) Has been cancelled
This commit is contained in:
@@ -346,6 +346,52 @@ inline void createParallelInsertSliceIntoBatchOutput(mlir::PatternRewriter& rewr
|
||||
mlir::tensor::ParallelInsertSliceOp::create(rewriter, loc, source, dest, offsets, sizes, strides);
|
||||
}
|
||||
|
||||
inline void publishGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
|
||||
mlir::Location loc,
|
||||
mlir::Value fragment,
|
||||
mlir::Value output,
|
||||
mlir::Value physicalSlot) {
|
||||
auto fragmentType = mlir::cast<mlir::RankedTensorType>(fragment.getType());
|
||||
mlir::SmallVector<mlir::OpFoldResult> offsets {physicalSlot};
|
||||
mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
|
||||
mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
|
||||
for (int64_t dim : fragmentType.getShape()) {
|
||||
offsets.push_back(rewriter.getIndexAttr(0));
|
||||
sizes.push_back(rewriter.getIndexAttr(dim));
|
||||
strides.push_back(rewriter.getIndexAttr(1));
|
||||
}
|
||||
createParallelInsertSliceIntoBatchOutput(rewriter, loc, fragment, output, offsets, sizes, strides);
|
||||
}
|
||||
|
||||
inline mlir::FailureOr<mlir::Value>
|
||||
extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
|
||||
mlir::Location loc,
|
||||
mlir::Value physicalBatch,
|
||||
mlir::OpFoldResult slot,
|
||||
mlir::RankedTensorType fragmentType) {
|
||||
if (fragmentType.getRank() == 0)
|
||||
return mlir::failure();
|
||||
auto physicalType = mlir::dyn_cast<mlir::RankedTensorType>(physicalBatch.getType());
|
||||
if (!physicalType || physicalType.getRank() != fragmentType.getRank() + 1)
|
||||
return mlir::failure();
|
||||
mlir::SmallVector<int64_t> selectedShape {1};
|
||||
llvm::append_range(selectedShape, fragmentType.getShape());
|
||||
auto selectedType = mlir::RankedTensorType::get(selectedShape, fragmentType.getElementType(), fragmentType.getEncoding());
|
||||
mlir::SmallVector<mlir::OpFoldResult> offsets {slot};
|
||||
mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
|
||||
mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
|
||||
for (int64_t dim : fragmentType.getShape()) {
|
||||
offsets.push_back(rewriter.getIndexAttr(0));
|
||||
sizes.push_back(rewriter.getIndexAttr(dim));
|
||||
strides.push_back(rewriter.getIndexAttr(1));
|
||||
}
|
||||
mlir::Value selected = mlir::tensor::ExtractSliceOp::create(rewriter, loc, selectedType, physicalBatch, offsets, sizes, strides);
|
||||
mlir::SmallVector<mlir::ReassociationIndices> reassociation {{0, 1}};
|
||||
for (int64_t dim = 2; dim <= fragmentType.getRank(); ++dim)
|
||||
reassociation.push_back({dim});
|
||||
return mlir::tensor::CollapseShapeOp::create(rewriter, loc, fragmentType, selected, reassociation).getResult();
|
||||
}
|
||||
|
||||
template <typename BodyFn>
|
||||
mlir::Value materializeOrComputeUnary(mlir::Value input,
|
||||
mlir::RankedTensorType resultType,
|
||||
|
||||
Reference in New Issue
Block a user