Bug 1135042: Inline SIMD stores in Ion

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

@@ -0,0 +1,62 @@
+load(libdir + 'simd.js');
+
+setJitCompilerOption("ion.warmup.trigger", 40);
+
+function f() {
+    var f32 = new Float32Array(16);
+    for (var i = 0; i < 16; i++)
+        f32[i] = i + 1;
+
+    var f64 = new Float64Array(f32.buffer);
+    var i32 = new Int32Array(f32.buffer);
+    var u32 = new Uint32Array(f32.buffer);
+    var i16 = new Int16Array(f32.buffer);
+    var u16 = new Uint16Array(f32.buffer);
+    var i8  = new Int8Array(f32.buffer);
+    var u8  = new Uint8Array(f32.buffer);
+
+    var f4 = SIMD.float32x4(42, 43, 44, 45);
+
+    function check() {
+        assertEq(f32[0], 42);
+        assertEq(f32[1], 43);
+        assertEq(f32[2], 44);
+        assertEq(f32[3], 45);
+
+        f32[0] = 1;
+        f32[1] = 2;
+        f32[2] = 3;
+        f32[3] = 4;
+    }
+
+    for (var i = 0; i < 150; i++) {
+        SIMD.float32x4.store(f64, 0, f4);
+        check();
+        SIMD.float32x4.store(f32, 0, f4);
+        check();
+        SIMD.float32x4.store(i32, 0, f4);
+        check();
+        SIMD.float32x4.store(u32, 0, f4);
+        check();
+        SIMD.float32x4.store(i16, 0, f4);
+        check();
+        SIMD.float32x4.store(u16, 0, f4);
+        check();
+        SIMD.float32x4.store(i8, 0, f4);
+        check();
+        SIMD.float32x4.store(u8, 0, f4);
+        check();
+
+        var caught = false;
+        try {
+            SIMD.float32x4.store(i8, (i < 149) ? 0 : (16 << 2) - (4 << 2) + 1, f4);
+            check();
+        } catch (e) {
+            caught = true;
+        }
+        assertEq(i < 149 || caught, true);
+    }
+}
+
+f();
+

js/src/jit/CodeGenerator.cpp

@@ -8663,15 +8663,18 @@ CodeGenerator::visitLoadTypedArrayElementHole(LLoadTypedArrayElementHole* lir)
 
 template <typename T>
 static inline void
-StoreToTypedArray(MacroAssembler &masm, Scalar::Type arrayType, const LAllocation *value, const T &dest)
+StoreToTypedArray(MacroAssembler &masm, Scalar::Type writeType, const LAllocation *value, const T &dest)
 {
-    if (arrayType == Scalar::Float32 || arrayType == Scalar::Float64) {
-        masm.storeToTypedFloatArray(arrayType, ToFloatRegister(value), dest);
+    if (Scalar::isSimdType(writeType) ||
+        writeType == Scalar::Float32 ||
+        writeType == Scalar::Float64)
+    {
+        masm.storeToTypedFloatArray(writeType, ToFloatRegister(value), dest);
     } else {
         if (value->isConstant())
-            masm.storeToTypedIntArray(arrayType, Imm32(ToInt32(value)), dest);
+            masm.storeToTypedIntArray(writeType, Imm32(ToInt32(value)), dest);
         else
-            masm.storeToTypedIntArray(arrayType, ToRegister(value), dest);
+            masm.storeToTypedIntArray(writeType, ToRegister(value), dest);
     }
 }
 
@@ -8681,16 +8684,16 @@ CodeGenerator::visitStoreTypedArrayElement(LStoreTypedArrayElement *lir)
     Register elements = ToRegister(lir->elements());
     const LAllocation *value = lir->value();
 
-    Scalar::Type arrayType = lir->mir()->arrayType();
-    int width = Scalar::byteSize(arrayType);
+    Scalar::Type writeType = lir->mir()->writeType();
+    int width = Scalar::byteSize(lir->mir()->arrayType());
 
     if (lir->index()->isConstant()) {
         Address dest(elements, ToInt32(lir->index()) * width + lir->mir()->offsetAdjustment());
-        StoreToTypedArray(masm, arrayType, value, dest);
+        StoreToTypedArray(masm, writeType, value, dest);
     } else {
         BaseIndex dest(elements, ToRegister(lir->index()), ScaleFromElemWidth(width),
                        lir->mir()->offsetAdjustment());
-        StoreToTypedArray(masm, arrayType, value, dest);
+        StoreToTypedArray(masm, writeType, value, dest);
     }
 }
 

js/src/jit/IonBuilder.h

@@ -838,7 +838,11 @@ class IonBuilder
                                      SimdTypeDescr::Type from, SimdTypeDescr::Type to);
     InliningStatus inlineSimdSelect(CallInfo &callInfo, JSNative native, bool isElementWise,
                                     SimdTypeDescr::Type type);
+
+    bool prepareForSimdLoadStore(CallInfo &callInfo, Scalar::Type simdType, MInstruction **elements,
+                                 MDefinition **index, Scalar::Type *arrayType);
     InliningStatus inlineSimdLoad(CallInfo &callInfo, JSNative native, SimdTypeDescr::Type type);
+    InliningStatus inlineSimdStore(CallInfo &callInfo, JSNative native, SimdTypeDescr::Type type);
 
     // Utility intrinsics.
     InliningStatus inlineIsCallable(CallInfo &callInfo);

js/src/jit/IonTypes.h

@@ -553,6 +553,32 @@ SimdTypeToLength(MIRType type)
     return 1 << ((type >> VECTOR_SCALE_SHIFT) & VECTOR_SCALE_MASK);
 }
 
+static inline MIRType
+ScalarTypeToMIRType(Scalar::Type type)
+{
+    switch (type) {
+      case Scalar::Int8:
+      case Scalar::Uint8:
+      case Scalar::Int16:
+      case Scalar::Uint16:
+      case Scalar::Int32:
+      case Scalar::Uint32:
+      case Scalar::Uint8Clamped:
+        return MIRType_Int32;
+      case Scalar::Float32:
+        return MIRType_Float32;
+      case Scalar::Float64:
+        return MIRType_Double;
+      case Scalar::Float32x4:
+        return MIRType_Float32x4;
+      case Scalar::Int32x4:
+        return MIRType_Int32x4;
+      case Scalar::MaxTypedArrayViewType:
+        break;
+    }
+    MOZ_CRASH("unexpected SIMD kind");
+}
+
 static inline unsigned
 ScalarTypeToLength(Scalar::Type type)
 {

js/src/jit/Lowering.cpp

@@ -2963,7 +2963,10 @@ LIRGenerator::visitStoreTypedArrayElement(MStoreTypedArrayElement* ins)
     MOZ_ASSERT(IsValidElementsType(ins->elements(), ins->offsetAdjustment()));
     MOZ_ASSERT(ins->index()->type() == MIRType_Int32);
 
-    if (ins->isFloatArray()) {
+    if (ins->isSimdWrite()) {
+        MOZ_ASSERT_IF(ins->writeType() == Scalar::Float32x4, ins->value()->type() == MIRType_Float32x4);
+        MOZ_ASSERT_IF(ins->writeType() == Scalar::Int32x4, ins->value()->type() == MIRType_Int32x4);
+    } else if (ins->isFloatArray()) {
         MOZ_ASSERT_IF(ins->arrayType() == Scalar::Float32, ins->value()->type() == MIRType_Float32);
         MOZ_ASSERT_IF(ins->arrayType() == Scalar::Float64, ins->value()->type() == MIRType_Double);
     } else {
@@ -2975,7 +2978,7 @@ LIRGenerator::visitStoreTypedArrayElement(MStoreTypedArrayElement* ins)
     LAllocation value;
 
     // For byte arrays, the value has to be in a byte register on x86.
-    if (ins->isByteArray())
+    if (ins->isByteArray() && !ins->isSimdWrite())
         value = useByteOpRegisterOrNonDoubleConstant(ins->value());
     else
         value = useRegisterOrNonDoubleConstant(ins->value());

js/src/jit/MCallOptimize.cpp

@@ -358,6 +358,12 @@ IonBuilder::inlineNativeCall(CallInfo& callInfo, JSFunction* target)
     if (native == js::simd_float32x4_load)
         return inlineSimdLoad(callInfo, native, SimdTypeDescr::TYPE_FLOAT32);
 
+
+    if (native == js::simd_int32x4_store)
+        return inlineSimdStore(callInfo, native, SimdTypeDescr::TYPE_INT32);
+    if (native == js::simd_float32x4_store)
+        return inlineSimdStore(callInfo, native, SimdTypeDescr::TYPE_FLOAT32);
+
     return InliningStatus_NotInlined;
 }
 
@@ -3155,54 +3161,97 @@ SimdTypeToScalarType(SimdTypeDescr::Type type)
     MOZ_CRASH("unexpected simd type");
 }
 
-IonBuilder::InliningStatus
-IonBuilder::inlineSimdLoad(CallInfo &callInfo, JSNative native, SimdTypeDescr::Type type)
+bool
+IonBuilder::prepareForSimdLoadStore(CallInfo &callInfo, Scalar::Type simdType, MInstruction **elements,
+                                    MDefinition **index, Scalar::Type *arrayType)
 {
-    InlineTypedObject *templateObj = nullptr;
-    if (!checkInlineSimd(callInfo, native, type, 2, &templateObj))
-        return InliningStatus_NotInlined;
-
     MDefinition *array = callInfo.getArg(0);
-    MDefinition *index = callInfo.getArg(1);
+    *index = callInfo.getArg(1);
 
-    Scalar::Type arrayType;
-    if (!ElementAccessIsAnyTypedArray(constraints(), array, index, &arrayType))
-        return InliningStatus_NotInlined;
+    if (!ElementAccessIsAnyTypedArray(constraints(), array, *index, arrayType))
+        return false;
 
-    MInstruction *indexAsInt32 = MToInt32::New(alloc(), index);
+    MInstruction *indexAsInt32 = MToInt32::New(alloc(), *index);
     current->add(indexAsInt32);
-    index = indexAsInt32;
+    *index = indexAsInt32;
 
-    MDefinition *indexForBoundsCheck = index;
+    MDefinition *indexForBoundsCheck = *index;
 
     // Artificially make sure the index is in bounds by adding the difference
     // number of slots needed (e.g. reading from Float32Array we need to make
     // sure to be in bounds for 4 slots, so add 3, etc.).
-    MOZ_ASSERT(Simd128DataSize % Scalar::byteSize(arrayType) == 0);
-    int32_t suppSlotsNeeded = Simd128DataSize / Scalar::byteSize(arrayType) - 1;
+    MOZ_ASSERT(Scalar::byteSize(simdType) % Scalar::byteSize(*arrayType) == 0);
+    int32_t suppSlotsNeeded = Scalar::byteSize(simdType) / Scalar::byteSize(*arrayType) - 1;
     if (suppSlotsNeeded) {
         MConstant *suppSlots = constant(Int32Value(suppSlotsNeeded));
-        MAdd *addedIndex = MAdd::New(alloc(), index, suppSlots);
-        // Even if this addition overflows, we're fine because the code generated
-        // for the bounds check uses uint32 arithmetic
+        MAdd *addedIndex = MAdd::New(alloc(), *index, suppSlots);
+        // We're fine even with the add overflows, as long as the generated code
+        // for the bounds check uses an unsigned comparison.
         addedIndex->setInt32();
         current->add(addedIndex);
         indexForBoundsCheck = addedIndex;
     }
 
     MInstruction *length;
-    MInstruction *elements;
-    addTypedArrayLengthAndData(array, SkipBoundsCheck, &index, &length, &elements);
+    addTypedArrayLengthAndData(array, SkipBoundsCheck, index, &length, elements);
 
     MInstruction *check = MBoundsCheck::New(alloc(), indexForBoundsCheck, length);
     current->add(check);
+    return true;
+}
+
+IonBuilder::InliningStatus
+IonBuilder::inlineSimdLoad(CallInfo &callInfo, JSNative native, SimdTypeDescr::Type type)
+{
+    InlineTypedObject *templateObj = nullptr;
+    if (!checkInlineSimd(callInfo, native, type, 2, &templateObj))
+        return InliningStatus_NotInlined;
+
+    Scalar::Type simdType = SimdTypeToScalarType(type);
+
+    MDefinition *index = nullptr;
+    MInstruction *elements = nullptr;
+    Scalar::Type arrayType;
+    if (!prepareForSimdLoadStore(callInfo, simdType, &elements, &index, &arrayType))
+        return InliningStatus_NotInlined;
 
     MLoadTypedArrayElement *load = MLoadTypedArrayElement::New(alloc(), elements, index, arrayType);
     load->setResultType(SimdTypeDescrToMIRType(type));
-    load->setReadType(SimdTypeToScalarType(type));
+    load->setReadType(simdType);
 
     return boxSimd(callInfo, load, templateObj);
 }
 
+IonBuilder::InliningStatus
+IonBuilder::inlineSimdStore(CallInfo &callInfo, JSNative native, SimdTypeDescr::Type type)
+{
+    InlineTypedObject *templateObj = nullptr;
+    if (!checkInlineSimd(callInfo, native, type, 3, &templateObj))
+        return InliningStatus_NotInlined;
+
+    Scalar::Type simdType = SimdTypeToScalarType(type);
+
+    MDefinition *index = nullptr;
+    MInstruction *elements = nullptr;
+    Scalar::Type arrayType;
+    if (!prepareForSimdLoadStore(callInfo, simdType, &elements, &index, &arrayType))
+        return InliningStatus_NotInlined;
+
+    MDefinition *valueToWrite = callInfo.getArg(2);
+    MStoreTypedArrayElement *store = MStoreTypedArrayElement::New(alloc(), elements, index,
+                                                                  valueToWrite, arrayType);
+    store->setWriteType(simdType);
+
+    current->add(store);
+    current->push(valueToWrite);
+
+    callInfo.setImplicitlyUsedUnchecked();
+
+    if (!resumeAfter(store))
+        return InliningStatus_Error;
+
+    return InliningStatus_Inlined;
+}
+
 } // namespace jit
 } // namespace js

js/src/jit/MIR.h

@@ -9047,6 +9047,7 @@ class MStoreTypedArrayElement
     public StoreTypedArrayPolicy::Data
 {
     Scalar::Type arrayType_;
+    Scalar::Type writeType_;
     bool requiresBarrier_;
     int32_t offsetAdjustment_;
 
@@ -9058,6 +9059,7 @@ class MStoreTypedArrayElement
                             int32_t offsetAdjustment)
       : MTernaryInstruction(elements, index, value),
         arrayType_(arrayType),
+        writeType_(arrayType),
         requiresBarrier_(requiresBarrier == DoesRequireMemoryBarrier),
         offsetAdjustment_(offsetAdjustment),
         racy_(false)
@@ -9083,9 +9085,18 @@ class MStoreTypedArrayElement
                                                   requiresBarrier, offsetAdjustment);
     }
 
+    void setWriteType(Scalar::Type type) {
+        writeType_ = type;
+    }
+    Scalar::Type writeType() const {
+        return writeType_;
+    }
     Scalar::Type arrayType() const {
         return arrayType_;
     }
+    bool isSimdWrite() const {
+        return Scalar::isSimdType(writeType());
+    }
     bool isByteArray() const {
         return arrayType_ == Scalar::Int8 ||
                arrayType_ == Scalar::Uint8 ||

js/src/jit/MacroAssembler.cpp

@@ -293,6 +293,12 @@ StoreToTypedFloatArray(MacroAssembler& masm, int arrayType, const S& value, cons
 #endif
         masm.storeDouble(value, dest);
         break;
+      case Scalar::Float32x4:
+        masm.storeUnalignedFloat32x4(value, dest);
+        break;
+      case Scalar::Int32x4:
+        masm.storeUnalignedInt32x4(value, dest);
+        break;
       default:
         MOZ_CRASH("Invalid typed array type");
     }

js/src/jit/RangeAnalysis.cpp

@@ -2637,7 +2637,7 @@ MDefinition::TruncateKind
 MStoreTypedArrayElement::operandTruncateKind(size_t index) const
 {
     // An integer store truncates the stored value.
-    return index == 2 && !isFloatArray() ? Truncate : NoTruncate;
+    return index == 2 && !isFloatArray() && !isSimdWrite() ? Truncate : NoTruncate;
 }
 
 MDefinition::TruncateKind

js/src/jit/TypePolicy.cpp

@@ -879,9 +879,13 @@ InstanceOfPolicy::adjustInputs(TempAllocator& alloc, MInstruction* def)
 }
 
 bool
-StoreTypedArrayPolicy::adjustValueInput(TempAllocator &alloc, MInstruction *ins, int arrayType,
-                                        MDefinition *value, int valueOperand)
+StoreTypedArrayPolicy::adjustValueInput(TempAllocator &alloc, MInstruction *ins,
+                                        Scalar::Type writeType, MDefinition *value, int valueOperand)
 {
+    // Storing a SIMD value just implies that we might need a SimdUnbox.
+    if (Scalar::isSimdType(writeType))
+        return MaybeSimdUnbox(alloc, ins, ScalarTypeToMIRType(writeType), valueOperand);
+
     MDefinition *curValue = value;
     // First, ensure the value is int32, boolean, double or Value.
     // The conversion is based on TypedArrayObjectTemplate::setElementTail.
@@ -922,7 +926,7 @@ StoreTypedArrayPolicy::adjustValueInput(TempAllocator &alloc, MInstruction *ins,
                value->type() == MIRType_Float32 ||
                value->type() == MIRType_Value);
 
-    switch (arrayType) {
+    switch (writeType) {
       case Scalar::Int8:
       case Scalar::Uint8:
       case Scalar::Int16:
@@ -969,7 +973,7 @@ StoreTypedArrayPolicy::adjustInputs(TempAllocator& alloc, MInstruction* ins)
     MOZ_ASSERT(IsValidElementsType(store->elements(), store->offsetAdjustment()));
     MOZ_ASSERT(store->index()->type() == MIRType_Int32);
 
-    return adjustValueInput(alloc, ins, store->arrayType(), store->value(), 2);
+    return adjustValueInput(alloc, store, store->writeType(), store->value(), 2);
 }
 
 bool