[WebAssembly] Fold constant i8x16.swizzle and i8x16.relaxed.swizzle to shufflevector (#169110)

Resolves #169058.

This adds ~~an InstCombine pass~~ a TTI hook to the WebAssembly backend
that folds `i8x16.swizzle` and `i8x16.relaxed.swizzle` operations to
`shufflevector` operations if their mask operands are constant.

This is mainly useful for abstractions over the raw intrinsics--for
instance, in architecture-generic SIMD code that may not be able to
expose the constant shuffles due to type system limitations.

I took most of this from the x86 backend (in particular,
`simplifyX86vpermilvar` in `X86InstCombineIntrinsic`), and adapted it
for the WebAssembly backend. There wasn't any previous
`instCombineIntrinsic` method on the WebAssembly `TargetTransformInfo`,
so I added it. Right now, this swizzle optimization is the only one it
performs.

As I noted in the transform itself, the "relaxed" swizzle actually has
stricter preconditions than the non-relaxed one. If a non-negative but
still out-of-bounds index is provided, the "relaxed" swizzle can choose
between returning 0 and the lane at the index modulo 16. However, it
must make the same choice every time, and we don't know which choice the
runtime will make, so we can't constant-fold it.

The regression tests were mostly generated by Claude and adapted a bit
by me (I tried to follow the [InstCombine contributor
guide](https://llvm.org/docs/InstCombineContributorGuide.html#tests)).
There was previously no WebAssembly subdirectory within the InstCombine
tests, so I created that too; as of now, the swizzle fold test is the
only file in it. Everything else was written by myself (well, partly
copy-pasted from the x86 backend).

I'm not sure how to write an Alive2 test for this; I can't find any
examples where the input is an arbitrary constant.

Author: valadaptive

llvm/lib/Target/WebAssembly/WebAssemblyTargetTransformInfo.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,87 @@ 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;
+    }
+
+    if (IsRelaxed && cast<ConstantInt>(COp)->getSExtValue() >= NumElts) {
+      // The relaxed_swizzle operation always returns 0 if the lane index is
+      // less than 0 when interpreted as a signed value. For lane indices above
+      // 15, however, it 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;
+    }
+
+    uint64_t Index = cast<ConstantInt>(COp)->getZExtValue();
+    if (Index >= NumElts) {
+      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.h

@@ -103,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/test/Transforms/InstCombine/WebAssembly/fold-swizzle.ll

@@ -0,0 +1,113 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt < %s -passes=instcombine -mtriple=wasm32-unknown-unknown -S | FileCheck %s
+
+; swizzle with a constant operand should be optimized to a shufflevector.
+
+; Identity swizzle pattern
+define <16 x i8> @swizzle_identity(<16 x i8> %v) {
+; CHECK-LABEL: define <16 x i8> @swizzle_identity(
+; CHECK-SAME: <16 x i8> [[V:%.*]]) {
+; CHECK-NEXT:    ret <16 x i8> [[V]]
+;
+  %result = tail call <16 x i8> @llvm.wasm.swizzle(<16 x i8> %v, <16 x i8> <i8 0, i8 1, i8 2, i8 3, i8 4, i8 5, i8 6, i8 7, i8 8, i8 9, i8 10, i8 11, i8 12, i8 13, i8 14, i8 15>)
+  ret <16 x i8> %result
+}
+
+; Reverse swizzle pattern
+define <16 x i8> @swizzle_reverse(<16 x i8> %v) {
+; CHECK-LABEL: define <16 x i8> @swizzle_reverse(
+; CHECK-SAME: <16 x i8> [[V:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = shufflevector <16 x i8> [[V]], <16 x i8> poison, <16 x i32> <i32 15, i32 14, i32 13, i32 12, i32 11, i32 10, i32 9, i32 8, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0>
+; CHECK-NEXT:    ret <16 x i8> [[RESULT]]
+;
+  %result = tail call <16 x i8> @llvm.wasm.swizzle(<16 x i8> %v, <16 x i8> <i8 15, i8 14, i8 13, i8 12, i8 11, i8 10, i8 9, i8 8, i8 7, i8 6, i8 5, i8 4, i8 3, i8 2, i8 1, i8 0>)
+  ret <16 x i8> %result
+}
+
+; poison elements
+define <16 x i8> @swizzle_with_poison(<16 x i8> %v) {
+; CHECK-LABEL: define <16 x i8> @swizzle_with_poison(
+; CHECK-SAME: <16 x i8> [[V:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = shufflevector <16 x i8> [[V]], <16 x i8> poison, <16 x i32> <i32 0, i32 poison, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+; CHECK-NEXT:    ret <16 x i8> [[RESULT]]
+;
+  %result = tail call <16 x i8> @llvm.wasm.swizzle(<16 x i8> %v, <16 x i8> <i8 0, i8 poison, i8 2, i8 3, i8 4, i8 5, i8 6, i8 7, i8 8, i8 9, i8 10, i8 11, i8 12, i8 13, i8 14, i8 15>)
+  ret <16 x i8> %result
+}
+
+; Negative test: non-constant operand
+define <16 x i8> @swizzle_non_constant(<16 x i8> %v, <16 x i8> %mask) {
+; CHECK-LABEL: define <16 x i8> @swizzle_non_constant(
+; CHECK-SAME: <16 x i8> [[V:%.*]], <16 x i8> [[MASK:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = tail call <16 x i8> @llvm.wasm.swizzle(<16 x i8> [[V]], <16 x i8> [[MASK]])
+; CHECK-NEXT:    ret <16 x i8> [[RESULT]]
+;
+  %result = tail call <16 x i8> @llvm.wasm.swizzle(<16 x i8> %v, <16 x i8> %mask)
+  ret <16 x i8> %result
+}
+
+; Out-of-bounds index, otherwise identity pattern
+define <16 x i8> @swizzle_out_of_bounds_1(<16 x i8> %v) {
+; CHECK-LABEL: define <16 x i8> @swizzle_out_of_bounds_1(
+; CHECK-SAME: <16 x i8> [[V:%.*]]) {
+; CHECK-NEXT:    [[RESULT1:%.*]] = insertelement <16 x i8> [[V]], i8 0, i64 15
+; CHECK-NEXT:    ret <16 x i8> [[RESULT1]]
+;
+  %result = tail call <16 x i8> @llvm.wasm.swizzle(<16 x i8> %v, <16 x i8> <i8 0, i8 1, i8 2, i8 3, i8 4, i8 5, i8 6, i8 7, i8 8, i8 9, i8 10, i8 11, i8 12, i8 13, i8 14, i8 16>)
+  ret <16 x i8> %result
+}
+
+; Out-of-bounds indices, both negative and positive
+define <16 x i8> @swizzle_out_of_bounds_2(<16 x i8> %v) {
+; CHECK-LABEL: define <16 x i8> @swizzle_out_of_bounds_2(
+; CHECK-SAME: <16 x i8> [[V:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = shufflevector <16 x i8> [[V]], <16 x i8> <i8 0, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison>, <16 x i32> <i32 16, i32 16, i32 13, i32 12, i32 11, i32 10, i32 9, i32 8, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0>
+; CHECK-NEXT:    ret <16 x i8> [[RESULT]]
+;
+  %result = tail call <16 x i8> @llvm.wasm.swizzle(<16 x i8> %v, <16 x i8> <i8 99, i8 -1, i8 13, i8 12, i8 11, i8 10, i8 9, i8 8, i8 7, i8 6, i8 5, i8 4, i8 3, i8 2, i8 1, i8 0>)
+  ret <16 x i8> %result
+}
+
+; Identity swizzle pattern (relaxed_swizzle)
+define <16 x i8> @relaxed_swizzle_identity(<16 x i8> %v) {
+; CHECK-LABEL: define <16 x i8> @relaxed_swizzle_identity(
+; CHECK-SAME: <16 x i8> [[V:%.*]]) {
+; CHECK-NEXT:    ret <16 x i8> [[V]]
+;
+  %result = tail call <16 x i8> @llvm.wasm.relaxed.swizzle(<16 x i8> %v, <16 x i8> <i8 0, i8 1, i8 2, i8 3, i8 4, i8 5, i8 6, i8 7, i8 8, i8 9, i8 10, i8 11, i8 12, i8 13, i8 14, i8 15>)
+  ret <16 x i8> %result
+}
+
+; Reverse swizzle pattern (relaxed_swizzle)
+define <16 x i8> @relaxed_swizzle_reverse(<16 x i8> %v) {
+; CHECK-LABEL: define <16 x i8> @relaxed_swizzle_reverse(
+; CHECK-SAME: <16 x i8> [[V:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = shufflevector <16 x i8> [[V]], <16 x i8> poison, <16 x i32> <i32 15, i32 14, i32 13, i32 12, i32 11, i32 10, i32 9, i32 8, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0>
+; CHECK-NEXT:    ret <16 x i8> [[RESULT]]
+;
+  %result = tail call <16 x i8> @llvm.wasm.relaxed.swizzle(<16 x i8> %v, <16 x i8> <i8 15, i8 14, i8 13, i8 12, i8 11, i8 10, i8 9, i8 8, i8 7, i8 6, i8 5, i8 4, i8 3, i8 2, i8 1, i8 0>)
+  ret <16 x i8> %result
+}
+
+; Out-of-bounds index, only negative (relaxed_swizzle)
+define <16 x i8> @relaxed_swizzle_out_of_bounds(<16 x i8> %v) {
+; CHECK-LABEL: define <16 x i8> @relaxed_swizzle_out_of_bounds(
+; CHECK-SAME: <16 x i8> [[V:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = shufflevector <16 x i8> [[V]], <16 x i8> <i8 0, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison, i8 poison>, <16 x i32> <i32 16, i32 16, i32 13, i32 12, i32 11, i32 10, i32 9, i32 8, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0>
+; CHECK-NEXT:    ret <16 x i8> [[RESULT]]
+;
+  %result = tail call <16 x i8> @llvm.wasm.relaxed.swizzle(<16 x i8> %v, <16 x i8> <i8 -99, i8 -1, i8 13, i8 12, i8 11, i8 10, i8 9, i8 8, i8 7, i8 6, i8 5, i8 4, i8 3, i8 2, i8 1, i8 0>)
+  ret <16 x i8> %result
+}
+
+; Negative test: out-of-bounds index, both positive and negative (relaxed_swizzle)
+; The choice between different relaxed semantics can only be made at runtime, since it must be consistent.
+define <16 x i8> @relaxed_swizzle_out_of_bounds_positive(<16 x i8> %v) {
+; CHECK-LABEL: define <16 x i8> @relaxed_swizzle_out_of_bounds_positive(
+; CHECK-SAME: <16 x i8> [[V:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = tail call <16 x i8> @llvm.wasm.relaxed.swizzle(<16 x i8> [[V]], <16 x i8> <i8 99, i8 -1, i8 13, i8 12, i8 11, i8 10, i8 9, i8 8, i8 7, i8 6, i8 5, i8 4, i8 3, i8 2, i8 1, i8 0>)
+; CHECK-NEXT:    ret <16 x i8> [[RESULT]]
+;
+  %result = tail call <16 x i8> @llvm.wasm.relaxed.swizzle(<16 x i8> %v, <16 x i8> <i8 99, i8 -1, i8 13, i8 12, i8 11, i8 10, i8 9, i8 8, i8 7, i8 6, i8 5, i8 4, i8 3, i8 2, i8 1, i8 0>)
+  ret <16 x i8> %result
+}