[InstCombine] Fold extractelt with select of constants

[tsymalla]

Summary

An extractelt with a constant index which extracts an element from the
two vector operands of a select can be directly folded into a select.

extractelt (select %x, %vec1, %vec2), %const ->
select %x, %vec1[%const], %vec2[%const]

Note: the implementation currently only works for constant vector operands.

[foad]

extractelt (select %cond, <vec1>, <vec2>), %c1, %c2 ->
select %cond, <vec1>[c1], <vec2>[c2]

What are %c1 and %c2? Surely these is only one index?

Also your description does not make it clear that you are only doing this for constant vectors. For general vectors it could be: extractelt (select %cond, <vec1>, <vec2>), %c -> select %cond, (extractelt <vec1>, %c), (extractelt <vec2>, %c)

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

396

Why do you need to check this? Generally it is OK when the top-level thing has multiple uses.

403

In most case you should use an early "return nullptr" to reduce indentation.

[nikic]

Can we extend FoldOpIntoSelect to handle this instead?

[tsymalla]

In D137934#3924785, @nikic wrote:

Can we extend FoldOpIntoSelect to handle this instead?

Okay.

[tsymalla]

In D137934#3924774, @foad wrote:

extractelt (select %cond, <vec1>, <vec2>), %c1, %c2 ->
select %cond, <vec1>[c1], <vec2>[c2]

What are %c1 and %c2? Surely these is only one index?

Also your description does not make it clear that you are only doing this for constant vectors. For general vectors it could be: extractelt (select %cond, <vec1>, <vec2>), %c -> select %cond, (extractelt <vec1>, %c), (extractelt <vec2>, %c)

You are right, there is only one constant. For now, I'm only handling constant vectors, therefore I'll adjust the comment.

[tsymalla]

• Move logic to FoldOpIntoSelect.
• Remove superfluous use check for the extractelt instruction.
• Make use of early returns.
• Add more / fix existing comments.

[foad]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1107

Why do these have to be ConstantInt? Surely it should work with floats too (should have a test for this) or any other kind of Constant?

[spatel]

This is missing at least one critical pattern-matching constraint (and regression test - because this will crash):

define i32 @extelt_vecselect_const_operand_vector(<3 x i1> %c) {
  %s = select <3 x i1> %c, <3 x i32> <i32 2, i32 3, i32 4>, <3 x i32> <i32 5, i32 6, i32 7>
  %r = extractelement <3 x i32> %s, i32 2
  ret i32 %r
}

[tsymalla]

In D137934#3925091, @spatel wrote:

This is missing at least one critical pattern-matching constraint (and regression test - because this will crash):

define i32 @extelt_vecselect_const_operand_vector(<3 x i1> %c) {
  %s = select <3 x i1> %c, <3 x i32> <i32 2, i32 3, i32 4>, <3 x i32> <i32 5, i32 6, i32 7>
  %r = extractelement <3 x i32> %s, i32 2
  ret i32 %r
}

Thanks!

[tsymalla]

Add additional test to show folding for other types.
Use Constants instead of ConstantInt.
Conditions must be Integer types, no vector types.

[tsymalla]

Update lit tests.

[arsenm]

llvm/test/Transforms/InstCombine/extractelement.ll

816

Needs a test with an out of bounds vector index

[tsymalla]

llvm/test/Transforms/InstCombine/extractelement.ll

816

Hmm, I don't know if that makes sense. A trivial case with an extractelement with OOB behavior will not be caught by my code, as far as I can see, so it will be folded to a poison value.
I'll double-check and probably remove the OOB check in my code.

[tsymalla]

llvm/test/Transforms/InstCombine/extractelement.ll

816

Found a possible different issue with the OOB behavior:

define i32 @extelt_select_const_operand_vector_oob(i1 %c) {
  %oob = select i1 %c, i32 3, i32 4
  %s = select i1 %c, <3 x i32> <i32 2, i32 3, i32 4>, <3 x i32> <i32 5, i32 6, i32 7>
  %r = extractelement <3 x i32> %s, i32 %oob
  ret i32 %r
}

will make opt crash during constant folding.

[tsymalla]

llvm/test/Transforms/InstCombine/extractelement.ll

816

Simple OOB extractelements get optimized away before InstCombine reaches that code, and above snippet is part of a different issue, so I'll remove the OOB logic for now.

[tsymalla]

Remove OOB logic. Simple cases will be handled earlier, so the relevant code
path will not be visited.

[foad]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1101

Don't need a cast. getAggregateElement already returns Constant.

[foad]

You don't need your new function FoldExtractElementSelectConstVector at all. All you should need to do is to call FoldOpIntoSelect from visitExtractElementInst, and add (trivial) support for ExtractElementInst in foldOperationIntoSelectOperand. For example: https://reviews.llvm.org/differential/diff/475368/

[tsymalla]

Remove superfluous dyn_cast for getAggregateElement.

[tsymalla]

In D137934#3926911, @foad wrote:

You don't need your new function FoldExtractElementSelectConstVector at all. All you should need to do is to call FoldOpIntoSelect from visitExtractElementInst, and add (trivial) support for ExtractElementInst in foldOperationIntoSelectOperand. For example: https://reviews.llvm.org/differential/diff/475368/

Thanks. It needs to be done earlier, in constantFoldOperationIntoSelectOperand, but should nevertheless work.

[tsymalla]

Removed the whole custom implementation as InstCombine is already able
to do the constant folding, but the Visitor for extractelement needs an
additional check to prevent folding the constants into select operands
when the select uses a condition vector.

[foad]

This version fails assertions during check-llvm-transforms-instcombine.

[tsymalla]

Fix assertion failures

[spatel]

This patch is more general than where it started. I'm not opposed to that, but it needs more tests.

This crashes:

define i32 @extelt_select_const_operand_vector_var_index(i1 %c, i32 %e) {
  %s = select i1 %c, <3 x i32> <i32 2, i32 3, i32 4>, <3 x i32> <i32 5, i32 6, i32 7>
  %r = extractelement <3 x i32> %s, i32 %e
  ret i32 %r
}

We can fold with a single constant select operand (the srem diff show this, but we should have minimal tests too):

define i32 @extelt_select_var_const_operand_vector(i1 %c, <3 x i32> %v) {
  %s = select i1 %c, <3 x i32> %v, <3 x i32> <i32 5, i32 6, i32 7>
  %r = extractelement <3 x i32> %s, i32 1
  ret i32 %r
}

define i32 @extelt_select_const_var_operand_vector(i1 %c, <3 x i32> %v) {
  %s = select i1 %c, <3 x i32> <i32 5, i32 6, i32 7>, <3 x i32> %v
  %r = extractelement <3 x i32> %s, i32 0
  ret i32 %r
}

Does this behave as expected when the select has another use besides the extract?

declare void @use(<3 x i32>)

define i32 @extelt_select_const_var_operands_vector_extra_use(i1 %c, <3 x i32> %x) {
  %s = select i1 %c, <3 x i32> <i32 42, i32 5, i32 4>, <3 x i32> %x
  call void @use(<3 x i32> %s)
  %r = extractelement <3 x i32> %s, i64 0
  ret i32 %r
}

define i32 @extelt_select_const_operands_vector_extra_use(i1 %c, <3 x i32> %x) {
  %s = select i1 %c, <3 x i32> <i32 42, i32 5, i32 4>, <3 x i32> <i32 5, i32 6, i32 7>
  call void @use(<3 x i32> %s)
  %r = extractelement <3 x i32> %s, i64 0
  ret i32 %r
}

[tsymalla]

@spatel Thanks for adding these tests. These work in general, except the case where the second operand of the ExtractElement instruction is not a constant. The reason is that currently constantFoldOperationIntoSelectOperand expects both operands to be either the original select or a constant, but not an additional select, so a new extractelement instruction can be created to extract the actual value. If one of the operands is another select, it could be obvious that the sequence shows OOB behavior:

%e = select i1 %c, i32 3, i32 4
%s = select i1 %c, <3 x i32> <i32 2, i32 3, i32 4>, <3 x i32> <i32 5, i32 6, i32 7>
%r = extractelement <3 x i32> %s, i32 %e
ret i32 %r

Even if I know how to make the regular case work, I'll wrap my head around that particular issue.

[spatel]

In D137934#3928765, @tsymalla wrote:

@spatel Thanks for adding these tests. These work in general, except the case where the second operand of the ExtractElement instruction is not a constant. The reason is that currently constantFoldOperationIntoSelectOperand expects both operands to be either the original select or a constant, but not an additional select, so a new extractelement instruction can be created to extract the actual value. If one of the operands is another select, it could be obvious that the sequence shows OOB behavior:

%e = select i1 %c, i32 3, i32 4
%s = select i1 %c, <3 x i32> <i32 2, i32 3, i32 4>, <3 x i32> <i32 5, i32 6, i32 7>
%r = extractelement <3 x i32> %s, i32 %e
ret i32 %r

Even if I know how to make the regular case work, I'll wrap my head around that particular issue.

Just bail out if the extract index value is not an immediate constant? The case where the index value is a select-of-constants that can be reduced might be another patch, but that doesn't seem like a common pattern.

[tsymalla]

In D137934#3936962, @spatel wrote:

In D137934#3928765, @tsymalla wrote:

@spatel Thanks for adding these tests. These work in general, except the case where the second operand of the ExtractElement instruction is not a constant. The reason is that currently constantFoldOperationIntoSelectOperand expects both operands to be either the original select or a constant, but not an additional select, so a new extractelement instruction can be created to extract the actual value. If one of the operands is another select, it could be obvious that the sequence shows OOB behavior:

%e = select i1 %c, i32 3, i32 4
%s = select i1 %c, <3 x i32> <i32 2, i32 3, i32 4>, <3 x i32> <i32 5, i32 6, i32 7>
%r = extractelement <3 x i32> %s, i32 %e
ret i32 %r

Even if I know how to make the regular case work, I'll wrap my head around that particular issue.

Just bail out if the extract index value is not an immediate constant? The case where the index value is a select-of-constants that can be reduced might be another patch, but that doesn't seem like a common pattern.

Yeah, I have that working, but it's not going to optimize such cases like in the shufflevector tests for now. I'll go with bailing out for now.

[tsymalla]

Don't access unreachable path when trying to constant-fold an extractelement with a select as index operand.

[spatel]

This seems close to complete now. Please pre-commit the tests with baseline CHECKs as an NFC patch (no pre-commit review needed). That way, we'll just show test diffs in this patch.

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

401–403

Omit this description of a negative pattern. The formula under this line describes what we expect to happen.
One minor adjustment - c generally refers to a constant value in this code, so use some other letter for the condition:

// extractelt (select %x, <vec1>, <vec2>), %const ->
// select %x, <vec1>[%const], <vec2>[%const]

594

Remove whitespace changes, so the patch is minimal.
You could do a more extensive formatting cleanup as an NFC patch if you want.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1072

Tab settings are off. Adjust your editor to conform with existing code or use clang-format.

1198

Avoid dyn_cast if we're not using the captured value:

if (!isa<Constant>(EI->getIndexOperand()))

[tsymalla]

Addressed review comments.

[foad]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1198

This whole check belongs in visitExtractElementInst, not here.

[tsymalla]

Moving constant check to visitExtractElementInst

[foad]

The code LGTM.

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

401–402

Why do you write <vec1> instead of %vec1? I found it confusing to have two different syntaxes for variables in the same line.

403–406

I don't really understand this comment. "FIXME" suggests that something is broken. What is broken? What would you like to fold extractelt (select %c, %<vec1>, %<vec2>), (select %c, %c1, %c2) into?

[tsymalla]

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

401–402

I wanted to use the vector-notation <v_1, ..., v_n, but I can adjust it.

403–406

Currently, cases like

extractelt (select %c, %<vec1>, %<vec2>), (select %c, 3, 4) are not handled. If such instruction sequence leads to OOB behavior at compile time, this could be caught here.
I'm going to adjust the comment.

[tsymalla]

Made comments more concise.

[spatel]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

409

nit: use "auto *" with dyn_cast to be consistent with the surrounding code.