24 rewriteWarpOpToScfFor(RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,  in rewriteWarpOpToScfFor()  argument
32   OpBuilder::InsertionGuard g(rewriter);  in rewriteWarpOpToScfFor()
33   rewriter.setInsertionPoint(warpOp);  in rewriteWarpOpToScfFor()
36   Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);  in rewriteWarpOpToScfFor()
37   Value isLane0 = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,  in rewriteWarpOpToScfFor()
39   auto ifOp = rewriter.create<scf::IfOp>(loc, isLane0,  in rewriteWarpOpToScfFor()
41   rewriter.eraseOp(ifOp.thenBlock()->getTerminator());  in rewriteWarpOpToScfFor()
50     rewriter.setInsertionPoint(ifOp);  in rewriteWarpOpToScfFor()
52         options.warpAllocationFn(loc, rewriter, warpOp, bbArg.getType());  in rewriteWarpOpToScfFor()
55     rewriter.setInsertionPoint(ifOp);  in rewriteWarpOpToScfFor()
58     Value storeOffset = rewriter.create<arith::MulIOp>(  in rewriteWarpOpToScfFor()
60         rewriter.create<arith::ConstantIndexOp>(loc, storeSize));  in rewriteWarpOpToScfFor()
61     rewriter.create<vector::StoreOp>(loc, val, buffer, storeOffset);  in rewriteWarpOpToScfFor()
64     rewriter.setInsertionPointToStart(ifOp.thenBlock());  in rewriteWarpOpToScfFor()
66     Value loadOp = rewriter.create<vector::LoadOp>(loc, bbArgType, buffer, c0);  in rewriteWarpOpToScfFor()
72     rewriter.setInsertionPoint(ifOp);  in rewriteWarpOpToScfFor()
73     options.warpSyncronizationFn(loc, rewriter, warpOp);  in rewriteWarpOpToScfFor()
77   rewriter.mergeBlocks(warpOpBody, ifOp.thenBlock(), bbArgReplacements);  in rewriteWarpOpToScfFor()
86     rewriter.setInsertionPoint(ifOp);  in rewriteWarpOpToScfFor()
88         options.warpAllocationFn(loc, rewriter, warpOp, val.getType());  in rewriteWarpOpToScfFor()
91     rewriter.setInsertionPoint(yieldOp);  in rewriteWarpOpToScfFor()
93       rewriter.create<vector::StoreOp>(yieldLoc, val, buffer, c0);  in rewriteWarpOpToScfFor()
95       rewriter.create<memref::StoreOp>(yieldLoc, val, buffer, c0);  in rewriteWarpOpToScfFor()
98     rewriter.setInsertionPointAfter(ifOp);  in rewriteWarpOpToScfFor()
107       Value loadOp = rewriter.create<memref::LoadOp>(loc, buffer, c0);  in rewriteWarpOpToScfFor()
114       Value loadOffset = rewriter.create<arith::MulIOp>(  in rewriteWarpOpToScfFor()
116           rewriter.create<arith::ConstantIndexOp>(loc, loadSize));  in rewriteWarpOpToScfFor()
117       Value loadOp = rewriter.create<vector::LoadOp>(loc, loadedVectorType,  in rewriteWarpOpToScfFor()
125     rewriter.setInsertionPointAfter(ifOp);  in rewriteWarpOpToScfFor()
126     options.warpSyncronizationFn(loc, rewriter, warpOp);  in rewriteWarpOpToScfFor()
130   rewriter.eraseOp(yieldOp);  in rewriteWarpOpToScfFor()
131   rewriter.setInsertionPointToEnd(ifOp.thenBlock());  in rewriteWarpOpToScfFor()
132   rewriter.create<scf::YieldOp>(yieldLoc);  in rewriteWarpOpToScfFor()
135   rewriter.replaceOp(warpOp, replacements);  in rewriteWarpOpToScfFor()
142     RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,  in moveRegionToNewWarpOpAndReplaceReturns()  argument
145   OpBuilder::InsertionGuard g(rewriter);  in moveRegionToNewWarpOpAndReplaceReturns()
146   rewriter.setInsertionPoint(warpOp);  in moveRegionToNewWarpOpAndReplaceReturns()
147   auto newWarpOp = rewriter.create<WarpExecuteOnLane0Op>(  in moveRegionToNewWarpOpAndReplaceReturns()
154   rewriter.inlineRegionBefore(opBody, newOpBody, newOpBody.begin());  in moveRegionToNewWarpOpAndReplaceReturns()
155   rewriter.eraseBlock(&newOpFirstBlock);  in moveRegionToNewWarpOpAndReplaceReturns()
162   rewriter.updateRootInPlace(  in moveRegionToNewWarpOpAndReplaceReturns()
170     RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,  in moveRegionToNewWarpOpAndAppendReturns()  argument
195       rewriter, warpOp, yieldValues.getArrayRef(), types);  in moveRegionToNewWarpOpAndAppendReturns()
196   rewriter.replaceOp(warpOp,  in moveRegionToNewWarpOpAndAppendReturns()
227 static Operation *cloneOpWithOperandsAndTypes(RewriterBase &rewriter,  in cloneOpWithOperandsAndTypes()  argument
233   return rewriter.create(res);  in cloneOpWithOperandsAndTypes()
272                                 PatternRewriter &rewriter) const override {  in matchAndRewrite()
273     return rewriteWarpOpToScfFor(rewriter, warpOp, options);  in matchAndRewrite()
284 static vector::TransferWriteOp cloneWriteOp(RewriterBase &rewriter,  in cloneWriteOp()  argument
290   OpBuilder::InsertionGuard g(rewriter);  in cloneWriteOp()
293       rewriter, warpOp, ValueRange{{writeOp.getVector()}},  in cloneWriteOp()
295   rewriter.setInsertionPointAfter(newWarpOp);  in cloneWriteOp()
297       cast<vector::TransferWriteOp>(rewriter.clone(*writeOp.getOperation()));  in cloneWriteOp()
298   rewriter.eraseOp(writeOp);  in cloneWriteOp()
328   LogicalResult tryDistributeOp(RewriterBase &rewriter,  in tryDistributeOp()
358         cloneWriteOp(rewriter, warpOp, writeOp, targetType);  in tryDistributeOp()
363     rewriter.setInsertionPoint(newWriteOp);  in tryDistributeOp()
376       auto scale = rewriter.getAffineConstantExpr(targetShape[vectorPos]);  in tryDistributeOp()
378           makeComposedAffineApply(rewriter, loc, d0 + scale * d1,  in tryDistributeOp()
387   LogicalResult tryExtractOp(RewriterBase &rewriter,  in tryExtractOp()
408         rewriter, warpOp, yieldValues, retTypes, newRetIndices);  in tryExtractOp()
409     rewriter.setInsertionPointAfter(newWarpOp);  in tryExtractOp()
412     auto secondWarpOp = rewriter.create<WarpExecuteOnLane0Op>(  in tryExtractOp()
415     rewriter.setInsertionPointToStart(&body);  in tryExtractOp()
417         cast<vector::TransferWriteOp>(rewriter.clone(*writeOp.getOperation()));  in tryExtractOp()
419     rewriter.eraseOp(writeOp);  in tryExtractOp()
420     rewriter.create<vector::YieldOp>(newWarpOp.getLoc());  in tryExtractOp()
425                                 PatternRewriter &rewriter) const override {  in matchAndRewrite()
446     if (succeeded(tryDistributeOp(rewriter, writeOp, warpOp)))  in matchAndRewrite()
449     if (succeeded(tryExtractOp(rewriter, writeOp, warpOp)))  in matchAndRewrite()
480                                 PatternRewriter &rewriter) const override {  in matchAndRewrite()
510         rewriter, warpOp, yieldValues, retTypes, newRetIndices);  in matchAndRewrite()
511     rewriter.setInsertionPointAfter(newWarpOp);  in matchAndRewrite()
517     OpBuilder::InsertionGuard g(rewriter);  in matchAndRewrite()
518     rewriter.setInsertionPointAfter(newWarpOp);  in matchAndRewrite()
520         rewriter, loc, elementWise, newOperands,  in matchAndRewrite()
544                                 PatternRewriter &rewriter) const override {  in matchAndRewrite()
558     rewriter.setInsertionPointAfter(warpOp);  in matchAndRewrite()
559     Value distConstant = rewriter.create<arith::ConstantOp>(loc, newAttr);  in matchAndRewrite()
586                                 PatternRewriter &rewriter) const override {  in matchAndRewrite()
599     OpBuilder::InsertionGuard g(rewriter);  in matchAndRewrite()
600     rewriter.setInsertionPointAfter(warpOp);  in matchAndRewrite()
612           makeComposedAffineApply(rewriter, read.getLoc(), d0 + scale * d1,  in matchAndRewrite()
615     Value newRead = rewriter.create<vector::TransferReadOp>(  in matchAndRewrite()
629                                 PatternRewriter &rewriter) const override {  in matchAndRewrite()
643         rewriter, warpOp, yieldValues, resultTypes);  in matchAndRewrite()
650     rewriter.eraseOp(warpOp);  in matchAndRewrite()
660                                 PatternRewriter &rewriter) const override {  in matchAndRewrite()
700                                 PatternRewriter &rewriter) const override {  in matchAndRewrite()
712         rewriter, warpOp, {broadcastOp.getSource()},  in matchAndRewrite()
714     rewriter.setInsertionPointAfter(newWarpOp);  in matchAndRewrite()
715     Value broadcasted = rewriter.create<vector::BroadcastOp>(  in matchAndRewrite()
727                                 PatternRewriter &rewriter) const override {  in matchAndRewrite()
739         rewriter, warpOp, {extractOp.getVector()}, {extractOp.getVectorType()},  in matchAndRewrite()
741     rewriter.setInsertionPointAfter(newWarpOp);  in matchAndRewrite()
742     Value newExtract = rewriter.create<vector::ExtractOp>(  in matchAndRewrite()
784                                 PatternRewriter &rewriter) const override {  in matchAndRewrite()
803     OpBuilder::InsertionGuard g(rewriter);  in matchAndRewrite()
804     rewriter.setInsertionPointAfter(warpOp);  in matchAndRewrite()
807     auto newForOp = rewriter.create<scf::ForOp>(  in matchAndRewrite()
810     rewriter.setInsertionPoint(newForOp.getBody(), newForOp.getBody()->begin());  in matchAndRewrite()
811     auto innerWarp = rewriter.create<WarpExecuteOnLane0Op>(  in matchAndRewrite()
824     rewriter.eraseOp(forOp.getBody()->getTerminator());  in matchAndRewrite()
825     rewriter.mergeBlocks(forOp.getBody(), innerWarp.getBody(), argMapping);  in matchAndRewrite()
826     rewriter.setInsertionPoint(innerWarp.getBody(), innerWarp.getBody()->end());  in matchAndRewrite()
827     rewriter.create<vector::YieldOp>(innerWarp.getLoc(), yieldOperands);  in matchAndRewrite()
828     rewriter.setInsertionPointAfter(innerWarp);  in matchAndRewrite()
830       rewriter.create<scf::YieldOp>(forOp.getLoc(), innerWarp.getResults());  in matchAndRewrite()
831     rewriter.eraseOp(forOp);  in matchAndRewrite()
869                                 PatternRewriter &rewriter) const override {  in matchAndRewrite()
880       return rewriter.notifyMatchFailure(  in matchAndRewrite()
884       return rewriter.notifyMatchFailure(  in matchAndRewrite()
889       return rewriter.notifyMatchFailure(  in matchAndRewrite()
905         rewriter, warpOp, yieldValues, retTypes, newRetIndices);  in matchAndRewrite()
906     rewriter.setInsertionPointAfter(newWarpOp);  in matchAndRewrite()
910     Value perLaneReduction = rewriter.create<vector::ReductionOp>(  in matchAndRewrite()
914         distributedReductionFn(reductionOp.getLoc(), rewriter, perLaneReduction,  in matchAndRewrite()
918           rewriter, reductionOp.getLoc(), reductionOp.getKind(), fullReduce,  in matchAndRewrite()