yet another gelu

[ngc92]

more complicated Packet128 for cleaner kernels

[ngc92]

This is how it would look like if we moved all the casting into the load/store functions:

template<class ElementType>
struct alignas(16) Packed128 {
    __device__ ElementType& operator[](int index) {
        return payload[index];
    }
    __device__ const ElementType& operator[](int index) const {
        return payload[index];
    }
    __device__ float fp32(int index) {
        return static_cast<float>(payload[index]);
    }
    static constexpr const size_t size = sizeof(int4) / sizeof(ElementType);

    ElementType payload[size];
};

// use this function to load a Packet128 from an aligned memory address
template<class ElementType, ELoadMode Mode=ELoadMode::CA>
__device__ __forceinline__ Packed128<ElementType> load_aligned(const ElementType* address, load_mode_t<Mode> mode = {}) {
    int4 bits = generic_load(reinterpret_cast<const int4*>(address), mode);
    Packed128<ElementType> result;
    static_assert(sizeof(bits) == sizeof(result), "Size mismatch.");
    memcpy(&result, &bits, sizeof(bits));
    return result;
}

// use this function to store a Packet128 to an aligned memory address
template<class ElementType, EStoreMode Mode=EStoreMode::WB>
__device__ void store_aligned(ElementType* target, Packed128<ElementType> value, store_mode_t<Mode> mode = {}) {
    int4 bits;
    static_assert(sizeof(bits) == sizeof(value), "Size mismatch.");
    memcpy(&bits, &value, sizeof(bits));
    generic_store(reinterpret_cast<int4*>(target), bits, mode);
}

[karpathy]

any reason we loadcs but store without cs?

[ngc92]

for this kernel, we read exactly once and store once, so these hints don't gain us anything. But by not keeping the input in cache, but the output there, maybe the next kernel can be a bit faster. This is just guesswork though, I haven't actually measured this.

[karpathy]

got it! that makes sense actually

[ngc92]

close by #298