Skip to content

Commit abfeb88

Browse files
author
ssjia
committed
[ET-VK][ez] Implement helper functions to get fastest moving dim
Pull Request resolved: #17107 Add C++ and GLSL helpers to query the fastest moving dimension (the dimension with stride 1 in buffer layout). This is useful for optimizing memory access patterns in shaders, as iterating along the fastest moving dimension maximizes cache locality. The C++ `fastest_whcn_dim()` method accounts for block-transposed layouts by returning `outer_packed_dim` instead of `packed_dim` when applicable. A corresponding GLSL macro extracts this info from the hashed layout. ghstack-source-id: 338638547 @exported-using-ghexport Differential Revision: [D92061369](https://our.internmc.facebook.com/intern/diff/D92061369/)
1 parent 1781382 commit abfeb88

3 files changed

Lines changed: 16 additions & 0 deletions

File tree

backends/vulkan/runtime/api/containers/Tensor.h

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -523,6 +523,16 @@ class vTensor final {
523523
return packed_dim_info_.packed_dim;
524524
}
525525

526+
/*
527+
* Returns the WHCN index of the fastest moving dimension (dim_order[0]).
528+
* This is the dimension with stride 1 in the buffer layout.
529+
* Note: dim_order_ is in NCHW order, so we convert to WHCN (3 - nchw_dim).
530+
*/
531+
inline int32_t fastest_whcn_dim() const {
532+
return packed_dim_info_.block_transposed ? packed_dim_info_.outer_packed_dim
533+
: packed_dim_info_.packed_dim;
534+
}
535+
526536
inline const PackedDimInfo& packed_dim_info() const {
527537
return packed_dim_info_;
528538
}

backends/vulkan/runtime/graph/ComputeGraph.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -456,6 +456,10 @@ class ComputeGraph final {
456456
return values_.at(idx).toConstTensor().packed_dim();
457457
}
458458

459+
inline int32_t fastest_whcn_dim_of(const ValueRef idx) const {
460+
return values_.at(idx).toConstTensor().fastest_whcn_dim();
461+
}
462+
459463
inline const api::PackedDimInfo& packed_dim_info_of(
460464
const ValueRef idx) const {
461465
return values_.at(idx).toConstTensor().packed_dim_info();

backends/vulkan/runtime/graph/ops/glsl/indexing.glslh

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,8 @@ bool is_channels_last(const int hashed_layout) {
7070
return layout_id(hashed_layout) == CHANNELS_LAST_BUFFER_LAYOUT_ID;
7171
}
7272

73+
#define get_fastest_dim(layout) ((layout) & 0xF)
74+
7375
// Extract packed dim info from hashed_layout (bits 16-31)
7476
// These match the format created by create_hashed_layout() in Tensor.cpp
7577
#define get_packed_dim(layout) (((layout) >> 16) & 0xF)

0 commit comments

Comments
 (0)