#include "dsl/ir/pass/pass.hpp"
#include "dsl/ir/pass/trace.hpp"
#include "dsl/utils/logging.hpp"
#include "dsl/utils/utils.hpp"
#include "gemmstone/config.hpp"
#include "gemmstone/dsl/dsl.hpp"
#include "gemmstone/strategy.hpp"
#include "generator_dsl/kernel_desc.hpp"
#include "gpu/intel/jit/utils/type_bridge.hpp"
GEMMSTONE_NAMESPACE_START
using namespace dsl;
inline type_t into_ir(Type t, int elems = 1) {
using namespace ir;
switch (t) {
case Type::invalid: return type_t::undef();
case Type::f4_e3m0: return type_t::f4_e3m0(elems);
case Type::f4_e2m1: return type_t::f4_e2m1(elems);
case Type::bf8: return type_t::bf8(elems);
case Type::hf8: return type_t::hf8(elems);
case Type::bf16: return type_t::bf16(elems);
case Type::f16: return type_t::f16(elems);
case Type::tf32: return type_t::tf32(elems);
case Type::f32: return type_t::f32(elems);
case Type::f64: return type_t::f64(elems);
case Type::u4: return type_t::u4(elems);
case Type::s4: return type_t::s4(elems);
case Type::u8: return type_t::u8(elems);
case Type::s8: return type_t::s8(elems);
case Type::u16: return type_t::u16(elems);
case Type::s16: return type_t::s16(elems);
case Type::u32: return type_t::u32(elems);
case Type::s32: return type_t::s32(elems);
case Type::u64: return type_t::u64(elems);
case Type::s64: return type_t::s64(elems);
default: stub(); return type_t::undef();
}
}
struct transform_t {
enum class kind_t { none, block, vnni, transpose_vnni };
transform_t() = default;
transform_t(kind_t t_kind, int pack_size, ngen::CacheSettingsLSC cache_hint,
std::array<idx_t, 2> dims)
: kind(t_kind)
, pack_size(pack_size)
, cache_hint(to_ir(cache_hint))
, dims(std::move(dims)) {}
layout_t get_layout(const tile_t &sizes, type_t type) const {
auto col_var = dims[0];
auto col = sizes[dims[0]];
auto row_var = dims[1];
auto row = sizes[dims[1]];
auto t = type.size();
auto normalized = kind;
if (normalized == kind_t::transpose_vnni) {
std::swap(col_var, row_var);
std::swap(col, row);
normalized = kind_t::vnni;
}
if (normalized == kind_t::vnni && t >= 4) normalized = kind_t::block;
int col_inner = pack_size ? pack_size : grf_size();
if (normalized == kind_t::block && col <= col_inner)
normalized = kind_t::none;
switch (normalized) {
case kind_t::none:
return layout_t(type, {{col_var, col, 1}, {row_var, row, col}});
case kind_t::block: {
int col_outer = (int)(col / col_inner);
return layout_t(type,
{{col_var, col_inner, 1}, {row_var, row, col_inner},
{col_var, col_outer, row * col_inner}});
}
case kind_t::vnni: {
int row_inner = 4 / t;
int row_outer = (int)(row / row_inner);
int col_outer = (int)(col / col_inner);
return layout_t(type,
{{row_var, row_inner, 1},
{col_var, col_inner, row_inner},
{row_var, row_outer, col_inner * row_inner},
{col_var, col_outer,
row_outer * col_inner * row_inner}});
}
case kind_t::transpose_vnni:
default: stub(); return {};
}
}
static send_cache_hint_t to_ir(ngen::CacheSettingsLSC hint) {
switch (hint) {
case ngen::CacheSettingsLSC::L1C_L3C:
return send_cache_hint_t::load_once;
case ngen::CacheSettingsLSC::Default:
return send_cache_hint_t::hw_default;
default: stub(); return send_cache_hint_t::undef;
}
}
kind_t kind = kind_t::none;
int pack_size = 0;
send_cache_hint_t cache_hint = send_cache_hint_t::undef;
std::array<idx_t, 2> dims = {};
};
static const idx_t m_var("m");
static const idx_t n_var("n");
static const idx_t k_var("k");
struct kloop_iterator_t {
virtual const global_tensor_t &A_prefetch() const = 0;
virtual const global_tensor_t &A_load() const = 0;
virtual const global_tensor_t &B_prefetch() const = 0;
virtual const global_tensor_t &B_load() const = 0;
virtual const global_tensor_t &C_store() const = 0;
virtual void A_prefetch_inc(int64_t k_block) = 0;
virtual void A_load_inc(int64_t k_block) = 0;
virtual void B_prefetch_inc(int64_t k_block) = 0;
virtual void B_load_inc(int64_t k_block) = 0;
virtual void kloop_inc(int64_t k_block) = 0;
virtual expr_t update_C() const = 0;
virtual expr_t is_inbounds(int64_t increment) const = 0;
};
transform_t get_transform(const MatrixAddressingStrategy &matrix_strategy,
std::array<idx_t, 2> dims, bool is_prefetch = false) {
switch (matrix_strategy.accessType) {
case AccessType::Scattered:
if (is_prefetch)
return transform_t(transform_t::kind_t::none, 0,
matrix_strategy.cachingR, dims);
return transform_t(transform_t::kind_t::transpose_vnni,
matrix_strategy.tileR, matrix_strategy.cachingR, dims);
case AccessType::ChannelScattered: stub(); return {};
case AccessType::Block2DTranspose:
return transform_t(transform_t::kind_t::transpose_vnni,
matrix_strategy.tileR, matrix_strategy.cachingR, dims);
case AccessType::Block:
case AccessType::PseudoBlock:
return transform_t(transform_t::kind_t::block,
matrix_strategy.tileR, matrix_strategy.cachingR, dims);
case AccessType::Block2D: {
return transform_t(transform_t::kind_t::block,
matrix_strategy.tileR, matrix_strategy.cachingR, dims);
};
case AccessType::Block2DVNNI: {
return transform_t(transform_t::kind_t::vnni, matrix_strategy.tileR,
matrix_strategy.cachingR, dims);
}
default: stub(); return {};
}
}
idx_map_t<expr_t> get_strides(
MatrixLayout layout, std::array<idx_t, 2> pvars, expr_t ld) {
switch (layout) {
case MatrixLayout::N: return {{pvars[0], 1}, {pvars[1], ld}};
case MatrixLayout::T: return {{pvars[0], ld}, {pvars[1], 1}};
default: stub(); return {};
};
}
struct tensor_config_t {
tensor_config_t(const global_tensor_t &g, transform_t t, int copies)
: transform(t) {
tile = g.tile;
layout = t.get_layout(g.tile, g.type);
layout = layout.with_block({k_var, copies});
}
tile_t tile;
layout_t layout;
transform_t transform;
};
void apply_post_ops(const dnnl::impl::gpu::intel::gpu_post_ops_t &ops,
const tensor_t &C, const std::vector<expr_t> &idxs,
const std::vector<idx_t> &dims) {
for (size_t i = 0; i < ops.len(); i++) {
if (ops[i].is_eltwise()) {
stub();
} else if (ops[i].is_sum()) {
stub();
} else if (ops[i].is_binary()) {
auto &e = ops[i].as_binary();
std::string i_s = std::to_string(i);
std::string stride_prefix = "binary" + i_s + "_stride";
int ndims = (int)dims.size();
idx_map_t<int> dim_to_md;
for (int i = 0; i < ndims; i++) {
dim_to_md[dims[i]] = i;
};
std::vector<expr_t> strides;
strides.reserve(idxs.size());
for (unsigned int j = 0; j < idxs.size(); j++) {
if (e.src1_desc.is_inner_dim(j, ndims)) {
strides.emplace_back(1);
} else if (e.src1_desc.is_broadcast(j, ndims)) {
strides.emplace_back(0);
} else {
strides.emplace_back(
arg(stride_prefix + std::to_string(j), true));
}
}
auto src_g = [&]() -> global_tensor_t {
expr_t src_g_offset = ir::simplify(arg("offset_binary" + i_s)
+ e.src1_desc.get_offset(idxs, strides));
idx_map_t<expr_t> g_strides;
idx_map_t<expr_t> g_sizes;
for (int i = 0; i < ndims; i++) {
g_strides[dims[i]] = strides[i];
g_sizes[dims[i]] = e.src1_desc.is_broadcast(i, ndims)
? expr_t(1)
: expr_t(0); }
return {arg("binary" + i_s),
dnnl::impl::gpu::intel::jit::to_ir(e.src1_desc.dt),
src_g_offset, coord_t(), g_sizes, g_strides, {}};
}();
layout_t src_layout = {src_g.type};
for (auto &b : C.layout.blocks()) {
if (!e.src1_desc.is_broadcast(dim_to_md[b.idx], ndims)) {
src_layout = src_layout.with_block({b.idx, b.size});
} else {
src_layout = src_layout.with_block({b.idx, 1});
}
}
tensor_t src = def("binary" + i_s + "_blk", src_layout);
std::cout << "src_g: " << src_g.str() << "\n";
std::cout << "src: " << src.str() << "\n";
load(src, src_g);
switch (e.alg) {
case dnnl::impl::alg_kind::binary_add:
binary(ir::op_kind_t::_add, C, C, src);
break;
default: stub();
}
} else {
stub();
}
}
}
struct basic_iterator_t : kloop_iterator_t {
basic_iterator_t(const global_tensor_t &A, int A_prefetch_k_blk,
int A_load_k_blk, const global_tensor_t &B, int B_prefetch_k_blk,
int B_load_k_blk, const global_tensor_t &C)
: m_idx_ {C.coord[m_var]}
, m_(C.sizes[m_var])
, n_idx_ {C.coord[n_var]}
, n_(C.sizes[n_var])
, k_idx_ {A.coord[k_var]}
, k_ {A.sizes[k_var]}
, A_prefetch_ {A.buf, A.type, A.base_offset, A.coord, A.sizes,
A.strides,
tile_t {{m_var, C.tile[m_var]}, {k_var, A_prefetch_k_blk}}}
, A_load_ {A.buf, A.type, A.base_offset, A.coord, A.sizes, A.strides,
tile_t {{m_var, C.tile[m_var]}, {k_var, A_load_k_blk}}}
, B_prefetch_ {B.buf, B.type, B.base_offset, B.coord, B.sizes,
B.strides,
tile_t {{k_var, B_prefetch_k_blk}, {n_var, C.tile[n_var]}}}
, B_load_ {B.buf, B.type, B.base_offset, B.coord, B.sizes, B.strides,
tile_t {{k_var, B_load_k_blk}, {n_var, C.tile[n_var]}}}
, C_store_ {C}
{
assume(m_idx_ % C.tile[m_var] == 0);
assume(n_idx_ % C.tile[n_var] == 0);
assume(m_idx_ >= 0);
assume(n_idx_ >= 0);
assume(k_idx_ >= 0);
}
const global_tensor_t &A_prefetch() const override { return A_prefetch_; }
const global_tensor_t &A_load() const override { return A_load_; }
const global_tensor_t &B_prefetch() const override { return B_prefetch_; }
const global_tensor_t &B_load() const override { return B_load_; }
const global_tensor_t &C_store() const override { return C_store_; }
void A_prefetch_inc(int64_t k_block) override {
A_prefetch_off += k_block;
A_prefetch_.coord[k_var] = k_idx_ + A_prefetch_off;
}
void A_load_inc(int64_t k_block) override {
A_load_off += k_block;
A_load_.coord[k_var] = k_idx_ + A_load_off;
}
void B_prefetch_inc(int64_t k_block) override {
B_prefetch_off += k_block;
B_prefetch_.coord[k_var] = k_idx_ + B_prefetch_off;
}
void B_load_inc(int64_t k_block) override {
B_load_off += k_block;
B_load_.coord[k_var] = k_idx_ + B_load_off;
}
void kloop_inc(int64_t k_block) override {
A_prefetch_inc(-k_block);
B_prefetch_inc(-k_block);
A_load_inc(-k_block);
B_load_inc(-k_block);
assign(k_idx_, k_idx_ + k_block);
}
expr_t update_C() const override { return false; }
expr_t is_inbounds(int64_t increment) const override {
return (m_idx_ < m_) & (n_idx_ < n_) & (k_idx_ < k_ - increment);
}
private:
expr_t m_idx_;
expr_t m_;
expr_t n_idx_;
expr_t n_;
expr_t k_idx_;
expr_t k_;
int64_t A_prefetch_off = 0;
int64_t A_load_off = 0;
int64_t B_prefetch_off = 0;
int64_t B_load_off = 0;
global_tensor_t A_prefetch_;
global_tensor_t A_load_;
global_tensor_t B_prefetch_;
global_tensor_t B_load_;
global_tensor_t C_store_;
};
struct generator_dsl_t {
generator_dsl_t(const generator_dsl_desc_t &desc)
: problem(desc.problem), strategy(desc.strategy) {}
kernel_t build(kernel::iface_t iface, ir::ir_context_t &ctx) {
if (strategy.kParallel || strategy.kParallelLocal) {
dsl_warning() << "kParallel support is unimplemented";
return {};
}
if (strategy.persistentLoop()) {
dsl_warning() << "persistentLoop support is unimplemented";
return {};
}
if (strategy.slmA || strategy.slmB) {
dsl_warning() << "slm copy support is unimplemented, disabling "
"slm copy";
}
if (strategy.wgPadFactor > 1) {
dsl_warning() << "work group padding is unimplemented";
return {};
}
if (strategy.cWalkOrder != WalkOrder::HW2D) {
dsl_warning() << "Unsupported walk order";
return {};
}
if (problem.Ta != problem.Ta_ext || problem.Tb != problem.Tb_ext
|| problem.Tc != problem.Tc_ext) {
dsl_warning() << "Type conversion support is unimplemented";
return {};
}
if (problem.batch != BatchMode::None
&& problem.batch != BatchMode::Strided) {
dsl_warning() << "Batch mode is unimplemented";
return {};
}
declare_kernel(iface, ctx);
const auto m = arg("m");
const auto n = arg("n");
const auto k = arg("k");
auto m_blk = strategy.unroll[LoopM];
auto n_blk = strategy.unroll[LoopN];
auto k_blk = strategy.unroll[LoopK];
std::array<idx_t, 2> A_vars = {m_var, k_var};
std::array<idx_t, 2> B_vars = {k_var, n_var};
std::array<idx_t, 2> C_vars = {m_var, n_var};
auto A_prefetch_transform
= get_transform(strategy.A_prefetch, A_vars, true);
auto A_load_transform = get_transform(strategy.A, A_vars);
auto B_prefetch_transform
= get_transform(strategy.B_prefetch, B_vars, true);
auto B_load_transform = get_transform(strategy.B, B_vars);
tile_t C_dims {{{m_var, m_blk}, {n_var, n_blk}}};
auto C_store_transform = get_transform(strategy.C, C_vars);
tensor_t C = def("C_blk",
C_store_transform.get_layout(C_dims, into_ir(problem.Tc)), 0);
idx_t subgroup_dim = C.layout[0].idx;
int m_group_idx = strategy.loopOrder[0] == LoopM ? 0 : 1;
auto m_idx = let("m_idx",
(group_id(m_group_idx) * local_size(m_group_idx)
+ local_id(m_group_idx))
* (subgroup_dim == m_var ? m_blk / strategy.subgroupSize
: m_blk));
int n_group_idx = strategy.loopOrder[0] == LoopN ? 0 : 1;
auto n_idx = let("n_idx",
(group_id(n_group_idx) * local_size(n_group_idx)
+ local_id(n_group_idx))
* (subgroup_dim == n_var ? n_blk / strategy.subgroupSize
: n_blk));
auto k_idx = def("k_idx", k.type(), 0);
auto offset_A = arg("offset_A");
auto offset_B = arg("offset_B");
auto offset_C = arg("offset_C");
std::vector<expr_t> C_idxs = {m_idx, n_idx};
if (problem.batch == BatchMode::Strided) {
struct info_t {
info_t(expr_t size, expr_t idiv_magic)
: size(std::move(size))
, idiv_magic(std::move(idiv_magic)) {}
expr_t size;
expr_t idiv_magic;
};
auto info = [&]() {
std::vector<info_t> ret;
ret.reserve(problem.batchDims - 1);
for (int i = 0; i < problem.batchDims - 1; i++) {
std::string i_s = std::to_string(i);
ret.emplace_back(
arg("batch_size" + i_s), arg("batch_magic" + i_s));
}
return ret;
}();
auto id = let("batch_id" + std::to_string(problem.batchDims - 1),
group_id(2) * local_size(2) + local_id(2));
for (int i = problem.batchDims - 1; i >= 0; i--) {
std::string i_s = std::to_string(i);
auto idx = let("batch_idx" + i_s, [&]() {
if (i == 0) return id;
auto id_next = let("batch_id" + std::to_string(i - 1),
ternary_idiv(id, info[i - 1].size,
info[i - 1].idiv_magic));
auto ret = id - info[i - 1].size * id_next;
id = id_next;
return ret;
}());
C_idxs.emplace_back(idx);
offset_A = offset_A + idx * arg("stride_A" + i_s);
offset_B = offset_B + idx * arg("stride_B" + i_s);
offset_C = offset_C + idx * arg("stride_C" + i_s);
}
}
global_tensor_t A_base {arg("A"), into_ir(problem.Ta_ext), offset_A,
{{m_var, m_idx}, {k_var, k_idx}}, {{m_var, m}, {k_var, k}},
get_strides(problem.A.layout, A_vars, arg("lda")), {}};
global_tensor_t B_base {arg("B"), into_ir(problem.Tb_ext), offset_B,
{{k_var, k_idx}, {n_var, n_idx}}, {{k_var, k}, {n_var, n}},
get_strides(problem.B.layout, B_vars, arg("ldb")), {}};
global_tensor_t C_base {arg("C"), into_ir(problem.Tc_ext), offset_B,
{{m_var, m_idx}, {n_var, n_idx}}, {{m_var, m}, {n_var, n}},
get_strides(problem.C.layout, C_vars, arg("ldc")),
{{m_var, m_blk}, {n_var, n_blk}}};
basic_iterator_t kloop_it(A_base, strategy.ka_prefetch,
strategy.ka_load, B_base, strategy.kb_prefetch,
strategy.kb_load, C_base);
auto store_C = [&]() {
apply_post_ops(problem.postOps.ops, C, C_idxs, {m_var, n_var});
store(kloop_it.C_store(), C, {}, {C_store_transform.cache_hint});
};
tensor_config_t A_load(
kloop_it.A_load(), A_load_transform, strategy.A_copies);
tensor_config_t B_load(
kloop_it.B_load(), B_load_transform, strategy.B_copies);
auto prefetchA = strategy.prefetchA
? round_down(strategy.prefetchA, strategy.ka_prefetch)
: 0;
if (prefetchA != strategy.prefetchA)
dsl_warning() << "Unimplemented partial A tile prefetch, modifying "
"prefetch distance "
<< strategy.prefetchA << " -> " << prefetchA;
auto prefetchB = strategy.prefetchB
? round_down(strategy.prefetchB, strategy.kb_prefetch)
: 0;
if (prefetchB != strategy.prefetchB)
dsl_warning() << "Unimplemented partial B tile prefetch, modifying "
"prefetch distance "
<< strategy.prefetchB << " -> " << prefetchB;
k_loop_config_t k_loop_main {k_blk, prefetchA, prefetchB, kloop_it,
std::move(A_load), std::move(B_load), A_prefetch_transform,
B_prefetch_transform, C};
dsl_assert(k_loop_main.A_load_warmup() % kloop_it.A_load().tile[k_var]
== 0);
dsl_assert(k_loop_main.B_load_warmup() % kloop_it.B_load().tile[k_var]
== 0);
tensor_config_t A_load_short(kloop_it.A_load(), A_load_transform, 1);
tensor_config_t B_load_short(kloop_it.B_load(), B_load_transform, 1);
auto k_blk_short
= (int)lcm(A_load_short.tile[k_var], B_load_short.tile[k_var]);
k_loop_config_t k_loop_short {k_blk_short, 0, 0, kloop_it,
std::move(A_load_short), std::move(B_load_short),
A_prefetch_transform, B_prefetch_transform, std::move(C)};
dsl_assert(k_loop_short.k_warmup() == 0);
if (problem.A.alignment) {
assume(arg("lda") % (problem.A.alignment / problem.Ta_ext) == 0);
}
if (problem.B.alignment) {
assume(arg("ldb") % (problem.B.alignment / problem.Tb_ext) == 0);
}
if (problem.C.alignment) {
assume(arg("ldc") % (problem.C.alignment / problem.Tc_ext) == 0);
}
_if(kloop_it.is_inbounds(0), [&]() {
_if(k >= k_loop_main.k_warmup(), [&]() {
build_k_loop(k_loop_main);
}, [&]() { build_k_loop(k_loop_short); });
store_C();
});
return end_kernel();
}
struct k_loop_config_t {
int64_t k_blk;
int64_t A_prefetch_warmup; int64_t B_prefetch_warmup; basic_iterator_t kloop_it;
tensor_config_t A_load;
tensor_config_t B_load;
transform_t A_prefetch_transform;
transform_t B_prefetch_transform;
tensor_t C;
int64_t A_load_warmup() const {
return A_load.layout.elems(k_var) - A_load.tile[k_var];
}
int64_t B_load_warmup() const {
return B_load.layout.elems(k_var) - B_load.tile[k_var];
}
int64_t k_warmup() const {
return std::max({A_load_warmup(), B_load_warmup(),
A_prefetch_warmup, B_prefetch_warmup});
}
};
void build_k_loop(const k_loop_config_t &cfg) {
auto k_blk = cfg.k_blk;
auto kloop_it = cfg.kloop_it;
auto &C = cfg.C;
tensor_t A = def("A_blk", cfg.A_load.layout);
tensor_t B = def("B_blk", cfg.B_load.layout);
auto mma_k_blk
= std::min(cfg.A_load.tile[k_var], cfg.B_load.tile[k_var]);
auto pipeline_idx
= [&](int64_t loop_idx, int64_t warmup_size, int64_t period) {
return (loop_idx + warmup_size) % period;
};
auto A_prefetch_blk
= cfg.A_prefetch_warmup ? kloop_it.A_prefetch().tile[k_var] : 0;
auto A_prefetch = [&](int64_t k_unroll_idx) {
if (cfg.A_prefetch_warmup == 0) return;
auto idx = pipeline_idx(
k_unroll_idx, cfg.A_prefetch_warmup, A_prefetch_blk);
if (idx % A_prefetch_blk != 0) return;
prefetch(kloop_it.A_prefetch(), {{k_var, 0}},
{cfg.A_prefetch_transform.cache_hint});
kloop_it.A_prefetch_inc(A_prefetch_blk);
};
auto A_load_blk = cfg.A_load.tile[k_var];
auto A_load = [&](int64_t k_unroll_idx) {
auto idx = pipeline_idx(k_unroll_idx, cfg.A_load_warmup(),
cfg.A_load.layout.elems(k_var));
if (idx % A_load_blk != 0) return;
load(A.sub(cfg.A_load.tile, {{k_var, idx}}), kloop_it.A_load(),
{{k_var, 0}}, {cfg.A_load.transform.cache_hint});
kloop_it.A_load_inc(A_load_blk);
};
auto B_prefetch_blk
= cfg.B_prefetch_warmup ? kloop_it.B_prefetch().tile[k_var] : 0;
auto B_prefetch = [&](int64_t k_unroll_idx) {
if (cfg.B_prefetch_warmup == 0) return;
auto idx = pipeline_idx(
k_unroll_idx, cfg.B_prefetch_warmup, B_prefetch_blk);
if (idx % B_prefetch_blk != 0) return;
prefetch(kloop_it.B_prefetch(), {{k_var, 0}},
{cfg.B_prefetch_transform.cache_hint});
kloop_it.B_prefetch_inc(B_prefetch_blk);
};
auto B_load_blk = cfg.B_load.tile[k_var];
auto B_load = [&](int64_t k_unroll_idx) {
auto idx = pipeline_idx(k_unroll_idx, cfg.B_load_warmup(),
cfg.B_load.layout.elems(k_var));
if (idx % B_load_blk != 0) return;
load(B.sub(cfg.B_load.tile, {{k_var, idx}}), kloop_it.B_load(),
{{k_var, 0}}, {cfg.B_load.transform.cache_hint});
kloop_it.B_load_inc(B_load_blk);
};
auto k_unroll_blk = [&]() {
auto ret = k_blk;
for (auto v :
{A_prefetch_blk, A_load_blk, B_prefetch_blk, B_load_blk}) {
ret = gcd(ret, v);
}
return ret;
}();
auto k_body
= [&](int64_t k_offset, bool do_A_prefetch, bool do_B_prefetch,
bool do_A_load, bool do_B_load, bool do_mma) {
if (do_A_prefetch) { A_prefetch(k_offset); }
if (do_B_prefetch) { B_prefetch(k_offset); }
if (do_A_load) { A_load(k_offset); }
if (do_B_load) { B_load(k_offset); }
if (do_mma) {
if (k_offset % mma_k_blk == 0) {
tile_t tile = C.layout.tile();
tile[k_var] = mma_k_blk;
mma(C, A, B, tile, {{k_var, k_offset}}, strategy.systolic);
}
}
};
auto warmup = cfg.k_warmup();
for (auto k_unroll_idx = -warmup; k_unroll_idx < 0;
k_unroll_idx += k_unroll_blk) {
bool A_prefetch = k_unroll_idx + cfg.A_prefetch_warmup >= 0;
bool B_prefetch = k_unroll_idx + cfg.B_prefetch_warmup >= 0;
bool A_load = k_unroll_idx + cfg.A_load_warmup() >= 0;
bool B_load = k_unroll_idx + cfg.B_load_warmup() >= 0;
bool do_mma = false;
k_body(k_unroll_idx, A_prefetch, B_prefetch, A_load, B_load,
do_mma);
}
_while(kloop_it.is_inbounds(warmup), [&]() {
for (int64_t k_unroll_idx = 0; k_unroll_idx < k_blk;
k_unroll_idx += k_unroll_blk) {
k_body(k_unroll_idx, cfg.A_prefetch_warmup,
cfg.B_prefetch_warmup, true, true, true);
}
kloop_it.kloop_inc(k_blk);
});
auto tail_end = round_up(warmup, k_blk);
for (int64_t k_unroll_idx = 0; k_unroll_idx < tail_end;
k_unroll_idx += k_unroll_blk) {
bool A_prefetch = k_unroll_idx + cfg.A_prefetch_warmup < tail_end;
bool B_prefetch = k_unroll_idx + cfg.B_prefetch_warmup < tail_end;
bool A_load = k_unroll_idx + cfg.A_load_warmup() < tail_end;
bool B_load = k_unroll_idx + cfg.B_load_warmup() < tail_end;
k_body(k_unroll_idx, A_prefetch, B_prefetch, A_load, B_load, true);
}
}
const GEMMProblem &problem;
const GEMMStrategy &strategy;
};
kernel_t make_kernel(const generator_dsl_desc_t &desc) {
ir::constraint_set_t cset;
ir::ir_context_t ctx(desc.options, cset);
ir::trace_start();
auto k = generator_dsl_t(desc).build(desc.kernel_iface(), ctx);
ir::trace_pass("build generator_dsl_t", k.body, ctx);
k.body = ir::simplify(k.body, ctx);
k.body = ir::inject_send(k.body, ctx);
k.body = ir::fixup_if_conditions(k.body, ctx);
k.body = ir::eliminate_common_subexprs(
k.body, ctx, desc.strategy.GRFs * ctx.hw().grf_size());
return k;
}
GEMMSTONE_NAMESPACE_END