[PATCH] D122233: [AggressiveInstCombine] Add `shufflevector` instr support to `TruncInstCombine`

[Anton Afanasyev via Phabricator via llvm-commits]

anton-afanasyev created this revision.
anton-afanasyev added reviewers: lebedev.ri, RKSimon.
Herald added a subscriber: hiraditya.
Herald added a project: All.
anton-afanasyev requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

  Add `shufflevector` instruction to the expression graph post-dominated by `trunc`,
  allowing `TruncInstCombine` to reduce bitwidth of expressions containing these
  instructions.
  
  Fixes #54149


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D122233

Files:
  llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp
  llvm/test/Transforms/AggressiveInstCombine/trunc_vector_instrs.ll


Index: llvm/test/Transforms/AggressiveInstCombine/trunc_vector_instrs.ll
===================================================================
--- llvm/test/Transforms/AggressiveInstCombine/trunc_vector_instrs.ll
+++ llvm/test/Transforms/AggressiveInstCombine/trunc_vector_instrs.ll
@@ -3,11 +3,10 @@
 
 define <4 x i16> @shuffle(<2 x i8> %a, <2 x i8> %b) {
 ; CHECK-LABEL: @shuffle(
-; CHECK-NEXT:    [[ZEXTA:%.*]] = zext <2 x i8> [[A:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[ZEXTB:%.*]] = zext <2 x i8> [[B:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <2 x i32> [[ZEXTA]], <2 x i32> [[ZEXTB]], <4 x i32> <i32 3, i32 2, i32 1, i32 0>
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <4 x i32> [[SHUF]] to <4 x i16>
-; CHECK-NEXT:    ret <4 x i16> [[TRUNC]]
+; CHECK-NEXT:    [[ZEXTA:%.*]] = zext <2 x i8> [[A:%.*]] to <2 x i16>
+; CHECK-NEXT:    [[ZEXTB:%.*]] = zext <2 x i8> [[B:%.*]] to <2 x i16>
+; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <2 x i16> [[ZEXTA]], <2 x i16> [[ZEXTB]], <4 x i32> <i32 3, i32 2, i32 1, i32 0>
+; CHECK-NEXT:    ret <4 x i16> [[SHUF]]
 ;
   %zexta = zext <2 x i8> %a to <2 x i32>
   %zextb = zext <2 x i8> %b to <2 x i32>
@@ -18,10 +17,9 @@
 
 define <2 x i16> @unary_shuffle(<2 x i8> %a) {
 ; CHECK-LABEL: @unary_shuffle(
-; CHECK-NEXT:    [[ZEXTA:%.*]] = zext <2 x i8> [[A:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <2 x i32> [[ZEXTA]], <2 x i32> undef, <2 x i32> <i32 1, i32 0>
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <2 x i32> [[SHUF]] to <2 x i16>
-; CHECK-NEXT:    ret <2 x i16> [[TRUNC]]
+; CHECK-NEXT:    [[ZEXTA:%.*]] = zext <2 x i8> [[A:%.*]] to <2 x i16>
+; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <2 x i16> [[ZEXTA]], <2 x i16> undef, <2 x i32> <i32 1, i32 0>
+; CHECK-NEXT:    ret <2 x i16> [[SHUF]]
 ;
   %zexta = zext <2 x i8> %a to <2 x i32>
   %shuf = shufflevector <2 x i32> %zexta, <2 x i32> undef, <2 x i32> <i32 1, i32 0>
@@ -31,9 +29,7 @@
 
 define <4 x i16> @const_shuffle() {
 ; CHECK-LABEL: @const_shuffle(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <2 x i32> <i32 1, i32 2>, <2 x i32> <i32 3, i32 7>, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <4 x i32> [[SHUF]] to <4 x i16>
-; CHECK-NEXT:    ret <4 x i16> [[TRUNC]]
+; CHECK-NEXT:    ret <4 x i16> <i16 7, i16 3, i16 2, i16 1>
 ;
   %shuf = shufflevector <2 x i32> <i32 1, i32 2>, <2 x i32> <i32 3, i32 7>, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
   %trunc = trunc <4 x i32> %shuf to <4 x i16>
Index: llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp
===================================================================
--- llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp
+++ llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp
@@ -66,6 +66,7 @@
   case Instruction::UDiv:
   case Instruction::URem:
   case Instruction::InsertElement:
+  case Instruction::ShuffleVector:
     Ops.push_back(I->getOperand(0));
     Ops.push_back(I->getOperand(1));
     break;
@@ -145,6 +146,7 @@
     case Instruction::UDiv:
     case Instruction::URem:
     case Instruction::InsertElement:
+    case Instruction::ShuffleVector:
     case Instruction::ExtractElement:
     case Instruction::Select: {
       SmallVector<Value *, 2> Operands;
@@ -162,10 +164,7 @@
       break;
     }
     default:
-      // TODO: Can handle more cases here:
-      // 1. shufflevector
-      // 2. sdiv, srem
-      // ...
+      // TODO: Can handle more cases here: sdiv, srem, ...
       return false;
     }
   }
@@ -456,6 +455,13 @@
       Res = Builder.CreateInsertElement(Vec, NewElt, Idx);
       break;
     }
+    case Instruction::ShuffleVector: {
+      Value *Vec1 = getReducedOperand(I->getOperand(0), SclTy);
+      Value *Vec2 = getReducedOperand(I->getOperand(1), SclTy);
+      ArrayRef<int> Mask = cast<ShuffleVectorInst>(I)->getShuffleMask();
+      Res = Builder.CreateShuffleVector(Vec1, Vec2, Mask);
+      break;
+    }
     case Instruction::Select: {
       Value *Op0 = I->getOperand(0);
       Value *LHS = getReducedOperand(I->getOperand(1), SclTy);


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D122233.417301.patch
Type: text/x-patch
Size: 4071 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20220322/b8f9a3af/attachment.bin>