大数跨境
0
0

NDArray高级内存操作:跨步视图与跨步复制机制

NDArray高级内存操作:跨步视图与跨步复制机制 ai算法芯片与系统
2025-12-06
1
导读:本文解析 NDArray 的 as_strided 和 strided_copy 操作,揭示内存管理技巧。介绍内存视图与复制概念,分析核心数据结构,对比两者实现原理与使用场景,给出性能优化建议。

 

摘要

本文深入解析NDArray中的两个核心操作:as_strided(原地视图)和strided_copy(跨步复制)。这两个操作体现了现代数值计算库在内存管理上的高级技巧。通过详细分析它们的实现原理、参数差异和使用场景,揭示如何在保持相同多维索引接口的情况下实现不同的内存操作策略。

本文首先介绍内存视图与复制的基本概念,然后详细阐述NDArray的核心数据结构,接着分别深入分析as_strided和strided_copy的实现原理和应用示例,最后对两者进行对比分析并给出性能优化建议。

目录

  1. 1. 内存视图与复制基础
  2. 2. NDArray核心数据结构
  3. 3. as_strided:原地视图操作
  4. 4. strided_copy:跨步复制操作
  5. 5. 操作对比与应用场景
  6. 6. 性能分析与优化
  7. 7. 总结与展望

1. 内存视图与复制基础

内存视图和数据复制是现代数值计算中两种根本不同的内存操作策略。理解它们的差异是高效使用现代数值计算库的关键。

视图操作通过重新解释现有数据创建新的逻辑结构,而不实际复制数据;复制操作则创建数据的独立副本,确保数据隔离。

概念
定义
核心特点
内存视图(View)
相同数据的逻辑重新解释
零拷贝
,共享存储
数据复制(Copy)
创建数据的独立副本
实际复制
,隔离存储
跨步(Stride)
维度间的内存跳跃距离
控制内存访问模式
存储偏移(Offset)
数据在存储中的起始位置
支持子张量操作

核心思想:通过调整形状跨步偏移这三个参数,可以从同一块内存创建出逻辑上完全不同的多个数组视图,而无需复制任何数据。这种机制是高性能数值计算的基础。

内存布局示例

考虑一个 矩阵,按行优先存储(即C风格)。假设内存中连续存储12个元素,从0到11。那么,形状为 的数组的跨步为 。这意味着沿着第一维(行)移动一步需要跳过4个元素,沿着第二维(列)移动一步需要跳过1个元素。

用数学公式表示,对于形状为 的数组,行优先布局的跨步为 。元素 在内存中的偏移为:

其中:

  • •   是第一维(行)索引,范围 
  • •   是第二维(列)索引,范围 
  • •   是第一维跨步,值为 
  • •   是第二维跨步,值为 

这个公式是多维数组内存访问的基础,通过简单的乘加运算将多维索引映射到一维内存地址。

视图与复制对比

操作类型
内存关系
特点
as_strided(视图)
原始数据与视图数据共享同一内存块
零拷贝
,修改视图会影响原始数据
strided_copy(复制)
原始数据与复制数据位于两个独立的内存块
实际复制
,修改复制数据不会影响原始数据

视图操作的最大优势是零拷贝,可以极大减少内存开销;而复制操作的优势是数据隔离,确保修改不会影响原始数据。


2. NDArray核心数据结构

在深入分析as_strided和strided_copy之前,我们需要先定义支持这些操作的NDArray类。这个类封装了多维数组的核心特性。

NDArray类通过五个核心成员变量管理多维数组:存储指针、存储大小、偏移量、形状向量和跨步向量。通过所有权标志区分拥有存储的数组和共享存储的视图。


   
    
   template<typename T>
class
 NDArray {
private
:
    T* storage_;                      // 底层存储指针
    size_t
 storage_size_;             // 存储总容量(元素个数)
    size_t
 offset_;                   // 存储偏移(元素偏移)
    std::vector<size_t> shape_;       // 形状 [dim0, dim1, ...]
    std::vector<size_t> strides_;     // 跨步(元素个数,非字节)
    bool
 owns_storage_;               // 是否拥有存储所有权
    
public
:
    // 构造函数:分配新存储

    NDArray
(const std::vector<size_t>& shape) 
        : shape_(shape), offset_(0), owns_storage_(true) {
        
        // 计算总元素数并分配内存

        size_t
 total_elements = 1;
        for
 (auto s : shape) total_elements *= s;
        storage_ = new T[total_elements];
        storage_size_ = total_elements;
        
        // 计算紧凑布局跨步(行优先)

        compute_compact_strides
();
    }
    
    // 构造函数:从现有存储创建视图

    NDArray
(T* storage, size_t storage_size, 
            const
 std::vector<size_t>& shape,
            const
 std::vector<size_t>& strides,
            size_t
 offset = 0)
        : storage_(storage), storage_size_(storage_size),
          offset_
(offset), shape_(shape), strides_(strides),
          owns_storage_
(false) {}
    
    // 析构函数:根据所有权释放内存

    ~NDArray() {
        if
 (owns_storage_ && storage_ != nullptr) {
            delete
[] storage_;
        }
    }
    
    // 计算紧凑布局跨步

    void compute_compact_strides() 
{
        strides_.resize(shape_.size());
        if
 (shape_.empty()) return;
        
        strides_.back() = 1;  // 最后一维跨步为1(元素个数)
        for
 (int i = shape_.size() - 2; i >= 0; --i) {
            strides_[i] = strides_[i + 1] * shape_[i + 1];
        }
    }
    
    // 多维下标访问运算符

    T& operator[](std::initializer_list<size_t> indices) {
        return
 storage_[offset_ + compute_offset(indices)];
    }
    
    // const版本访问

    const
 T& operator[](std::initializer_list<size_t> indices) const {
        return
 storage_[offset_ + compute_offset(indices)];
    }
    
    // 计算物理偏移(核心函数)

    size_t compute_offset(std::initializer_list<size_t> indices) const 
{
        assert
(indices.size() == shape_.size());
        
        size_t
 offset = 0;
        size_t
 dim = 0;
        for
 (auto idx : indices) {
            // 验证索引有效性

            assert
(idx < shape_[dim]);
            offset += idx * strides_[dim];
            dim++;
        }
        return
 offset;
    }
    
    // 获取基本信息

    const std::vector<size_t>& shape() const 
{ return shape_; }
    const std::vector<size_t>& strides() const 
{ return strides_; }
    size_t offset() const 
{ return offset_; }
    T* storage() const 
{ return storage_; }
    size_t storage_size() const 
{ return storage_size_; }
    bool owns_storage() const 
{ return owns_storage_; }
    
    // 获取元素总数

    size_t numel() const 
{
        size_t
 total = 1;
        for
 (auto s : shape_) total *= s;
        return
 total;
    }
};

NDArray类设计说明:

  • • 存储管理:通过owns_storage_标志区分视图和拥有存储的数组,确保正确的内存管理
  • • 多维索引:通过重载operator[]支持[{i,j,k}]语法,内部使用compute_offset函数将多维索引映射到一维偏移
  • • 跨步计算compute_compact_strides计算行优先的紧凑布局跨步,即最后一维跨步为1,向前逐维乘以形状

NDArray类提供了统一的多维数组接口,隐藏了底层复杂的内存映射细节。通过所有权机制,可以安全地管理内存生命周期,避免内存泄漏和悬垂指针。

多维索引映射公式

对于形状为 维数组,其跨步为 ,索引为 的元素在内存中的偏移为:

其中 是第 维的索引,满足 是第 维的跨步。

对于行优先(C风格)紧凑布局,跨步的计算公式为:

这个递推公式确保相邻元素在内存中连续存储,对于缓存友好的访问模式非常重要。

示例:对于形状为 的矩阵,按行优先布局,其跨步为 。那么,元素 在内存中的偏移为:

具体来说:

索引 
偏移计算
内存位置
0
1
4
11

这种偏移计算方式使得多维数组的访问非常高效,只需简单的乘加运算即可定位到内存地址。


3. as_strided:原地视图操作

as_strided是一个原地操作,它在原始张量的存储偏移处创建一个新视图,不复制任何数据。这是实现零拷贝变换的核心函数。

as_strided操作的核心思想是"重新解释"现有数据。通过调整形状、跨步和偏移这三个参数,可以从同一块内存创建出逻辑上完全不同的数组视图,而不需要移动任何数据。

3.1 函数定义与参数说明


   
    
   // as_strided:创建原地视图
NDArray<float> as_strided(
    const
 NDArray<float>& input,               // 输入张量
    const
 std::vector<size_t>& new_shape,      // 新形状
    const
 std::vector<size_t>& new_strides,    // 新跨步
    size_t
 storage_offset = 0                  // 存储偏移(相对输入存储)
)
 
{
    // 验证参数有效性

    assert
(storage_offset < input.storage_size());
    
    // 计算新视图需要的存储空间

    size_t
 required_storage = storage_offset;
    for
 (size_t i = 0; i < new_shape.size(); ++i) {
        // 计算每个维度可能的最大偏移

        size_t
 max_dim_offset = (new_shape[i] - 1) * new_strides[i];
        required_storage += max_dim_offset;
    }
    
    assert
(required_storage < input.storage_size());
    
    // 创建并返回新视图(共享底层存储)

    return
 NDArray<float>(
        input.storage(),      // 共享相同存储
        input.storage_size(), // 相同存储大小
        new_shape,           // 新形状
        new_strides,         // 新跨步
        input.offset() + storage_offset // 新偏移
    );
}

as_strided函数详细说明:

功能:在现有张量的存储上创建新的逻辑视图,不复制任何数据。通过调整形状、跨步和偏移参数,重新解释现有数据。

参数说明

  1. 1. input:输入张量,提供原始存储
  2. 2. new_shape:新视图的形状,元素个数可以与原始不同
  3. 3. new_strides:新视图的跨步,决定新视图中各维度间的内存跳跃距离
  4. 4. storage_offset:相对于输入存储的偏移量,用于从存储的中间位置开始创建视图

数学原理
给定输入张量 ,其存储起始位置为 ,偏移为 ,跨步为 ,新视图 的偏移计算公式为:

其中:

  • •  是新视图的维度数
  • •  是新视图第 维的索引
  • •  是新视图第 维的跨步

存储边界验证:为确保访问安全,需要验证新视图的最大偏移不超过存储边界:

这个验证确保不会访问到未分配的内存区域,防止内存越界错误。

as_strided函数的核心优势是零拷贝,它只创建新的视图元数据(形状、跨步、偏移),而不移动任何实际数据。这使得操作非常高效,特别适合大数组操作。

3.2 原地转置函数与示例

函数定义:原地转置


   
    
   // 原地转置函数:对二维矩阵进行转置
NDArray<float> transpose(const NDArray<float>& matrix) 
{
    // 获取原始矩阵的形状

    const
 std::vector<size_t>& shape = matrix.shape();
    assert
(shape.size() == 2);  // 确保是二维矩阵
    
    // 转置后的形状:交换行和列

    std::vector<size_t> transposed_shape = {shape[1], shape[0]};
    
    // 转置后的跨步:原始跨步为(shape[1], 1),转置后为(1, shape[1])

    // 注意:原始矩阵的跨步不一定是紧凑的,所以这里使用原始跨步进行计算

    const
 std::vector<size_t>& strides = matrix.strides();
    std::vector<size_t> transposed_strides = {1, strides[0]};
    
    // 创建转置视图

    return
 as_strided(matrix, transposed_shape, transposed_strides, 0);
}

转置操作的数学原理
对于二维矩阵 ,其形状为 ,跨步为 。转置矩阵 的形状为 ,跨步应为

对于紧凑行优先布局(默认情况), ,因此转置后的跨步为

索引映射关系

  • • 原始矩阵 中元素 的偏移:
  • • 转置视图 中元素 的偏移:
  • • 由于紧凑布局中 ,因此

转置操作的关键在于交换跨步。通过将行跨步变为列跨步,列跨步变为行跨步,实现了逻辑上的转置,而物理存储保持不变。

调用示例:原地转置


   
    
   // 使用示例
void transpose_example() 
{
    // 创建一个3x4的矩阵

    std::vector<size_t> shape = {3, 4};
    NDArray<float> matrix(shape)
;
    
    // 初始化矩阵数据

    for
 (size_t i = 0; i < 3; ++i) {
        for
 (size_t j = 0; j < 4; ++j) {
            matrix[{i, j}] = i * 10 + j;
        }
    }
    
    // 创建转置视图

    NDArray<float> transposed_view = transpose(matrix);
    
    // 验证转置效果

    // 原始矩阵元素 (1,2) 应该等于转置视图元素 (2,1)

    assert
(matrix[{1, 2}] == transposed_view[{2, 1}]);
    
    // 修改转置视图会影响原始矩阵

    transposed_view[{2, 1}] = 999;
    assert
(matrix[{1, 2}] == 999);
    
    // 打印结果

    std::cout << "原始矩阵形状: " << shape[0] << "x" << shape[1] << std::endl;
    std::cout << "转置视图形状: " << transposed_view.shape()[0] << "x" 
              << transposed_view.shape()[1] << std::endl;
    std::cout << "A[1,2] = " << matrix[{1, 2}] << std::endl;
    std::cout << "A^T[2,1] = " << transposed_view[{2, 1}] << std::endl;
}

转置操作的数据关系

原始矩阵A
转置视图
偏移计算
[0,0]
[0,0]
0
[0,1]
[1,0]
1
[0,2]
[2,0]
2
[0,3]
[3,0]
3
[1,0]
[0,1]
10
[1,1]
[1,1]
11
[1,2]
[2,1]
12
[1,3]
[3,1]
13
...
...
...
...
[2,3]
[3,2]
23

转置操作展示了视图操作的核心优势:通过重新解释现有数据来满足不同需求,无需额外的内存分配和数据移动。这对于大矩阵操作特别高效,避免了数据复制的开销。


4. strided_copy:跨步复制操作

strided_copy是一个数据复制操作,它将源张量中特定视图的数据复制到目标张量的指定位置。与as_strided不同,它执行实际的数据复制,创建数据的独立副本。

strided_copy操作的核心思想是"选择性复制"。它允许从源张量的任意位置(通过偏移和跨步定义)复制数据到目标张量的任意位置,支持复杂的张量重组操作。

4.1 函数定义与参数说明


   
    
   // strided_copy:跨步复制操作
void strided_copy(
    const
 NDArray<float>& input,               // 输入张量
    const
 std::vector<size_t>& new_input_shape,    // 输入张量的新视图形状
    const
 std::vector<size_t>& new_input_strides,  // 输入张量的新视图跨步
    size_t
 input_storage_offset,               // 输入存储偏移
    NDArray<float>& output,                    // 输出张量
    size_t
 output_storage_offset               // 输出存储偏移
)
 
{
    // 验证输入视图形状与输出张量形状匹配

    assert
(new_input_shape.size() == output.shape().size());
    for
 (size_t i = 0; i < new_input_shape.size(); ++i) {
        assert
(new_input_shape[i] == output.shape()[i]);
    }
    
    // 验证存储偏移有效性

    assert
(input_storage_offset < input.storage_size());
    assert
(output_storage_offset < output.storage_size());
    
    const
 size_t ndim = new_input_shape.size();
    
    // 创建输入视图(临时对象,不拥有存储)

    NDArray<float> input_view(
        input.storage(),
        input.storage_size(),
        new_input_shape,
        new_input_strides,
        input.offset() + input_storage_offset
    )
;
    
    // 创建输出视图(临时对象,不拥有存储)

    NDArray<float> output_view(
        output.storage(),
        output.storage_size(),
        output.shape(),      // 使用输出张量原有的形状
        output.strides(),    // 使用输出张量原有的跨步
        output.offset() + output_storage_offset
    )
;
    
    // 使用递归函数遍历所有索引并复制

    std::vector<size_t> indices(ndim, 0)
;
    strided_copy_recursive
(input_view, output_view, indices, 0);
}

strided_copy函数详细说明:

功能:将源张量中特定视图的数据实际复制到目标张量的指定位置。创建数据的独立副本,不共享存储。

参数说明

  1. 1. input:源张量,提供源数据
  2. 2. new_input_shape:源张量上临时视图的形状
  3. 3. new_input_strides:源张量上临时视图的跨步
  4. 4. input_storage_offset:源张量存储的偏移,用于从存储的特定位置开始复制
  5. 5. output:目标张量,接收复制数据
  6. 6. output_storage_offset:目标张量存储的偏移,用于指定复制到的位置

数学原理
对于输入视图 中的每个索引 ,执行复制操作:

输入视图的偏移计算公式:

【声明】内容源于网络
0
0
ai算法芯片与系统
长期关注ai领域,算法,芯片,软件(系统,框架,编译器,算子库)等联合设计
内容 196
粉丝 0
ai算法芯片与系统 长期关注ai领域,算法,芯片,软件(系统,框架,编译器,算子库)等联合设计
总阅读90
粉丝0
内容196