Introduction to Layout

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

A Layout consists of two core components:

• Shape: defines the logical shape of the data, such as the number of rows and columns of a two-dimensional matrix or the size of each dimension of a multi-dimensional tensor.
• Stride: defines the stride of each dimension in the memory, that is, the interval between adjacent elements in the same dimension in the memory. The unit of the interval is element, which 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), then:

• The matrix has four rows and two columns.
• The stride in the column direction is 1, that is, the interval between adjacent elements in the same row in the memory is one element. The stride in the row direction is 4, that is, the interval between the start addresses of adjacent rows is four 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 struct construction method
template <typename... Ts>
__aicore__ inline constexpr Shape<Ts...> MakeShape(const Ts&... t)

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

// Layout struct 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;