ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-15 11:34:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add last steps: activations functions

Browse Source
This commit is contained in:
Damien Lejeune
2026-01-29 08:30:45 -05:00
parent da895cdd88
commit 6ea40157f1
4 changed files with 190 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
include/ck_tile/host/reference/reference_mhc.hpp
View File

@@ -7,6 +7,7 @@
#include <thread>
#include "ck_tile/core.hpp"
#include "ck_tile/host/host_tensor.hpp"
#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp"
namespace ck_tile {
@@ -14,17 +15,19 @@ namespace ck_tile {
template <typename XDataType,
typename PhiDataType,
typename YDataType,
typename ComputeDataType = float>
typename ComputeDataType = float,
typename Activation = element_wise::Sigmoid>
CK_TILE_HOST void reference_mhc(const HostTensor<XDataType>& x_b_nc, // [B, nC]
const HostTensor<PhiDataType>& phi_nc_out, // [nC, 2n+n²]
HostTensor<YDataType>& output_b_out, // [B, 2n+n²]
const HostTensor<PhiDataType>& phi_nc_out, // [nC, 2n+n^2]
HostTensor<YDataType>& output_b_out, // [B, 2n+n^2]
int n, // expansion factor
int C, // channels per stream
[[maybe_unused]] float r = 1.0f,
[[maybe_unused]] float alpha_pre = 1.0f,
[[maybe_unused]] float alpha_post = 1.0f,
[[maybe_unused]] float alpha_res = 1.0f,
[[maybe_unused]] float bias = 0.0f)
[[maybe_unused]] float bias = 0.0f,
Activation activation = Activation{})
{
const int B = x_b_nc.get_length(0);
const int nC = n * C;
@@ -43,7 +46,7 @@ CK_TILE_HOST void reference_mhc(const HostTensor<XDataType>& x_b_nc, // [B
// Step 2 & 3: Perform GEMM and apply elementwise operations
// Process H^{pre}: x * phi[:, 0:n] -> output[:, 0:n]
// Process H^{pre}: x * phi[:, 0:n] -> sigma(output[:, 0:n])
for(int out_idx = 0; out_idx < n; out_idx++)
{
ComputeDataType sum = 0.0f;
@@ -52,11 +55,14 @@ CK_TILE_HOST void reference_mhc(const HostTensor<XDataType>& x_b_nc, // [B
sum += type_convert<ComputeDataType>(x_b_nc(b, k)) *
type_convert<ComputeDataType>(phi_nc_out(k, out_idx));
}
// Apply: 1/r * alpha_pre * sum + bias
output_b_out(b, out_idx) = type_convert<YDataType>((alpha_pre / r) * sum + bias);
// Step 4: Apply activation σ(H^{pre})
ComputeDataType activated_value;
activation(activated_value, sum);
output_b_out(b, out_idx) =
type_convert<YDataType>((alpha_pre / r) * activated_value + bias);
}
// Process H^{post}: x * phi[:, n:2n] -> output[:, n:2n]
// Process H^{post}: x * phi[:, n:2n] -> 2*sigma(output[:, n:2n])
for(int out_idx = 0; out_idx < n; out_idx++)
{
ComputeDataType sum = 0.0f;
@@ -65,11 +71,14 @@ CK_TILE_HOST void reference_mhc(const HostTensor<XDataType>& x_b_nc, // [B
sum += type_convert<ComputeDataType>(x_b_nc(b, k)) *
type_convert<ComputeDataType>(phi_nc_out(k, n + out_idx));
}
// Apply: 1/r * alpha_post * sum + bias
output_b_out(b, n + out_idx) = type_convert<YDataType>((alpha_post / r) * sum + bias);
// Step 5: Apply 2*σ(H^{post})
ComputeDataType activated_value;
activation(activated_value, sum);
output_b_out(b, n + out_idx) =
type_convert<YDataType>((alpha_post / r) * 2.0f * activated_value + bias);
}
// Process H^{res}: x * phi[:, 2n:2n+n²] -> output[:, 2n:2n+n²]
// Process H^{res}: x * phi[:, 2n:2n+n^2] -> output[:, 2n:2n+n^2]
int n_squared = n * n;
for(int out_idx = 0; out_idx < n_squared; out_idx++)
{

49
include/ck_tile/ops/mhc/kernel/mhc_kernel_tile_v2.hpp
View File

@@ -8,6 +8,7 @@
#include "ck_tile/ops/mhc/pipeline/mhc_problem.hpp"
#include "ck_tile/ops/mhc/pipeline/mhc_default_policy.hpp"
#include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1.hpp"
#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp"
// Manifold Constrained Hyper Connection Kernel (True CK Tile Version):
// =====================================================================
@@ -20,15 +21,17 @@
namespace ck_tile {
template <typename Problem_,
typename Policy_ = MHCDefaultPolicy,
index_t B_ = 16, // Batch size (compile-time)
index_t N_ = 4, // Expansion factor (compile-time)
index_t C_ = 64, // Channels per stream (compile-time)
index_t KTile_ = 256> // K-tile size for shared memory (compile-time)
typename Policy_ = MHCDefaultPolicy,
index_t B_ = 16, // Batch size (compile-time)
index_t N_ = 4, // Expansion factor (compile-time)
index_t C_ = 64, // Channels per stream (compile-time)
index_t KTile_ = 256, // K-tile size for shared memory (compile-time)
typename Activation_ = element_wise::Sigmoid> // Activation function (compile-time)
struct ManifoldConstrainedHyperConnectionTiled
{
using Problem = ck_tile::remove_cvref_t<Problem_>;
using Policy = ck_tile::remove_cvref_t<Policy_>;
using Activation = ck_tile::remove_cvref_t<Activation_>;
using Problem = ck_tile::remove_cvref_t<Problem_>;
using Policy = ck_tile::remove_cvref_t<Policy_>;
using XDataType = ck_tile::remove_cvref_t<typename Problem::XDataType>;
using ComputeDataType = ck_tile::remove_cvref_t<typename Problem::ComputeDataType>;
@@ -241,16 +244,32 @@ struct ManifoldConstrainedHyperConnectionTiled
if(global_batch < B && global_col < output_dim)
{
// Determine alpha based on the actual output column
float alpha = (global_col < kN) ? alpha_pre
: (global_col < 2 * kN) ? alpha_post
: alpha_res;
constexpr auto i_j_idx = make_tuple(idx0, idx1);
ComputeDataType value = result_tile[i_j_idx];
// Step 4 & 5: Apply activation functions based on output section
if(global_col < kN)
{
// H^{pre}: Apply sigma(H^{pre})
ComputeDataType activated_value;
Activation{}(activated_value, value);
value = (alpha_pre / r) * activated_value + bias;
}
else if(global_col < 2 * kN)
{
// H^{post}: Apply 2*sigma(H^{post})
ComputeDataType activated_value;
Activation{}(activated_value, value);
value = (alpha_post / r) * 2.0f * activated_value + bias;
}
else
{
// H^{res}: No activation (will be Sinkhorn-Knopp later)
value = (alpha_res / r) * value + bias;
}
// Apply scaling and bias, then store: result = (alpha / r) * result + bias
constexpr auto i_j_idx = make_tuple(idx0, idx1);
const index_t global_idx = global_batch * output_dim + global_col;
p_output[global_idx] =
type_convert<YDataType>((alpha / r) * result_tile[i_j_idx] + bias);
p_output[global_idx] = type_convert<YDataType>(value);
}
});
});

16
test/ck_tile/mhc/test_mhc.cpp
View File

@@ -63,6 +63,22 @@ TYPED_TEST(TestCkTileMHC, TestBatchSize2N4C1024) { this->template RunBatchSizeTe
TYPED_TEST(TestCkTileMHC, TestBatchSize2N4C4096) { this->template RunBatchSizeTest<2, 4, 4096>(); }
// Test with different activation functions
TYPED_TEST(TestCkTileMHC, TestBatchSize16WithTanh)
{
this->template RunBatchSizeTestWithActivation<16, 4, 64, ck_tile::element_wise::TanH>();
}
TYPED_TEST(TestCkTileMHC, TestBatchSize16WithRelu)
{
this->template RunBatchSizeTestWithActivation<16, 4, 64, ck_tile::element_wise::Relu>();
}
TYPED_TEST(TestCkTileMHC, TestBatchSize16WithSilu)
{
this->template RunBatchSizeTestWithActivation<16, 4, 64, ck_tile::element_wise::Silu>();
}
TYPED_TEST(TestCkTileMHC, TestBatchSize16N4C4096)
{
this->template RunBatchSizeTest<16, 4, 4096>();

120
test/ck_tile/mhc/test_mhc_impl.hpp
View File

@@ -377,6 +377,126 @@ class TestCkTileMHC : public ::testing::Test
EXPECT_TRUE(pass);
}
// Test with specific batch size and custom activation function
template <int B = 16,
int n = 4,
int C = 64,
typename ActivationFunc = ck_tile::element_wise::Sigmoid>
void RunBatchSizeTestWithActivation()
{
const int nC = n * C; // Total input dimension
const int output_dim = 2 * n + n * n; // 2n + n^2
std::cout << "\n--- Testing batch size B=" << B << " (n=" << n << ", C=" << C
<< ") with activation: " << ActivationFunc::name << " ---" << std::endl;
// Allocate host tensors
ck_tile::HostTensor<float> h_x({B, nC});
ck_tile::HostTensor<float> h_phi({nC, output_dim});
ck_tile::HostTensor<float> h_output({B, output_dim});
// Initialize with random data
ck_tile::FillUniformDistribution<float>{-1.0f, 1.0f}(h_x);
ck_tile::FillUniformDistribution<float>{-0.5f, 0.5f}(h_phi);
h_output.SetZero();
// Allocate device memory
ck_tile::DeviceMem d_x_mem(h_x.get_element_space_size_in_bytes());
ck_tile::DeviceMem d_phi_mem(h_phi.get_element_space_size_in_bytes());
ck_tile::DeviceMem d_output_mem(h_output.get_element_space_size_in_bytes());
// Copy data to device
d_x_mem.ToDevice(h_x.data());
d_phi_mem.ToDevice(h_phi.data());
d_output_mem.ToDevice(h_output.data());
// Define block shape
using BlockShape = ck_tile::Generic2dBlockShape<ck_tile::sequence<1, 256>,
ck_tile::sequence<1, 256>,
ck_tile::sequence<1, 1>>;
using Problem = ck_tile::MHCProblem<float, float, float, BlockShape>;
// Use template parameters for B, n, C, and Activation (compile-time)
using KernelExpansionParallel =
ck_tile::ManifoldConstrainedHyperConnectionTiled<Problem,
ck_tile::MHCDefaultPolicy,
B,
n,
C,
256,
ActivationFunc>;
const ck_tile::index_t kBlockSize = KernelExpansionParallel::BlockSize();
const ck_tile::index_t kGridSize = (output_dim + 15) / 16;
constexpr ck_tile::index_t kBlockPerCu = 1;
const float r = 2.0f, alpha_pre = 1.5f, alpha_post = 2.5f, alpha_res = 3.5f, bias = 1.5f;
// Launch kernel
ck_tile::launch_kernel(
ck_tile::stream_config{nullptr, false, 0},
ck_tile::make_kernel<kBlockPerCu>(KernelExpansionParallel{},
kGridSize,
kBlockSize,
0,
static_cast<float*>(d_x_mem.GetDeviceBuffer()),
static_cast<float*>(d_phi_mem.GetDeviceBuffer()),
static_cast<float*>(d_output_mem.GetDeviceBuffer()),
r,
alpha_pre,
alpha_post,
alpha_res,
bias));
d_output_mem.FromDevice(h_output.data());
// Compute reference with the same activation function
ck_tile::HostTensor<float> h_output_ref({B, output_dim});
h_output_ref.SetZero();
ck_tile::reference_mhc<float, float, float, float, ActivationFunc>(h_x,
h_phi,
h_output_ref,
n,
C,
r,
alpha_pre,
alpha_post,
alpha_res,
bias,
ActivationFunc{});
// Validate
bool pass = ck_tile::check_err(
h_output, h_output_ref, "Error: Activation function mismatch!", 1e-3f, 1e-3f);
std::cout << " Result: " << (pass ? "PASS" : "FAIL") << std::endl;
if(!pass)
{
// Print first few values for debugging
std::cout << " First batch kernel output: [";
for(int i = 0; i < std::min(8, output_dim); i++)
{
std::cout << h_output(0, i);
if(i < std::min(8, output_dim) - 1)
std::cout << ", ";
}
std::cout << " ...]" << std::endl;
std::cout << " First batch reference: [";
for(int i = 0; i < std::min(8, output_dim); i++)
{
std::cout << h_output_ref(0, i);
if(i < std::min(8, output_dim) - 1)
std::cout << ", ";
}
std::cout << " ...]" << std::endl;
}
EXPECT_TRUE(pass);
}
// Test with multiple arbitrary batch sizes
void RunArbitraryBatchSizeTest()
{