mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-02 04:31:25 +00:00
b661eab573
[CK TILE] Refactor MX FLATMM example Refactor the MX FLATMM example to support more pipelines across different architectures. This work facilitates the NPI team roadmap.
35 lines
987 B
C++
35 lines
987 B
C++
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
#pragma once
|
|
|
|
#include <string>
|
|
|
|
#include "ck_tile/core.hpp"
|
|
#include "ck_tile/host/kernel_launch.hpp"
|
|
#include "ck_tile/ops/epilogue.hpp"
|
|
#include "ck_tile/ops/flatmm.hpp"
|
|
#include "ck_tile/ops/gemm.hpp"
|
|
|
|
#include "mx_flatmm_arch_traits.hpp"
|
|
|
|
template <typename MXFlatmmArchTraits,
|
|
typename ADataType,
|
|
typename BDataType,
|
|
typename DsDatatype,
|
|
typename AccDataType,
|
|
typename CDataType,
|
|
typename ALayout,
|
|
typename BLayout,
|
|
typename DsLayout,
|
|
typename ELayout,
|
|
typename ScaleM,
|
|
typename ScaleN,
|
|
bool persistent,
|
|
typename CDEElementWise,
|
|
bool Splitk,
|
|
bool HasHotLoop,
|
|
ck_tile::TailNumber TailNum>
|
|
float mx_flatmm_calc(const ck_tile::ScaleFlatmmHostArgs<ScaleM, ScaleN>& args,
|
|
const ck_tile::stream_config& s);
|