derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
335 lines
14 KiB
C++
335 lines
14 KiB
C++
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
|
// SPDX-License-Identifier: MIT
|
|
/*
|
|
* helper_rocke.helpers.grid.c -- C99 port of rocke.helpers.grid.
|
|
*
|
|
* Ports chiplet_aware_super_tile_dynamic and the two dynamic helpers it
|
|
* composes. The builder-call sequence is byte-identical to the Python so the
|
|
* emitted IR op stream matches exactly.
|
|
*/
|
|
#include "rocke/helper_rocke.helpers.grid.h"
|
|
|
|
#include "rocke/ir_internal.h" /* rocke_i_set_err for Python-ValueError parity */
|
|
|
|
/* ------------------------------------------------------------------------
|
|
* chiplet_transform_chunked_dynamic (file-local; Python helper)
|
|
*
|
|
* block = num_xcds * chunk_size
|
|
* c_num_xcds = b.const_i32(int(num_xcds))
|
|
* c_chunk_size = b.const_i32(int(chunk_size))
|
|
* c_block = b.const_i32(int(block))
|
|
* limit = b.mul(b.div(num_wgs, c_block), c_block)
|
|
* xcd = b.mod(wgid, c_num_xcds)
|
|
* local_pid = b.div(wgid, c_num_xcds)
|
|
* chunk_idx = b.div(local_pid, c_chunk_size)
|
|
* pos_in_chunk = b.mod(local_pid, c_chunk_size)
|
|
* new_wgid = b.add(b.add(b.mul(chunk_idx, c_block),
|
|
* b.mul(xcd, c_chunk_size)),
|
|
* pos_in_chunk)
|
|
* in_full_block = b.cmp_lt(wgid, limit)
|
|
* return b.select(in_full_block, new_wgid, wgid)
|
|
* ------------------------------------------------------------------------ */
|
|
static rocke_value_t* rocke_i_chiplet_transform_chunked_dynamic(rocke_ir_builder_t* b,
|
|
rocke_value_t* wgid,
|
|
rocke_value_t* num_wgs,
|
|
int num_xcds,
|
|
int chunk_size)
|
|
{
|
|
int block;
|
|
rocke_value_t* c_num_xcds;
|
|
rocke_value_t* c_chunk_size;
|
|
rocke_value_t* c_block;
|
|
rocke_value_t* limit;
|
|
rocke_value_t* xcd;
|
|
rocke_value_t* local_pid;
|
|
rocke_value_t* chunk_idx;
|
|
rocke_value_t* pos_in_chunk;
|
|
rocke_value_t* new_wgid;
|
|
rocke_value_t* in_full_block;
|
|
|
|
if(num_xcds <= 0 || chunk_size <= 0)
|
|
{
|
|
return (rocke_value_t*)rocke_i_set_err(
|
|
b,
|
|
ROCKE_ERR_VALUE,
|
|
"chiplet_transform_chunked_dynamic: invalid args num_xcds=%d chunk_size=%d",
|
|
num_xcds,
|
|
chunk_size);
|
|
}
|
|
|
|
block = num_xcds * chunk_size;
|
|
|
|
c_num_xcds = rocke_b_const_i32(b, (int64_t)num_xcds);
|
|
c_chunk_size = rocke_b_const_i32(b, (int64_t)chunk_size);
|
|
c_block = rocke_b_const_i32(b, (int64_t)block);
|
|
|
|
/* limit = (num_wgs / block) * block (largest full-block boundary) */
|
|
limit = rocke_b_mul(b, rocke_b_div(b, num_wgs, c_block), c_block);
|
|
|
|
xcd = rocke_b_mod(b, wgid, c_num_xcds);
|
|
local_pid = rocke_b_div(b, wgid, c_num_xcds);
|
|
chunk_idx = rocke_b_div(b, local_pid, c_chunk_size);
|
|
pos_in_chunk = rocke_b_mod(b, local_pid, c_chunk_size);
|
|
|
|
/* Python: b.add(b.add(b.mul(chunk_idx, c_block), b.mul(xcd, c_chunk_size)),
|
|
* pos_in_chunk). Python evaluates the inner add's first arg (mul chunk_idx)
|
|
* before its second (mul xcd), so the SSA order is mul(chunk_idx,block),
|
|
* mul(xcd,chunk_size), add, add. C's argument evaluation order is
|
|
* unspecified -- pin it with explicit temporaries. */
|
|
{
|
|
rocke_value_t* mul_chunk = rocke_b_mul(b, chunk_idx, c_block);
|
|
rocke_value_t* mul_xcd = rocke_b_mul(b, xcd, c_chunk_size);
|
|
new_wgid = rocke_b_add(b, rocke_b_add(b, mul_chunk, mul_xcd), pos_in_chunk);
|
|
}
|
|
|
|
in_full_block = rocke_b_cmp_lt(b, wgid, limit);
|
|
return rocke_b_select(b, in_full_block, new_wgid, wgid);
|
|
}
|
|
|
|
/* ------------------------------------------------------------------------
|
|
* super_tile_swizzle_dynamic (file-local; Python helper)
|
|
*
|
|
* if wgm <= 0: raise ValueError
|
|
* c_wgm = b.const_i32(int(wgm))
|
|
* num_wgid_in_group = b.mul(c_wgm, num_pid_n)
|
|
* group_id = b.div(wgid, num_wgid_in_group)
|
|
* first_pid_m = b.mul(group_id, c_wgm)
|
|
* rem = b.sub(num_pid_m, first_pid_m)
|
|
* use_wgm = b.cmp_lt(c_wgm, rem)
|
|
* group_size_m = b.select(use_wgm, c_wgm, rem)
|
|
* local_id = b.mod(wgid, num_wgid_in_group)
|
|
* pid_m = b.add(first_pid_m, b.mod(local_id, group_size_m))
|
|
* pid_n = b.div(local_id, group_size_m)
|
|
* return SuperTileSwizzleResult(row=pid_m, col=pid_n)
|
|
* ------------------------------------------------------------------------ */
|
|
static rocke_super_tile_swizzle_result_t
|
|
rocke_i_super_tile_swizzle_dynamic(rocke_ir_builder_t* b,
|
|
rocke_value_t* wgid,
|
|
rocke_value_t* num_pid_m,
|
|
rocke_value_t* num_pid_n,
|
|
int wgm)
|
|
{
|
|
rocke_super_tile_swizzle_result_t res;
|
|
rocke_value_t* c_wgm;
|
|
rocke_value_t* num_wgid_in_group;
|
|
rocke_value_t* group_id;
|
|
rocke_value_t* first_pid_m;
|
|
rocke_value_t* rem;
|
|
rocke_value_t* use_wgm;
|
|
rocke_value_t* group_size_m;
|
|
rocke_value_t* local_id;
|
|
|
|
res.row = NULL;
|
|
res.col = NULL;
|
|
|
|
if(wgm <= 0)
|
|
{
|
|
rocke_i_set_err(b, ROCKE_ERR_VALUE, "super_tile_swizzle_dynamic: wgm=%d must be > 0", wgm);
|
|
return res;
|
|
}
|
|
|
|
c_wgm = rocke_b_const_i32(b, (int64_t)wgm);
|
|
num_wgid_in_group = rocke_b_mul(b, c_wgm, num_pid_n);
|
|
|
|
group_id = rocke_b_div(b, wgid, num_wgid_in_group);
|
|
first_pid_m = rocke_b_mul(b, group_id, c_wgm);
|
|
rem = rocke_b_sub(b, num_pid_m, first_pid_m);
|
|
use_wgm = rocke_b_cmp_lt(b, c_wgm, rem);
|
|
group_size_m = rocke_b_select(b, use_wgm, c_wgm, rem);
|
|
|
|
local_id = rocke_b_mod(b, wgid, num_wgid_in_group);
|
|
res.row = rocke_b_add(b, first_pid_m, rocke_b_mod(b, local_id, group_size_m));
|
|
res.col = rocke_b_div(b, local_id, group_size_m);
|
|
|
|
return res;
|
|
}
|
|
|
|
/* ------------------------------------------------------------------------
|
|
* chiplet_aware_super_tile_dynamic (public)
|
|
*
|
|
* num_wgs = b.mul(num_pid_m, num_pid_n)
|
|
* remapped = chiplet_transform_chunked_dynamic(
|
|
* b, wgid, num_wgs=num_wgs,
|
|
* num_xcds=num_xcds, chunk_size=chunk_size)
|
|
* return super_tile_swizzle_dynamic(
|
|
* b, remapped, num_pid_m=num_pid_m,
|
|
* num_pid_n=num_pid_n, wgm=wgm)
|
|
* ------------------------------------------------------------------------ */
|
|
rocke_super_tile_swizzle_result_t rocke_chiplet_aware_super_tile_dynamic(rocke_ir_builder_t* b,
|
|
rocke_value_t* wgid,
|
|
rocke_value_t* num_pid_m,
|
|
rocke_value_t* num_pid_n,
|
|
int wgm,
|
|
int num_xcds,
|
|
int chunk_size)
|
|
{
|
|
rocke_value_t* num_wgs;
|
|
rocke_value_t* remapped;
|
|
|
|
num_wgs = rocke_b_mul(b, num_pid_m, num_pid_n);
|
|
remapped = rocke_i_chiplet_transform_chunked_dynamic(b, wgid, num_wgs, num_xcds, chunk_size);
|
|
return rocke_i_super_tile_swizzle_dynamic(b, remapped, num_pid_m, num_pid_n, wgm);
|
|
}
|
|
|
|
/* ------------------------------------------------------------------------
|
|
* chiplet_transform_chunked (file-local; Python COMPILE-TIME helper)
|
|
*
|
|
* Differs from the *_dynamic peer only in that num_wgs (hence `limit`) is a
|
|
* compile-time int, so `c_limit` is a folded const and `in_full_block` compares
|
|
* against it directly (no div/mul IR for the limit). Python (grid.py 84-110):
|
|
* block = num_xcds * chunk_size
|
|
* limit = (num_wgs // block) * block
|
|
* c_num_xcds = b.const_i32(num_xcds)
|
|
* c_chunk_size = b.const_i32(chunk_size)
|
|
* c_block = b.const_i32(block)
|
|
* c_limit = b.const_i32(limit)
|
|
* xcd = b.mod(wgid, c_num_xcds)
|
|
* local_pid = b.div(wgid, c_num_xcds)
|
|
* chunk_idx = b.div(local_pid, c_chunk_size)
|
|
* pos_in_chunk = b.mod(local_pid, c_chunk_size)
|
|
* new_wgid = b.add(b.add(b.mul(chunk_idx, c_block),
|
|
* b.mul(xcd, c_chunk_size)),
|
|
* pos_in_chunk)
|
|
* in_full_block = b.cmp_lt(wgid, c_limit)
|
|
* return b.select(in_full_block, new_wgid, wgid)
|
|
* ------------------------------------------------------------------------ */
|
|
static rocke_value_t* rocke_i_chiplet_transform_chunked(
|
|
rocke_ir_builder_t* b, rocke_value_t* wgid, int num_wgs, int num_xcds, int chunk_size)
|
|
{
|
|
int block;
|
|
int limit;
|
|
rocke_value_t* c_num_xcds;
|
|
rocke_value_t* c_chunk_size;
|
|
rocke_value_t* c_block;
|
|
rocke_value_t* c_limit;
|
|
rocke_value_t* xcd;
|
|
rocke_value_t* local_pid;
|
|
rocke_value_t* chunk_idx;
|
|
rocke_value_t* pos_in_chunk;
|
|
rocke_value_t* new_wgid;
|
|
rocke_value_t* in_full_block;
|
|
|
|
if(num_xcds <= 0 || chunk_size <= 0 || num_wgs <= 0)
|
|
{
|
|
return (rocke_value_t*)rocke_i_set_err(
|
|
b,
|
|
ROCKE_ERR_VALUE,
|
|
"chiplet_transform_chunked: invalid args num_wgs=%d num_xcds=%d chunk_size=%d",
|
|
num_wgs,
|
|
num_xcds,
|
|
chunk_size);
|
|
}
|
|
|
|
block = num_xcds * chunk_size;
|
|
limit = (num_wgs / block) * block;
|
|
|
|
c_num_xcds = rocke_b_const_i32(b, (int64_t)num_xcds);
|
|
c_chunk_size = rocke_b_const_i32(b, (int64_t)chunk_size);
|
|
c_block = rocke_b_const_i32(b, (int64_t)block);
|
|
c_limit = rocke_b_const_i32(b, (int64_t)limit);
|
|
|
|
xcd = rocke_b_mod(b, wgid, c_num_xcds);
|
|
local_pid = rocke_b_div(b, wgid, c_num_xcds);
|
|
chunk_idx = rocke_b_div(b, local_pid, c_chunk_size);
|
|
pos_in_chunk = rocke_b_mod(b, local_pid, c_chunk_size);
|
|
|
|
{
|
|
rocke_value_t* mul_chunk = rocke_b_mul(b, chunk_idx, c_block);
|
|
rocke_value_t* mul_xcd = rocke_b_mul(b, xcd, c_chunk_size);
|
|
new_wgid = rocke_b_add(b, rocke_b_add(b, mul_chunk, mul_xcd), pos_in_chunk);
|
|
}
|
|
|
|
in_full_block = rocke_b_cmp_lt(b, wgid, c_limit);
|
|
return rocke_b_select(b, in_full_block, new_wgid, wgid);
|
|
}
|
|
|
|
/* ------------------------------------------------------------------------
|
|
* super_tile_swizzle (file-local; Python COMPILE-TIME helper)
|
|
*
|
|
* Differs from the *_dynamic peer in that num_pid_m / num_pid_n are
|
|
* compile-time ints: num_wgid_in_group (= wgm * num_pid_n) and c_num_pid_m are
|
|
* folded consts, not IR muls. Python (grid.py 159-181):
|
|
* num_wgid_in_group = wgm * num_pid_n
|
|
* c_wgm = b.const_i32(wgm)
|
|
* c_num_pid_m = b.const_i32(num_pid_m)
|
|
* c_num_wgid_in_group = b.const_i32(num_wgid_in_group)
|
|
* group_id = b.div(wgid, c_num_wgid_in_group)
|
|
* first_pid_m = b.mul(group_id, c_wgm)
|
|
* rem = b.sub(c_num_pid_m, first_pid_m)
|
|
* use_wgm = b.cmp_lt(c_wgm, rem)
|
|
* group_size_m = b.select(use_wgm, c_wgm, rem)
|
|
* local_id = b.mod(wgid, c_num_wgid_in_group)
|
|
* pid_m = b.add(first_pid_m, b.mod(local_id, group_size_m))
|
|
* pid_n = b.div(local_id, group_size_m)
|
|
* ------------------------------------------------------------------------ */
|
|
static rocke_super_tile_swizzle_result_t rocke_i_super_tile_swizzle(
|
|
rocke_ir_builder_t* b, rocke_value_t* wgid, int num_pid_m, int num_pid_n, int wgm)
|
|
{
|
|
rocke_super_tile_swizzle_result_t res;
|
|
int num_wgid_in_group;
|
|
rocke_value_t* c_wgm;
|
|
rocke_value_t* c_num_pid_m;
|
|
rocke_value_t* c_num_wgid_in_group;
|
|
rocke_value_t* group_id;
|
|
rocke_value_t* first_pid_m;
|
|
rocke_value_t* rem;
|
|
rocke_value_t* use_wgm;
|
|
rocke_value_t* group_size_m;
|
|
rocke_value_t* local_id;
|
|
|
|
res.row = NULL;
|
|
res.col = NULL;
|
|
|
|
if(num_pid_m <= 0 || num_pid_n <= 0 || wgm <= 0)
|
|
{
|
|
rocke_i_set_err(b,
|
|
ROCKE_ERR_VALUE,
|
|
"super_tile_swizzle: invalid args num_pid_m=%d num_pid_n=%d wgm=%d",
|
|
num_pid_m,
|
|
num_pid_n,
|
|
wgm);
|
|
return res;
|
|
}
|
|
|
|
num_wgid_in_group = wgm * num_pid_n;
|
|
|
|
c_wgm = rocke_b_const_i32(b, (int64_t)wgm);
|
|
c_num_pid_m = rocke_b_const_i32(b, (int64_t)num_pid_m);
|
|
c_num_wgid_in_group = rocke_b_const_i32(b, (int64_t)num_wgid_in_group);
|
|
|
|
group_id = rocke_b_div(b, wgid, c_num_wgid_in_group);
|
|
first_pid_m = rocke_b_mul(b, group_id, c_wgm);
|
|
rem = rocke_b_sub(b, c_num_pid_m, first_pid_m);
|
|
use_wgm = rocke_b_cmp_lt(b, c_wgm, rem);
|
|
group_size_m = rocke_b_select(b, use_wgm, c_wgm, rem);
|
|
|
|
local_id = rocke_b_mod(b, wgid, c_num_wgid_in_group);
|
|
res.row = rocke_b_add(b, first_pid_m, rocke_b_mod(b, local_id, group_size_m));
|
|
res.col = rocke_b_div(b, local_id, group_size_m);
|
|
|
|
return res;
|
|
}
|
|
|
|
/* ------------------------------------------------------------------------
|
|
* chiplet_aware_super_tile (public; Python COMPILE-TIME composition)
|
|
*
|
|
* num_wgs = int(num_pid_m) * int(num_pid_n)
|
|
* remapped = chiplet_transform_chunked(b, wgid, num_wgs=num_wgs,
|
|
* num_xcds=num_xcds, chunk_size=chunk_size)
|
|
* return super_tile_swizzle(b, remapped, num_pid_m=num_pid_m,
|
|
* num_pid_n=num_pid_n, wgm=wgm)
|
|
* ------------------------------------------------------------------------ */
|
|
rocke_super_tile_swizzle_result_t rocke_chiplet_aware_super_tile(rocke_ir_builder_t* b,
|
|
rocke_value_t* wgid,
|
|
int num_pid_m,
|
|
int num_pid_n,
|
|
int wgm,
|
|
int num_xcds,
|
|
int chunk_size)
|
|
{
|
|
int num_wgs = num_pid_m * num_pid_n;
|
|
rocke_value_t* remapped
|
|
= rocke_i_chiplet_transform_chunked(b, wgid, num_wgs, num_xcds, chunk_size);
|
|
return rocke_i_super_tile_swizzle(b, remapped, num_pid_m, num_pid_n, wgm);
|
|
}
|