[Mlir-commits] [mlir] Lower shuffle to single-result form if possible. (PR #84321)
Johannes Reifferscheid
llvmlistbot at llvm.org
Thu Mar 21 02:33:21 PDT 2024
================
@@ -176,14 +174,26 @@ struct GPUShuffleOpLowering : public ConvertOpToLLVMPattern<gpu::ShuffleOp> {
rewriter.create<LLVM::SubOp>(loc, int32Type, adaptor.getWidth(), one);
}
- auto returnValueAndIsValidAttr = rewriter.getUnitAttr();
+ bool predIsUsed = !op->getResult(1).use_empty();
+ UnitAttr returnValueAndIsValidAttr = nullptr;
+ Type resultTy = valueTy;
+ if (predIsUsed) {
+ returnValueAndIsValidAttr = rewriter.getUnitAttr();
+ resultTy = LLVM::LLVMStructType::getLiteral(rewriter.getContext(),
+ {valueTy, predTy});
+ }
Value shfl = rewriter.create<NVVM::ShflOp>(
loc, resultTy, activeMask, adaptor.getValue(), adaptor.getOffset(),
maskAndClamp, convertShflKind(op.getMode()), returnValueAndIsValidAttr);
- Value shflValue = rewriter.create<LLVM::ExtractValueOp>(loc, shfl, 0);
- Value isActiveSrcLane = rewriter.create<LLVM::ExtractValueOp>(loc, shfl, 1);
-
- rewriter.replaceOp(op, {shflValue, isActiveSrcLane});
+ if (predIsUsed) {
+ Value shflValue = rewriter.create<LLVM::ExtractValueOp>(loc, shfl, 0);
+ Value isActiveSrcLane =
+ rewriter.create<LLVM::ExtractValueOp>(loc, shfl, 1);
+ rewriter.replaceOp(op, {shflValue, isActiveSrcLane});
+ } else {
+ Value falseCst = rewriter.create<LLVM::ConstantOp>(loc, predTy, 0);
----------------
jreiffers wrote:
Yep, looks like it was just me. Thanks for checking!
https://github.com/llvm/llvm-project/pull/84321
More information about the Mlir-commits
mailing list