TheRock/rocm-libraries/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/loads.cpp

// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
/*
 * C99 port of rocke/helpers/loads.py: CoalescedTileLoader, AsyncTileLoader
 * (and the AsyncTileLoaderSlot that AsyncTileLoader.bind returns).
 *
 * See the header for the original Python and the contract. The builder-emitting
 * methods (load / bind / issue) reproduce the const_i32 / arith / load / store
 * call sequence byte-faithfully so the downstream IR is identical to the Python.
 */
#include "rocke/helper_rocke.helpers.loads.h"

#include "rocke/ir_internal.h" /* rocke_i_set_err */

/* ============================================================================
 * CoalescedTileLoader
 * ========================================================================== */

rocke_status_t rocke_coalesced_tile_loader_choose_vec(
    int tile_rows, int tile_cols, int block_size, int max_vec, int* out_vec)
{
    int v;

    /* Python:
     *   v = max_vec
     *   while v >= 1:
     *       if (tile_cols % v == 0
     *           and (tile_rows*tile_cols)//v >= block_size
     *           and ((tile_rows*tile_cols)//v) % block_size == 0):
     *           return v
     *       v //= 2
     *   raise ValueError(...)
     */
    v = max_vec;
    while(v >= 1)
    {
        if(tile_cols % v == 0 && (tile_rows * tile_cols) / v >= block_size
           && ((tile_rows * tile_cols) / v) % block_size == 0)
        {
            if(out_vec != NULL)
            {
                *out_vec = v;
            }
            return ROCKE_OK;
        }
        v /= 2;
    }
    return ROCKE_ERR_VALUE;
}

rocke_status_t rocke_coalesced_tile_loader_from_tile(int tile_rows,
                                                     int tile_cols,
                                                     int block_size,
                                                     int max_vec,
                                                     bool use_buffer_rsrc,
                                                     rocke_coalesced_tile_loader_t* out)
{
    int vec;
    rocke_status_t st;

    /* Python:
     *   vec = cls.choose_vec(...)
     *   return cls(tile_rows=..., tile_cols=..., block_size=...,
     *              load_vec=vec, use_buffer_rsrc=use_buffer_rsrc)
     */
    st = rocke_coalesced_tile_loader_choose_vec(tile_rows, tile_cols, block_size, max_vec, &vec);
    if(st != ROCKE_OK)
    {
        return st;
    }
    if(out != NULL)
    {
        out->tile_rows = tile_rows;
        out->tile_cols = tile_cols;
        out->block_size = block_size;
        out->load_vec = vec;
        out->use_buffer_rsrc = use_buffer_rsrc;
        /* Python (1 << 31) - 1 == 2147483647 (arbitrary-precision ints);
         * spell it as the literal to avoid the C int shift-into-sign overflow. */
        out->oob_sentinel = 2147483647; /* dataclass default */
        out->has_inner_dim = false; /* inner_dim default None */
        out->inner_dim = 0;
    }
    return ROCKE_OK;
}

rocke_status_t
    rocke_coalesced_tile_loader_vecs_per_thread(const rocke_coalesced_tile_loader_t* self, int* out)
{
    int total_vecs;

    if(self == NULL)
    {
        return ROCKE_ERR_VALUE;
    }
    /* Python:
     *   total_vecs = (tile_rows * tile_cols) // load_vec
     *   if total_vecs % block_size: raise ValueError(...)
     *   return total_vecs // block_size
     */
    total_vecs = (self->tile_rows * self->tile_cols) / self->load_vec;
    if(total_vecs % self->block_size)
    {
        return ROCKE_ERR_VALUE;
    }
    if(out != NULL)
    {
        *out = total_vecs / self->block_size;
    }
    return ROCKE_OK;
}

int rocke_coalesced_tile_loader_cols_per_vec(const rocke_coalesced_tile_loader_t* self)
{
    /* Python: return tile_cols // load_vec */
    if(self == NULL)
    {
        return 0;
    }
    return self->tile_cols / self->load_vec;
}

void rocke_coalesced_tile_loader_load(rocke_ir_builder_t* b,
                                      const rocke_coalesced_tile_loader_t* self,
                                      rocke_value_t* tid,
                                      rocke_value_t* smem_dst,
                                      rocke_loads_descriptor_fn descriptor,
                                      void* descriptor_user,
                                      rocke_value_t* rsrc,
                                      rocke_value_t* ptr)
{
    rocke_value_t* c_threads;
    rocke_value_t* c_load_vec;
    rocke_value_t* c_cols_per_vec;
    rocke_value_t* c_half_bytes;
    rocke_value_t* c0;
    rocke_value_t* c_oob;
    int vecs_per_thread;
    int e;

    if(b != NULL && b->status != ROCKE_OK)
    {
        return; /* already in error: no-op */
    }
    if(self == NULL)
    {
        rocke_i_set_err(b, ROCKE_ERR_VALUE, "CoalescedTileLoader.load: NULL loader");
        return;
    }

    /* Python:
     *   if self.use_buffer_rsrc and rsrc is None:
     *       raise ValueError("... use_buffer_rsrc=True requires rsrc")
     *   if not self.use_buffer_rsrc and ptr is None:
     *       raise ValueError("... use_buffer_rsrc=False requires ptr")
     */
    if(self->use_buffer_rsrc && rsrc == NULL)
    {
        rocke_i_set_err(
            b, ROCKE_ERR_VALUE, "CoalescedTileLoader: use_buffer_rsrc=True requires rsrc");
        return;
    }
    if(!self->use_buffer_rsrc && ptr == NULL)
    {
        rocke_i_set_err(
            b, ROCKE_ERR_VALUE, "CoalescedTileLoader: use_buffer_rsrc=False requires ptr");
        return;
    }

    /* Python:
     *   c_threads      = b.const_i32(self.block_size)
     *   c_load_vec     = b.const_i32(self.load_vec)
     *   c_cols_per_vec = b.const_i32(self.cols_per_vec)
     *   c_half_bytes   = b.const_i32(2)
     *   c0             = b.const_i32(0)
     *   c_oob          = b.const_i32(self.oob_sentinel)
     */
    c_threads = rocke_b_const_i32(b, self->block_size);
    c_load_vec = rocke_b_const_i32(b, self->load_vec);
    c_cols_per_vec = rocke_b_const_i32(b, rocke_coalesced_tile_loader_cols_per_vec(self));
    c_half_bytes = rocke_b_const_i32(b, 2);
    c0 = rocke_b_const_i32(b, 0);
    c_oob = rocke_b_const_i32(b, self->oob_sentinel);

    /* self.vecs_per_thread (Python property; raises ValueError on bad divide). */
    if(rocke_coalesced_tile_loader_vecs_per_thread(self, &vecs_per_thread) != ROCKE_OK)
    {
        rocke_i_set_err(b,
                        ROCKE_ERR_VALUE,
                        "CoalescedTileLoader: tile %dx%d / %d not divisible by "
                        "block_size %d",
                        self->tile_rows,
                        self->tile_cols,
                        self->load_vec,
                        self->block_size);
        return;
    }

    for(e = 0; e < vecs_per_thread; ++e)
    {
        rocke_value_t* vec_idx;
        rocke_value_t* row;
        rocke_value_t* col_v;
        rocke_value_t* col;
        rocke_value_t* off_elems;
        rocke_value_t* valid;

        /* Python:
         *   vec_idx = b.add(b.mul(b.const_i32(e), c_threads), tid)
         *   row     = b.div(vec_idx, c_cols_per_vec)
         *   col_v   = b.mod(vec_idx, c_cols_per_vec)
         *   col     = b.mul(col_v, c_load_vec) if self.load_vec > 1 else col_v
         */
        vec_idx = rocke_b_add(b, rocke_b_mul(b, rocke_b_const_i32(b, e), c_threads), tid);
        row = rocke_b_div(b, vec_idx, c_cols_per_vec);
        col_v = rocke_b_mod(b, vec_idx, c_cols_per_vec);
        col = (self->load_vec > 1) ? rocke_b_mul(b, col_v, c_load_vec) : col_v;

        /* Python: off_elems, valid = descriptor(b, row, col) */
        valid = NULL; /* default => "None" if callback leaves it NULL */
        off_elems = descriptor(b, row, col, &valid, descriptor_user);

        if(self->use_buffer_rsrc)
        {
            rocke_value_t* off_bytes;
            rocke_value_t* safe;
            rocke_value_t* indices[2];

            /* Python:
             *   off_bytes = b.mul(off_elems, c_half_bytes)
             *   safe = b.select(valid, off_bytes, c_oob) if valid is not None
             *          else off_bytes
             */
            off_bytes = rocke_b_mul(b, off_elems, c_half_bytes);
            if(valid != NULL)
            {
                safe = rocke_b_select(b, valid, off_bytes, c_oob);
            }
            else
            {
                safe = off_bytes;
            }

            indices[0] = row;
            indices[1] = col;
            if(self->load_vec == 1)
            {
                /* Python:
                 *   v = b.buffer_load_f16(rsrc, safe, c0)
                 *   b.smem_store_f16(smem_dst, [row, col], v)
                 */
                rocke_value_t* v = rocke_b_buffer_load_f16(b, rsrc, safe, c0);
                rocke_b_smem_store_f16(b, smem_dst, indices, 2, v);
            }
            else
            {
                /* Python:
                 *   dwords = self.load_vec // 2
                 *   v = b.buffer_load_vN_f16(rsrc, safe, c0, dwords)
                 *   b.smem_store_vN_f16(smem_dst, [row, col], v, self.load_vec)
                 */
                int dwords = self->load_vec / 2;
                rocke_value_t* v = rocke_b_buffer_load_vN_f16(b, rsrc, safe, c0, dwords);
                rocke_b_smem_store_vN_f16(b, smem_dst, indices, 2, v, self->load_vec);
            }
        }
        else
        {
            rocke_value_t* indices[2];
            indices[0] = row;
            indices[1] = col;
            if(self->load_vec == 1)
            {
                /* Python:
                 *   v = b.global_load_f16(ptr, off_elems)
                 *   b.smem_store_f16(smem_dst, [row, col], v)
                 */
                rocke_value_t* v = rocke_b_global_load_f16(b, ptr, off_elems, 0);
                rocke_b_smem_store_f16(b, smem_dst, indices, 2, v);
            }
            else
            {
                /* Python:
                 *   v = b.global_load_vN_f16(ptr, off_elems, self.load_vec)
                 *   b.smem_store_vN_f16(smem_dst, [row, col], v, self.load_vec)
                 */
                rocke_value_t* v = rocke_b_global_load_vN_f16(b, ptr, off_elems, self->load_vec, 0);
                rocke_b_smem_store_vN_f16(b, smem_dst, indices, 2, v, self->load_vec);
            }
        }
    }
}

/* ============================================================================
 * AsyncTileLoader / AsyncTileLoaderSlot
 * ========================================================================== */

rocke_status_t rocke_async_tile_loader_choose_dwords(
    int tile_rows, int tile_cols, int block_size, int max_dwords, int* out)
{
    /* Python:
     *   if max_dwords > 4: max_dwords = 4
     *   for d in (4, 3, 1):
     *       if d > max_dwords: continue
     *       halves = d * 2
     *       if tile_cols % halves != 0: continue
     *       chunks = (tile_rows*tile_cols) // halves
     *       if chunks < block_size: continue
     *       return d
     *   raise ValueError(...)
     */
    static const int kCandidates[3] = {4, 3, 1};
    int i;

    if(max_dwords > 4)
    {
        max_dwords = 4;
    }
    for(i = 0; i < 3; ++i)
    {
        int d = kCandidates[i];
        int halves;
        int chunks;
        if(d > max_dwords)
        {
            continue;
        }
        halves = d * 2;
        if(tile_cols % halves != 0)
        {
            continue;
        }
        chunks = (tile_rows * tile_cols) / halves;
        if(chunks < block_size)
        {
            continue;
        }
        if(out != NULL)
        {
            *out = d;
        }
        return ROCKE_OK;
    }
    return ROCKE_ERR_VALUE;
}

rocke_status_t rocke_async_tile_loader_from_tile(int tile_rows,
                                                 int tile_cols,
                                                 int block_size,
                                                 int wave_size,
                                                 int max_dwords,
                                                 rocke_async_tile_loader_t* out)
{
    int d;
    int halves;
    int chunks;
    int passes;
    rocke_status_t st;

    /* Python:
     *   d = cls.choose_dwords(...)
     *   halves = d * 2
     *   chunks = (tile_rows * tile_cols) // halves
     *   passes = (chunks + block_size - 1) // block_size
     *   return cls(... dwords=d, chunks_total=chunks,
     *              chunks_per_pass=block_size, passes=passes)
     */
    st = rocke_async_tile_loader_choose_dwords(tile_rows, tile_cols, block_size, max_dwords, &d);
    if(st != ROCKE_OK)
    {
        return st;
    }
    halves = d * 2;
    chunks = (tile_rows * tile_cols) / halves;
    passes = (chunks + block_size - 1) / block_size;
    if(out != NULL)
    {
        out->tile_rows = tile_rows;
        out->tile_cols = tile_cols;
        out->block_size = block_size;
        out->wave_size = wave_size;
        out->dwords = d;
        out->chunks_total = chunks;
        out->chunks_per_pass = block_size;
        out->passes = passes;
    }
    return ROCKE_OK;
}

int rocke_async_tile_loader_halves_per_chunk(const rocke_async_tile_loader_t* self)
{
    /* Python: return dwords * 2 */
    return (self == NULL) ? 0 : self->dwords * 2;
}

int rocke_async_tile_loader_bytes_per_chunk(const rocke_async_tile_loader_t* self)
{
    /* Python: return dwords * 4 */
    return (self == NULL) ? 0 : self->dwords * 4;
}

int rocke_async_tile_loader_cols_per_chunk(const rocke_async_tile_loader_t* self)
{
    /* Python: return self.halves_per_chunk */
    return rocke_async_tile_loader_halves_per_chunk(self);
}

int rocke_async_tile_loader_wave_bytes(const rocke_async_tile_loader_t* self)
{
    /* Python: return wave_size * self.bytes_per_chunk */
    if(self == NULL)
    {
        return 0;
    }
    return self->wave_size * rocke_async_tile_loader_bytes_per_chunk(self);
}

int rocke_async_tile_loader_pass_bytes(const rocke_async_tile_loader_t* self)
{
    /* Python: return block_size * self.bytes_per_chunk */
    if(self == NULL)
    {
        return 0;
    }
    return self->block_size * rocke_async_tile_loader_bytes_per_chunk(self);
}

rocke_status_t rocke_async_tile_loader_bind(rocke_ir_builder_t* b,
                                            const rocke_async_tile_loader_t* self,
                                            rocke_value_t* smem_dst,
                                            rocke_value_t* wave_id,
                                            rocke_async_tile_loader_slot_t* out_slot)
{
    rocke_value_t* lds_base;
    rocke_value_t* wave_byte_off_i32;
    rocke_value_t* wave_byte_off_i64;
    rocke_value_t* per_wave_lds;

    if(b != NULL && b->status != ROCKE_OK)
    {
        return b->status; /* already in error: no-op */
    }
    if(self == NULL)
    {
        rocke_i_set_err(b, ROCKE_ERR_VALUE, "AsyncTileLoader.bind: NULL loader");
        return ROCKE_ERR_VALUE;
    }

    /* Python:
     *   lds_base = b.smem_addr_of(smem_dst)                       # i64
     *   wave_byte_off_i32 = b.mul(wave_id, b.const_i32(self.wave_bytes))
     *   wave_byte_off_i32 = b.to_sgpr_u32(wave_byte_off_i32)
     *   wave_byte_off_i64 = b.zext(wave_byte_off_i32, I64)
     *   per_wave_lds = b.smem_ptr_add(lds_base, wave_byte_off_i64)
     *   return AsyncTileLoaderSlot(loader=self, smem_dst=smem_dst,
     *                              per_wave_lds_base=per_wave_lds)
     */
    lds_base = rocke_b_smem_addr_of(b, smem_dst);
    wave_byte_off_i32
        = rocke_b_mul(b, wave_id, rocke_b_const_i32(b, rocke_async_tile_loader_wave_bytes(self)));
    wave_byte_off_i32 = rocke_b_to_sgpr_u32(b, wave_byte_off_i32);
    wave_byte_off_i64 = rocke_b_zext(b, wave_byte_off_i32, rocke_i64());
    per_wave_lds = rocke_b_smem_ptr_add(b, lds_base, wave_byte_off_i64);

    if(out_slot != NULL)
    {
        out_slot->loader = *self;
        out_slot->smem_dst = smem_dst;
        out_slot->per_wave_lds_base = per_wave_lds;
    }
    return (b != NULL) ? b->status : ROCKE_OK;
}

void rocke_async_tile_loader_slot_issue(rocke_ir_builder_t* b,
                                        const rocke_async_tile_loader_slot_t* self,
                                        rocke_value_t* tid,
                                        rocke_value_t* rsrc,
                                        rocke_loads_descriptor_fn descriptor,
                                        void* descriptor_user,
                                        int oob_sentinel,
                                        int coherency)
{
    const rocke_async_tile_loader_t* L;
    rocke_value_t* c_half_bytes;
    rocke_value_t* c_oob;
    rocke_value_t* c0;
    rocke_value_t* c_cols_per_chunk;
    int p;

    if(b != NULL && b->status != ROCKE_OK)
    {
        return; /* already in error: no-op */
    }
    if(self == NULL)
    {
        rocke_i_set_err(b, ROCKE_ERR_VALUE, "AsyncTileLoaderSlot.issue: NULL slot");
        return;
    }

    /* Python:
     *   L = self.loader
     *   c_half_bytes     = b.const_i32(2)
     *   c_oob            = b.const_i32(oob_sentinel)
     *   c0               = b.const_i32(0)
     *   c_cols_per_chunk = b.const_i32(L.cols_per_chunk)
     */
    L = &self->loader;
    c_half_bytes = rocke_b_const_i32(b, 2);
    c_oob = rocke_b_const_i32(b, oob_sentinel);
    c0 = rocke_b_const_i32(b, 0);
    c_cols_per_chunk = rocke_b_const_i32(b, rocke_async_tile_loader_cols_per_chunk(L));

    for(p = 0; p < L->passes; ++p)
    {
        int pass_byte_off;
        rocke_value_t* pass_base;
        rocke_value_t* chunk_idx;
        rocke_value_t* row;
        rocke_value_t* col_v;
        rocke_value_t* col;
        rocke_value_t* off_elems;
        rocke_value_t* valid;
        rocke_value_t* off_bytes;
        rocke_value_t* in_pass;
        rocke_value_t* valid_final;
        rocke_value_t* safe;

        /* Python:
         *   pass_byte_off = p * L.pass_bytes
         *   pass_base = (b.smem_ptr_add(self.per_wave_lds_base,
         *                               b.zext(b.const_i32(pass_byte_off), I64))
         *                if p > 0 else self.per_wave_lds_base)
         */
        pass_byte_off = p * rocke_async_tile_loader_pass_bytes(L);
        if(p > 0)
        {
            pass_base = rocke_b_smem_ptr_add(
                b,
                self->per_wave_lds_base,
                rocke_b_zext(b, rocke_b_const_i32(b, pass_byte_off), rocke_i64()));
        }
        else
        {
            pass_base = self->per_wave_lds_base;
        }

        /* Python:
         *   chunk_idx = b.add(tid, b.const_i32(p * L.block_size))
         *   row   = b.div(chunk_idx, c_cols_per_chunk)
         *   col_v = b.mod(chunk_idx, c_cols_per_chunk)
         *   col   = b.mul(col_v, b.const_i32(L.halves_per_chunk))
         */
        chunk_idx = rocke_b_add(b, tid, rocke_b_const_i32(b, p * L->block_size));
        row = rocke_b_div(b, chunk_idx, c_cols_per_chunk);
        col_v = rocke_b_mod(b, chunk_idx, c_cols_per_chunk);
        col = rocke_b_mul(
            b, col_v, rocke_b_const_i32(b, rocke_async_tile_loader_halves_per_chunk(L)));

        /* Python:
         *   off_elems, valid = descriptor(b, row, col)
         *   off_bytes = b.mul(off_elems, c_half_bytes)
         *   in_pass = b.cmp_lt(chunk_idx, b.const_i32(L.chunks_total))
         *   valid_final = b.land(valid, in_pass) if valid is not None else in_pass
         *   safe = b.select(valid_final, off_bytes, c_oob)
         *   b.async_buffer_load_lds_addr(rsrc, pass_base, safe, c0, L.dwords,
         *                                coherency=coherency)
         */
        valid = NULL;
        off_elems = descriptor(b, row, col, &valid, descriptor_user);
        off_bytes = rocke_b_mul(b, off_elems, c_half_bytes);
        in_pass = rocke_b_cmp_lt(b, chunk_idx, rocke_b_const_i32(b, L->chunks_total));
        if(valid != NULL)
        {
            valid_final = rocke_b_land(b, valid, in_pass);
        }
        else
        {
            valid_final = in_pass;
        }
        safe = rocke_b_select(b, valid_final, off_bytes, c_oob);
        rocke_b_async_buffer_load_lds_addr(b, rsrc, pass_base, safe, c0, L->dwords, coherency);
    }
}

int rocke_async_tile_loader_slot_required_lds_bytes(const rocke_async_tile_loader_slot_t* self)
{
    /* Python: return self.loader.passes * self.loader.pass_bytes */
    if(self == NULL)
    {
        return 0;
    }
    return self->loader.passes * rocke_async_tile_loader_pass_bytes(&self->loader);
}