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

Update to example_hstu_attention_fwd.cpp

Browse Source
This commit is contained in:
Qianfeng Zhang
2026-06-22 09:22:10 +00:00
parent e349a39f3d
commit db5cb5d19d
2 changed files with 128 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files

120
example/ck_tile/18_hstu_attention/example_helper.hpp Normal file
View File

@@ -0,0 +1,120 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include <iostream>
#include <stdexcept>
template <typename T>
static void dumpBufferToFile(const char* fileName, T* data, size_t dataNumItems)
{
std::ofstream outFile(fileName, std::ios::binary);
if(outFile)
{
outFile.write(reinterpret_cast<char*>(data), dataNumItems * sizeof(T));
outFile.close();
printf("Write output to file %s\n", fileName);
}
else
{
printf("Could not open file %s for writing\n", fileName);
}
}
template <typename T>
static void readDataToBufferFromFile(T* data, size_t dataNumItems, const std::string& fileName)
{
std::ifstream infile(fileName, std::ios::binary);
if(infile)
{
try
{
infile.read(reinterpret_cast<char*>(data), dataNumItems * sizeof(T));
infile.close();
}
catch(const std::runtime_error& e)
{
throw e;
};
}
else
{
throw std::runtime_error("could not open the file for reading");
}
}
static std::vector<int> get_integers_from_string(std::string srcStr)
{
std::vector<int> integers;
std::size_t pos = 0;
std::size_t new_pos;
new_pos = srcStr.find(',', pos);
while(new_pos != std::string::npos)
{
std::string sliceStr = srcStr.substr(pos, new_pos - pos);
int len = std::stoi(sliceStr);
integers.push_back(len);
pos = new_pos + 1;
new_pos = srcStr.find(',', pos);
};
std::string sliceStr = srcStr.substr(pos);
if(!sliceStr.empty())
{
int len = std::stoi(sliceStr);
integers.push_back(len);
};
return (integers);
};
static std::vector<float> get_floats_from_string(std::string srcStr)
{
std::vector<float> values;
std::size_t pos = 0;
std::size_t new_pos;
new_pos = srcStr.find(',', pos);
while(new_pos != std::string::npos)
{
std::string sliceStr = srcStr.substr(pos, new_pos - pos);
float val = std::stof(sliceStr);
values.push_back(val);
pos = new_pos + 1;
new_pos = srcStr.find(',', pos);
};
std::string sliceStr = srcStr.substr(pos);
if(!sliceStr.empty())
{
float val = std::stof(sliceStr);
values.push_back(val);
};
return (values);
};
template <typename T>
static void supplement_array_by_last_element(std::vector<T>& arr, int target_num_elements)
{
if(static_cast<int>(arr.size()) < target_num_elements)
{
T last_val = arr.back();
for(int i = arr.size(); i < target_num_elements; i++)
arr.push_back(last_val);
};
};

128
example/ck_tile/18_hstu_attention/example_hstu_attention_fwd.cpp
View File

@@ -2,9 +2,7 @@
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
#include <array>
#include <cstring>
#include <functional>
#include <numeric>
#include <string>
#include <ostream>
#include <string>
#include <tuple>
@@ -30,43 +28,7 @@
#include "hstu_attention_host_util.hpp"
#include "hstu_attention_api.hpp"
template <typename T>
void dumpBufferToFile(const char* fileName, T* data, size_t dataNumItems)
{
std::ofstream outFile(fileName, std::ios::binary);
if(outFile)
{
outFile.write(reinterpret_cast<char*>(data), dataNumItems * sizeof(T));
outFile.close();
printf("Write output to file %s\n", fileName);
}
else
{
printf("Could not open file %s for writing\n", fileName);
}
}
template <typename T>
void readDataToBufferFromFile(T* data, size_t dataNumItems, const std::string& fileName)
{
std::ifstream infile(fileName, std::ios::binary);
if(infile)
{
try
{
infile.read(reinterpret_cast<char*>(data), dataNumItems * sizeof(T));
infile.close();
}
catch(const std::runtime_error& e)
{
throw e;
};
}
else
{
throw std::runtime_error("could not open the file for reading");
}
}
#include "example_helper.hpp"
template <typename T>
std::ostream& operator<<(std::ostream& os, const std::vector<T>& v)
@@ -130,80 +92,6 @@ auto create_args(int argc, char* argv[])
return std::make_tuple(result, arg_parser);
}
static std::vector<int> get_integers_from_string(std::string srcStr)
{
std::vector<int> integers;
std::size_t pos = 0;
std::size_t new_pos;
new_pos = srcStr.find(',', pos);
while(new_pos != std::string::npos)
{
std::string sliceStr = srcStr.substr(pos, new_pos - pos);
int len = std::stoi(sliceStr);
integers.push_back(len);
pos = new_pos + 1;
new_pos = srcStr.find(',', pos);
};
std::string sliceStr = srcStr.substr(pos);
if(!sliceStr.empty())
{
int len = std::stoi(sliceStr);
integers.push_back(len);
};
return (integers);
};
static std::vector<float> get_floats_from_string(std::string srcStr)
{
std::vector<float> values;
std::size_t pos = 0;
std::size_t new_pos;
new_pos = srcStr.find(',', pos);
while(new_pos != std::string::npos)
{
std::string sliceStr = srcStr.substr(pos, new_pos - pos);
float val = std::stof(sliceStr);
values.push_back(val);
pos = new_pos + 1;
new_pos = srcStr.find(',', pos);
};
std::string sliceStr = srcStr.substr(pos);
if(!sliceStr.empty())
{
float val = std::stof(sliceStr);
values.push_back(val);
};
return (values);
};
template <typename T>
void supplement_array_by_last_element(std::vector<T>& arr, int target_num_elements)
{
if(static_cast<int>(arr.size()) < target_num_elements)
{
T last_val = arr.back();
for(int i = arr.size(); i < target_num_elements; i++)
arr.push_back(last_val);
};
};
// threshold for different dtypes
template <typename DataType>
auto get_elimit()
@@ -223,7 +111,7 @@ auto get_elimit<ck_tile::bf16_t>()
}
template <typename InOutDataType>
bool run_no_group_hstu(const ck_tile::ArgParser& arg_parser, bool is_jagged)
bool run_no_group_hstu_forward(const ck_tile::ArgParser& arg_parser, bool is_jagged)
{
using CompDataType = typename HstuAttentionFwdTypeConfig<InOutDataType>::CompDataType;
@@ -690,7 +578,7 @@ bool run_no_group_hstu(const ck_tile::ArgParser& arg_parser, bool is_jagged)
}
template <typename InOutDataType>
bool run_group_hstu(const ck_tile::ArgParser& arg_parser, int num_group)
bool run_group_hstu_forward(const ck_tile::ArgParser& arg_parser, int num_group)
{
using CompDataType = typename HstuAttentionFwdTypeConfig<InOutDataType>::CompDataType;
@@ -1180,11 +1068,11 @@ int main(int argc, char* argv[])
{
if(data_type == "fp16")
{
return run_group_hstu<ck_tile::half_t>(arg_parser, num_group) ? 0 : -2;
return run_group_hstu_forward<ck_tile::half_t>(arg_parser, num_group) ? 0 : -2;
}
else if(data_type == "bf16")
{
return run_group_hstu<ck_tile::bf16_t>(arg_parser, num_group) ? 0 : -2;
return run_group_hstu_forward<ck_tile::bf16_t>(arg_parser, num_group) ? 0 : -2;
}
}
else
@@ -1193,11 +1081,11 @@ int main(int argc, char* argv[])
if(data_type == "fp16")
{
return run_no_group_hstu<ck_tile::half_t>(arg_parser, is_jagged) ? 0 : -2;
return run_no_group_hstu_forward<ck_tile::half_t>(arg_parser, is_jagged) ? 0 : -2;
}
else if(data_type == "bf16")
{
return run_no_group_hstu<ck_tile::bf16_t>(arg_parser, is_jagged) ? 0 : -2;
return run_no_group_hstu_forward<ck_tile::bf16_t>(arg_parser, is_jagged) ? 0 : -2;
}
};