[Mlir-commits] [mlir] Lower shuffle to single-result form if possible. (PR #84321)

[Mehdi Amini]

================
@@ -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);
----------------
joker-eph wrote:

That reproduced for me FYI:
```
# | #12 0x00005640f7453cc8 decltype(auto) llvm::cast<mlir::LLVM::LLVMArrayType, mlir::Type>(mlir::Type&) /home/mamini/projects/llvm-project2/llvm/include/llvm/Support/Casting.h:573:37
# | #13 0x00005640f71e5735 getInsertExtractValueElementType(mlir::Type, llvm::ArrayRef<long>) /home/mamini/projects/llvm-project2/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp:1590:18
# | #14 0x00005640f71e55ca mlir::LLVM::ExtractValueOp::build(mlir::OpBuilder&, mlir::OperationState&, mlir::Value, llvm::ArrayRef<long>) /home/mamini/projects/llvm-project2/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp:1645:9
# | #15 0x00005640f99a78e9 mlir::LLVM::ExtractValueOp mlir::OpBuilder::create<mlir::LLVM::ExtractValueOp, mlir::Value&, int>(mlir::Location, mlir::Value&, int&&) /home/mamini/projects/llvm-project2/mlir/include/mlir/IR/Builders.h:511:5
# | #16 0x00005640f9be1283 (anonymous namespace)::GPUShuffleOpLowering::matchAndRewrite(mlir::gpu::ShuffleOp, mlir::gpu::ShuffleOpAdaptor, mlir::ConversionPatternRewriter&) const /home/mamini/projects/llvm-project2/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp:191:34
```

https://github.com/llvm/llvm-project/pull/84321