Overview

The Layout<Shape, Stride> data structure is a basic template class that describes the memory layout of multi-dimensional tensors. It implements the mapping from the logical coordinate space to the one-dimensional memory address space based on the shape and stride information at compile time, providing fundamental support for complex tensor operations and hardware optimization. With the template metaprogramming technique, this class completes computation and code generation at compile time, thereby reducing runtime overhead.

A layout consists of two core components:

• Shape: defines the logical shape of data, such as the number of rows and columns in a two-dimensional matrix or the size of each dimension in a multi-dimensional tensor.
• Stride: defines the stride size of each dimension in memory, that is, the interval between adjacent elements in the same dimension in memory. The interval is measured in elements and corresponds to the dimension information of the shape.

For example, if the shape of a two-dimensional matrix is (4, 2) and the stride is (4, 1), it indicates that:

• The matrix has 4 rows and 2 columns.
• The stride in the column direction is 1, meaning that the interval between adjacent elements in each row is 1 element. The stride in the row direction is 4, meaning that the interval between the start addresses of adjacent rows is 4 elements.

Table 1 shows the one-dimensional memory address space view, and Table 2 shows the logical view of the two-dimensional matrix.

1

#include "kernel_operator_layout.h"

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

template <typename ShapeType, typename StrideType>
struct Layout : private Std::tuple<ShapeType, StrideType> {
    __aicore__ inline constexpr Layout(const ShapeType& shape  = {}, const StrideType& stride = {}) : Std::tuple<ShapeType, StrideType>(shape, stride) {}

    __aicore__ inline constexpr decltype(auto) layout() {}
    __aicore__ inline constexpr decltype(auto) layout() const {}
    
    __aicore__ inline constexpr decltype(auto) GetShape() {}   
    __aicore__ inline constexpr decltype(auto) GetShape() const {}
    
    __aicore__ inline constexpr decltype(auto) GetStride() {}    
    __aicore__ inline constexpr decltype(auto) GetStride() const {}
    
    template <typename CoordType>
    __aicore__ inline constexpr auto operator()(const CoordType& coord) const {}
}

__aicore__ inline constexpr Layout(const ShapeType& shape  = {}, const StrideType& stride = {}) : Std::tuple<ShapeType, StrideType>(shape, stride) 
__aicore__ inline constexpr decltype(auto) layout()
__aicore__ inline constexpr decltype(auto) layout() const
__aicore__ inline constexpr decltype(auto) GetShape()  
__aicore__ inline constexpr decltype(auto) GetShape() const
__aicore__ inline constexpr decltype(auto) GetStride()    
__aicore__ inline constexpr decltype(auto) GetStride() const
template <typename CoordType> __aicore__ inline constexpr auto operator()(const CoordType& coord) const {}

// Shape construction method
template <typename... Ts>
__aicore__ inline constexpr Shape<Ts...> MakeShape(const Ts&... t)

// Stride construction method
template <typename... Ts>
__aicore__ inline constexpr Stride<Ts...> MakeStride(const Ts&... t)

// Layout construction method
template <typename ShapeType, typename StrideType>
__aicore__ inline constexpr auto MakeLayout(const ShapeType& shape, const StrideType& stride)

// is_layout prototype definition
template <T>
struct is_layout;