How to convert the CUDA PTX codes in faiss library to the same effection codes?

[dwwcqu]

When I working on convert faiss library from CUDA-gpu backend to ROCm-gpu backend, I got the obstacles about how to convert the CUDA PTX codes in faiss library to the same effect code. Because ROCm-gpu does not support CUDA-PTX code. Here lists a part codes I have converted:

// faiss/gpu/impl/PQCodeLoad.h
template <>
struct LoadCode32<2> {
    static inline __device__ void load(
            unsigned int code32[1],
            uint8_t* p,
            int offset) {
        p += offset * 2;
        #ifdef __HIP_PLATFORM_NVIDIA__
            asm("ld.global.cs.u16 {%0}, [%1];" : "=r"(code32[0]) : "l"(p));
        #else
            uint16_t *temp_p = (uint16_t *)p;
            code32[0] = (uint32_t)*temp_p;
        #endif
    }
};

template <>
struct LoadCode32<3> {
    static inline __device__ void load(
            unsigned int code32[1],
            uint8_t* p,
            int offset) {
        p += offset * 3;
        unsigned int a;
        unsigned int b;
        unsigned int c;
        #ifdef __HIP_PLATFORM_NVIDIA__
            asm("ld.global.cs.u8 {%0}, [%1 + 0];" : "=r"(a) : "l"(p));
            asm("ld.global.cs.u8 {%0}, [%1 + 1];" : "=r"(b) : "l"(p));
            asm("ld.global.cs.u8 {%0}, [%1 + 2];" : "=r"(c) : "l"(p));
        #else
        // FIXME: this is a non-coalesced, unaligned, non-vectorized load
        // unfortunately need to reorganize memory layout by warp
            uint8_t *temp_p = p;
            a = (unsigned int)*temp_p;   ++temp_p;
            b = (unsigned int)*temp_p;   ++temp_p;
            c = (unsigned int)*temp_p;
        #endif
        // FIXME: this is also slow, since we have to recover the
        // individual bytes loaded
        code32[0] = (c << 16) | (b << 8) | a;
    }
};

template <>
struct LoadCode32<4> {
    static inline __device__ void load(
            unsigned int code32[1],
            uint8_t* p,
            int offset) {
        p += offset * 4;
        #ifdef __HIP_PLATFORM_NVIDIA__
            asm("ld.global.cs.u32 {%0}, [%1];" : "=r"(code32[0]) : "l"(p));
        #else
            uint32_t *temp_p = (uint32_t *)p;
            code32[0] = (unsigned int)*temp_p;
        #endif
    }
};

template <>
struct LoadCode32<8> {
    static inline __device__ void load(
            unsigned int code32[2],
            uint8_t* p,
            int offset) {
        p += offset * 8;
        #ifdef __HIP_PLATFORM_NVIDIA__
            asm("ld.global.cs.v2.u32 {%0, %1}, [%2];"
                : "=r"(code32[0]), "=r"(code32[1])
                : "l"(p));
        #else
        uint32_t *temp_p = (uint32_t *)p;
        for(int sz = 0; sz < 2; ++sz){
            code32[sz] = *temp_p;
            ++temp_p;
        }
        #endif
    }
};

// faiss/gpu/utils/LoadStoreOperators.h
struct Half4 {
    half2 a;
    half2 b;
};
template <>
struct LoadStore<Half4> {
    static inline __device__ Half4 load(void* p) {
        Half4 out;
#ifdef __HIP_PLATFORM_NVIDIA__
    #if CUDA_VERSION >= 9000
            asm("ld.global.v2.u32 {%0, %1}, [%2];"
                : "=r"(__HALF2_TO_UI(out.a)), "=r"(__HALF2_TO_UI(out.b))
                : "l"(p));
    #else
            asm("ld.global.v2.u32 {%0, %1}, [%2];"
                : "=r"(out.a.x), "=r"(out.b.x)
                : "l"(p));
    #endif
#else
    uint16_t *ptr = (uint16_t *)p;
    __half2 temp_a;
    __half2 temp_b;
    temp_a.x = *ptr; ++ptr;
    temp_a.y = *ptr; ++ptr;
    temp_b.x = *ptr; ++ptr;
    temp_b.y = *ptr;

    out.a = temp_a;
    out.b = temp_b;
#endif
    return out;
}

// faiss/gpu/utils/PtxUtils.h
#ifdef __HIP_PLATFORM_NVIDIA__
__device__ __forceinline__ unsigned int getBitfield(
        unsigned int val,
        int pos,
        int len) {
    unsigned int ret;
    asm("bfe.u32 %0, %1, %2, %3;" : "=r"(ret) : "r"(val), "r"(pos), "r"(len));
    return ret;
}

__device__ __forceinline__ uint64_t
getBitfield(uint64_t val, int pos, int len) {
    uint64_t ret;
    asm("bfe.u64 %0, %1, %2, %3;" : "=l"(ret) : "l"(val), "r"(pos), "r"(len));
    return ret;
}

__device__ __forceinline__ unsigned int setBitfield(
        unsigned int val,
        unsigned int toInsert,
        int pos,
        int len) {
    unsigned int ret;
    asm("bfi.b32 %0, %1, %2, %3, %4;"
        : "=r"(ret)
        : "r"(toInsert), "r"(val), "r"(pos), "r"(len));
    return ret;
}
#else
__device__ __forceinline__ unsigned int getBitfield(
        unsigned int val,
        int pos,
        int len) {
    unsigned int ret;
    ret = (val >> pos) & ((1 << len) - 1);
    return ret;
}

__device__ __forceinline__ uint64_t
getBitfield(uint64_t val, int pos, int len) {
    uint64_t ret;
    ret = (val >> pos) & ((1 << len) - 1);
    return ret;
}

__device__ __forceinline__ unsigned int setBitfield(
        unsigned int val,
        unsigned int toInsert,
        int pos,
        int len) {
    unsigned int ret;
    unsigned int mask = ((1 << len) - 1) << pos;
    ret = (val | mask) & (toInsert | (~mask));
    return ret;
}
#endif

As the code show above, the PTX codes locate in #ifdef __HIP_PLATFORM_NVIDIA__ scope. And the corresponding same effection codes locate in #else scope.

Because I don't familiar with the function of PTX code in faiss library. If you have any suggestions about the codes above. Thanks for your help!

[matrixji]

IMO, It's not an easy thing for porting the faiss to ROCm platform, I've done some related work. For porting faiss to ROCm, for now, I see blow works(maybe more, hope below tips could helps U):

• First for your question, about the PTX code, you have mentioned above, the code in PQCodeLoad.h, you could look at the __ldcs functions in ROCm. Actually, CUDA DO have Load/Store functions for implementing ld/st PTX code, I guess the reason is Faiss is written for years, and some code is reused from fbcuda
• Second part is the most difficult part, CUDA uses 32 warp size, ROCm uses 64 wavefronts, and there are a lot of codes in faiss that are hard-coded with 32 parallels, you need to port it to 64. I see the K-selection and some implementations in IVF are designed for CUDA's 32 warp size.
• Still faiss's implementation use cublas as BLAS, I guess the ROCm has similar implementations, you need to check it if could be porting to ROCm smoothly.
• For the latest code on master, the faiss could use rapidsai/raft, it depends on many Nvidia opensource projects like rmm, cutlass, I guess it is impossible for ROCm right now if you're trying porting the latest master code, you may need to disable raft.

[dwwcqu]

Thanks for your suggestions! Now, I concentrating on the version 1.7.3 of faiss library and don't find the dependent on rapidsai/raft. The problems I facing are the CUDA-PTX code. I don't sure the correction of my convertion. Following your suggestions, I will check the corrections again. If you have any others suggestions about the the converted codes above, thanks for your helps.