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

save tmp

Browse Source
This commit is contained in:
Max Podkorytov
2025-11-10 13:11:13 -06:00
parent 76c4c12f59
commit 1e2dac15a1
7 changed files with 365 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
include/ck_tile/core/algorithm/coordinate_transform.hpp
View File

@@ -1326,8 +1326,21 @@ struct xor_t : public base_transform<2, 2>
idx_low(number<0>{}) = idx_up[number<0>{}];
// original (1.0 rate):
// idx_low(number<1>{}) =
// idx_up[number<1>{}] ^ (idx_up[number<0>{}] % up_lengths_[number<1>{}]);
// 0.5 rate
// idx_low(number<1>{}) =
// idx_up[number<1>{}] ^ (2 * idx_up[number<0>{}]) % up_lengths_[number<1>{}];
idx_low(number<1>{}) =
idx_up[number<1>{}] ^ (idx_up[number<0>{}] % up_lengths_[number<1>{}]);
idx_up[number<1>{}] ^ (2 * idx_up[number<0>{}] % up_lengths_[number<1>{}]);
// if (threadIdx.x < 64)
// {
// printf("lane: %u | idx_low: (%d, %d) | idx_up: (%d, %d)\n",
// __lane_id(),
// idx_low[number<0>{}], idx_low[number<1>{}],
// idx_up[number<0>{}], idx_up[number<1>{}]);
// }
}
template <typename LowIdxDiff, typename UpIdxDiff, typename LowIdx, typename UpIdx>

164
include/ck_tile/core/tensor/distributed_lds_tensor.hpp Normal file
View File

@@ -0,0 +1,164 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "ck_tile/core/config.hpp"
#include "ck_tile/core/numeric/integer.hpp"
#include "ck_tile/core/numeric/integral_constant.hpp"
#include "ck_tile/core/numeric/numeric.hpp"
#include "ck_tile/core/utility/type_traits.hpp"
#include "ck_tile/core/tensor/tensor_descriptor.hpp"
#include "ck_tile/core/tensor/tile_distribution.hpp"
#include "ck_tile/core/arch/arch.hpp"
namespace ck_tile
{
// A static distributed tensor backed by shared memory (LDS),
// with indexing defined by a naive packed tensor descriptor of the Y-dimensions
// of the provided StaticTileDistribution.
//
// Notes:
// - This does NOT allocate LDS. Callers must provide a __shared__-allocated buffer
// of size get_lds_buffer_size() elements of DataType.
// - Index mapping uses a naive packed descriptor derived from the Y lengths of the
// distribution to compute a linear offset into LDS.
// - Distributed indexing API mirrors static_distributed_tensor (operator[]/operator()).
//
// Example (inside __global__ kernel):
// using Dist = decltype(make_static_tile_distribution(my_encoding{}));
// using Tensor = ck_tile::static_distributed_lds_tensor<float, Dist>;
//
// // Compile-time LDS size in elements
// __shared__ float lds[Tensor::get_lds_buffer_size()];
//
// Tensor t{lds};
// // write using distributed indices (compile-time indices)
// t(ck_tile::detail::make_tile_distributed_index(ck_tile::sequence<0>{},
// ck_tile::sequence<0>{})) = 1.0f;
// ck_tile::block_sync_lds();
//
// // read back similarly
// auto v = t[ck_tile::detail::make_tile_distributed_index(ck_tile::sequence<0>{},
// ck_tile::sequence<0>{})];
//
template <typename DataType_, typename StaticTileDistribution_>
struct static_distributed_lds_tensor
{
using DataType = remove_cvref_t<DataType_>;
using StaticTileDistribution = remove_cvref_t<StaticTileDistribution_>;
static_assert(StaticTileDistribution::is_static(),
"StaticTileDistribution must be fully known at compile time");
// Build a naive packed Y descriptor based on the distribution's Y lengths.
// This mirrors the usage pattern from static_distributed_tensor where offsets are
// computed through the Ys-to-D descriptor, but here we explicitly construct a naive
// packed descriptor to satisfy the requirement.
using YsDescriptorOriginal =
remove_cvref_t<decltype(StaticTileDistribution{}.get_ys_to_d_descriptor())>;
using YsNaivePackedDescriptor =
remove_cvref_t<decltype(make_naive_tensor_descriptor_packed(YsDescriptorOriginal{}.get_lengths()))>;
using ThreadTensorDesc = YsNaivePackedDescriptor;
static constexpr index_t PackedSize =
ck_tile::numeric_traits<remove_cvref_t<DataType>>::PackedSize;
static constexpr index_t kLdsElementSpaceSize = ThreadTensorDesc{}.get_element_space_size();
static_assert(kLdsElementSpaceSize > 0, "Make sure tile distribution is valid for LDS");
// Number of DataType elements required in LDS buffer (considering PackedSize for DataType)
CK_TILE_HOST_DEVICE static constexpr index_t get_lds_buffer_size()
{
return kLdsElementSpaceSize / PackedSize;
}
CK_TILE_HOST_DEVICE static constexpr auto get_num_of_dimension()
{
return StaticTileDistribution::get_num_of_dimension_x();
}
CK_TILE_HOST_DEVICE static constexpr auto get_lengths()
{
return StaticTileDistribution::get_lengths();
}
CK_TILE_HOST_DEVICE static constexpr auto get_tile_distribution()
{
return StaticTileDistribution{};
}
CK_TILE_HOST_DEVICE static constexpr auto get_distributed_spans()
{
return StaticTileDistribution::get_distributed_spans();
}
// Construct with an LDS pointer (caller-provided __shared__ storage).
CK_TILE_HOST_DEVICE explicit constexpr static_distributed_lds_tensor(DataType* lds_ptr)
: p_lds_{lds_ptr}
{
}
// Bind or rebind LDS pointer at runtime if needed.
CK_TILE_DEVICE void bind(DataType* lds_ptr) { p_lds_ = lds_ptr; }
CK_TILE_DEVICE constexpr const DataType* data() const { return p_lds_; }
CK_TILE_DEVICE constexpr DataType* data() { return p_lds_; }
// Distributed indexing: const access
template <typename TileDistributedIndices>
CK_TILE_HOST_DEVICE constexpr const DataType& operator[](TileDistributedIndices) const
{
static_assert(is_static_v<TileDistributedIndices>,
"TileDistributedIndices must be static");
constexpr auto y_idx = get_tile_distribution().get_y_indices_from_distributed_indices(
TileDistributedIndices{});
constexpr index_t linear = ThreadTensorDesc{}.calculate_offset(y_idx) / PackedSize;
return p_lds_[linear];
}
// Distributed indexing: mutable access
template <typename TileDistributedIndices>
CK_TILE_HOST_DEVICE constexpr DataType& operator()(TileDistributedIndices)
{
static_assert(is_static_v<TileDistributedIndices>,
"TileDistributedIndices must be static");
constexpr auto y_idx = get_tile_distribution().get_y_indices_from_distributed_indices(
TileDistributedIndices{});
constexpr index_t linear = ThreadTensorDesc{}.calculate_offset(y_idx) / PackedSize;
return p_lds_[linear];
}
// Returns the naive Y-descriptor (type-only, constexpr-friendly)
CK_TILE_HOST_DEVICE static constexpr auto get_y_naive_packed_descriptor()
{
return ThreadTensorDesc{};
}
private:
DataType* p_lds_ = nullptr; // shared (LDS) pointer provided by caller
};
// Helper: make function
template <typename DataType, typename StaticTileDistribution>
CK_TILE_HOST_DEVICE constexpr auto
make_static_distributed_lds_tensor(DataType* lds_ptr, const StaticTileDistribution&)
{
return static_distributed_lds_tensor<remove_cvref_t<DataType>,
remove_cvref_t<StaticTileDistribution>>{lds_ptr};
}
// Utility to query required LDS size (in elements of DataType) for a given distribution
template <typename DataType, typename StaticTileDistribution>
CK_TILE_HOST_DEVICE constexpr index_t
get_required_lds_size_elems(const StaticTileDistribution&)
{
using tensor_t = static_distributed_lds_tensor<remove_cvref_t<DataType>,
remove_cvref_t<StaticTileDistribution>>;
return tensor_t::get_lds_buffer_size();
}
} // namespace ck_tile

18
include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp
View File

@@ -59,6 +59,10 @@ struct GemmPipelineAgBgCrImplBase
const ElementFunction& element_func) const
{
const auto block_tile_tmp = tile_elementwise_in(element_func, src_block_tile);
ignore = lds_tile_window;
ignore = src_block_tile;
ignore = element_func;
ignore = block_tile_tmp;
store_tile(lds_tile_window, block_tile_tmp);
}
@@ -66,6 +70,10 @@ struct GemmPipelineAgBgCrImplBase
CK_TILE_DEVICE void LocalPrefill(DstTileWindow& lds_tile_window,
const SrcBlockTile& src_block_tile) const
{
ignore = lds_tile_window;
ignore = src_block_tile;
// CK_PRINT<typename remove_cvref_t<decltype(src_block_tile)>::StaticTileDistribution::DstrEncode>();
store_tile(lds_tile_window, src_block_tile);
}
@@ -74,10 +82,12 @@ struct GemmPipelineAgBgCrImplBase
const SrcTileWindow& lds_tile_window,
bool_constant<LoadTranspose> = {}) const
{
if constexpr(LoadTranspose)
dst_block_tile = load_tile_transpose(lds_tile_window);
else
load_tile(dst_block_tile, lds_tile_window);
// if constexpr(LoadTranspose)
// dst_block_tile = load_tile_transpose(lds_tile_window);
// else
// load_tile(dst_block_tile, lds_tile_window);
ignore = dst_block_tile;
ignore = lds_tile_window;
}
CK_TILE_DEVICE auto GetABLdsTensorViews(void* p_smem) const

8
include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp
View File

@@ -444,12 +444,16 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3<Problem>
constexpr auto b_lds_load_tile_distr =
make_static_tile_distribution(BlockGemm::MakeBBlockDistributionEncode());
// CK_PRINT<typename decltype(a_lds_load_tile_distr)::DstrEncode>();
// A DRAM tile window for load
// A LDS tile window for store
// A LDS tile for block GEMM
auto&& [a_copy_dram_window, a_copy_lds_window, a_lds_gemm_window] =
Base::GetAWindows(a_dram_block_window_tmp, a_lds_block, a_lds_load_tile_distr);
// CK_PRINT<typename remove_cvref_t<decltype(a_copy_dram_window[number<0>{}])>::Base::TileDstr::DstrEncode>();
// B DRAM tile window for load
// B LDS tile window for store
// B LDS tile for block GEMM
@@ -478,6 +482,10 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3<Problem>
auto elementwise_As_res =
load_tile_with_elementwise(a_copy_dram_window, a_element_func);
// CK_PRINT<typename decltype(elementwise_As_res)::StaticTileDistribution::DstrEncode>();
// CK_PRINT<typename remove_cvref_t<decltype(a_copy_dram_window[number<0>{}])>::Base::TileDstr::DstrEncode>();
// Move each A — the enhanced function move_tile_window is executed, which takes a tuple
// as input.
move_tile_window(a_copy_dram_window, a_dram_tile_window_step);

2
include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp
View File

@@ -545,6 +545,8 @@ struct UniversalGemmBasePolicy
VecLoadSize,
getATileAccessPattern(),
NumWaveGroups>;
CK_PRINT<TileEncodingPattern>();
CK_PRINT<typename remove_cvref_t<decltype(TileEncodingPattern::make_2d_static_tile_distribution())>::DstrEncode>();
return TileEncodingPattern::make_2d_static_tile_distribution();
}
// Tile: KPerBlock X MPerBlock