[WebAssembly] Rearrange new TTI hook code

Author: valadaptive

llvm/lib/Target/WebAssembly/CMakeLists.txt

@@ -32,7 +32,6 @@ add_llvm_target(WebAssemblyCodeGen
   WebAssemblyFixIrreducibleControlFlow.cpp
   WebAssemblyFixFunctionBitcasts.cpp
   WebAssemblyFrameLowering.cpp
-  WebAssemblyInstCombineIntrinsic.cpp
   WebAssemblyISelDAGToDAG.cpp
   WebAssemblyISelLowering.cpp
   WebAssemblyInstrInfo.cpp

llvm/lib/Target/WebAssembly/WebAssemblyInstCombineIntrinsic.cpp

@@ -13,6 +13,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "WebAssemblyTargetTransformInfo.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/IntrinsicsWebAssembly.h"
+#include "llvm/Transforms/InstCombine/InstCombiner.h"
 
 #include "llvm/CodeGen/CostTable.h"
 using namespace llvm;
@@ -493,3 +496,86 @@ bool WebAssemblyTTIImpl::isProfitableToSinkOperands(
 
   return false;
 }
+
+/// Attempt to convert [relaxed_]swizzle to shufflevector if the mask is
+/// constant.
+static Value *simplifyWasmSwizzle(const IntrinsicInst &II,
+                                  InstCombiner::BuilderTy &Builder,
+                                  bool IsRelaxed) {
+  auto *V = dyn_cast<Constant>(II.getArgOperand(1));
+  if (!V)
+    return nullptr;
+
+  auto *VecTy = cast<FixedVectorType>(II.getType());
+  unsigned NumElts = VecTy->getNumElements();
+  assert(NumElts == 16);
+
+  // Construct a shuffle mask from constant integers or UNDEFs.
+  int Indexes[16];
+  bool AnyOutOfBounds = false;
+
+  for (unsigned I = 0; I < NumElts; ++I) {
+    Constant *COp = V->getAggregateElement(I);
+    if (!COp || (!isa<UndefValue>(COp) && !isa<ConstantInt>(COp)))
+      return nullptr;
+
+    if (isa<UndefValue>(COp)) {
+      Indexes[I] = -1;
+      continue;
+    }
+
+    int64_t Index = cast<ConstantInt>(COp)->getSExtValue();
+
+    if (Index >= NumElts && IsRelaxed) {
+      // For lane indices above 15, the relaxed_swizzle operation can choose
+      // between returning 0 or the lane at `Index % 16`. However, the choice
+      // must be made consistently. As the WebAssembly spec states:
+      //
+      // "The result of relaxed operators are implementation-dependent, because
+      // the set of possible results may depend on properties of the host
+      // environment, such as its hardware. Technically, their behaviour is
+      // controlled by a set of global parameters to the semantics that an
+      // implementation can instantiate in different ways. These choices are
+      // fixed, that is, parameters are constant during the execution of any
+      // given program."
+      //
+      // The WebAssembly runtime may choose differently from us, so we can't
+      // optimize a relaxed swizzle with lane indices above 15.
+      return nullptr;
+    }
+
+    if (Index >= NumElts || Index < 0) {
+      AnyOutOfBounds = true;
+      // If there are out-of-bounds indices, the swizzle instruction returns
+      // zeroes in those lanes. We'll provide an all-zeroes vector as the
+      // second argument to shufflevector and read the first element from it.
+      Indexes[I] = NumElts;
+      continue;
+    }
+
+    Indexes[I] = Index;
+  }
+
+  auto *V1 = II.getArgOperand(0);
+  auto *V2 =
+      AnyOutOfBounds ? Constant::getNullValue(VecTy) : PoisonValue::get(VecTy);
+
+  return Builder.CreateShuffleVector(V1, V2, ArrayRef(Indexes, NumElts));
+}
+
+std::optional<Instruction *>
+WebAssemblyTTIImpl::instCombineIntrinsic(InstCombiner &IC,
+                                         IntrinsicInst &II) const {
+  Intrinsic::ID IID = II.getIntrinsicID();
+  switch (IID) {
+  case Intrinsic::wasm_swizzle:
+  case Intrinsic::wasm_relaxed_swizzle:
+    if (Value *V = simplifyWasmSwizzle(
+            II, IC.Builder, IID == Intrinsic::wasm_relaxed_swizzle)) {
+      return IC.replaceInstUsesWith(II, V);
+    }
+    break;
+  }
+
+  return std::nullopt;
+}

llvm/lib/Target/WebAssembly/WebAssemblyTargetTransformInfo.cpp

@@ -90,8 +90,6 @@ class WebAssemblyTTIImpl final : public BasicTTIImplBase<WebAssemblyTTIImpl> {
                                      TTI::TargetCostKind CostKind,
                                      unsigned Index, const Value *Op0,
                                      const Value *Op1) const override;
-  std::optional<Instruction *>
-  instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const override;
   InstructionCost getPartialReductionCost(
       unsigned Opcode, Type *InputTypeA, Type *InputTypeB, Type *AccumType,
       ElementCount VF, TTI::PartialReductionExtendKind OpAExtend,
@@ -105,6 +103,9 @@ class WebAssemblyTTIImpl final : public BasicTTIImplBase<WebAssemblyTTIImpl> {
   bool isProfitableToSinkOperands(Instruction *I,
                                   SmallVectorImpl<Use *> &Ops) const override;
 
+  std::optional<Instruction *>
+  instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const override;
+
   /// @}
 };
 

llvm/lib/Target/WebAssembly/WebAssemblyTargetTransformInfo.h

@@ -54,7 +54,6 @@ static_library("LLVMWebAssemblyCodeGen") {
     "WebAssemblyFixFunctionBitcasts.cpp",
     "WebAssemblyFixIrreducibleControlFlow.cpp",
     "WebAssemblyFrameLowering.cpp",
-    "WebAssemblyInstCombineIntrinsic.cpp",
     "WebAssemblyISelDAGToDAG.cpp",
     "WebAssemblyISelLowering.cpp",
     "WebAssemblyInstrInfo.cpp",