From 421c02e5c6b376fbbd8cc61c20feb35f1cee7a5a Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev@redking.me.uk>
Date: Thu, 2 Jul 2020 11:56:32 +0100
Subject: [PATCH] [InstCombine] Add some (vXi1 trunc(lshr(x,c))) ->
 icmp_eq(and(x,c')) tests for non-uniform vectors

As noticed on PR46531
---
 .../Transforms/InstCombine/apint-shift.ll     | 12 ++++++
 llvm/test/Transforms/InstCombine/icmp.ll      | 41 +++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/apint-shift.ll b/llvm/test/Transforms/InstCombine/apint-shift.ll
index 0f05f76d1a68..d6dd529c92f7 100644
--- a/llvm/test/Transforms/InstCombine/apint-shift.ll
+++ b/llvm/test/Transforms/InstCombine/apint-shift.ll
@@ -329,6 +329,18 @@ define <2 x i1> @test16vec(<2 x i84> %X) {
   ret <2 x i1> %cmp
 }
 
+define <2 x i1> @test16vec_nonuniform(<2 x i84> %X) {
+; CHECK-LABEL: @test16vec_nonuniform(
+; CHECK-NEXT:    [[SHR1:%.*]] = lshr <2 x i84> [[X:%.*]], <i84 4, i84 2>
+; CHECK-NEXT:    [[CMP:%.*]] = trunc <2 x i84> [[SHR1]] to <2 x i1>
+; CHECK-NEXT:    ret <2 x i1> [[CMP]]
+;
+  %shr = ashr <2 x i84> %X, <i84 4, i84 2>
+  %and = and <2 x i84> %shr, <i84 1, i84 1>
+  %cmp = icmp ne <2 x i84> %and, zeroinitializer
+  ret <2 x i1> %cmp
+}
+
 define i1 @test17(i106 %A) {
 ; CHECK-LABEL: @test17(
 ; CHECK-NEXT:    [[B_MASK:%.*]] = and i106 [[A:%.*]], -8
diff --git a/llvm/test/Transforms/InstCombine/icmp.ll b/llvm/test/Transforms/InstCombine/icmp.ll
index 67d254a3fe52..eb214f059a1d 100644
--- a/llvm/test/Transforms/InstCombine/icmp.ll
+++ b/llvm/test/Transforms/InstCombine/icmp.ll
@@ -2686,6 +2686,20 @@ define <2 x i1> @icmp_and_or_lshr_cst_vec(<2 x i32> %x) {
   ret <2 x i1> %ret
 }
 
+define <2 x i1> @icmp_and_or_lshr_cst_vec_nonuniform(<2 x i32> %x) {
+; CHECK-LABEL: @icmp_and_or_lshr_cst_vec_nonuniform(
+; CHECK-NEXT:    [[SHF:%.*]] = lshr <2 x i32> [[X:%.*]], <i32 1, i32 2>
+; CHECK-NEXT:    [[OR:%.*]] = or <2 x i32> [[SHF]], [[X]]
+; CHECK-NEXT:    [[RET:%.*]] = trunc <2 x i32> [[OR]] to <2 x i1>
+; CHECK-NEXT:    ret <2 x i1> [[RET]]
+;
+  %shf = lshr <2 x i32> %x, <i32 1, i32 2>
+  %or = or <2 x i32> %shf, %x
+  %and = and <2 x i32> %or, <i32 1, i32 1>
+  %ret = icmp ne <2 x i32> %and, zeroinitializer
+  ret <2 x i1> %ret
+}
+
 define <2 x i1> @icmp_and_or_lshr_cst_vec_commute(<2 x i32> %xp) {
 ; CHECK-LABEL: @icmp_and_or_lshr_cst_vec_commute(
 ; CHECK-NEXT:    [[X:%.*]] = srem <2 x i32> [[XP:%.*]], <i32 42, i32 42>
@@ -2701,6 +2715,22 @@ define <2 x i1> @icmp_and_or_lshr_cst_vec_commute(<2 x i32> %xp) {
   ret <2 x i1> %ret
 }
 
+define <2 x i1> @icmp_and_or_lshr_cst_vec_nonuniform_commute(<2 x i32> %xp) {
+; CHECK-LABEL: @icmp_and_or_lshr_cst_vec_nonuniform_commute(
+; CHECK-NEXT:    [[X:%.*]] = srem <2 x i32> [[XP:%.*]], <i32 42, i32 42>
+; CHECK-NEXT:    [[SHF:%.*]] = lshr <2 x i32> [[X]], <i32 1, i32 2>
+; CHECK-NEXT:    [[OR:%.*]] = or <2 x i32> [[X]], [[SHF]]
+; CHECK-NEXT:    [[RET:%.*]] = trunc <2 x i32> [[OR]] to <2 x i1>
+; CHECK-NEXT:    ret <2 x i1> [[RET]]
+;
+  %x = srem <2 x i32> %xp, <i32 42, i32 -42> ; prevent complexity-based canonicalization
+  %shf = lshr <2 x i32> %x, <i32 1, i32 2>
+  %or = or <2 x i32> %x, %shf
+  %and = and <2 x i32> %or, <i32 1, i32 1>
+  %ret = icmp ne <2 x i32> %and, zeroinitializer
+  ret <2 x i1> %ret
+}
+
 define i1 @shl_ap1_zero_ap2_non_zero_2(i32 %a) {
 ; CHECK-LABEL: @shl_ap1_zero_ap2_non_zero_2(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i32 [[A:%.*]], 29
@@ -2721,6 +2751,17 @@ define <2 x i1> @shl_ap1_zero_ap2_non_zero_2_vec(<2 x i32> %a) {
   ret <2 x i1> %cmp
 }
 
+define <2 x i1> @shl_ap1_zero_ap2_non_zero_2_vec_nonuniform(<2 x i32> %a) {
+; CHECK-LABEL: @shl_ap1_zero_ap2_non_zero_2_vec_nonuniform(
+; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i32> <i32 4, i32 5>, [[A:%.*]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <2 x i32> [[SHL]], zeroinitializer
+; CHECK-NEXT:    ret <2 x i1> [[CMP]]
+;
+  %shl = shl <2 x i32> <i32 4, i32 5>, %a
+  %cmp = icmp eq <2 x i32> %shl, zeroinitializer
+  ret <2 x i1> %cmp
+}
+
 define i1 @shl_ap1_zero_ap2_non_zero_4(i32 %a) {
 ; CHECK-LABEL: @shl_ap1_zero_ap2_non_zero_4(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i32 [[A:%.*]], 30