Bug 1135039: Implement general SIMD swizzle and use it

js/src/builtin/SIMD.h

@@ -241,11 +241,11 @@
     _(splat)                         \
     _(not)                           \
     _(neg)                           \
+    _(swizzle)                       \
     _(check)
 #define FOREACH_COMMONX4_SIMD_OP(_)  \
     ION_COMMONX4_SIMD_OP(_)          \
     COMP_COMMONX4_TO_INT32X4_SIMD_OP(_) \
-    _(swizzle)                       \
     _(shuffle)                       \
     _(load)                          \
     _(loadX)                         \

js/src/jit/IonBuilder.h

@@ -832,6 +832,7 @@ class IonBuilder
     InliningStatus inlineSimdWith(CallInfo &callInfo, JSNative native, SimdLane lane,
                                   SimdTypeDescr::Type type);
     InliningStatus inlineSimdSplat(CallInfo &callInfo, JSNative native, SimdTypeDescr::Type type);
+    InliningStatus inlineSimdSwizzle(CallInfo &callInfo, JSNative native, SimdTypeDescr::Type type);
     InliningStatus inlineSimdCheck(CallInfo &callInfo, JSNative native, SimdTypeDescr::Type type);
     InliningStatus inlineSimdConvert(CallInfo &callInfo, JSNative native, bool isCast,
                                      SimdTypeDescr::Type from, SimdTypeDescr::Type to);

js/src/jit/LIR-Common.h

@@ -375,6 +375,49 @@ class LSimdSwizzleF : public LSimdSwizzleBase
     {}
 };
 
+class LSimdGeneralSwizzleBase : public LInstructionHelper<1, 5, 1>
+{
+  public:
+    LSimdGeneralSwizzleBase(const LAllocation &base, const LAllocation lanes[4],
+                            const LDefinition &temp)
+    {
+        setOperand(0, base);
+        for (size_t i = 0; i < 4; i++)
+            setOperand(1 + i, lanes[i]);
+        setTemp(0, temp);
+    }
+
+    const LAllocation *base() {
+        return getOperand(0);
+    }
+    const LAllocation *lane(size_t i) {
+        return getOperand(1 + i);
+    }
+    const LDefinition *temp() {
+        return getTemp(0);
+    }
+};
+
+class LSimdGeneralSwizzleI : public LSimdGeneralSwizzleBase
+{
+  public:
+    LIR_HEADER(SimdGeneralSwizzleI);
+    LSimdGeneralSwizzleI(const LAllocation &base, const LAllocation lanes[4],
+                         const LDefinition &temp)
+      : LSimdGeneralSwizzleBase(base, lanes, temp)
+    {}
+};
+
+class LSimdGeneralSwizzleF : public LSimdGeneralSwizzleBase
+{
+  public:
+    LIR_HEADER(SimdGeneralSwizzleF);
+    LSimdGeneralSwizzleF(const LAllocation &base, const LAllocation lanes[4],
+                         const LDefinition &temp)
+      : LSimdGeneralSwizzleBase(base, lanes, temp)
+    {}
+};
+
 // Base class for both int32x4 and float32x4 shuffle instructions.
 class LSimdShuffle : public LInstructionHelper<1, 2, 1>
 {

js/src/jit/LOpcodes.h

@@ -28,6 +28,8 @@
     _(SimdInsertElementI)           \
     _(SimdInsertElementF)           \
     _(SimdSignMaskX4)               \
+    _(SimdGeneralSwizzleI)          \
+    _(SimdGeneralSwizzleF)          \
     _(SimdSwizzleI)                 \
     _(SimdSwizzleF)                 \
     _(SimdShuffle)                  \

js/src/jit/Lowering.cpp

@@ -3948,6 +3948,29 @@ LIRGenerator::visitSimdShuffle(MSimdShuffle* ins)
     lir->setTemp(0, temp);
 }
 
+void
+LIRGenerator::visitSimdGeneralSwizzle(MSimdGeneralSwizzle *ins)
+{
+    MOZ_ASSERT(IsSimdType(ins->input()->type()));
+    MOZ_ASSERT(IsSimdType(ins->type()));
+
+    LAllocation lanesUses[4];
+    for (size_t i = 0; i < 4; i++)
+        lanesUses[i] = use(ins->lane(i));
+
+    if (ins->input()->type() == MIRType_Int32x4) {
+        LSimdGeneralSwizzleI *lir = new (alloc()) LSimdGeneralSwizzleI(useRegister(ins->input()),
+                                                                       lanesUses, temp());
+        define(lir, ins);
+    } else if (ins->input()->type() == MIRType_Float32x4) {
+        LSimdGeneralSwizzleF *lir = new (alloc()) LSimdGeneralSwizzleF(useRegister(ins->input()),
+                                                                       lanesUses, temp());
+        define(lir, ins);
+    } else {
+        MOZ_CRASH("Unknown SIMD kind when getting lane");
+    }
+}
+
 void
 LIRGenerator::visitSimdUnaryArith(MSimdUnaryArith *ins)
 {

js/src/jit/Lowering.h

@@ -271,6 +271,7 @@ class LIRGenerator : public LIRGeneratorSpecific
     void visitSimdInsertElement(MSimdInsertElement *ins);
     void visitSimdSignMask(MSimdSignMask *ins);
     void visitSimdSwizzle(MSimdSwizzle *ins);
+    void visitSimdGeneralSwizzle(MSimdGeneralSwizzle *ins);
     void visitSimdShuffle(MSimdShuffle *ins);
     void visitSimdUnaryArith(MSimdUnaryArith *ins);
     void visitSimdBinaryComp(MSimdBinaryComp *ins);

js/src/jit/MCallOptimize.cpp

@@ -350,6 +350,9 @@ IonBuilder::inlineNativeCall(CallInfo& callInfo, JSFunction* target)
     if (native == js::simd_float32x4_bitselect)
         return inlineSimdSelect(callInfo, native, IsElementWise(false), SimdTypeDescr::TYPE_FLOAT32);
 
+    if (native == js::simd_int32x4_swizzle)
+        return inlineSimdSwizzle(callInfo, native, SimdTypeDescr::TYPE_INT32);
+
     return InliningStatus_NotInlined;
 }
 
@@ -3120,5 +3123,21 @@ IonBuilder::inlineSimdCheck(CallInfo &callInfo, JSNative native, SimdTypeDescr::
     return InliningStatus_Inlined;
 }
 
+IonBuilder::InliningStatus
+IonBuilder::inlineSimdSwizzle(CallInfo &callInfo, JSNative native, SimdTypeDescr::Type type)
+{
+    InlineTypedObject *templateObj = nullptr;
+    if (!checkInlineSimd(callInfo, native, type, 5, &templateObj))
+        return InliningStatus_NotInlined;
+
+    MDefinition *lanes[4];
+    for (size_t i = 0; i < 4; i++)
+        lanes[i] = callInfo.getArg(1 + i);
+
+    MIRType mirType = SimdTypeDescrToMIRType(type);
+    MSimdGeneralSwizzle *ins = MSimdGeneralSwizzle::New(alloc(), callInfo.getArg(0), lanes, mirType);
+    return boxSimd(callInfo, ins, templateObj);
+}
+
 } // namespace jit
 } // namespace js

js/src/jit/MIR.cpp

@@ -931,6 +931,20 @@ MSimdSwizzle::foldsTo(TempAllocator& alloc)
 }
 
 
+MDefinition *
+MSimdGeneralSwizzle::foldsTo(TempAllocator &alloc)
+{
+    int32_t lanes[4];
+    for (size_t i = 0; i < 4; i++) {
+        if (!lane(i)->isConstant() || lane(i)->type() != MIRType_Int32)
+            return this;
+        lanes[i] = lane(i)->toConstant()->value().toInt32();
+        if (lanes[i] < 0 || lanes[i] >= 4)
+            return this;
+    }
+    return MSimdSwizzle::New(alloc, input(), type(), lanes[0], lanes[1], lanes[2], lanes[3]);
+}
+
 template <typename T>
 static void
 PrintOpcodeOperation(T *mir, FILE *fp)

js/src/jit/MIR.h

@@ -1814,6 +1814,12 @@ class MSimdSwizzle
         return new(alloc) MSimdSwizzle(obj, type, laneX, laneY, laneZ, laneW);
     }
 
+    static MSimdSwizzle *New(TempAllocator &alloc, MDefinition *obj, MIRType type,
+                             uint32_t laneX, uint32_t laneY, uint32_t laneZ, uint32_t laneW)
+    {
+        return new(alloc) MSimdSwizzle(obj, type, laneX, laneY, laneZ, laneW);
+    }
+
     bool congruentTo(const MDefinition *ins) const MOZ_OVERRIDE {
         if (!ins->isSimdSwizzle())
             return false;
@@ -1830,6 +1836,57 @@ class MSimdSwizzle
     ALLOW_CLONE(MSimdSwizzle)
 };
 
+// A "general swizzle" is a swizzle with non-constant lane indices.  This is the
+// one that Ion inlines and it can be folded into a MSimdSwizzle if lane indices
+// are constant. Performance of general swizzle does not really matter, as we
+// expect to always get constant indices.
+class MSimdGeneralSwizzle :
+    public MAryInstruction<5>,
+    public SimdSwizzlePolicy::Data
+{
+  protected:
+    MSimdGeneralSwizzle(MDefinition *vec, MDefinition *lanes[4], MIRType type)
+    {
+        MOZ_ASSERT(IsSimdType(type));
+        MOZ_ASSERT(SimdTypeToLength(type) == 4);
+
+        initOperand(0, vec);
+        for (unsigned i = 0; i < 4; i++)
+            initOperand(1 + i, lanes[i]);
+
+        setResultType(type);
+        specialization_ = type;
+        setMovable();
+    }
+
+  public:
+    INSTRUCTION_HEADER(SimdGeneralSwizzle);
+    ALLOW_CLONE(MSimdGeneralSwizzle);
+
+    static MSimdGeneralSwizzle *New(TempAllocator &alloc, MDefinition *vec, MDefinition *lanes[4],
+                                    MIRType type)
+    {
+        return new(alloc) MSimdGeneralSwizzle(vec, lanes, type);
+    }
+
+    MDefinition *input() const {
+        return getOperand(0);
+    }
+    MDefinition *lane(size_t i) const {
+        return getOperand(1 + i);
+    }
+
+    bool congruentTo(const MDefinition *ins) const MOZ_OVERRIDE {
+        return congruentIfOperandsEqual(ins);
+    }
+
+    MDefinition *foldsTo(TempAllocator &alloc) MOZ_OVERRIDE;
+
+    AliasSet getAliasSet() const MOZ_OVERRIDE {
+        return AliasSet::None();
+    }
+};
+
 // Applies a shuffle operation to the inputs, selecting the 2 first lanes of the
 // output from lanes of the first input, and the 2 last lanes of the output from
 // lanes of the second input.

js/src/jit/MOpcodes.h

@@ -24,6 +24,7 @@ namespace jit {
     _(SimdInsertElement)                                                    \
     _(SimdSignMask)                                                         \
     _(SimdSwizzle)                                                          \
+    _(SimdGeneralSwizzle)                                                   \
     _(SimdShuffle)                                                          \
     _(SimdUnaryArith)                                                       \
     _(SimdBinaryComp)                                                       \

js/src/jit/TypePolicy.cpp

@@ -788,6 +788,31 @@ SimdPolicy<Op>::adjustInputs(TempAllocator &alloc, MInstruction *ins)
 template bool
 SimdPolicy<0>::adjustInputs(TempAllocator &alloc, MInstruction *ins);
 
+bool
+SimdSwizzlePolicy::adjustInputs(TempAllocator &alloc, MInstruction *ins)
+{
+    MIRType specialization = ins->typePolicySpecialization();
+
+    // First input is the vector input.
+    if (!MaybeSimdUnbox(alloc, ins, specialization, 0))
+        return false;
+
+    // Next inputs are the lanes, which need to be int32
+    for (unsigned i = 0; i < 4; i++) {
+        MDefinition *in = ins->getOperand(i + 1);
+        if (in->type() == MIRType_Int32)
+            continue;
+
+        MInstruction *replace = MTruncateToInt32::New(alloc, in);
+        ins->block()->insertBefore(ins, replace);
+        ins->replaceOperand(i + 1, replace);
+        if (!replace->typePolicy()->adjustInputs(alloc, replace))
+            return false;
+    }
+
+    return true;
+}
+
 bool
 SimdSelectPolicy::adjustInputs(TempAllocator &alloc, MInstruction *ins)
 {
@@ -1088,6 +1113,7 @@ FilterTypeSetPolicy::adjustInputs(TempAllocator& alloc, MInstruction* ins)
     _(PowPolicy)                                \
     _(SimdAllPolicy)                            \
     _(SimdSelectPolicy)                         \
+    _(SimdSwizzlePolicy)                        \
     _(StoreTypedArrayElementStaticPolicy)       \
     _(StoreTypedArrayHolePolicy)                \
     _(StoreTypedArrayPolicy)                    \

js/src/jit/TypePolicy.h

@@ -310,22 +310,29 @@ class SimdScalarPolicy final : public TypePolicy
 
 class SimdAllPolicy MOZ_FINAL : public TypePolicy
 {
-    public:
+  public:
     SPECIALIZATION_DATA_;
     virtual bool adjustInputs(TempAllocator &alloc, MInstruction *ins) MOZ_OVERRIDE;
 };
 
 template <unsigned Op>
 class SimdPolicy MOZ_FINAL : public TypePolicy
 {
-    public:
+  public:
     SPECIALIZATION_DATA_;
     virtual bool adjustInputs(TempAllocator &alloc, MInstruction *ins) MOZ_OVERRIDE;
 };
 
 class SimdSelectPolicy MOZ_FINAL : public TypePolicy
 {
-    public:
+  public:
+    SPECIALIZATION_DATA_;
+    virtual bool adjustInputs(TempAllocator &alloc, MInstruction *ins) MOZ_OVERRIDE;
+};
+
+class SimdSwizzlePolicy MOZ_FINAL : public TypePolicy
+{
+  public:
     SPECIALIZATION_DATA_;
     virtual bool adjustInputs(TempAllocator &alloc, MInstruction *ins) MOZ_OVERRIDE;
 };

js/src/jit/arm/CodeGenerator-arm.h

@@ -244,6 +244,8 @@ class CodeGeneratorARM : public CodeGeneratorShared
     void visitSimdExtractElementI(LSimdExtractElementI *ins) { MOZ_CRASH("NYI"); }
     void visitSimdExtractElementF(LSimdExtractElementF *ins) { MOZ_CRASH("NYI"); }
     void visitSimdSignMaskX4(LSimdSignMaskX4 *ins) { MOZ_CRASH("NYI"); }
+    void visitSimdGeneralSwizzleI(LSimdGeneralSwizzleI *lir) { MOZ_CRASH("NYI"); }
+    void visitSimdGeneralSwizzleF(LSimdGeneralSwizzleF *lir) { MOZ_CRASH("NYI"); }
     void visitSimdSwizzleI(LSimdSwizzleI *lir) { MOZ_CRASH("NYI"); }
     void visitSimdSwizzleF(LSimdSwizzleF *lir) { MOZ_CRASH("NYI"); }
     void visitSimdBinaryCompIx4(LSimdBinaryCompIx4 *lir) { MOZ_CRASH("NYI"); }

js/src/jit/mips/CodeGenerator-mips.h

@@ -290,6 +290,8 @@ class CodeGeneratorMIPS : public CodeGeneratorShared
     void visitSimdBinaryArithIx4(LSimdBinaryArithIx4 *lir) { MOZ_CRASH("NYI"); }
     void visitSimdBinaryArithFx4(LSimdBinaryArithFx4 *lir) { MOZ_CRASH("NYI"); }
     void visitSimdBinaryBitwiseX4(LSimdBinaryBitwiseX4 *lir) { MOZ_CRASH("NYI"); }
+    void visitSimdGeneralSwizzleI(LSimdGeneralSwizzleI *lir) { MOZ_CRASH("NYI"); }
+    void visitSimdGeneralSwizzleF(LSimdGeneralSwizzleF *lir) { MOZ_CRASH("NYI"); }
 };
 
 typedef CodeGeneratorMIPS CodeGeneratorSpecific;