Squeeze

产品支持情况

• Ascend 950PR/Ascend 950DT：支持

• Atlas A3 训练系列产品/Atlas A3 推理系列产品：不支持

• Atlas A2 训练系列产品/Atlas A2 推理系列产品：不支持

• Atlas 训练系列产品：不支持

• Atlas 推理系列产品AI Core：不支持

• Atlas 推理系列产品Vector Core：不支持

• Atlas 200I/500 A2 推理产品：不支持

功能说明

需要包含的头文件为：#include "tensor_api/tensor.h"。

Squeeze用于删除Layout的Shape中大小为1的指定维度，并同步删除Stride中的对应维度。输入为Tensor时，接口更新Tensor的Layout，Tensor的Engine、存储位置和数据起始地址保持不变。

接口支持以下两种方式指定待删除维度：

• 通过模板参数指定Shape顶层维度的索引。
• 通过与Shape结构相同的Pattern指定维度。Pattern中使用_1标记待删除维度，使用_标记保留维度，嵌套Pattern会递归处理嵌套Shape。

返回的Layout或Tensor保留输入Layout携带的LayoutPattern和LayoutTrait信息。

函数原型

• 通过维度索引删除大小为1的维度。

  C++

  template <size_t... SqueezeDims, typename T>
  __aicore__ inline constexpr auto Squeeze(const T& x)
  

• 通过Pattern删除大小为1的维度。

  C++

  template <typename Pattern, typename T>
  __aicore__ inline constexpr auto Squeeze(const T& x, const Pattern& pattern)
  

参数说明

表1 模板参数说明

参数名	描述
SqueezeDims	待删除维度在Shape顶层结构中的索引。支持同时指定多个索引，索引顺序不影响结果。
T	输入对象类型，需要为Layout类型或Tensor类型。
Pattern	Pattern类型，需要为与输入Layout的Shape结构相同的元组类型。

表2 参数说明

参数名	输入/输出	描述
x	输入	待处理的Layout或Tensor对象。
pattern	输入	维度选择Pattern。_1表示删除对应的大小为1的维度，_表示保留对应维度。Pattern可由MakeCoord构造。

返回值说明

输入为Layout时，返回删除指定维度后的Layout对象；输入为Tensor时，返回使用新Layout描述的Tensor对象。返回对象保留原Layout的LayoutPattern和LayoutTrait信息。

约束说明

• SqueezeDims至少需要包含一个维度索引，并指向输入Shape的顶层维度。
• pattern需要与输入Layout的Shape具有相同的嵌套结构。
• 对于编译期常量维度，只有大小为1的维度会被删除；被指定但大小不为1的维度会被保留。
• 对于运行时维度，接口不会校验对应Shape值是否为1。用户需要保证被指定删除的运行时维度大小为1，否则返回的Layout或Tensor不能正确描述原始数据。
• 删除维度后，如果当前层只保留一个嵌套元组，接口会展开该层。例如，Shape从(1, (M, N))转换为(M, N)。
• 删除维度只改变Shape和Stride的结构，不移动或修改内存中的数据。

调用示例

按维度索引删除

C++

#include "tensor_api/tensor.h"

using namespace AscendC::Te;

auto layout = MakeLayout(MakeShape(_1{}, _4{}, _1{}), MakeStride(_4{}, _1{}, _1{}));
auto squeezedLayout = Squeeze<0, 2>(layout);

// squeezedLayout的Shape为(4)，Stride为(1)。

按Pattern删除

C++

#include "tensor_api/tensor.h"

using namespace AscendC::Te;

auto layout = MakeLayout(MakeShape(_1{}, _4{}, _1{}), MakeStride(_4{}, _1{}, _1{}));
auto pattern = MakeCoord(_1{}, _, _1{});
auto squeezedLayout = Squeeze(layout, pattern);

// squeezedLayout的Shape为(4)，Stride为(1)。

删除Tensor的Batch维度

C++

#include "tensor_api/tensor.h"

using namespace AscendC::Te;

__aicore__ inline void SqueezeTensorExample()
{
    constexpr uint32_t batch = 1;
    constexpr uint32_t m = 32;
    constexpr uint32_t n = 64;

    __cbuf__ half l1Buffer[batch * m * n];
    auto batchLayout = MakeFrameLayout<NZLayoutPtn, half>(batch, m, n);
    auto batchTensor = MakeTensor(MakeMemPtr(l1Buffer), batchLayout);

    auto matrixTensor = Squeeze<0>(batchTensor);
    // matrixTensor与batchTensor指向相同的数据，Layout中不再包含Batch维度。
}