[llvm] [InstCombine] Fold zext(sub(0, trunc(X))) to and(sub(0, X), mask) (PR #207564)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Ayaan (def3r)

<details>
<summary>Changes</summary>

Problem: vector rotate and funnel shift fails to fold for vectors > 16 bytes on AVX-512. This is because of the `trunc` and `zext` instructions.

Example:
```llvm
define dso_local <8 x i64> @<!-- -->baz(<8 x i64> %0, <8 x i64> %1) local_unnamed_addr {
Entry:
  %2 = trunc <8 x i64> %1 to <8 x i6>
  %3 = sub <8 x i6> zeroinitializer, %2
  %4 = zext <8 x i6> %3 to <8 x i64>
  %5 = shl <8 x i64> %0, %4
  %6 = and <8 x i64> %1, splat (i64 63)
  %7 = lshr <8 x i64> %0, %6
  %8 = or <8 x i64> %5, %7
  ret <8 x i64> %8
}
```   

Solution; Canonicalize the transformation `zext(sub(0, trunc(X))) -> and(sub(0, X), mask)` for scalars and vectors.

Closes #<!-- -->165306

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


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp (+9) 
- (added) llvm/test/Transforms/InstCombine/zext-sub-trunc.ll (+55) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
index 7e747ddb9013e..84c143897c8d4 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
@@ -1686,6 +1686,15 @@ Instruction *InstCombinerImpl::visitZExt(ZExtInst &Zext) {
     return BinaryOperator::CreateXor(Builder.CreateAnd(X, ZC), ZC);
   }
 
+  // zext(sub(0, trunc(X))) -> and(sub(0, X), mask)
+  if (match(Src, m_Sub(m_Zero(), m_Trunc(m_Value(X)))) &&
+      X->getType() == DestTy) {
+    APInt Mask = APInt::getLowBitsSet(DestTy->getScalarSizeInBits(),
+                                      SrcTy->getScalarSizeInBits());
+    Value *Neg = Builder.CreateSub(ConstantInt::get(DestTy, 0), X);
+    return BinaryOperator::CreateAnd(Neg, ConstantInt::get(DestTy, Mask));
+  }
+
   // If we are truncating, masking, and then zexting back to the original type,
   // that's just a mask. This is not handled by canEvaluateZextd if the
   // intermediate values have extra uses. This could be generalized further for
diff --git a/llvm/test/Transforms/InstCombine/zext-sub-trunc.ll b/llvm/test/Transforms/InstCombine/zext-sub-trunc.ll
new file mode 100644
index 0000000000000..47c31acc800c5
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/zext-sub-trunc.ll
@@ -0,0 +1,55 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define i64 @test_scalar_mask(i64 %a) {
+; CHECK-LABEL: define i64 @test_scalar_mask(
+; CHECK-SAME: i64 [[A:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = sub i64 0, [[A]]
+; CHECK-NEXT:    [[ZEXT:%.*]] = and i64 [[TMP1]], 63
+; CHECK-NEXT:    ret i64 [[ZEXT]]
+;
+  %trunc = trunc i64 %a to i6
+  %neg = sub i6 0, %trunc
+  %zext = zext i6 %neg to i64
+  ret i64 %zext
+}
+
+define <8 x i64> @test_vector_mask_v8i64(<8 x i64> %a0) {
+; CHECK-LABEL: define <8 x i64> @test_vector_mask_v8i64(
+; CHECK-SAME: <8 x i64> [[A0:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = sub <8 x i64> zeroinitializer, [[A0]]
+; CHECK-NEXT:    [[ZEXT:%.*]] = and <8 x i64> [[TMP1]], splat (i64 63)
+; CHECK-NEXT:    ret <8 x i64> [[ZEXT]]
+;
+  %trunc = trunc <8 x i64> %a0 to <8 x i6>
+  %neg = sub <8 x i6> zeroinitializer, %trunc
+  %zext = zext <8 x i6> %neg to <8 x i64>
+  ret <8 x i64> %zext
+}
+
+define <8 x i32> @test_vector_mask_v8i32(<8 x i32> %a0) {
+; CHECK-LABEL: define <8 x i32> @test_vector_mask_v8i32(
+; CHECK-SAME: <8 x i32> [[A0:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = sub <8 x i32> zeroinitializer, [[A0]]
+; CHECK-NEXT:    [[ZEXT:%.*]] = and <8 x i32> [[TMP1]], splat (i32 15)
+; CHECK-NEXT:    ret <8 x i32> [[ZEXT]]
+;
+  %trunc = trunc <8 x i32> %a0 to <8 x i4>
+  %neg = sub <8 x i4> zeroinitializer, %trunc
+  %zext = zext <8 x i4> %neg to <8 x i32>
+  ret <8 x i32> %zext
+}
+
+define <8 x i64> @test_negative_vector_mask_v8i32(<8 x i32> %a0) {
+; CHECK-LABEL: define <8 x i64> @test_negative_vector_mask_v8i32(
+; CHECK-SAME: <8 x i32> [[A0:%.*]]) {
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <8 x i32> [[A0]] to <8 x i4>
+; CHECK-NEXT:    [[NEG:%.*]] = sub <8 x i4> zeroinitializer, [[TRUNC]]
+; CHECK-NEXT:    [[ZEXT:%.*]] = zext <8 x i4> [[NEG]] to <8 x i64>
+; CHECK-NEXT:    ret <8 x i64> [[ZEXT]]
+;
+  %trunc = trunc <8 x i32> %a0 to <8 x i4>
+  %neg = sub <8 x i4> zeroinitializer, %trunc
+  %zext = zext <8 x i4> %neg to <8 x i64>
+  ret <8 x i64> %zext
+}

``````````

</details>


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