@@ -2558,7 +2558,8 @@ FailureOr<TilingResult> OnlineAttentionOp::tileToPartialReduction(
2558
2558
OpBuilder &b, Location loc, ReductionTilingStrategy strategy,
2559
2559
ValueRange init, ArrayRef<OpFoldResult> offsets,
2560
2560
ArrayRef<OpFoldResult> sizes,
2561
- const llvm::SetVector<unsigned > &reductionDims) {
2561
+ const llvm::SetVector<unsigned > &reductionDims,
2562
+ ArrayRef<OpFoldResult> splitReductionIvs) {
2562
2563
FailureOr<AttentionOpDetail> maybeOpInfo = AttentionOpDetail::get (
2563
2564
getQueryMap (), getKeyMap (), getValueMap (), getOutputMap ());
2564
2565
if (failed (maybeOpInfo)) {
@@ -2753,9 +2754,10 @@ FailureOr<MergeResult> OnlineAttentionOp::mergeReductions(
2753
2754
}
2754
2755
2755
2756
LogicalResult OnlineAttentionOp::getPartialResultTilePosition (
2756
- OpBuilder &b, unsigned resultNumber, ArrayRef<OpFoldResult> offsets ,
2757
- ArrayRef<OpFoldResult> sizes,
2757
+ OpBuilder &b, unsigned resultNumber, ReductionTilingStrategy tilingStrategy ,
2758
+ ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
2758
2759
const llvm::SetVector<unsigned > &reductionDims,
2760
+ ArrayRef<OpFoldResult> splitReductionIvs,
2759
2761
SmallVector<OpFoldResult> &resultOffsets,
2760
2762
SmallVector<OpFoldResult> &resultSizes) {
2761
2763
0 commit comments