Bug 1240796 - Implement Uint32x4 extractLane in Ion. r=nbp

Since Uint32 can't be represented in a MIRType_Int32, this function should
return a MIRType_Double.

Allow MSimdExtractElement(Uint32x4) to return a MIRType_Int32 too. It will work
like the double version followed by MTruncateToInt32 which bitcasts the Uint32
value range into the Int32 value range.

js/src/jit-test/tests/SIMD/getters.js

@@ -4,6 +4,7 @@ setJitCompilerOption("ion.warmup.trigger", 50);
 
 function f() {
     var i4 = SIMD.Int32x4(1, -2, 3, -4);
+    var u4 = SIMD.Uint32x4(1, -2, 3, 0x88000000);
     var b4 = SIMD.Bool32x4(true, true, false, true);
 
 
@@ -19,6 +20,11 @@ function f() {
         assertEq(SIMD.Int32x4.extractLane(i4, 2), 3);
         assertEq(SIMD.Int32x4.extractLane(i4, 3), -4);
 
+        assertEq(SIMD.Uint32x4.extractLane(u4, 0), 1);
+        assertEq(SIMD.Uint32x4.extractLane(u4, 1), -2 >>> 0);
+        assertEq(SIMD.Uint32x4.extractLane(u4, 2), 3);
+        assertEq(SIMD.Uint32x4.extractLane(u4, 3), 0x88000000);
+
         assertEq(SIMD.Float32x4.extractLane(f4, 0), v);
         assertEq(SIMD.Float32x4.extractLane(f4, 1), NaN);
         assertEq(SIMD.Float32x4.extractLane(f4, 2), Infinity);

js/src/jit/Lowering.cpp

@@ -4149,18 +4149,27 @@ LIRGenerator::visitSimdExtractElement(MSimdExtractElement* ins)
 
     switch (ins->input()->type()) {
       case MIRType_Int32x4: {
+        MOZ_ASSERT(ins->signedness() != SimdSign::NotApplicable);
         // Note: there could be int16x8 in the future, which doesn't use the
         // same instruction. We either need to pass the arity or create new LIns.
         LUse use = useRegisterAtStart(ins->input());
-        define(new(alloc()) LSimdExtractElementI(use), ins);
+        if (ins->type() == MIRType_Double) {
+            // Extract an Uint32 lane into a double.
+            MOZ_ASSERT(ins->signedness() == SimdSign::Unsigned);
+            define(new (alloc()) LSimdExtractElementU2D(use, temp()), ins);
+        } else {
+            define(new (alloc()) LSimdExtractElementI(use), ins);
+        }
         break;
       }
       case MIRType_Float32x4: {
+        MOZ_ASSERT(ins->signedness() == SimdSign::NotApplicable);
         LUse use = useRegisterAtStart(ins->input());
         define(new(alloc()) LSimdExtractElementF(use), ins);
         break;
       }
       case MIRType_Bool32x4: {
+        MOZ_ASSERT(ins->signedness() == SimdSign::NotApplicable);
         LUse use = useRegisterAtStart(ins->input());
         define(new(alloc()) LSimdExtractElementB(use), ins);
         break;

js/src/jit/MCallOptimize.cpp

@@ -3468,10 +3468,6 @@ IonBuilder::inlineSimdExtractLane(CallInfo& callInfo, JSNative native,
         return InliningStatus_NotInlined;
     }
 
-    // TODO JSO: Implement unsigned integer lane values.
-    if (sign == SimdSign::Unsigned)
-        return InliningStatus_NotInlined;
-
     MDefinition* arg = callInfo.getArg(1);
     if (!arg->isConstantValue() || arg->type() != MIRType_Int32)
         return InliningStatus_NotInlined;
@@ -3481,8 +3477,13 @@ IonBuilder::inlineSimdExtractLane(CallInfo& callInfo, JSNative native,
 
     // See comment in inlineSimdBinary
     MIRType laneType = SimdTypeToLaneType(vecType);
+
+    // An Uint32 lane can't be represented in MIRType_Int32. Get it as a double.
+    if (sign == SimdSign::Unsigned && vecType == MIRType_Int32x4)
+        laneType = MIRType_Double;
+
     MSimdExtractElement* ins = MSimdExtractElement::New(alloc(), callInfo.getArg(0),
-                                                        vecType, laneType, SimdLane(lane));
+                                                        vecType, laneType, SimdLane(lane), sign);
     current->add(ins);
     current->push(ins);
     callInfo.setImplicitlyUsedUnchecked();

js/src/jit/MIR.h

@@ -1642,23 +1642,38 @@ class MSimdReinterpretCast
 };
 
 // Extracts a lane element from a given vector type, given by its lane symbol.
+//
+// For integer SIMD types, a SimdSign must be provided so the lane value can be
+// converted to a scalar correctly.
 class MSimdExtractElement
   : public MUnaryInstruction,
     public SimdPolicy<0>::Data
 {
   protected:
     SimdLane lane_;
+    SimdSign sign_;
 
-    MSimdExtractElement(MDefinition* obj, MIRType vecType, MIRType laneType, SimdLane lane)
-      : MUnaryInstruction(obj), lane_(lane)
+    MSimdExtractElement(MDefinition* obj, MIRType vecType, MIRType laneType, SimdLane lane,
+                        SimdSign sign)
+      : MUnaryInstruction(obj), lane_(lane), sign_(sign)
     {
         MOZ_ASSERT(IsSimdType(vecType));
         MOZ_ASSERT(uint32_t(lane) < SimdTypeToLength(vecType));
         MOZ_ASSERT(!IsSimdType(laneType));
+        MOZ_ASSERT((sign != SimdSign::NotApplicable) == IsIntegerSimdType(vecType),
+                   "Signedness must be specified for integer SIMD extractLanes");
         // The resulting type should match the lane type.
         // Allow extracting boolean lanes directly into an Int32 (for asm.js).
+        // Allow extracting Uint32 lanes into a double.
+        //
+        // We also allow extracting Uint32 lanes into a MIRType_Int32. This is
+        // equivalent to extracting the Uint32 lane to a double and then
+        // applying MTruncateToInt32, but it bypasses the conversion to/from
+        // double.
         MOZ_ASSERT(SimdTypeToLaneType(vecType) == laneType ||
-                   (IsBooleanSimdType(vecType) && laneType == MIRType_Int32));
+                   (IsBooleanSimdType(vecType) && laneType == MIRType_Int32) ||
+                   (vecType == MIRType_Int32x4 && laneType == MIRType_Double &&
+                    sign == SimdSign::Unsigned));
 
         setMovable();
         specialization_ = vecType;
@@ -1671,19 +1686,26 @@ class MSimdExtractElement
     static MSimdExtractElement* NewAsmJS(TempAllocator& alloc, MDefinition* obj, MIRType type,
                                          SimdLane lane)
     {
-        return new(alloc) MSimdExtractElement(obj, obj->type(), type, lane);
+        // Only signed integer types in AsmJS so far.
+        SimdSign sign =
+          IsIntegerSimdType(obj->type()) ? SimdSign::Signed : SimdSign::NotApplicable;
+        return new (alloc) MSimdExtractElement(obj, obj->type(), type, lane, sign);
     }
 
     static MSimdExtractElement* New(TempAllocator& alloc, MDefinition* obj, MIRType vecType,
-                                    MIRType scalarType, SimdLane lane)
+                                    MIRType scalarType, SimdLane lane, SimdSign sign)
     {
-        return new(alloc) MSimdExtractElement(obj, vecType, scalarType, lane);
+        return new(alloc) MSimdExtractElement(obj, vecType, scalarType, lane, sign);
     }
 
     SimdLane lane() const {
         return lane_;
     }
 
+    SimdSign signedness() const {
+        return sign_;
+    }
+
     AliasSet getAliasSet() const override {
         return AliasSet::None();
     }
@@ -1691,7 +1713,7 @@ class MSimdExtractElement
         if (!ins->isSimdExtractElement())
             return false;
         const MSimdExtractElement* other = ins->toSimdExtractElement();
-        if (other->lane_ != lane_)
+        if (other->lane_ != lane_ || other->sign_ != sign_)
             return false;
         return congruentIfOperandsEqual(other);
     }

js/src/jit/shared/LIR-shared.h

@@ -280,6 +280,24 @@ class LSimdExtractElementF : public LSimdExtractElementBase
     {}
 };
 
+// Extracts an element from an Uint32x4 SIMD vector, converts to double.
+class LSimdExtractElementU2D : public LInstructionHelper<1, 1, 1>
+{
+  public:
+    LIR_HEADER(SimdExtractElementU2D);
+    explicit LSimdExtractElementU2D(const LAllocation& base, const LDefinition& temp) {
+        setOperand(0, base);
+        setTemp(0, temp);
+    }
+    SimdLane lane() const {
+        return mir_->toSimdExtractElement()->lane();
+    }
+    const LDefinition* temp() {
+        return getTemp(0);
+    }
+};
+
+
 class LSimdInsertElementBase : public LInstructionHelper<1, 2, 0>
 {
   protected:

js/src/jit/shared/LOpcodes-shared.h

@@ -26,6 +26,7 @@
     _(SimdAnyTrue)                  \
     _(SimdReinterpretCast)          \
     _(SimdExtractElementI)          \
+    _(SimdExtractElementU2D)        \
     _(SimdExtractElementB)          \
     _(SimdExtractElementF)          \
     _(SimdInsertElementI)           \

js/src/jit/x86-shared/CodeGenerator-x86-shared.cpp

@@ -2484,13 +2484,10 @@ CodeGeneratorX86Shared::visitSimdReinterpretCast(LSimdReinterpretCast* ins)
     }
 }
 
+// Extract an integer lane from the vector register |input| and place it in |output|.
 void
-CodeGeneratorX86Shared::visitSimdExtractElementB(LSimdExtractElementB* ins)
+CodeGeneratorX86Shared::emitSimdExtractLane(FloatRegister input, Register output, unsigned lane)
 {
-    FloatRegister input = ToFloatRegister(ins->input());
-    Register output = ToRegister(ins->output());
-
-    SimdLane lane = ins->lane();
     if (lane == LaneX) {
         // The value we want to extract is in the low double-word
         masm.moveLowInt32(input, output);
@@ -2501,6 +2498,15 @@ CodeGeneratorX86Shared::visitSimdExtractElementB(LSimdExtractElementB* ins)
         masm.shuffleInt32(mask, input, ScratchSimd128Reg);
         masm.moveLowInt32(ScratchSimd128Reg, output);
     }
+}
+
+void
+CodeGeneratorX86Shared::visitSimdExtractElementB(LSimdExtractElementB* ins)
+{
+    FloatRegister input = ToFloatRegister(ins->input());
+    Register output = ToRegister(ins->output());
+
+    emitSimdExtractLane(input, output, ins->lane());
 
     // We need to generate a 0/1 value. We have 0/-1.
     masm.and32(Imm32(1), output);
@@ -2512,18 +2518,18 @@ CodeGeneratorX86Shared::visitSimdExtractElementI(LSimdExtractElementI* ins)
     FloatRegister input = ToFloatRegister(ins->input());
     Register output = ToRegister(ins->output());
 
-    SimdLane lane = ins->lane();
-    if (lane == LaneX) {
-        // The value we want to extract is in the low double-word
-        masm.moveLowInt32(input, output);
-    } else if (AssemblerX86Shared::HasSSE41()) {
-        masm.vpextrd(lane, input, output);
-    } else {
-        uint32_t mask = MacroAssembler::ComputeShuffleMask(lane);
-        ScratchSimd128Scope scratch(masm);
-        masm.shuffleInt32(mask, input, scratch);
-        masm.moveLowInt32(scratch, output);
-    }
+    emitSimdExtractLane(input, output, ins->lane());
+}
+
+void
+CodeGeneratorX86Shared::visitSimdExtractElementU2D(LSimdExtractElementU2D* ins)
+{
+    FloatRegister input = ToFloatRegister(ins->input());
+    FloatRegister output = ToFloatRegister(ins->output());
+    Register temp = ToRegister(ins->temp());
+
+    emitSimdExtractLane(input, temp, ins->lane());
+    masm.convertUInt32ToDouble(temp, output);
 }
 
 void

js/src/jit/x86-shared/CodeGenerator-x86-shared.h

@@ -185,6 +185,8 @@ class CodeGeneratorX86Shared : public CodeGeneratorShared
 
     void emitTableSwitchDispatch(MTableSwitch* mir, Register index, Register base);
 
+    void emitSimdExtractLane(FloatRegister input, Register output, unsigned lane);
+
   public:
     CodeGeneratorX86Shared(MIRGenerator* gen, LIRGraph* graph, MacroAssembler* masm);
 
@@ -265,6 +267,7 @@ class CodeGeneratorX86Shared : public CodeGeneratorShared
     void visitSimdReinterpretCast(LSimdReinterpretCast* lir);
     void visitSimdExtractElementB(LSimdExtractElementB* lir);
     void visitSimdExtractElementI(LSimdExtractElementI* lir);
+    void visitSimdExtractElementU2D(LSimdExtractElementU2D* lir);
     void visitSimdExtractElementF(LSimdExtractElementF* lir);
     void visitSimdInsertElementI(LSimdInsertElementI* lir);
     void visitSimdInsertElementF(LSimdInsertElementF* lir);