[InstCombine] Canonicalize extractvalue + select

[dtcxzyw]

This patch canonicalizes extractvalue (select Cond, TV, FV) into select Cond, (extractvalue TV), (extractvalue FV). The latter form may enable more optimizations.

Looks like this patch doesn't affect the codegen: https://godbolt.org/z/vPEGd8vWf

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

This patch canonicalizes extractvalue (select Cond, TV, FV) into select Cond, (extractvalue TV), (extractvalue FV). The latter form may enable more optimizations.

Looks like this patch doesn't affect the codegen: https://godbolt.org/z/vPEGd8vWf

---

Full diff: https://github.com/llvm/llvm-project/pull/84686.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstructionCombining.cpp (+10)
• (added) llvm/test/Transforms/InstCombine/extract-select-agg.ll (+60)

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 1a831805dc72a0..caf4355a0b1ca5 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -3834,6 +3834,16 @@ Instruction *InstCombinerImpl::visitExtractValueInst(ExtractValueInst &EV) {
     if (Instruction *Res = foldOpIntoPhi(EV, PN))
       return Res;
 
+  // Canonicalize extract (select Cond, TV, FV)
+  // -> select cond, (extract TV), (extract FV)
+  if (auto *SI = dyn_cast<SelectInst>(Agg)) {
+    Value *Cond = SI->getCondition();
+    Value *TV = Builder.CreateExtractValue(SI->getTrueValue(), EV.getIndices());
+    Value *FV =
+        Builder.CreateExtractValue(SI->getFalseValue(), EV.getIndices());
+    return SelectInst::Create(Cond, TV, FV);
+  }
+
   // We could simplify extracts from other values. Note that nested extracts may
   // already be simplified implicitly by the above: extract (extract (insert) )
   // will be translated into extract ( insert ( extract ) ) first and then just
diff --git a/llvm/test/Transforms/InstCombine/extract-select-agg.ll b/llvm/test/Transforms/InstCombine/extract-select-agg.ll
new file mode 100644
index 00000000000000..385fec6d4f06f5
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/extract-select-agg.ll
@@ -0,0 +1,60 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define i64 @test_select_agg_constant(i64 %val, i1 %cond) {
+; CHECK-LABEL: define i64 @test_select_agg_constant(
+; CHECK-SAME: i64 [[VAL:%.*]], i1 [[COND:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[RET:%.*]] = zext i1 [[COND]] to i64
+; CHECK-NEXT:    ret i64 [[RET]]
+;
+entry:
+  %a = insertvalue { i64, i64 } { i64 1, i64 poison }, i64 %val, 1
+  %b = insertvalue { i64, i64 } { i64 0, i64 poison }, i64 %val, 1
+  %sel = select i1 %cond, { i64, i64 } %a, { i64, i64 } %b
+  %ret = extractvalue { i64, i64 } %sel, 0
+  ret i64 %ret
+}
+
+define void @test_select_agg_multiuse(i1 %cond, i64 %v1, i64 %v2, i64 %v3, i64 %v4) {
+; CHECK-LABEL: define void @test_select_agg_multiuse(
+; CHECK-SAME: i1 [[COND:%.*]], i64 [[V1:%.*]], i64 [[V2:%.*]], i64 [[V3:%.*]], i64 [[V4:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[X:%.*]] = select i1 [[COND]], i64 [[V1]], i64 [[V3]]
+; CHECK-NEXT:    call void @use(i64 [[X]])
+; CHECK-NEXT:    [[Y:%.*]] = select i1 [[COND]], i64 [[V2]], i64 [[V4]]
+; CHECK-NEXT:    call void @use(i64 [[Y]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %a0 = insertvalue { i64, i64 } poison, i64 %v1, 0
+  %a1 = insertvalue { i64, i64 } %a0, i64 %v2, 1
+  %b0 = insertvalue { i64, i64 } poison, i64 %v3, 0
+  %b1 = insertvalue { i64, i64 } %b0, i64 %v4, 1
+  %sel = select i1 %cond, { i64, i64 } %a1, { i64, i64 } %b1
+  %x = extractvalue { i64, i64 } %sel, 0
+  call void @use(i64 %x)
+  %y = extractvalue { i64, i64 } %sel, 1
+  call void @use(i64 %y)
+  ret void
+}
+
+define i64 @test_select_agg_simplified(i1 %cond, i64 %v1, i64 %v2, i64 %v3, i64 %v4) {
+; CHECK-LABEL: define i64 @test_select_agg_simplified(
+; CHECK-SAME: i1 [[COND:%.*]], i64 [[V1:%.*]], i64 [[V2:%.*]], i64 [[V3:%.*]], i64 [[V4:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[RET:%.*]] = select i1 [[COND]], i64 [[V2]], i64 0
+; CHECK-NEXT:    ret i64 [[RET]]
+;
+entry:
+  %a0 = insertvalue { i64, i64 } poison, i64 %v1, 0
+  %a1 = insertvalue { i64, i64 } %a0, i64 %v2, 1
+  %b0 = insertvalue { i64, i64 } poison, i64 %v3, 0
+  %b1 = insertvalue { i64, i64 } %b0, i64 %v4, 1
+  %sel = select i1 %cond, { i64, i64 } %a1, { i64, i64 } %b1
+  %y = extractvalue { i64, i64 } %sel, 1
+  %ret = select i1 %cond, i64 %y, i64 0
+  ret i64 %ret
+}
+
+declare void @use(i64)

[nikic]

I think the proper way to do this would be to call FoldOpIntoSelect for extractvalue.

[dtcxzyw]

I think the proper way to do this would be to call FoldOpIntoSelect for extractvalue.

llvm-project/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

Lines 1658 to 1661 in 193b3d6

	Value *TV = SI->getTrueValue();
	Value *FV = SI->getFalseValue();
	if (!(isa<Constant>(TV) || isa<Constant>(FV)))
	return nullptr;

FoldOpIntoSelect doesn't handle non-constant select arms.

[nikic]

@dtcxzyw If it is viable from a compile-time perspective, we could change the fold to use InstSimplify instead of ConstantFolding...

[dtcxzyw]

@dtcxzyw If it is viable from a compile-time perspective, we could change the fold to use InstSimplify instead of ConstantFolding...

Compile-time impact: https://llvm-compile-time-tracker.com/compare.php?from=3ec1f25f3cf9346590892397ec9ddd859397d363&to=0370e494c190b02a4548f85d480c437ddafd44d8&stat=instructions:u

[dtcxzyw]

I think the proper way to do this would be to call FoldOpIntoSelect for extractvalue.

Done.

[goldsteinn]

LGTM, wait on nikic

[nikic]

LGTM