Star462
269
代码
代码
Issues46
Pull Requests71
Actions
项目讨论
Wiki
分析
Star462
269
已合并
TLA增加originShape，新增TileView、MakeTensorLike、Tensor的operator()和增加strided_batched_matmul_tla样例 #469
sinceseto创建于 2月13日
TLA增加originShape，新增TileView、MakeTensorLike、Tensor的operator()和增加strided_batched_matmul_tla样例 #469
已合并
sinceseto创建于 2月13日
从sinceseto/catlass_tlatla_modify_pr合入到cann/catlassmaster
共 30 个文件变更+2688-327
Mdocs/2_Design/02_tla/01_layout.md+267-95
@@ -1,170 +1,342 @@
1# TLA Layouts1# TLA Layouts
2 2 
3-这篇文档描述了CATLASS的TLA(Tensor Layout Abstraction)下的`Layout`数据结构，它提供了多维坐标与内存的映射关系。3+本文介绍 TLA（Tensor Layout Abstraction）中的 `Layout`。
4 4 
5-`Layout`为多维数组的访问提供了一个通用接口，它将数组元素在内存中的组织细节抽象化。这使得用户能够编写通用的多维数组访问算法，从而在布局发生变化时，无需修改用户代码。5+如果把 Tensor 看成“逻辑上的多维数组”，那么 `Layout` 负责回答以下问题：
6 6 
7-## 基础类型与概念7+- 一个逻辑坐标 `(i, j, ...)` 对应到哪一个线性地址。
8+- 这块 Tensor 在逻辑上有多大。
9+- 当底层存在分块、对齐或填充时，哪些位置是逻辑有效数据。
10+ 
11+因此，`Layout` 可以理解为“逻辑坐标到内存地址的映射规则”。算法通常依赖这套规则访问数据，而不直接依赖底层物理排布。这样，同一段计算逻辑就可以适配普通 ND 布局、行优先、列优先以及 `zN`、`nZ` 等分形布局。
12+ 
13+## 先建立三个基本概念
14+ 
15+### 逻辑坐标 coord
16+ 
17+`coord` 表示元素在 Tensor 逻辑空间中的位置，约定如下：
18+ 
19+- 坐标从 0 开始计数。
20+- 坐标单位是“元素”，不是字节，也不是 tile 编号。
21+- `coord` 的 rank 必须与 Tensor 或 Layout 的逻辑维度一致。
22+- 即使底层采用 `zN`、`nZ` 这类嵌套布局，`coord` 仍然描述逻辑上的行列位置，例如 `(row, col)`。
23+ 
24+例如，对一个逻辑形状为 `(8, 16)` 的矩阵，`coord = (2, 4)` 表示第 2 行、第 4 列的元素。它不关心这块数据在内存中是按行连续、按列连续，还是按分形块组织。
25+ 
26+### 逻辑形状与内存布局
27+ 
28+在 TLA 中，这两个概念被刻意分离：
29+ 
30+- 逻辑形状：从使用者视角看，Tensor 有多少行、多少列。
31+- 内存布局：这些逻辑元素在内存中如何排布，跨一个维度移动时需要跳过多少位置。
32+ 
33+`Layout` 的核心价值，就是把“逻辑上多大”和“内存里怎样排”同时表达清楚。
34+ 
35+### Tail tile
36+ 
37+当矩阵尺寸不是 tile 大小的整数倍时，边界 tile 往往只包含部分有效元素。这类边界 tile 通常称为 tail tile。
38+ 
39+TLA 使用 `originShape` 表达“逻辑上实际有效的范围”。因此，用户通常不需要手工推导每个边缘 tile 的真实尺寸。
40+ 
41+## 基础类型
8 42 
9### Tuple43### Tuple
10 44 
11-TLA以元组[`tla::tuple`](../../../include/tla/tuple.hpp)为起始，tla::tuple 包含了若干个元素组成的有限序列元组，其行为与 std::tuple 类似，但引入了一些 C++ template arguments 的限制，并削减了部分实现以提升性能。45+TLA 以 [`tla::tuple`](../../include/tla/tuple.hpp) 为基础。它与 `std::tuple` 的用途相似，都是表达定长元素序列；不同之处在于，TLA 对模板元编程和高性能场景做了定制。
12 46 
13### IntTuple47### IntTuple
14 48 
15-TLA 还定义了[`IntTuple`](../../../include/tla/int_tuple.hpp)概念。`IntTuple`既可作为一个整数，也可作为一个`Tuple`类型。这个递归定义允许我们构建任意嵌套的 Layout。49+[`IntTuple`](../../include/tla/int_tuple.hpp) 是 TLA 中最常用的基础概念之一。它可以是：
16 50 
17-以下任何一个都是 `IntTuple` 的有效模板参数：51+- 一个整数，例如 `int{2}`、`size_t{16}`。
52+- 一个编译期整数，例如 `Int<3>{}` 或别名 `_3`。
53+- 一个由以上元素递归组成的 tuple，例如 `make_tuple(int{2}, Int<3>{})`。
18 54 
19-* `int{2}`: 运行时整数，也称为动态整数，就是 C++ 的正常整数类型比如`int`/`size_t`等等，只要是`std::is_integral<T>`的都是。55+因此，`IntTuple` 既可以表示一维尺寸，也可以表示带层次结构的嵌套尺寸。
20 56 
21-* `Int<3>{}`： 编译期整数，或称之为静态整数。TLA 通过 `tla::C<Value>` 来定义兼容的静态整数类型，使得这些整数的计算能在编译期内完成。TLA 将别名 _1、_2、_3等定义为`Int<1>`、`Int<2>`、`Int<3>`等类型。更多信息可查看[`integral_constant`](../../../include/tla/numeric/integral_constant.hpp)。57+常用操作如下：
22 58 
23-* 带有任何模板参数的 IntTuple，例如 `make_tuple(int{2}, Int<3>{})`。59+- `rank(IntTuple)`：返回元素个数。
60+- `get<I>(IntTuple)`：返回第 `I` 个元素。
61+- `depth(IntTuple)`：返回嵌套层数；普通整数的 `depth` 为 0。
24 62 
25-TLA 不仅将 `IntTuple` 用在了`Layout`上，还会在很多其他的地方比如 `Shape` 和 `Stride` 等用到它，详见63+`IntTuple` 不仅用于 `Layout`，也用于 `Shape`、`Stride` 等类型，定义见 [`include/tla/layout.hpp`](../../include/tla/layout.hpp)。
26-[`include/tla/layout.hpp`](../../../include/tla/layout.hpp)。
27 64 
28-`IntTuple` 的相关 API 操作：65+## Layout 由什么组成
29 66 
30-* `rank(IntTuple)`: 返回 `IntTuple` 的元素个数。67+`Layout` 本质上由三个 `IntTuple` 组成：`Shape`、`Stride` 和 `OriginShape`。
31 68 
32-* `get<I>(IntTuple)`: 返回 `IntTuple` 的第 `I` 个元素。69+| 字段 | 作用 | 关注点 |
70+| --- | --- | --- |
71+| `Shape` | 用于内存布局计算的尺寸描述 | 决定布局结构，不一定等于逻辑实际尺寸 |
72+| `Stride` | 各维度上的步长 | 决定坐标如何映射到线性地址 |
73+| `OriginShape` | Tensor 的逻辑实际尺寸 | 决定哪些元素在逻辑上有效 |
33 74 
34-* `depth(IntTuple)`: 返回 `IntTuple` 的嵌套层数，整数为 0。75+可以先把它们理解成：
35 76 
36-### Layout77+- `Shape` 说明“内存按什么结构排”。
78+- `Stride` 说明“每跨一步跳多远”。
79+- `OriginShape` 说明“逻辑上到底有多少有效数据”。
37 80 
38-`Layout` 本质上就是由2个 `IntTuple` 组成， `Shape` 和 `Stride` 。 `Shape` 定义了 `Tensor` 的形状，`Stride` 定义了元素间的距离。81+这里最容易混淆的是 `Shape` 和 `OriginShape`。两者并不重复：
39 82 
40-## Layout 使用83+- `Shape` 面向布局计算，允许包含对齐、分块和填充后的结构。
84+- `OriginShape` 面向逻辑语义，只描述真实有效的数据范围。
41 85 
42-`Layout` 也有许多与 `IntTuple` 类似的操作：86+![Origin_Shape-layout_1.png](https://raw.gitcode.com/user-images/assets/7631999/936388d2-81b6-400d-afe2-788eaf140f81/Origin_Shape-layout_1.png 'Origin_Shape-layout_1.png')
43 87 
44-* `rank(Layout)`: `Layout` 的维度，等同于 `Stride` 的 rank(IntTuple)。88+`OriginShape` 用于把“内存怎样排”与“逻辑上哪些数据有效”区分开。
45 89 
46-* `get<I>(Layout)`: 返回 `Layout` 的第 `Ith` 个元素，`I < rank`。90+- `Shape`：服务于布局计算，可能包含对齐、分块或填充后的尺寸。
91+- `OriginShape`：服务于逻辑语义，描述真实有效的数据范围。
47 92 
48-* `depth(Layout)`: 返回 `Layout` 的嵌套层数，整数为0。93+例如，一个逻辑大小为 `100 x 100` 的矩阵采用 `zN` 布局时，可能出现：
49 94 
50-* `shape(Layout)`: 返回 `Layout` 的 `Shape` 。95+- `originShape = (100, 100)`
96+- `shape = ((16, 7), (16, 7))`
51 97 
52-* `stride(Layout)`: 返回 `Layout` 的 `Stride` 。98+原因是：
53 99 
54-此外，为了便于操作，还定义了下述一些函数：100+- `16 * 7 = 112`，说明底层内存按 `112 x 112` 的块化结构组织。
101+- 但逻辑上只有 `100 x 100` 是有效元素。
55 102 
56-* `get<I0,I1,...,IN>(x) := get<IN>(...(get<I1>(get<I0>(x)))...)`： 获取第 `I0` 个单元的第 `I1` 个单元的 ... 的第 `IN` 个单元。103+这也是 TLA 能自动处理 tail tile 的基础。用户在 block 层和 kernel 层通常只需要按 tile 编程，边界有效范围由 `originShape` 传递和裁剪，无需每一层都手动判断尾块。
57 104 
58-* `rank<I...>(x) := rank(get<I...>(x))`： 获取第 `I...` 个单元维度。105+## Layout 的常用接口
59 106 
60-* `depth<I...>(x) := depth(get<I...>(x))`： 获取第 `I...` 个单元的嵌套层数。107+`Layout` 提供了一组与 `IntTuple` 风格一致的访问接口：
61 108 
62-* `shape<I...>(x) := shape(get<I...>(x))`： 获取第 `I...` 个单元的形状。109+- `rank(Layout)`：布局的逻辑维度。
110+- `get<I>(Layout)`：取出第 `I` 个分量。
111+- `depth(Layout)`：布局的嵌套层数。
112+- `shape(Layout)`：返回 `Shape`。
113+- `stride(Layout)`：返回 `Stride`。
114+- `originShape(Layout)`：返回 `OriginShape`。
63 115 
64-* ...116+另外还提供递归版本的辅助接口，例如：
65 117 
66-### Layout 构造118+- `get<I0, I1, ..., IN>(x)`：逐层向下取子单元。
119+- `rank<I...>(x)`：查看某个子单元的 rank。
120+- `depth<I...>(x)`：查看某个子单元的 depth。
121+- `shape<I...>(x)`：查看某个子单元的 shape。
122+- `originShape<I...>(x)`：查看某个子单元的 origin shape。
67 123 
68-`Layout` 有多种构造方式，可以是静态整数和动态整数的任意结合，可以定义任意维度。124+## Layout 构造
69-**注：在昇腾CUBE核内部，存在 `zN` 、 `nZ` 、 `zZ` 、 `nN` 格式，因此目前昇腾算子模板库中只定义与使用`行优先`、`列优先`和前述4种格式**。125+ 
126+`Layout` 支持静态整数、动态整数及其混合构造，也支持普通矩阵布局和 Ascend 常用内部布局。
127+ 
128+在昇腾 CUBE 核内部，常见内部格式包括 `zN`、`nZ`、`zZ`、`nN`、`L0C` 等；在 GEMV、Scale、Bias 等场景中，也会使用一维 `VectorLayout`。
70 129 
71```c++130```c++
72using namespace tla;131using namespace tla;
132+ 
133+// 1. 直接给 shape 和 stride，originShape 由系统推导
73Layout w2xh4 = MakeLayout(MakeShape(Int<2>{}, 4),134Layout w2xh4 = MakeLayout(MakeShape(Int<2>{}, 4),
74 MakeStride(Int<12>{}, Int<1>{}));135 MakeStride(Int<12>{}, Int<1>{}));
75 136 
76-Layout w32xh48 = MakeLayout(MakeShape(MakeShape(16,2), MakeShape(16,3)),137+// 2. 嵌套布局，originShape 隐式推导为 (16*2, 16*3) = (32, 48)
77- MakeStride(MakeStride(16,256), MakeStride(1,512)));138+Layout w32xh48 = MakeLayout(MakeShape(MakeShape(16, 2), MakeShape(16, 3)),
139+ MakeStride(MakeStride(16, 256), MakeStride(1, 512)));
140+ 
141+// 3. 显式指定 originShape
142+Layout w2xh4_explicit = MakeLayout(MakeShape(Int<2>{}, 4),
143+ MakeStride(Int<12>{}, Int<1>{}),
144+ MakeShape(2, 4));
145+ 
146+Layout w32xh48_explicit = MakeLayout(MakeShape(MakeShape(16, 2), MakeShape(16, 3)),
147+ MakeStride(MakeStride(16, 256), MakeStride(1, 512)),
148+ MakeShape(32, 48));
149+ 
150+// 4. rank=2 时，也可以用 LayoutTag + (rows, cols) 构造
151+auto rm = MakeLayout<float, Catlass::layout::RowMajor>(2, 4);
152+ 
153+// 5. 一维 VectorLayout
154+auto vec = MakeLayout(128);
78```155```
79 156 
80- `MakeLayout` 函数返回 `Layout`。 `MakeShape` 函数返回 `Shape`。 `MakeStride` 函数返回 `Stride`。157+其中：
81 158 
82-上述Layout格式如下159+- `MakeLayout` 返回 `Layout`。
160+- `MakeShape` 返回 `Shape`。
161+- `MakeStride` 返回 `Stride`。
83 162 
84-```163+上面的布局可写成：
85-w2xh4 : (_2,4):(_12,_1)164+ 
86-w32xh48 : ((16,2),(16,3)):((16,256),(1,512))165+```text
166+w2xh4 : (_2, 4):(_12, _1)
167+w32xh48 : ((16, 2), (16, 3)):((16, 256), (1, 512))
87```168```
88 169 
89-`(_2,4):(_12,_1)` 中前面的括号表示 `Tensor` 的形状，后面的括号表示在不同维度下的 `Stride`。170+读法如下：
90 171 
91-### Matrix examples172+- 前一部分是 `Shape`。
173+- 后一部分是 `Stride`。
174+- 如果省略 `OriginShape`，表示它可由 `Shape` 推导，或与逻辑尺寸一致。
92 175 
93-可以定义一个matrix的 `layout` 如下几种类型：176+## 从直观例子理解 Shape 与 Stride
94 177 
95-2x3 `行优先` layout178+### 2x3 行优先
96 179 
97-```180+```text
98-(2,3):(3,1)181+shape = (2, 3)
99- 0 1 2182+stride = (3, 1)
100- +---+---+---+
101- 0 | 0 | 1 | 2 |
102- +---+---+---+
103- 1 | 3 | 4 | 5 |
104- +---+---+---+
105```183```
106 184 
107-定义了一个 2x3 的 tensor，2 行 3 列。至于 stride，在前一维度（行维度），stride=3，表示映射到一维空间中，按行方向递增时，内存跨度为3；在后一维度（列维度），stride=1，表示映射到一维空间中，按列方向递增时，内存跨度为1。185+含义是：
108 186 
109-2x3 `列优先` layout187+- 行维度前进一步，线性地址增加 3。
188+- 列维度前进一步，线性地址增加 1。
110 189 
111-```190+因此线性地址顺序为：
112-(2,3):(1,2)191+ 
113- 0 1 2192+| 逻辑坐标 | 线性地址 |
114- +---+---+---+193+| --- | --- |
115- 0 | 0 | 2 | 4 |194+| `(0, 0)` | `0` |
116- +---+---+---+195+| `(0, 1)` | `1` |
117- 1 | 1 | 3 | 5 |196+| `(0, 2)` | `2` |
118- +---+---+---+197+| `(1, 0)` | `3` |
198+| `(1, 1)` | `4` |
199+| `(1, 2)` | `5` |
200+ 
201+### 2x3 列优先
202+ 
203+```text
204+shape = (2, 3)
205+stride = (1, 2)
119```206```
120 207 
121-定义了一个 2x3 的 tensor，2 行 3 列。至于 stride，在前一维度（行维度），stride=1，表示映射到一维空间中，按行方向递增时，内存跨度为1；在后一维度（列维度），stride=2，表示映射到一维空间中，按列方向递增时，内存跨度为2。208+含义是：
122 209 
123-`zN` layout（ `示例展示为4x4块，实际核内为 16x(32/sizeof(ElemType))` ），其余3种内部格式类似210+- 行维度前进一步，线性地址增加 1。
211+- 列维度前进一步，线性地址增加 2。
124 212 
125-```213+因此线性地址顺序为：
126-((4,2),(4,3)):((4,16),(1,32))214+ 
127- 0 1 2 3 4 5 6 7 8 9 10 11215+| 逻辑坐标 | 线性地址 |
128- +----+----+----+----+----+----+----+----+----+----+----+----+216+| --- | --- |
129- 0 | 0 | 1 | 2 | 3 | 32 | 33 | 34 | 35 | 64 | 65 | 66 | 67 |217+| `(0, 0)` | `0` |
130- +----+----+----+----+----+----+----+----+----+----+----+----+218+| `(1, 0)` | `1` |
131- 1 | 4 | 5 | 6 | 7 | 36 | 37 | 38 | 39 | 68 | 69 | 70 | 71 |219+| `(0, 1)` | `2` |
132- +----+----+----+----+----+----+----+----+----+----+----+----+220+| `(1, 1)` | `3` |
133- 2 | 8 | 9 | 10 | 11 | 40 | 41 | 42 | 43 | 72 | 73 | 74 | 75 |221+| `(0, 2)` | `4` |
134- +----+----+----+----+----+----+----+----+----+----+----+----+222+| `(1, 2)` | `5` |
135- 3 | 12 | 13 | 14 | 15 | 44 | 45 | 46 | 47 | 76 | 77 | 78 | 79 |223+ 
136- +----+----+----+----+----+----+----+----+----+----+----+----+224+### 以 `zN` 为例理解嵌套布局
137- 4 | 16 | 17 | 18 | 19 | 48 | 49 | 50 | 51 | 80 | 81 | 82 | 83 |225+ 
138- +----+----+----+----+----+----+----+----+----+----+----+----+226+示例布局：
139- 5 | 20 | 21 | 22 | 23 | 52 | 53 | 54 | 55 | 84 | 85 | 86 | 87 |227+ 
140- +----+----+----+----+----+----+----+----+----+----+----+----+228+```text
141- 6 | 24 | 25 | 26 | 27 | 56 | 57 | 58 | 59 | 88 | 89 | 90 | 91 |229+shape = ((4, 2), (4, 3))
142- +----+----+----+----+----+----+----+----+----+----+----+----+230+stride = ((4, 16), (1, 32))
143- 7 | 28 | 29 | 30 | 31 | 60 | 61 | 62 | 63 | 92 | 93 | 94 | 95 |
144- +----+----+----+----+----+----+----+----+----+----+----+----+
145```231```
146 232 
147-在行维度上，我们有个子 tensor ，该子 tensor 有四行（此为第一个 shape 的第一个4），行之间的 stride 为 4（所以第一个 stride 第一个数为 4）；然后该子 tensor 在整个大 tensor 行中会重复两次（此为第一个 shape 的第二个 2），相对应地，子 tensor 间的 stride 为 16（此为第一个 stride 的第二个 16）。233+可以理解为：
148-同理，在列维度上，我们有个子 tensor ，该子 tensor 有四列（此为第二个 shape 第一个4），列之间的 stride 为 1（所以第二个 stride 第一个数为 1）；然后该子 tensor 在整个大 tensor 列中会重复三次（此为第二个 shape 的第二个 3），相对应地，子 tensor 间的 stride 为 32（此为第二个 stride 的第二个 32）。
149 234 
150-### Layout 坐标与索引235+- 行方向先以 4 为一个内层块，再沿行方向重复 2 次。
236+- 列方向先以 4 为一个内层块，再沿列方向重复 3 次。
237+- 子块内部如何走、子块之间如何跳，分别由嵌套 `Stride` 给出。
151 238 
152-在 TLA 中，可使用 `tla::crd2offset(c, shape, stride)` 将坐标转换到索引，目前坐标需为二维。239+关键点不在于记住每个数字，而在于理解：TLA 用嵌套 `Shape` 和 `Stride` 显式表达分块布局的结构层次，而不是把这类格式硬编码进算法。
240+ 
241+## 坐标如何映射为索引
242+ 
243+在 TLA 中，可以使用 `tla::crd2offset(coord, shape, stride)` 将逻辑坐标转换为线性索引。
244+ 
245+约束如下：
246+ 
247+- `coord`、`shape`、`stride` 的 rank 必须一致。
248+- `coord` 表示逻辑元素坐标，而不是字节偏移。
153 249 
154```cpp250```cpp
155-auto shape = Shape<Shape<_4,_2>,Shape<_4,_3>>{};251+auto shape = Shape<Shape<_4, _2>, Shape<_4, _3>>{};
156-auto stride = Stride<Stride<_4,_16>,Stride<_1,_32>>{};252+auto stride = Stride<Stride<_4, _16>, Stride<_1, _32>>{};
253+ 
157print(crd2offset(tla::MakeCoord(1, 5), shape, stride)); // 37254print(crd2offset(tla::MakeCoord(1, 5), shape, stride)); // 37
158```255```
159 256 
160-### 获取 Tilelayout257+这段代码表示：在一个逻辑大小为 `(8, 12)`、底层按分形格式排布的矩阵中，逻辑坐标 `(1, 5)` 对应的线性索引为 `37`。
161 258 
162-Tilelayout 可以用下列方式获取：259+## 获取 TileLayout
260+ 
261+TileLayout 可以通过 `GetTileLayout` 获取：
163 262 
164```cpp263```cpp
264+template <class Layout, class TileShape, class Coord>
265+auto GetTileLayout(Layout const& layout,
266+ TileShape const& tileShape,
267+ Coord const& coord);
268+ 
165using namespace tla;269using namespace tla;
166-Layout a = Layout<Shape<Shape<_4,_2>,Shape<_4,_3>>, Stride<Stride<_4,_16>,Stride<_1,_32>>>{}; // ((4,2),(4,3)):((4,16),(1,32))270+Layout a = Layout<Shape<Shape<_4, _2>, Shape<_4, _3>>,
167-Layout a0 = MakeLayoutTile(a, MakeShape(4, 4)); // ((4,1),(4,1)):((4,16),(1,32))271+ Stride<Stride<_4, _16>, Stride<_1, _32>>,
272+ Shape<_8, _12>>{};
273+ 
274+Layout a0 = GetTileLayout(a, MakeShape(4, 4), MakeCoord(6, 10));
275+// 结果可理解为：stride 保持不变，逻辑有效范围裁剪为 (2, 2)
168```276```
169 277 
170-`MakeLayoutTile`接口改变了原本layout的shape，不改变stride。278+参数语义如下：
279+ 
280+- `tileShape`：期望取出的 tile 大小，单位是元素。
281+- `coord`：tile 左上角在父 layout 逻辑空间中的元素坐标，单位也是元素。
282+ 
283+也就是说，`coord = (6, 10)` 的含义是“从逻辑第 6 行、第 10 列开始取 tile”，而不是“第 6 个 tile、第 10 个 tile”。
284+ 
285+### `GetTileLayout` 的核心语义
286+ 
287+`GetTileLayout` 返回的是一个 tile 视图的 `Layout`，不会改变底层数据排布。它主要做三件事：
288+ 
289+1. 保留原有 `stride()`，因为底层内存布局没有变化。
290+2. 用 `tileShape` 构造 tile 的 `shape()`；当父布局带有嵌套结构时，返回结果会在需要时保持同样的结构层次。
291+3. 根据父 layout 的 `originShape()` 和起始 `coord`，计算 tile 的 `originShape()`。
292+ 
293+其中第 3 步最关键：
294+ 
295+$$
296+origin\_shape[d] = \min(tileShape[d], \max(origin\_base[d] - coord[d], 0))
297+$$
298+ 
299+它表示“从当前位置开始，在逻辑上还剩多少有效元素”。因此：
300+ 
301+- 中间区域的 tile，`originShape == tileShape`。
302+- 触边的 tail tile，`originShape` 会自动缩小。
303+ 
304+### “按父 layout 的结构转换成对应的 `shape()`”是什么意思
305+ 
306+这句话的含义是：当父布局本身是嵌套布局时，tile 的 `shape()` 也需要保持同样的结构层次，这样后续访问规则才能继续复用。
307+ 
308+例如，父布局的行和列都按 `16` 为内层块组织：
309+ 
310+```text
311+parent shape = ((16, 7), (16, 7))
312+parent originShape = (100, 100)
313+```
314+ 
315+如果希望取一个逻辑大小为 `(32, 48)` 的 tile，那么这个 tile 的逻辑尺寸可以直接写成 `(32, 48)`，但在父布局是 `zN` 的前提下，它对应的 `shape()` 会按父布局的结构表达成：
316+ 
317+```text
318+tile logical size = (32, 48)
319+tile shape = ((16, 2), (16, 3))
320+```
321+ 
322+这里发生的是“结构转换”，不是“重新排布数据”：
323+ 
324+- 逻辑上，tile 仍然是 `32 x 48`。
325+- 布局上，它被表达成“每维一个 16 的内层块，再乘以外层块个数”。
326+- `stride()` 仍继承自父布局，因此访问规则不变。
327+ 
328+这样做的目的，是保证父 layout 和 tile layout 在结构层次上保持一致。
329+ 
330+![Origin_Shape-layout_2.png](https://raw.gitcode.com/user-images/assets/7631999/649c84f3-981f-49eb-be77-6cbf6fd1e5b3/Origin_Shape-layout_2.png 'Origin_Shape-layout_2.png')
331+ 
332+### 参数约束
333+ 
334+- `tileShape` 与 `coord` 都必须是一层 tuple，即 `depth == 1`。
335+- `rank(coord) == rank(tileShape)`。
336+ 
337+### 不同布局下的行为
338+ 
339+- 如果父 layout 是普通 vector 或 matrix，返回 layout 的 `shape()` 通常就等于 `tileShape`。
340+- 如果父 layout 是嵌套或分形布局，例如 `zN`、`nZ`、`zZ`、`L0C`，当前实现仅支持 `rank == 2`，并会把 `(rows, cols)` 形式的 `tileShape` 转换成与父布局同结构的嵌套 `Shape`。
341+ 
342+ 
Mdocs/2_Design/02_tla/03_tensor.md+398-26
@@ -1,56 +1,428 @@
1# TLA Tensors1# TLA Tensors
2 2 
3-这篇文档描述了CATLASS的TLA(Tensor Layout Abstraction)下的`Tensor`。3+本文介绍 TLA 中的 `Tensor`。
4 4 
5-本质上，`Tensor` （张量）表示一个多维数组。Tensor 抽象了数组元素在内存中的组织方式与存储方式的细节。这使得用户能够编写通用的访问多维数组的算法，并可根据张量的特性（traits）对算法进行特化。如张量的depth、rank、layout、数据的类型、位置等。5+如果说 `Layout` 负责描述“逻辑坐标如何映射到内存”，那么 `Tensor` 就是在 `Layout` 的基础上，再绑定具体数据、当前视图起点和存储层级后的可访问对象。
6 6 
7-`Tensor` 包含4个模板参数: `BuiltinTensor`、 `Layout`、 `Coord`、 `Position`。7+在本文中，`Tensor` 一律指逻辑视图：
8-关于 `Layout` 的描述, 请参考 [Layout](./01_layout.md)。
9 8 
10-## BuiltinTensor 和 Position9+- `MakeTensor` 创建的是视图，不发生数据拷贝。
10+- `operator()` 的切片结果是子视图，不发生数据拷贝。
11+- `GetTile` 与 `TileView` 返回的是 tile 视图，不发生数据拷贝。
12+- `MakeTensorLike` 只是把一块已有存储绑定成“与参考 Tensor 逻辑尺寸一致”的新视图，本身不执行数据搬运。
11 13 
12-`BuiltinTensor` 为AscendC内的 `GlobalTensor` 或者 `LocalTensor`，`Position` 为AscendC定义的各层级位置。相关使用参考AscendC文档。14+真正的数据移动应由显式的搬运或计算接口完成，而不是由这些视图构造接口隐式完成。
15+ 
16+关于 `Layout` 的基础定义，请先参考 [Layout](./01_layout.md)。
17+ 
18+## 先分清四个组成部分
19+ 
20+`Tensor` 的模板参数是 `BuiltinTensor`、`Layout`、`Coord`、`Position`。第一次接触时，建议先把这四部分分开理解。
21+ 
22+### BuiltinTensor
23+ 
24+`BuiltinTensor` 是 AscendC 提供的底层张量对象，例如 `GlobalTensor` 或 `LocalTensor`。它表示“底层存储对象本身”。
25+ 
26+### Layout
27+ 
28+`Layout` 描述逻辑坐标如何映射到内存，以及逻辑有效范围如何表达。
29+ 
30+### Coord
31+ 
32+`Coord` 是当前 `Tensor` 视图在BuiltinTensor所表达的父逻辑空间中的起点坐标。
33+ 
34+这里需要特别强调两点：
35+ 
36+- `coord` 的单位是元素，不是字节。
37+- `coord` 表示“这个视图从BuiltinTensor所表达的父逻辑空间的哪里开始看”，不是 tile 编号。
38+ 
39+例如，一个逻辑大小为 `(8, 16)` 的矩阵中，如果某个子 Tensor 的 `coord()` 是 `(2, 4)`，它表示“这个视图的左上角，对应父逻辑矩阵的第 2 行、第 4 列”。
40+ 
41+### Position
42+ 
43+`Position` 是 AscendC 中的位置标签，例如 `Arch::PositionGM{}`、`Arch::PositionL1{}`。它用于区分数据位于 GM、L1、L0 等哪一层存储。
13 44 
14## Tensor 构造45## Tensor 构造
15 46 
16-当前提供 `MakeTensor` 接口构造`Tensor`， 包含四个模板参数： `BuiltinTensor`、 `Layout`、 `Coord`、 `Position`。47+当前使用 `MakeTensor` 构造 `Tensor`。
17- 
18-有如下两种方式构造：
19 48 
20```cpp49```cpp
21using namespace tla;50using namespace tla;
22GlobalTensor<float> A = ...;51GlobalTensor<float> A = ...;
23 52 
24-// 缺省Coord, 默认为（0， 0）53+auto layout = tla::MakeLayout<float, Catlass::layout::RowMajor>(8, 16);
25-Layout w8xh16 = MakeLayout(MakeShape(8, Int<16>{}), MakeStride(Int<16>{},Int<1>{}));
26-Tensor tensor_8x16 = MakeTensor(A, w8xh16, Arch::PositionGM{});
27 54 
28-// 用户指定Coord55+// 1. 默认从逻辑坐标 (0, 0) 开始
29-Tensor tensor_8x16 = MakeTensor(A, w8xh16, tla::MakeCoord(1, 5), Arch::PositionGM{});56+auto tensorA = MakeTensor(A, layout, Arch::PositionGM{});
57+ 
58+// 2. 显式指定当前视图起点
59+auto tensorA_sub = MakeTensor(A, layout, tla::MakeCoord(1, 5), Arch::PositionGM{});
30```60```
31 61 
32-## Tensors 接口62+可以按下面的方式理解：
33 63 
34-TLA `Tensor` 提供获取相应特性的接口：64+- `layout` 决定“如何解释这块内存”。
65+- `coord` 决定“当前视图从BuiltinTensor所表达的父逻辑空间的哪里开始”。
35 66 
36-* `.data()`. 返回 `Tensor` 的内存。67+## Tensor 的常用接口
37 68 
38-* `.layout()`. 返回 `Tensor` 的 `layout`。69+TLA `Tensor` 提供以下常用接口：
39 70 
40-* `.coord()`. 返回 `Tensor` 的 `coord`。71+- `.data()`：返回底层内存对象。
72+- `.layout()`：返回布局。
73+- `.coord()`：返回当前视图起点。
74+- `.shape()`：返回 `layout.shape()`。
75+- `.stride()`：返回 `layout.stride()`。
76+- `.originShape()`：返回 `layout.originShape()`。
77+- `(coord0, coord1, ...)`：按坐标索引或切片。
41 78 
42-* `.shape()`. 返回 `Tensor` 的 `shape`。79+## 统一理解三类“坐标”
43 80 
44-* `.stride()`. 返回 `Tensor` 的 `stride`。81+TLA 文档中最容易混淆的是几类不同的“坐标”。下面给出统一约定。
45 82 
46-## 获取 TileTensor83+### 元素坐标 element coord
47 84 
48-提供一个 `GetTile` 接口获取 `Tensor` 的一片子tensor，会根据输入坐标对coord进行更新，并依据新的Tile的shape变换layout（只是逻辑层面的数据组织形式），底层的数据实体不变更。85+元素坐标表示“按元素计数的逻辑位置”，例如 `(row, col)`。`GetTile`、`crd2offset`、普通索引访问等接口使用的都是这种坐标。
86+ 
87+### tile 坐标 tile coord
88+ 
89+tile 坐标表示“第几个 tile”，不是第几个元素。例如在 `tileShape = (64, 128)` 时：
90+ 
91+- `tileCoord = (1, 2)` 表示第 1 个行 tile、第 2 个列 tile。
92+- 它对应的元素起点是 `(1 * 64, 2 * 128)`。
93+ 
94+### 视图起点 view coord
95+ 
96+`tensor.coord()` 表示当前 `Tensor` 视图在BuiltinTensor所表达的父逻辑空间中的起点。它由创建这个视图的操作决定，例如 `MakeTensor`、`GetTile`、`TileView` 或切片操作。
97+ 
98+可以用一句话概括：
99+ 
100+- `element coord` 是元素位置。
101+- `tile coord` 是 tile 编号。
102+- `tensor.coord()` 是当前视图的起点。
103+ 
104+## 用一个完整示例理解 `coord()`
105+ 
106+下面用同一个矩阵，串联 `MakeTensor`、`GetTile` 几种情形。
49 107 
50```cpp108```cpp
51using namespace tla;109using namespace tla;
52-Layout w8xh16 = MakeLayout(MakeShape(8, Int<16>{}), MakeStride(Int<16>{},Int< 1>{}));
53-Tensor tensor_8x16 = MakeTensor(A, w8xh16, Arch::PositionGM{});
54 110 
55-auto tensor_tile = GetTile(tensor_8x16, tla::MakeCoord(2, 4), MakeShape(4, 8)); // (4,8):(_16,_1)111+GlobalTensor<float> A = ...;
112+auto layout = tla::MakeLayout<float, Catlass::layout::RowMajor>(8, 16);
113+ 
114+auto tensorA = MakeTensor(A, layout, Arch::PositionGM{});
115+// tensorA.coord() == (0, 0)
116+ 
117+auto tensorA_sub = MakeTensor(A, layout, MakeCoord(1, 5), Arch::PositionGM{});
118+// tensorA_sub.coord() == (1, 5)
119+ 
120+auto tileA = GetTile(tensorA_sub, MakeCoord(2, 4), MakeShape(4, 8));
121+// tileA.coord() == (3, 9)
56```122```
123+ 
124+上面分别表示：
125+ 
126+1. `tensorA` 直接观察整块逻辑矩阵，因此起点是 `(0, 0)`。
127+2. `tensorA_sub` 从BuiltinTensor所表达的父逻辑空间的 `(1, 5)` 开始观察，因此起点变为 `(1, 5)`。
128+3. `tileA` 在 `tensorA_sub` 的基础上再取一个起点为 `(2, 4)` 的 tile，因此新视图起点是 `(1, 5) + (2, 4) = (3, 9)`。
129+ 
130+![Origin_Shape-tensor_1.png](https://raw.gitcode.com/user-images/assets/7631999/004cd08b-cda4-4c99-a5c1-63770d71f151/Origin_Shape-tensor_1.png 'Origin_Shape-tensor_1.png')
131+ 
132+## 使用 `operator()` 进行索引与切片
133+ 
134+TLA `Tensor` 支持使用 `operator()` 做索引，也支持使用 `tla::_` 表达整维切片，返回子 Tensor 视图。
135+ 
136+### 基本规则
137+ 
138+- 不带 `tla::_` 时，`tensor(i, j, ...)` 返回一个底层 `BuiltinTensor` 访问结果，本质上对应 `tensor.data()[offset]`。
139+- 带 `tla::_` 时，`tensor(..., tla::_, ...)` 返回子 Tensor 视图；被索引的维度会被固定，保留 `tla::_` 所在维度。
140+- 这里使用的坐标参数必须是一层 tuple，即每个维度都是标量或 `tla::_`，不支持嵌套 tuple。
141+ 
142+等价语义可写为：
143+ 
144+```cpp
145+tensor.data()[tensor.layout()(tensor.coord() + coord_arg)]
146+```
147+ 
148+### 输出 Tensor 的维度
149+ 
150+设输入张量 rank 为 $R$，`coord` 中出现 `tla::_` 的维度索引集合为 $\{d_0, d_1, ..., d_{k-1}\}$，则：
151+ 
152+- 输出 Tensor 的 rank 为 $k$。
153+- 输出 Tensor 的 `layout.shape()`、`layout.stride()`、`layout.originShape()` 是输入布局在这些维度上的投影。
154+- 输出 Tensor 的 `coord()` 会重新从全 0 开始，因为它已经成为新的局部视图。
155+ 
156+例如，对 3D 张量 `A(B, M, K)`：
157+ 
158+```cpp
159+auto A2 = A3(b, tla::_, tla::_); // 3D -> 2D，得到 (M, K) 视图
160+auto A1 = A2(r, tla::_) // 2D -> 1D，得到 (K)视图
161+```
162+ 
163+![Origin_Shape-tensor_2.png](https://raw.gitcode.com/user-images/assets/7631999/dd1a947a-2371-4e6e-9fef-51a5ceb1556a/Origin_Shape-tensor_2.png 'Origin_Shape-tensor_2.png')
164+ 
165+## 获取 TileTensor
166+ 
167+### GetTile
168+ 
169+`GetTile` 用于从父 Tensor 上切出一个 tile 视图，不拷贝数据。
170+ 
171+```cpp
172+template <class Tensor, class Coord, class Shape>
173+auto GetTile(Tensor const& tensor,
174+ Coord const& coord,
175+ Shape const& shape);
176+```
177+ 
178+参数语义如下：
179+ 
180+- `coord`：元素坐标，表示 tile 左上角在父 Tensor 逻辑空间中的起点。
181+- `shape`：tile 的期望尺寸，单位是元素。
182+ 
183+```cpp
184+using namespace tla;
185+ 
186+auto layout = tla::MakeLayout<float, Catlass::layout::RowMajor>(8, 16);
187+auto tensor = MakeTensor(A, layout, Arch::PositionGM{});
188+ 
189+// 从逻辑坐标 (2, 4) 开始，取一个 4 x 8 的 tile
190+auto tile = GetTile(tensor, tla::MakeCoord(2, 4), MakeShape(4, 8));
191+```
192+ 
193+返回结果可理解为：
194+ 
195+- `tile.coord()` = `tensor.coord()` + `(2, 4)`。
196+- `tile.layout().shape()` 表示期望 tile 尺寸或其与父布局结构一致的表达形式。
197+- `tile.layout().originShape()` 表示该 tile 真实有效的逻辑范围，触边时会自动裁剪。
198+ 
199+### 使用约束
200+ 
201+- 支持 `tensor.layout().depth == 1`。
202+- 若 `tensor.layout().depth > 1`，即分形或嵌套布局，当前 `GetTileLayout` 仅支持 `rank == 2`。
203+- `coord` 与 `shape` 都必须为一层 tuple，并满足 `rank(coord) == rank(shape) == Tensor::rank`。
204+ 
205+### 边界行为
206+ 
207+例如父 Tensor 的逻辑尺寸是 `(8, 16)`，执行：
208+ 
209+```cpp
210+auto tail = GetTile(tensor, tla::MakeCoord(6, 10), MakeShape(4, 8));
211+```
212+ 
213+那么：
214+ 
215+- 期望尺寸仍然是 `(4, 8)`。
216+- 但逻辑上只剩下 2 行、6 列有效数据。
217+- 因此 `tail.layout().originShape()` 会变成 `(2, 6)`。
218+ 
219+### TileView
220+ 
221+`TileView` 与 `GetTile` 的行为等价，区别只在于输入坐标的单位不同：
222+ 
223+- `GetTile` 接收元素坐标。
224+- `TileView` 接收 tile 坐标。
225+ 
226+```cpp
227+template <class TensorT, class TileCoord, class TileShape>
228+auto TileView(TensorT const& tensor,
229+ TileCoord const& tileCoord,
230+ TileShape const& tileShape);
231+```
232+ 
233+例如：
234+ 
235+```cpp
236+auto tensorTileA = tla::TileView(
237+ tensorA,
238+ tla::MakeCoord(0u, kLoopIdx),
239+ tla::MakeShape(Int<L1_TILE_M>{}, Int<L1_TILE_K>{})
240+);
241+```
242+ 
243+### 等价关系
244+ 
245+`TileView` 与 `GetTile` 可以直接按下面的等式理解：
246+ 
247+```cpp
248+TileView(t, tileCoord, tileShape) = GetTile(t, tileCoord ⊙ tileShape, tileShape)
249+```
250+ 
251+这里的 `⊙` 表示逐维相乘，例如：
252+ 
253+```cpp
254+(1, 2) ⊙ (64, 128) = (64, 256)
255+```
256+ 
257+这条等式表示：
258+ 
259+1. `TileView` 先把 tile 坐标转换为元素坐标。
260+2. 然后按 `GetTile` 的规则创建同一个 tile 视图。
261+ 
262+因此，两者的差别只在于调用者提供的是哪一种坐标单位，而不是返回结果的逻辑语义。
263+ 
264+![Origin_Shape-tensor_3.png](https://raw.gitcode.com/user-images/assets/7631999/b543ee2b-c938-4e19-87b8-40abd0c81b53/Origin_Shape-tensor_3.png 'Origin_Shape-tensor_3.png')
265+ 
266+### 为什么 `TileView` 更适合分块循环
267+ 
268+在实际 kernel 或 block 循环中，循环变量通常就是 tile 编号，而不是元素坐标。因此 `TileView` 往往更直接。
269+ 
270+下面用同一个按 K 维分块的例子做对比。
271+ 
272+#### 写法一：使用 `GetTile`
273+ 
274+```cpp
275+constexpr uint32_t tileM = 64;
276+constexpr uint32_t tileK = 128;
277+ 
278+for (uint32_t kTile = 0; kTile < kTiles; ++kTile) {
279+ auto coord = tla::MakeCoord(0u, kTile * tileK);
280+ auto shape = tla::MakeShape(tileM, tileK);
281+ auto tensorTileA = tla::GetTile(tensorA, coord, shape);
282+ // use tensorTileA
283+}
284+```
285+ 
286+#### 写法二：使用 `TileView`
287+ 
288+```cpp
289+constexpr uint32_t tileM = 64;
290+constexpr uint32_t tileK = 128;
291+ 
292+for (uint32_t kTile = 0; kTile < kTiles; ++kTile) {
293+ auto tensorTileA = tla::TileView(
294+ tensorA,
295+ tla::MakeCoord(0u, kTile),
296+ tla::MakeShape(tileM, tileK)
297+ );
298+ // use tensorTileA
299+}
300+```
301+ 
302+这两段代码的逻辑结果相同，但第二种写法直接使用 tile 坐标，更贴近分块循环本身的语义，也更不容易把“tile 坐标”和“元素坐标”混淆。
303+ 
304+## 创建类似的 Tensor
305+ 
306+### MakeTensorLike
307+ 
308+`MakeTensorLike` 用于创建一个“逻辑尺寸与 `likeTensor` 一致”的新 Tensor。最常见的用途是：从一个已有 tile 视图出发，在另一层内存中构造对应 Tensor，并自动继承其 `originShape()`。
309+ 
310+在未指定layoutbase时，行为为根据 LayoutTagDst 决定布局，从 LikeTensor::Element 推断 ElementDst，从 likeTensor 的 originShape 提取尺寸。调用MakeLayout<ElementDst, LayoutTagDst>(originShape())构造目标 layout（可能会因分型布局合法要求对shape进行以分型为粒度的向上取整）
311+ 
312+指定layoutbase时，使用MakeLayout(layoutBase.shape(), layoutBase.stride(), likeTensor.originShape())构造目标layout。
313+ 
314+这里仍然需要强调：`MakeTensorLike` 构造的是新视图，不执行数据搬运。它只是把用户传入的 `builtinTensor` 绑定成一个新的 TLA `Tensor`，并让这个新视图复用 `likeTensor` 的逻辑尺寸语义。
315+ 
316+当前 `MakeTensorLike` 仅支持 `likeTensor.rank <= 2`。
317+ 
318+接口分为三类典型场景。
319+ 
320+```cpp
321+// 1) 从 LikeTensor::Element 推断 ElementDst
322+template <class LayoutTagDst, class BuiltinTensor, class LikeTensor, class PositionType>
323+auto MakeTensorLike(BuiltinTensor const& builtinTensor,
324+ LikeTensor const& likeTensor,
325+ PositionType);
326+ 
327+// 2) 显式指定 ElementDst
328+template <class LayoutTagDst, class ElementDst, class BuiltinTensor, class LikeTensor, class PositionType>
329+auto MakeTensorLike(BuiltinTensor const& builtinTensor,
330+ LikeTensor const& likeTensor,
331+ PositionType);
332+ 
333+// 3) 提供 layoutBase
334+template <class LayoutTagDst, class BuiltinTensor, class LikeTensor, class PositionType, class LayoutBase>
335+auto MakeTensorLike(BuiltinTensor const& builtinTensor,
336+ LikeTensor const& likeTensor,
337+ PositionType,
338+ LayoutBase const& layoutBase);
339+ 
340+template <class LayoutTagDst, class ElementDst, class BuiltinTensor, class LikeTensor, class PositionType, class LayoutBase>
341+auto MakeTensorLike(BuiltinTensor const& builtinTensor,
342+ LikeTensor const& likeTensor,
343+ PositionType,
344+ LayoutBase const& layoutBase);
345+```
346+ 
347+### 场景一：源和目标元素类型相同
348+ 
349+这是最常见的场景。例如从 GM 中的一个 `half` tile 创建对应的 L1 Tensor，元素类型不变，只是存储层级改变。
350+ 
351+```cpp
352+auto tensorTileA = tla::TileView(
353+ tensorA,
354+ tla::MakeCoord(blockM, kTile),
355+ tla::MakeShape(L1_TILE_M, L1_TILE_K)
356+);
357+ 
358+auto tensorL1A = tla::MakeTensorLike<LayoutTagL1A>(
359+ l1ATensorList[l1ListId],
360+ tensorTileA,
361+ Arch::PositionL1{}
362+);
363+ 
364+// 结果：
365+// 1. tensorL1A 使用 L1 目标布局
366+// 2. tensorL1A 的 originShape 与 tensorTileA 相同
367+// 3. 元素类型从 likeTensor 自动推断
368+```
369+ 
370+### 场景二：目标元素类型不同
371+ 
372+当目标 Tensor 的元素类型与源 Tensor 不一致时，需要显式指定 `ElementDst`。例如：
373+ 
374+- L0C 中使用 accumulator 类型。
375+- 需要从 `half` 输入生成 `float` 累加视图。
376+- 目标内存对象的 `PrimType` 与 `LikeTensor::Element` 不同。
377+ 
378+```cpp
379+auto tensorL0C = tla::MakeTensorLike<LayoutTagL0C, float>(
380+ l0cTensor,
381+ tensorTileC,
382+ Arch::PositionL0C{}
383+);
384+ 
385+// 结果：
386+// 1. tensorL0C 的逻辑尺寸继承自 tensorTileC
387+// 2. 目标元素类型显式为 float
388+// 3. 适用于 accumulator 或类型提升场景
389+```
390+ 
391+### 场景三：目标布局需要额外控制
392+ 
393+有些场景下，仅指定 `LayoutTagDst` 还不够，因为目标布局的基础形状或步长需要用户显式给出。例如：
394+ 
395+- 目标 Tensor 采用特定分形布局。
396+- 需要固定某个 L1 的物理排布。注意：L0的排布由originShape唯一确定，因此定制L0上的非预期排布为不合法行为。
397+- 需要预先给出特殊的 `shape/stride` 结构，但逻辑有效范围仍要继承自 `likeTensor`。
398+ 
399+```cpp
400+auto layoutBaseL1A = tla::MakeLayout<half, LayoutTagL1A>(L1_TILE_M, L1_TILE_K);
401+ 
402+auto tensorL1A = tla::MakeTensorLike<LayoutTagL1A>(
403+ l1ATensor,
404+ tensorTileA,
405+ Arch::PositionL1A{},
406+ layoutBaseL1A
407+);
408+ 
409+// 结果：
410+// 1. tensorL1A 的 shape/stride 来自 layoutBaseL1A
411+// 2. tensorL1A 的 originShape 继承自GM上的 tensorTileA
412+// 3. 即使当前 tile 是尾块，逻辑有效范围也不会丢失
413+```
414+ 
415+如果既要控制目标布局，又要显式指定目标元素类型，可以使用同时带 `layoutBase` 和 `ElementDst` 的重载。
416+ 
417+## 实际使用模式
418+ 
419+在 block 层和 kernel 层，常见写法通常是两步：
420+ 
421+1. 用 `TileView` 从父 Tensor 得到 tile 视图，自动处理边界。
422+2. 用 `MakeTensorLike` 在目标内存层级构造对应 Tensor，自动继承 `originShape()`。
423+ 
424+这套模式的价值在于：
425+ 
426+- 主流程始终围绕 tile 编程。
427+- 尾块逻辑通过 `originShape` 自动传递。
428+- 数据搬运和计算阶段都能复用同一套逻辑尺寸语义，减少边界分支和歧义。
Mexamples/40_flash_attention_infer_tla/fai_kernel.cpp+25-3
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4 * CANN Open Software License Agreement Version 2.0 (the "License").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -225,8 +225,30 @@ class FAInferKernelTla {
225 225 
226 auto layoutQCube = tla::MakeLayout(MakeShape(qSBlockSize, qNBlockSize, embed), 226 auto layoutQCube = tla::MakeLayout(MakeShape(qSBlockSize, qNBlockSize, embed),
227 MakeStride((int64_t)qHeads * embed, (int64_t)embed, (int64_t)Int<1>{}));227 MakeStride((int64_t)qHeads * embed, (int64_t)embed, (int64_t)Int<1>{}));
228- auto layoutKCube = tla::MakeLayout<ElementK, LayoutK>(strideKV, blockStackNum * pagedBlockSize);228+ // Build kvHead-sliced physical views for K/V.
229- auto layoutVCube = tla::MakeLayout<ElementV, LayoutV>(blockStackNum * pagedBlockSize, strideKV);229+ //
230+ // Paged KV cache physical storage: [numBlocks, blockSize, kvHeads, embed]
231+ // For a fixed kvHeadIdx, a token's embedding starts at:
232+ // base = gmOffsetK/V + token_slot * strideKV, where token_slot = blockId * blockSize + blockOffset
233+ //
234+ // express this as:
235+ // - paged: token_slot in [0, numBlocks * blockSize)
236+ // - contig: token_slot in [0, kvSeqlen)
237+ //
238+ // K_head_view: shape [embed, token_slot_len], stride [1, strideKV]
239+ // V_head_view: shape [token_slot_len, embed], stride [strideKV, 1]
240+ uint32_t kvPhysTokenSlots = 0;
241+ if constexpr (PAGED_CACHE_FLAG) {
242+ kvPhysTokenSlots = fATilingData->numBlocks * pagedBlockSize;
243+ } else {
244+ kvPhysTokenSlots = kvSeqlen;
245+ }
246+ auto layoutKCube = tla::MakeLayout(
247+ MakeShape(embed, kvPhysTokenSlots),
248+ MakeStride(Int<1>{}, (int64_t)strideKV));
249+ auto layoutVCube = tla::MakeLayout(
250+ MakeShape(kvPhysTokenSlots, embed),
251+ MakeStride((int64_t)strideKV, Int<1>{}));
230 auto layoutOTmpCube = tla::MakeLayout<ElementOTmp, LayoutOTmp>(rowNum, embed);252 auto layoutOTmpCube = tla::MakeLayout<ElementOTmp, LayoutOTmp>(rowNum, embed);
231 auto tensorQ = tla::MakeTensor(gQ[gmQOffset], layoutQCube, Arch::PositionGM{});253 auto tensorQ = tla::MakeTensor(gQ[gmQOffset], layoutQCube, Arch::PositionGM{});
232 auto tensorK = tla::MakeTensor(gK[gmKOffset], layoutKCube, Arch::PositionGM{});254 auto tensorK = tla::MakeTensor(gK[gmKOffset], layoutKCube, Arch::PositionGM{});
Aexamples/45_strided_batched_matmul_tla/CMakeLists.txt+12-0
@@ -0,0 +1,12 @@
1+# -----------------------------------------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# -----------------------------------------------------------------------------------------------------------
10+ 
11+set_source_files_properties(strided_batched_matmul_tla.cpp PROPERTIES LANGUAGE ASCEND)
12+catlass_example_add_executable(45_strided_batched_matmul_tla cube strided_batched_matmul_tla.cpp)
Aexamples/45_strided_batched_matmul_tla/README.md+46-0
@@ -0,0 +1,46 @@
1+# StridedBatchedMatmulTla Example Readme
2+## 代码组织
3+```
4+├── 45_strided_batched_matmul_tla
5+│   ├── CMakeLists.txt # CMake编译文件
6+│   ├── README.md
7+│   └── strided_batched_matmul_tla.cpp # 主文件
8+```
9+## 使用示例
10+- 获取代码之后编译相应的算子可执行文件，可参考[quickstart](../../docs/quickstart.md#算子编译)
S
Ssunhao_hw3月24日

quickstart.md路径有更新

likedislike
S
sunhao_hw
3月24日 评论：
11+- 执行算子
12+```
13+# 编译指定用例
14+bash scripts/build.sh 45_strided_batched_matmul_tla
15+cd output/bin
16+# 基本用法：可执行文件名 batch轴|m轴|n轴|k轴|Device ID
17+# Device ID 可选，默认为0
18+./45_strided_batched_matmul_tla 5 256 512 1024 0
19+ 
20+# layout 定制（仅支持 row/col，大小写不敏感；可选，默认 row row）
21+# - layoutA: A(M,K) 的 layout
22+# - layoutB: B(K,N) 的 layout
23+# layout 是一个“可选的尾部分组”，可以追加在任意一种参数组合的末尾；
24+./45_strided_batched_matmul_tla 5 256 512 1024 row col
25+./45_strided_batched_matmul_tla 5 256 512 1024 0 row col
26+ 
27+# stride 定制（单位：elements）
28+# - lda/ldb/ldc：分别为 A(M,K)/B(K,N)/C(M,N) 的 leading dimension
29+# - A: row 时 lda>=K；col 时 lda>=M
30+# - B: row 时 ldb>=N；col 时 ldb>=K
31+# - C: 本示例固定为 row，因此 ldc>=N
32+# - strideA/strideB/strideC：batch 维度上相邻两矩阵的步长
33+#
34+# 只指定 lda/ldb/ldc（strideBatch 默认连续）
35+./45_strided_batched_matmul_tla 5 256 512 1024 0 1100 600 600
36+#
37+# 同时指定 batch stride（支持 batch 间 padding）
38+./45_strided_batched_matmul_tla 5 256 512 1024 0 1100 600 600 300000 400000 500000
39+ 
40+# layout + stride 混用（当传 layoutA/layoutB 时，必须放在命令行最后两个参数位置）
41+./45_strided_batched_matmul_tla 5 256 512 1024 0 1100 600 600 300000 400000 500000 col row
42+```
43+执行结果如下，说明精度比对成功。
44+```
45+Compare success.
46+```
Aexamples/45_strided_batched_matmul_tla/strided_batched_matmul_tla.cpp+399-0
@@ -0,0 +1,399 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+// By setting the K_MAX_SHAPE_DIM macro, the dimension of the AscendC Tensor's ShapeInfo is configured to 0,
12+// optimizing stack space. If you need to use the ShapeInfo of the AscendC Tensor, please undefine this macro.
13+#ifndef K_MAX_SHAPE_DIM
14+#define K_MAX_SHAPE_DIM 0
15+#endif
16+ 
17+#include "catlass/gemm/kernel/strided_batched_matmul_tla.hpp"
18+ 
19+#include "catlass/arch/arch.hpp"
20+#include "catlass/catlass.hpp"
21+#include "catlass/gemm/block/block_mmad.hpp"
22+#include "catlass/gemm/block/block_swizzle.hpp"
23+#include "catlass/gemm/device/device_gemm.hpp"
24+#include "catlass/gemm/dispatch_policy.hpp"
25+#include "catlass/gemm/gemm_type.hpp"
26+#include "catlass/layout/layout.hpp"
27+#include "catlass/status.hpp"
28+#include "tla/layout.hpp"
29+#include "tla/tensor.hpp"
30+ 
31+#include "golden.hpp"
32+#include "helper.hpp"
33+ 
34+ 
35+using namespace Catlass;
36+using namespace tla;
37+ 
38+struct BatchedMatmulTlaOptions {
39+ const std::string HELPER =
40+ "problem_count m n k [device_id] [lda ldb ldc] [strideA strideB strideC] [layoutA layoutB]\n"
41+ " layoutA/layoutB: row | col (case-insensitive)\n"
42+ " Note: C is row in this example (ldc>=n, strideC based on row).";
43+ 
44+ Catlass::GemmCoord problemShape{128, 128, 128};
45+ uint32_t problemCount{1};
46+ int32_t deviceId{0};
47+ 
48+ // Stride customization (unit: elements).
49+ // - lda: stride of A on M axis (A is [M,K] RowMajor)
50+ // - ldb: stride of B on K axis (B is [K,N] RowMajor)
51+ // - ldc: stride of C on M axis (C is [M,N] RowMajor)
52+ int64_t lda{-1};
53+ int64_t ldb{-1};
54+ int64_t ldc{-1};
55+ // stride between batches
56+ int64_t strideA{-1};
57+ int64_t strideB{-1};
58+ int64_t strideC{-1};
59+ 
60+ // Layout selection
61+ // - A: [M,K] (RowMajor or ColumnMajor)
62+ // - B: [K,N] (RowMajor or ColumnMajor)
63+ // - C: [M,N] (RowMajor only in this example)
64+ enum class MatrixLayout { RowMajor, ColumnMajor };
65+ MatrixLayout layoutA{MatrixLayout::RowMajor};
66+ MatrixLayout layoutB{MatrixLayout::RowMajor};
67+ 
68+ BatchedMatmulTlaOptions() = default;
69+ 
70+ static bool IsLayoutToken(const std::string &s) {
71+ if (s.empty()) {
72+ return false;
73+ }
74+ std::string t;
75+ t.resize(s.size());
76+ std::transform(s.begin(), s.end(), t.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
77+ return (t == "row" || t == "col");
78+ }
79+ 
80+ static MatrixLayout ParseLayoutToken(const std::string &s) {
81+ std::string t;
82+ t.resize(s.size());
83+ std::transform(s.begin(), s.end(), t.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
84+ if (t == "row") {
85+ return MatrixLayout::RowMajor;
86+ }
87+ return MatrixLayout::ColumnMajor;
88+ }
89+ 
90+ int Parse(int argc, const char **argv) {
91+ // Supported forms:
92+ // 1) name problem_count m n k
93+ // 2) name problem_count m n k device_id
94+ // 3) name problem_count m n k device_id lda ldb ldc
95+ // 4) name problem_count m n k device_id lda ldb ldc strideA strideB strideC
96+ // Each form may optionally append: layoutA layoutB (two tokens), e.g. "row col".
97+ int argcEffective = argc;
98+ if (argc >= 7) {
99+ std::string maybeA(argv[argc - 2]);
100+ std::string maybeB(argv[argc - 1]);
101+ if (IsLayoutToken(maybeA) && IsLayoutToken(maybeB)) {
102+ layoutA = ParseLayoutToken(maybeA);
103+ layoutB = ParseLayoutToken(maybeB);
104+ argcEffective -= 2;
105+ }
106+ }
107+ 
108+ if (!(argcEffective == 5 || argcEffective == 6 || argcEffective == 9 || argcEffective == 12)) {
109+ std::cerr << TOSTRING(CATLASS_EXAMPLE_NAME) << " " << HELPER << std::endl;
110+ return -1;
111+ }
112+ 
113+ 
114+ problemCount = std::atoi(argv[1]);
115+ problemShape.m() = std::atoi(argv[2]);
116+ problemShape.n() = std::atoi(argv[3]);
117+ problemShape.k() = std::atoi(argv[4]);
118+ 
119+ if (argcEffective >= 6) {
120+ deviceId = std::atoi(argv[5]);
121+ }
122+ 
123+ uint32_t m = problemShape.m();
124+ uint32_t n = problemShape.n();
125+ uint32_t k = problemShape.k();
126+ 
127+ // Default: contiguous per-matrix and contiguous between batches.
128+ // Interpret lda/ldb as the leading dimension in the chosen layout.
129+ lda = (layoutA == MatrixLayout::RowMajor) ? static_cast<int64_t>(k) : static_cast<int64_t>(m);
130+ ldb = (layoutB == MatrixLayout::RowMajor) ? static_cast<int64_t>(n) : static_cast<int64_t>(k);
131+ ldc = static_cast<int64_t>(n); // C is RowMajor in this example
132+ 
133+ strideA = (layoutA == MatrixLayout::RowMajor) ? static_cast<int64_t>(m) * lda : static_cast<int64_t>(k) * lda;
134+ strideB = (layoutB == MatrixLayout::RowMajor) ? static_cast<int64_t>(k) * ldb : static_cast<int64_t>(n) * ldb;
135+ strideC = static_cast<int64_t>(m) * ldc;
136+ 
137+ if (argcEffective >= 9) {
138+ lda = std::atoll(argv[6]);
139+ ldb = std::atoll(argv[7]);
140+ ldc = std::atoll(argv[8]);
141+ 
142+ strideA = (layoutA == MatrixLayout::RowMajor) ? static_cast<int64_t>(m) * lda : static_cast<int64_t>(k) * lda;
143+ strideB = (layoutB == MatrixLayout::RowMajor) ? static_cast<int64_t>(k) * ldb : static_cast<int64_t>(n) * ldb;
144+ strideC = static_cast<int64_t>(m) * ldc;
145+ }
146+ if (argcEffective == 12) {
147+ strideA = std::atoll(argv[9]);
148+ strideB = std::atoll(argv[10]);
149+ strideC = std::atoll(argv[11]);
150+ }
151+ 
152+ // Basic validation for ND layouts.
153+ int64_t minLda = (layoutA == MatrixLayout::RowMajor) ? static_cast<int64_t>(k) : static_cast<int64_t>(m);
154+ int64_t minLdb = (layoutB == MatrixLayout::RowMajor) ? static_cast<int64_t>(n) : static_cast<int64_t>(k);
155+ if (lda < minLda || ldb < minLdb || ldc < static_cast<int64_t>(n)) {
156+ std::cerr << "Invalid leading dimensions: require lda>=" << minLda
157+ << ", ldb>=" << minLdb
158+ << ", ldc>=" << n << "." << std::endl;
159+ return -1;
160+ }
161+ 
162+ int64_t minMatA = (layoutA == MatrixLayout::RowMajor)
163+ ? (static_cast<int64_t>(m - 1) * lda + static_cast<int64_t>(k))
164+ : (static_cast<int64_t>(k - 1) * lda + static_cast<int64_t>(m));
165+ int64_t minMatB = (layoutB == MatrixLayout::RowMajor)
166+ ? (static_cast<int64_t>(k - 1) * ldb + static_cast<int64_t>(n))
167+ : (static_cast<int64_t>(n - 1) * ldb + static_cast<int64_t>(k));
168+ int64_t minMatC = static_cast<int64_t>(m - 1) * ldc + static_cast<int64_t>(n);
169+ 
170+ if (strideA < minMatA || strideB < minMatB || strideC < minMatC) {
171+ std::cerr << "Invalid batch strides: require strideA/strideB/strideC large enough for one matrix."
172+ << std::endl;
173+ return -1;
174+ }
175+ 
176+ return 0;
177+ }
178+};
179+ 
180+using Options = BatchedMatmulTlaOptions;
181+ 
182+template <typename LayoutTagA>
183+static auto MakeTlaLayoutA(uint32_t batchCount, uint32_t m, uint32_t k, int64_t strideA, int64_t lda) {
184+ if constexpr (std::is_same_v<LayoutTagA, layout::RowMajor>) {
185+ return tla::MakeLayout(
186+ tla::MakeShape(batchCount, m, k),
187+ tla::MakeStride(strideA, lda, tla::Int<1>{})
188+ );
189+ } else {
190+ return tla::MakeLayout(
191+ tla::MakeShape(batchCount, m, k),
192+ tla::MakeStride(strideA, tla::Int<1>{}, lda)
193+ );
194+ }
195+}
196+ 
197+template <typename LayoutTagB>
198+static auto MakeTlaLayoutB(uint32_t batchCount, uint32_t k, uint32_t n, int64_t strideB, int64_t ldb) {
199+ if constexpr (std::is_same_v<LayoutTagB, layout::RowMajor>) {
200+ return tla::MakeLayout(
201+ tla::MakeShape(batchCount, k, n),
202+ tla::MakeStride(strideB, ldb, tla::Int<1>{})
203+ );
204+ } else {
205+ return tla::MakeLayout(
206+ tla::MakeShape(batchCount, k, n),
207+ tla::MakeStride(strideB, tla::Int<1>{}, ldb)
208+ );
209+ }
210+}
211+ 
212+template <typename LayoutTagA, typename LayoutTagB>
213+static void RunWithLayouts(const Options &options) {
214+ aclrtStream stream{nullptr};
215+ 
216+ ACL_CHECK(aclInit(nullptr));
217+ ACL_CHECK(aclrtSetDevice(options.deviceId));
218+ ACL_CHECK(aclrtCreateStream(&stream));
219+ 
220+ uint32_t batchCount = options.problemCount;
221+ uint32_t m = options.problemShape.m();
222+ uint32_t n = options.problemShape.n();
223+ uint32_t k = options.problemShape.k();
224+ 
225+ using ElementA = half;
226+ using ElementB = half;
227+ using ElementC = half;
228+ 
229+ using LayoutTagC = layout::RowMajor; // must be RowMajor
230+ LayoutTagA tagA{m, k, options.lda};
231+ LayoutTagB tagB{k, n, options.ldb};
232+ LayoutTagC tagC{m, n, options.ldc};
233+ 
234+ // Capacity in elements (last element offset + 1)
235+ int64_t capA = (static_cast<int64_t>(batchCount) - 1) * options.strideA +
236+ static_cast<int64_t>(tagA.GetOffset(MakeCoord(m - 1, k - 1))) + 1;
237+ int64_t capB = (static_cast<int64_t>(batchCount) - 1) * options.strideB +
238+ static_cast<int64_t>(tagB.GetOffset(MakeCoord(k - 1, n - 1))) + 1;
239+ int64_t capC = (static_cast<int64_t>(batchCount) - 1) * options.strideC +
240+ static_cast<int64_t>(tagC.GetOffset(MakeCoord(m - 1, n - 1))) + 1;
241+ 
242+ size_t lenA = static_cast<size_t>(capA);
243+ size_t lenB = static_cast<size_t>(capB);
244+ size_t lenC = static_cast<size_t>(capC);
245+ 
246+ size_t sizeA = lenA * sizeof(ElementA);
247+ size_t sizeB = lenB * sizeof(ElementB);
248+ size_t sizeC = lenC * sizeof(ElementC);
249+ 
250+ // allocate memory of A and copy to device side
251+ std::vector<ElementA> hostA(lenA, 1.0f);
252+ golden::FillRandomData<ElementA>(hostA, -5.0f, 5.0f);
253+ uint8_t *deviceA{nullptr};
254+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void **>(&deviceA), sizeA, ACL_MEM_MALLOC_HUGE_FIRST));
255+ ACL_CHECK(aclrtMemcpy(deviceA, sizeA, hostA.data(), sizeA, ACL_MEMCPY_HOST_TO_DEVICE));
256+ 
257+ // allocate memory of B and copy to device side
258+ std::vector<ElementB> hostB(lenB, 1.0f);
259+ golden::FillRandomData<ElementB>(hostB, -5.0f, 5.0f);
260+ uint8_t *deviceB{nullptr};
261+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void **>(&deviceB), sizeB, ACL_MEM_MALLOC_HUGE_FIRST));
262+ ACL_CHECK(aclrtMemcpy(deviceB, sizeB, hostB.data(), sizeB, ACL_MEMCPY_HOST_TO_DEVICE));
263+ 
264+ // allocate memory of C
265+ std::vector<ElementC> hostC(lenC);
266+ uint8_t *deviceC{nullptr};
267+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void **>(&deviceC), sizeC, ACL_MEM_MALLOC_HUGE_FIRST));
268+ 
269+ // Get the number of cube cores of the current hardware
270+ auto aicCoreNum = platform_ascendc::PlatformAscendCManager::GetInstance()->GetCoreNumAic();
271+ 
272+ using ArchTag = Arch::AtlasA2;
273+ using DispatchPolicy = Gemm::MmadPingpongTlaV2<ArchTag, true>;
274+ using L1TileShape = Shape<_128, _256, _256>;
275+ using L0TileShape = Shape<_128, _256, _64>;
276+ 
277+ using TileCopy =
278+ Gemm::Tile::PackedTileCopyTla<ArchTag, ElementA, LayoutTagA, ElementB, LayoutTagB, ElementC, LayoutTagC>;
279+ using BlockMmad = Gemm::Block::BlockMmadTla<
280+ DispatchPolicy, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, void, TileCopy>;
281+ using BlockEpilogue = void;
282+ 
283+ auto layoutA = MakeTlaLayoutA<LayoutTagA>(batchCount, m, k, options.strideA, options.lda);
284+ auto layoutB = MakeTlaLayoutB<LayoutTagB>(batchCount, k, n, options.strideB, options.ldb);
285+ auto layoutC = tla::MakeLayout(
286+ tla::MakeShape(batchCount, m, n),
287+ tla::MakeStride(options.strideC, options.ldc, tla::Int<1>{})
288+ );
289+ 
290+ if (options.problemShape.m() > options.problemShape.n()) {
291+ // Swizzle offset is 3 and direction is 0.
292+ using BlockScheduler = typename Gemm::Block::GemmIdentityBlockSwizzle<3, 0>;
293+ 
294+ // kernel level
295+ using MatmulKernel = Gemm::Kernel::StridedBatchedMatmulTla<BlockMmad, BlockEpilogue, BlockScheduler>;
296+ 
297+ using MatmulAdapter = Gemm::Device::DeviceGemm<MatmulKernel>;
298+ typename MatmulKernel::Arguments arguments{
299+ batchCount, options.problemShape,
300+ deviceA, layoutA,
301+ deviceB, layoutB,
302+ deviceC, layoutC
303+ };
304+ MatmulAdapter matmulOp;
305+ 
306+ uint8_t *deviceWorkspace{nullptr};
307+ matmulOp.CanImplement(arguments);
308+ matmulOp.Initialize(arguments, deviceWorkspace);
309+ matmulOp(stream, aicCoreNum);
310+ ACL_CHECK(aclrtSynchronizeStream(stream));
311+ 
312+ } else {
313+ // Swizzle offset is 3 and direction is 1.
314+ using BlockScheduler = typename Gemm::Block::GemmIdentityBlockSwizzle<3, 1>;
315+ 
316+ // kernel level
317+ using MatmulKernel = Gemm::Kernel::StridedBatchedMatmulTla<BlockMmad, BlockEpilogue, BlockScheduler>;
318+ 
319+ using MatmulAdapter = Gemm::Device::DeviceGemm<MatmulKernel>;
320+ typename MatmulKernel::Arguments arguments{
321+ batchCount, options.problemShape,
322+ deviceA, layoutA,
323+ deviceB, layoutB,
324+ deviceC, layoutC
325+ };
326+ MatmulAdapter matmulOp;
327+ 
328+ uint8_t *deviceWorkspace{nullptr};
329+ matmulOp.CanImplement(arguments);
330+ matmulOp.Initialize(arguments, deviceWorkspace);
331+ matmulOp(stream, aicCoreNum);
332+ ACL_CHECK(aclrtSynchronizeStream(stream));
333+ }
334+ 
335+ ACL_CHECK(aclrtMemcpy(hostC.data(), sizeC, deviceC, sizeC, ACL_MEMCPY_DEVICE_TO_HOST));
336+ 
337+ // comparison of precision with matmul computed on cpu
338+ size_t packedLenC = static_cast<size_t>(batchCount) * m * n;
339+ std::vector<ElementC> packedC(packedLenC);
340+ std::vector<float> packedGolden(packedLenC);
341+ 
342+ for (uint32_t b = 0; b < batchCount; ++b) {
343+ size_t basePacked = static_cast<size_t>(b) * m * n;
344+ size_t baseA = static_cast<size_t>(b) * static_cast<size_t>(options.strideA);
345+ size_t baseB = static_cast<size_t>(b) * static_cast<size_t>(options.strideB);
346+ size_t baseC = static_cast<size_t>(b) * static_cast<size_t>(options.strideC);
347+ for (uint32_t i = 0; i < m; ++i) {
348+ for (uint32_t j = 0; j < n; ++j) {
349+ size_t idxPacked = basePacked + static_cast<size_t>(i) * n + j;
350+ size_t offC = baseC + static_cast<size_t>(tagC.GetOffset(MakeCoord(i, j)));
351+ packedC[idxPacked] = hostC[offC];
352+ 
353+ float acc = 0.0f;
354+ for (uint32_t kk = 0; kk < k; ++kk) {
355+ size_t offA = baseA + static_cast<size_t>(tagA.GetOffset(MakeCoord(i, kk)));
356+ size_t offB = baseB + static_cast<size_t>(tagB.GetOffset(MakeCoord(kk, j)));
357+ acc += static_cast<float>(hostA[offA]) * static_cast<float>(hostB[offB]);
358+ }
359+ packedGolden[idxPacked] = acc;
360+ }
361+ }
362+ }
363+ 
364+ std::vector<uint64_t> errorIndices = golden::CompareData(packedC, packedGolden, k);
365+ if (errorIndices.empty()) {
366+ std::cout << "Compare success." << std::endl;
367+ } else {
368+ std::cerr << "Compare failed. Error count: " << errorIndices.size() << std::endl;
369+ }
370+ 
371+ ACL_CHECK(aclrtFree(deviceA));
372+ ACL_CHECK(aclrtFree(deviceB));
373+ ACL_CHECK(aclrtFree(deviceC));
374+ ACL_CHECK(aclrtDestroyStream(stream));
375+ ACL_CHECK(aclrtResetDevice(options.deviceId));
376+ ACL_CHECK(aclFinalize());
377+}
378+ 
379+static void Run(const Options &options) {
380+ using ML = Options::MatrixLayout;
381+ if (options.layoutA == ML::RowMajor && options.layoutB == ML::RowMajor) {
382+ RunWithLayouts<layout::RowMajor, layout::RowMajor>(options);
383+ } else if (options.layoutA == ML::RowMajor && options.layoutB == ML::ColumnMajor) {
384+ RunWithLayouts<layout::RowMajor, layout::ColumnMajor>(options);
385+ } else if (options.layoutA == ML::ColumnMajor && options.layoutB == ML::RowMajor) {
386+ RunWithLayouts<layout::ColumnMajor, layout::RowMajor>(options);
387+ } else {
388+ RunWithLayouts<layout::ColumnMajor, layout::ColumnMajor>(options);
389+ }
390+}
391+ 
392+int main(int argc, const char **argv) {
393+ Options options;
394+ if (options.Parse(argc, argv) != 0) {
395+ return -1;
396+ }
397+ Run(options);
398+ return 0;
399+}
Mexamples/CMakeLists.txt+1-0
@@ -139,6 +139,7 @@ set(EXAMPLE_ATLASA2
139 41_sparse_matmul_tla139 41_sparse_matmul_tla
140 42_quant_optimized_matmul_tla140 42_quant_optimized_matmul_tla
141 44_quant_matmul_full_loadA_tla141 44_quant_matmul_full_loadA_tla
142+ 45_strided_batched_matmul_tla
142 102_dynamic_optimized_matmul143 102_dynamic_optimized_matmul
143 103_dynamic_optimized_quant_matmul_per_token_basic144 103_dynamic_optimized_quant_matmul_per_token_basic
144)145)
Minclude/catlass/gemm/block/block_mmad.hpp+1-0
@@ -140,6 +140,7 @@ struct BlockMmadTla {
140 140 
141#include "catlass/gemm/block/block_mmad_pingpong_tla.hpp"141#include "catlass/gemm/block/block_mmad_pingpong_tla.hpp"
142#include "catlass/gemm/block/block_mmad_pingpong_dequant_tla.hpp"142#include "catlass/gemm/block/block_mmad_pingpong_dequant_tla.hpp"
143+#include "catlass/gemm/block/block_mmad_pingpong_tla_v2.hpp"
143#include "catlass/gemm/block/block_mmad_preload_tla.hpp"144#include "catlass/gemm/block/block_mmad_preload_tla.hpp"
144#include "catlass/gemm/block/block_mmad_preload_async_with_callback_tla.hpp"145#include "catlass/gemm/block/block_mmad_preload_async_with_callback_tla.hpp"
145 146 
Ainclude/catlass/gemm/block/block_mmad_pingpong_tla_v2.hpp+400-0
@@ -0,0 +1,400 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_V2_HPP
12+#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_V2_HPP
13+ 
14+#include "catlass/catlass.hpp"
15+#include "catlass/arch/resource.hpp"
16+#include "catlass/coord.hpp"
17+#include "catlass/gemm_coord.hpp"
18+#include "catlass/gemm/dispatch_policy.hpp"
19+#include "catlass/gemm/helper.hpp"
20+#include "catlass/gemm/tile/tile_copy.hpp"
21+#include "catlass/gemm/tile/tile_mmad.hpp"
22+#include "tla/layout.hpp"
23+#include "tla/tensor.hpp"
24+ 
25+namespace Catlass::Gemm::Block {
26+ 
27+template <
28+ class ArchTag_,
29+ bool ENABLE_UNIT_FLAG_,
30+ class L1TileShape_,
31+ class L0TileShape_,
32+ class ElementA_,
33+ class ElementB_,
34+ class ElementC_,
35+ class ElementBias_,
36+ class TileCopy_,
37+ class TileMmad_
38+>
39+struct BlockMmadTla <
40+ MmadPingpongTlaV2<ArchTag_, ENABLE_UNIT_FLAG_>,
41+ L1TileShape_,
42+ L0TileShape_,
43+ ElementA_,
44+ ElementB_,
45+ ElementC_,
46+ ElementBias_,
47+ TileCopy_,
48+ TileMmad_
49+> {
50+public:
51+ // Type Aliases
52+ using DispatchPolicy = MmadPingpongTlaV2<ArchTag_, ENABLE_UNIT_FLAG_>;
53+ using ArchTag = typename DispatchPolicy::ArchTag;
54+ using L1TileShape = L1TileShape_;
55+ using L0TileShape = L0TileShape_;
56+ using ElementA = ElementA_;
57+ using LayoutA = typename TileCopy_::LayoutA;
58+ using ElementB = ElementB_;
59+ using LayoutB = typename TileCopy_::LayoutB;
60+ using ElementC = ElementC_;
61+ using LayoutC = typename TileCopy_::LayoutC;
62+ 
63+ using TileMmad = TileMmad_;
64+ 
65+ using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A;
66+ using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B;
67+ 
68+ using ElementAccumulator = typename TileCopy_::ElementAccumulator;
69+ 
70+ using LayoutTagL1A = typename TileCopy_::LayoutTagL1A;
71+ using LayoutTagL1B = typename TileCopy_::LayoutTagL1B;
72+ using LayoutTagL0A = typename TileCopy_::LayoutTagL0A;
73+ using LayoutTagL0B = typename TileCopy_::LayoutTagL0B;
74+ using LayoutTagL0C = typename TileCopy_::LayoutTagL0C;
75+ 
76+ using L1AAlignHelper = typename TileCopy_::L1AAlignHelper;
77+ using L1BAlignHelper = typename TileCopy_::L1BAlignHelper;
78+ 
79+ static_assert(tla::is_tuple<L1TileShape>::value && tla::is_static<L1TileShape>::value,
80+ "L1TileShape must be tla::tuple and static!");
81+ static_assert(tla::is_tuple<L0TileShape>::value && tla::is_static<L0TileShape>::value,
82+ "L0TileShape must be tla::tuple and static!");
83+ 
84+ static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG;
85+ static constexpr uint32_t STAGES = DispatchPolicy::STAGES;
86+ static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{});
87+ static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{});
88+ static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{});
89+ static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{});
90+ static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{});
91+ static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{});
92+ 
93+ // L1 tile size
94+ static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA);
95+ static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB);
96+ // L0 tile size
97+ static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA);
98+ static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB);
99+ static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator);
100+ 
101+ // Check LayoutC
102+ static_assert(tla::detail::isRowMajor<LayoutC>::value ||
103+ ((std::is_same_v<ElementC, half> || std::is_same_v<ElementC, bfloat16_t> ||
104+ std::is_same_v<ElementC, float>) && tla::detail::iszN<ElementC, LayoutC>::value),
105+ "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!");
106+ 
107+ // Check L1TileShape
108+ static_assert((L1A_TILE_SIZE + L1B_TILE_SIZE) * STAGES <= ArchTag::L1_SIZE,
109+ "L1TileShape exceeding the L1 space!");
110+ 
111+ // Check L0TileShape
112+ static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!");
113+ static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!");
114+ static_assert(L0C_TILE_SIZE <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!");
115+ 
116+ static constexpr uint32_t _32B = 32*8; // in bits
117+ static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N,
118+ "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet");
119+ static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K");
120+ static_assert(L1_TILE_M * SizeOfBits<ElementA>::value % _32B == 0, "L1TileShape::M must be 32B aligned.");
121+ static_assert(L1_TILE_K * SizeOfBits<ElementA>::value % _32B == 0, "L1TileShape::K must be 32B aligned.");
122+ static_assert(L1_TILE_K * SizeOfBits<ElementB>::value % _32B == 0, "L1TileShape::K must be 32B aligned.");
123+ static_assert(L1_TILE_N * SizeOfBits<ElementB>::value % _32B == 0, "L1TileShape::N must be 32B aligned.");
124+ static_assert(L0_TILE_K * SizeOfBits<ElementB>::value % _32B == 0, "L0TileShape::K must be 32B aligned.");
125+ 
126+ static constexpr auto L1A_LAYOUT = tla::MakeLayout<ElementA, LayoutTagL1A>(tla::Int<L1_TILE_M>{}, tla::Int<L1_TILE_K>{});
127+ static constexpr auto L1B_LAYOUT = tla::MakeLayout<ElementB, LayoutTagL1B>(tla::Int<L1_TILE_K>{}, tla::Int<L1_TILE_N>{});
128+ 
129+ /// Construct
130+ CATLASS_DEVICE
131+ BlockMmadTla(Arch::Resource<ArchTag> &resource, uint32_t l1BufAddrStart = 0)
132+ {
133+ uint32_t l1AOffset = l1BufAddrStart;
134+ uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * STAGES;
135+ // Init buffers
136+ for (uint32_t i = 0; i < STAGES; i++) {
137+ // Assign L1/L0A/L0B space for each stages
138+ l1ATensorList[i] = resource.l1Buf.template GetBufferByByte<ElementA>(l1AOffset + L1A_TILE_SIZE * i);
139+ l1BTensorList[i] = resource.l1Buf.template GetBufferByByte<ElementB>(l1BOffset + L1B_TILE_SIZE * i);
140+ l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(L0A_TILE_SIZE * i);
141+ l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(L0B_TILE_SIZE * i);
142+ 
143+ // Assign event ID for each stages
144+ l1AEventList[i] = i;
145+ l1BEventList[i] = i + STAGES;
146+ l0AEventList[i] = i;
147+ l0BEventList[i] = i + STAGES;
148+ 
149+ // The event id that needs to be set before the loop
150+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[i]);
151+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[i]);
152+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[i]);
153+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[i]);
154+ }
155+ l0CTensor = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(0);
156+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
157+ }
158+ 
159+ /// Destructor
160+ CATLASS_DEVICE
161+ ~BlockMmadTla()
162+ {
163+ for (uint32_t i = 0; i < STAGES; i++) {
164+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[i]);
165+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[i]);
166+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[i]);
167+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[i]);
168+ }
169+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
170+ }
171+ 
172+ /// Perform a block-scoped matrix multiply-accumulate
173+ template <class TensorA, class TensorB, class TensorC>
174+ CATLASS_DEVICE
175+ void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC)
176+ {
177+ using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A<TensorA>;
178+ using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B<TensorB>;
179+ CopyGmToL1A copyGmToL1A;
180+ CopyGmToL1B copyGmToL1B;
181+#if (defined (CATLASS_ARCH) && CATLASS_ARCH == 2201)
182+ using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm<TensorC>;
183+ CopyL0CToGm copyL0CToDst;
184+#endif
185+#if (defined (CATLASS_ARCH) && CATLASS_ARCH == 3510)
186+ using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst<TensorC>;
187+ CopyL0CToDst copyL0CToDst;
188+#endif
189+ 
190+ // Create an accumulator tensor view on L0C buffer.
191+ // - Logical size comes from tensorC.layout().originShape() (tail-aware)
192+ // - Layout is constructed from LayoutTagL0C
193+ // - coord is initialized to (0, 0) (new buffer view)
194+ auto tensorL0C = tla::MakeTensorLike<LayoutTagL0C, ElementAccumulator>(l0CTensor, tensorC, Arch::PositionL0C{});
195+ 
196+ // load first matrix A tile from GM to L1
197+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[l1ListId]);
198+ // TileView: tileCoord is in tile units (not element units).
199+ // It internally converts tileCoord to elementOffset = tileCoord ⊙ tileShape and handles tail tiles via originShape.
200+ auto tensorTileA = tla::TileView(tensorA,
201+ tla::MakeCoord(0u, 0u), // (m_tile, k_tile)
202+ tla::MakeShape(tla::Int<L1_TILE_M>{}, tla::Int<L1_TILE_K>{}));
203+ auto tensorL1A = tla::MakeTensorLike<LayoutTagL1A>(l1ATensorList[l1ListId], tensorTileA, Arch::PositionL1{}, L1A_LAYOUT);
204+ copyGmToL1A(tensorL1A, tensorTileA);
205+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList[l1ListId]);
206+ 
207+ // load first matrix B tile from GM to L1
208+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[l1ListId]);
209+ auto tensorTileB = tla::TileView(tensorB,
210+ tla::MakeCoord(0u, 0u), // (k_tile, n_tile)
211+ tla::MakeShape(tla::Int<L1_TILE_K>{}, tla::Int<L1_TILE_N>{}));
212+ auto tensorL1B = tla::MakeTensorLike<LayoutTagL1B>(l1BTensorList[l1ListId], tensorTileB, Arch::PositionL1{}, L1B_LAYOUT);
213+ copyGmToL1B(tensorL1B, tensorTileB);
214+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList[l1ListId]);
215+ 
216+ if constexpr (!ENABLE_UNIT_FLAG) {
217+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
218+ }
219+ 
220+ uint32_t mPartLoop = CeilDiv<L0_TILE_M>(tla::get<0>(tensorL0C.originShape())); // dim 0 = M
221+ uint32_t nPartLoop = CeilDiv<L0_TILE_N>(tla::get<1>(tensorL0C.originShape())); // dim 1 = N
222+ 
223+ // main loop
224+ uint32_t kTileCount = CeilDiv<L1_TILE_K>(tla::get<1>(tensorA.originShape())); // dim 1 = K
225+ for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) {
226+ uint32_t l1ListIdNext = (l1ListId + 1 < STAGES) ? (l1ListId + 1) : 0;
227+ // preload next tile from GM to L1
228+ if (kLoopIdx < kTileCount - 1) {
229+ uint32_t kLoopIdxNext = kLoopIdx + 1;
230+ 
231+ // Get L1 tensor for next stage
232+ auto l1ATensor = l1ATensorList[l1ListIdNext];
233+ auto l1BTensor = l1BTensorList[l1ListIdNext];
234+ // Get GM tile for next stage
235+ auto tensorTileA = tla::TileView(tensorA,
236+ tla::MakeCoord(0u, kLoopIdxNext), // (m_tile, k_tile)
237+ tla::MakeShape(tla::Int<L1_TILE_M>{}, tla::Int<L1_TILE_K>{}));
238+ auto tensorTileB = tla::TileView(tensorB,
239+ tla::MakeCoord(kLoopIdxNext, 0u), // (k_tile, n_tile)
240+ tla::MakeShape(tla::Int<L1_TILE_K>{}, tla::Int<L1_TILE_N>{}));
241+ auto tensorL1A = tla::MakeTensorLike<LayoutTagL1A>(l1ATensor, tensorTileA, Arch::PositionL1{}, L1A_LAYOUT);
242+ auto tensorL1B = tla::MakeTensorLike<LayoutTagL1B>(l1BTensor, tensorTileB, Arch::PositionL1{}, L1B_LAYOUT);
243+ 
244+ // load next matrix A tile from GM to L1
245+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[l1ListIdNext]);
246+ copyGmToL1A(tensorL1A, tensorTileA);
247+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList[l1ListIdNext]);
248+ 
249+ // load next matrix B tile from GM to L1
250+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[l1ListIdNext]);
251+ copyGmToL1B(tensorL1B, tensorTileB);
252+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList[l1ListIdNext]);
253+ }
254+ 
255+ // Get L1 tensor for current stage
256+ auto l1ATensor = l1ATensorList[l1ListId];
257+ auto l1BTensor = l1BTensorList[l1ListId];
258+ // Create tile view for current K iteration
259+ auto tensorTileA = tla::TileView(tensorA,
260+ tla::MakeCoord(0u, kLoopIdx), // (m_tile, k_tile)
261+ tla::MakeShape(tla::Int<L1_TILE_M>{}, tla::Int<L1_TILE_K>{}));
262+ auto tensorTileB = tla::TileView(tensorB,
263+ tla::MakeCoord(kLoopIdx, 0u), // (k_tile, n_tile)
264+ tla::MakeShape(tla::Int<L1_TILE_K>{}, tla::Int<L1_TILE_N>{}));
265+ auto tensorL1A = tla::MakeTensorLike<LayoutTagL1A>(l1ATensor, tensorTileA, Arch::PositionL1{}, L1A_LAYOUT);
266+ auto tensorL1B = tla::MakeTensorLike<LayoutTagL1B>(l1BTensor, tensorTileB, Arch::PositionL1{}, L1B_LAYOUT);
267+ // Get the loop nums on L0 based on current L1 tile's actual K size
268+ uint32_t kPartLoop = CeilDiv<L0_TILE_K>(tla::get<1>(tensorL1A.originShape())); // dim 1 = K
269+ 
270+ for (int mPartIdx = 0; mPartIdx < mPartLoop; mPartIdx++) {
271+ for (int kPartIdx = 0; kPartIdx < kPartLoop; kPartIdx++) {
272+ 
273+ // Locate the current tile on L0A
274+ auto l0ATile = l0ATensorList[l0AListId];
275+ // Locate the current tile of matrix A on L1
276+ // Take a (L0_TILE_M, L0_TILE_K) tile from the current L1A tile (tile coordinates within L1 tile).
277+ auto tensorTileL1A = tla::TileView(
278+ tensorL1A,
279+ tla::MakeCoord(mPartIdx, kPartIdx),
280+ tla::MakeShape(tla::Int<L0_TILE_M>{}, tla::Int<L0_TILE_K>{})
281+ );
282+ auto tensorL0A = tla::MakeTensorLike<LayoutTagL0A>(l0ATile, tensorTileL1A, Arch::PositionL0A{});
283+ 
284+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[l0AListId]);
285+ if ((mPartIdx == 0) && (kPartIdx == 0)) {
286+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList[l1ListId]);
287+ }
288+ 
289+ // Load current tile from L1 to L0A
290+ copyL1ToL0A(tensorL0A, tensorTileL1A);
291+ 
292+ if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) {
293+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[l1ListId]);
294+ }
295+ 
296+ for (int nPartIdx = 0; nPartIdx < nPartLoop; nPartIdx++) {
297+ // Locate the current tile on L0B
298+ auto l0BTile = l0BTensorList[l0BListId];
299+ // Locate the current tile of matrix B on L1
300+ // Take a (L0_TILE_K, L0_TILE_N) tile from the current L1B tile (tile coordinates within L1 tile).
301+ auto tensorTileL1B = tla::TileView(
302+ tensorL1B,
303+ tla::MakeCoord(kPartIdx, nPartIdx),
304+ tla::MakeShape(tla::Int<L0_TILE_K>{}, tla::Int<L0_TILE_N>{})
305+ );
306+ auto tensorL0B = tla::MakeTensorLike<LayoutTagL0B>(l0BTile, tensorTileL1B, Arch::PositionL0B{});
307+ 
308+ // Wait for mmad finished
309+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[l0BListId]);
310+ // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1
311+ if ((kPartIdx == 0) && (nPartIdx == 0)) {
312+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList[l1ListId]);
313+ }
314+ 
315+ // Load current tile from L1 to L0B
316+ copyL1ToL0B(tensorL0B, tensorTileL1B);
317+ 
318+ // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1
319+ if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) {
320+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[l1ListId]);
321+ }
322+ // Notify to do mmad
323+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
324+ 
325+ // Locate the current tile on L0C
326+ // View into L0C accumulator tile (tile coordinates in (m_part, n_part)).
327+ auto tensorTileL0C = tla::TileView(
328+ tensorL0C,
329+ tla::MakeCoord(mPartIdx, nPartIdx),
330+ tla::MakeShape(tla::Int<L0_TILE_M>{}, tla::Int<L0_TILE_N>{})
331+ );
332+ 
333+ // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator
334+ // Wait for loading L0B
335+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
336+ 
337+ // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0
338+ bool initC = ((kLoopIdx == 0) && (kPartIdx == 0));
339+ // If the unit flag is enabled, the unit flag is set according to the calculation progress
340+ uint8_t unitFlag = 0b00;
341+ if constexpr (ENABLE_UNIT_FLAG) {
342+ if ((kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) &&
343+ (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) {
344+ unitFlag = 0b11;
345+ } else {
346+ unitFlag = 0b10;
347+ }
348+ }
349+ // Perform calculation operations
350+ tileMmad(tensorTileL0C, tensorL0A, tensorL0B, initC, unitFlag);
351+ 
352+ // Notify to move the next L0B tile
353+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[l0BListId]);
354+ l0BListId = (l0BListId + 1 < STAGES) ? (l0BListId + 1) : 0;
355+ }
356+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[l0AListId]);
357+ l0AListId = (l0AListId + 1 < STAGES) ? (l0AListId + 1) : 0;
358+ }
359+ }
360+ l1ListId = l1ListIdNext;
361+ }
362+ 
363+ // copy block out
364+ if constexpr (!ENABLE_UNIT_FLAG) {
365+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
366+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
367+ copyL0CToDst(tensorC, tensorL0C);
368+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
369+ } else {
370+ copyL0CToDst(tensorC, tensorL0C, 0b11);
371+ }
372+ }
373+ 
374+protected:
375+ // Multi-stage tensors list
376+ AscendC::LocalTensor<ElementA> l1ATensorList[STAGES];
377+ AscendC::LocalTensor<ElementB> l1BTensorList[STAGES];
378+ AscendC::LocalTensor<ElementA> l0ATensorList[STAGES];
379+ AscendC::LocalTensor<ElementB> l0BTensorList[STAGES];
380+ AscendC::LocalTensor<ElementAccumulator> l0CTensor;
381+ 
382+ // Multi-stage event id list
383+ int32_t l1AEventList[STAGES];
384+ int32_t l1BEventList[STAGES];
385+ int32_t l0AEventList[STAGES];
386+ int32_t l0BEventList[STAGES];
387+ 
388+ // The id of current stage
389+ uint32_t l1ListId{0};
390+ uint32_t l0AListId{0};
391+ uint32_t l0BListId{0};
392+ 
393+ TileMmad tileMmad;
394+ CopyL1ToL0A copyL1ToL0A;
395+ CopyL1ToL0B copyL1ToL0B;
396+};
397+ 
398+} // namespace Catlass::Gemm::Block
399+ 
400+#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_V2_HPP
Minclude/catlass/gemm/dispatch_policy.hpp+7-0
@@ -316,6 +316,13 @@ struct MmadPreloadAsyncWithCallback : public MmadBase<ArchTag_, true> {
316 static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_;316 static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_;
317};317};
318 318 
319+// 基于TLA提供block层不感知尾块逻辑，tile层感知originShape的流程。使用TileView与MakeTensorLike。
320+template <class ArchTag_, bool ENABLE_UNIT_FLAG_ = false>
tianxinghui
tianxinghuitianxinghui3月17日

加一下这个policy的注释，相比不带v2的差别

likedislike
321+struct MmadPingpongTlaV2 : public MmadBase<ArchTag_, ENABLE_UNIT_FLAG_> {
322+ static constexpr uint32_t STAGES = 2;
323+ static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
324+};
325+ 
319template <class ArchTag_, bool ENABLE_UNIT_FLAG_ = false>326template <class ArchTag_, bool ENABLE_UNIT_FLAG_ = false>
320struct SparseMatmulMultiBlockOnKAxis : public MmadBase<ArchTag_, false> {327struct SparseMatmulMultiBlockOnKAxis : public MmadBase<ArchTag_, false> {
321 static constexpr uint32_t STAGES = 2;328 static constexpr uint32_t STAGES = 2;
Ainclude/catlass/gemm/kernel/strided_batched_matmul_tla.hpp+169-0
@@ -0,0 +1,169 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef CATLASS_GEMM_KERNEL_STRIDED_BATCHED_MATMUL_TLA_HPP
12+#define CATLASS_GEMM_KERNEL_STRIDED_BATCHED_MATMUL_TLA_HPP
13+ 
14+#include "catlass/catlass.hpp"
15+#include "catlass/arch/resource.hpp"
16+#include "catlass/coord.hpp"
17+#include "catlass/gemm_coord.hpp"
18+#include "catlass/matrix_coord.hpp"
19+#include "tla/layout.hpp"
20+#include "tla/tensor.hpp"
21+ 
22+namespace Catlass::Gemm::Kernel {
23+ 
24+// Template for Strided Batched Matmul kernel. Compute strided batched C = A * B
25+template <
26+ class BlockMmad_,
27+ class BlockEpilogue_,
28+ class BlockScheduler_
29+>
30+class StridedBatchedMatmulTla {
31+public:
32+ using BlockMmad = BlockMmad_;
33+ using ArchTag = typename BlockMmad::ArchTag;
34+ using L1TileShape = typename BlockMmad::L1TileShape;
35+ using ElementA = typename BlockMmad::ElementA;
36+ using LayoutA2D = typename BlockMmad::LayoutA;
37+ using ElementB = typename BlockMmad::ElementB;
38+ using LayoutB2D = typename BlockMmad::LayoutB;
39+ using ElementC = typename BlockMmad::ElementC;
40+ using LayoutC2D = typename BlockMmad::LayoutC;
41+ using ElementAccumulator = typename BlockMmad::ElementAccumulator;
42+ 
43+ using LayoutA = tla::MakeBatchedLayout_t<LayoutA2D>;
44+ using LayoutB = tla::MakeBatchedLayout_t<LayoutB2D>;
45+ using LayoutC = tla::MakeBatchedLayout_t<LayoutC2D>;
46+ 
47+ using BlockScheduler = BlockScheduler_;
48+ 
49+ static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{});
50+ static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{});
51+ 
52+ /// Parameters structure
53+ struct Params {
54+ uint32_t batchCount;
55+ GemmCoord problemShape;
56+ GM_ADDR ptrA;
57+ LayoutA layoutA;
58+ GM_ADDR ptrB;
59+ LayoutB layoutB;
60+ GM_ADDR ptrC;
61+ LayoutC layoutC;
62+ 
63+ CATLASS_HOST_DEVICE
64+ Params() {}
65+ 
66+ CATLASS_HOST_DEVICE
67+ Params(uint32_t batchCount_, GemmCoord const &problemShape_,
68+ GM_ADDR ptrA_, LayoutA layoutA_,
69+ GM_ADDR ptrB_, LayoutB layoutB_,
70+ GM_ADDR ptrC_, LayoutC layoutC_)
71+ : batchCount(batchCount_), problemShape(problemShape_),
72+ ptrA(ptrA_), layoutA(layoutA_),
73+ ptrB(ptrB_), layoutB(layoutB_),
74+ ptrC(ptrC_), layoutC(layoutC_) {}
75+ };
76+ 
77+ struct Arguments {
78+ uint32_t batchCount;
79+ GemmCoord problemShape;
80+ uint8_t *ptrA; LayoutA layoutA;
81+ uint8_t *ptrB; LayoutB layoutB;
82+ uint8_t *ptrC; LayoutC layoutC;
83+ };
84+ 
85+ static bool CanImplement(const Arguments &args)
86+ {
87+ return true;
88+ }
89+ 
90+ static size_t GetWorkspaceSize(const Arguments &args)
91+ {
92+ return 0;
93+ }
94+ 
95+ static Params ToUnderlyingArguments(const Arguments &args, uint8_t *workspace)
96+ {
97+ Params params{args.batchCount, args.problemShape,
98+ args.ptrA, args.layoutA,
99+ args.ptrB, args.layoutB,
100+ args.ptrC, args.layoutC};
101+ return params;
102+ }
103+ 
104+ CATLASS_DEVICE
105+ StridedBatchedMatmulTla() {}
106+ 
107+ template <int32_t CORE_TYPE = g_coreType>
108+ CATLASS_DEVICE
109+ void operator()(Params const &params);
110+ 
111+ template <>
112+ CATLASS_DEVICE
113+ void operator()<AscendC::AIC>(Params const &params) {
114+ BlockScheduler matmulBlockScheduler(params.problemShape, MakeCoord(L1_TILE_M, L1_TILE_N));
115+ uint32_t coreLoops = params.batchCount * matmulBlockScheduler.GetCoreLoops();
116+ 
117+ Arch::Resource<ArchTag> resource;
118+ BlockMmad blockMmad(resource);
119+ 
120+ AscendC::GlobalTensor<ElementA> gmA;
121+ gmA.SetGlobalBuffer((__gm__ ElementA *)params.ptrA);
122+ AscendC::GlobalTensor<ElementB> gmB;
123+ gmB.SetGlobalBuffer((__gm__ ElementB *)params.ptrB);
124+ AscendC::GlobalTensor<ElementC> gmC;
125+ gmC.SetGlobalBuffer((__gm__ ElementC *)params.ptrC);
126+ 
127+ auto tensorA3 = tla::MakeTensor(gmA, params.layoutA, Arch::PositionGM{});
128+ auto tensorB3 = tla::MakeTensor(gmB, params.layoutB, Arch::PositionGM{});
129+ auto tensorC3 = tla::MakeTensor(gmC, params.layoutC, Arch::PositionGM{});
130+ 
131+ for (uint32_t loopIdx = AscendC::GetBlockIdx(); loopIdx < coreLoops; loopIdx += AscendC::GetBlockNum()) {
132+ uint32_t batchIdx = matmulBlockScheduler.GetBatchIdx(loopIdx);
133+ GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);
134+ 
135+ // Slice to rank-2 tensors: Tensor(batchIdx, _, _)
136+ auto tensorA = tensorA3(batchIdx, tla::_, tla::_);
137+ auto tensorB = tensorB3(batchIdx, tla::_, tla::_);
138+ auto tensorC = tensorC3(batchIdx, tla::_, tla::_);
139+ 
140+ auto tensorBlockA = tla::TileView(
141+ tensorA,
142+ tla::MakeCoord(blockCoord.m(), 0u),
143+ tla::MakeShape(L1_TILE_M, params.problemShape.k())
144+ );
145+ auto tensorBlockB = tla::TileView(
146+ tensorB,
147+ tla::MakeCoord(0u, blockCoord.n()),
148+ tla::MakeShape(params.problemShape.k(), L1_TILE_N)
149+ );
150+ auto tensorBlockC = tla::TileView(
151+ tensorC,
152+ tla::MakeCoord(blockCoord.m(), blockCoord.n()),
153+ tla::MakeShape(L1_TILE_M, L1_TILE_N)
154+ );
155+ 
156+ blockMmad(tensorBlockA, tensorBlockB, tensorBlockC);
157+ }
158+ 
159+ AscendC::PipeBarrier<PIPE_ALL>();
160+ }
161+ 
162+ template <>
163+ CATLASS_DEVICE
164+ void operator()<AscendC::AIV>(Params const &params) {}
165+};
166+ 
167+} // namespace Catlass::Gemm::Kernel
168+ 
169+#endif // CATLASS_GEMM_KERNEL_STRIDED_BATCHED_MATMUL_TLA_HPP
Minclude/catlass/gemm/tile/ascend950/copy_gm_to_l1.hpp+9-9
@@ -49,8 +49,8 @@ struct TileCopyTla<
49 "The input parameters do not match. TensorSrc must be GM and RowMajor, while TensorDst must be L1 and zN"49 "The input parameters do not match. TensorSrc must be GM and RowMajor, while TensorDst must be L1 and zN"
tianxinghui
tianxinghuitianxinghui3月11日

包含了gm2l1、l0c2gm，但没有提供l12l0的copy，在maketensorlike指定layoutbase时，会有精度问题。

likedislike
50 );50 );
51 51 
52- const uint32_t nValue = tla::get<0>(srcTensor.shape());52+ const uint32_t nValue = tla::get<0>(srcTensor.originShape());
53- const uint32_t dValue = tla::get<1>(srcTensor.shape());53+ const uint32_t dValue = tla::get<1>(srcTensor.originShape());
54 const uint32_t srcDValue = tla::get<0>(srcTensor.stride());54 const uint32_t srcDValue = tla::get<0>(srcTensor.stride());
55 const uint32_t dstInnerStrideRow = tla::get<0, 0>(dstTensor.stride());55 const uint32_t dstInnerStrideRow = tla::get<0, 0>(dstTensor.stride());
56 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());56 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
@@ -97,8 +97,8 @@ struct TileCopyTla<
97 "The input parameters do not match. TensorSrc must be GM and zN, while TensorDst must be L1 and zN"97 "The input parameters do not match. TensorSrc must be GM and zN, while TensorDst must be L1 and zN"
98 );98 );
99 99 
100- const uint32_t blockCount = tla::get<1, 1>(srcTensor.shape());100+ uint32_t blockCount = CeilDiv<ELE_NUM_PER_C0>(tla::get<1>(srcTensor.originShape()));
101- const uint32_t blockLen = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape());101+ uint32_t blockLen = RoundUp<C0_NUM_PER_FRACTAL>(tla::get<0>(srcTensor.originShape()));
102 102 
103 AscendC::DataCopyParams repeatParams;103 AscendC::DataCopyParams repeatParams;
104 104 
@@ -145,8 +145,8 @@ struct TileCopyTla<
145 "while TensorDst must be L1 and nZ"145 "while TensorDst must be L1 and nZ"
146 );146 );
147 147 
148- const uint32_t nValue = tla::get<1>(srcTensor.shape());148+ const uint32_t nValue = tla::get<1>(srcTensor.originShape());
149- const uint32_t dValue = tla::get<0>(srcTensor.shape());149+ const uint32_t dValue = tla::get<0>(srcTensor.originShape());
150 const uint32_t srcDValue = tla::get<1>(srcTensor.stride());150 const uint32_t srcDValue = tla::get<1>(srcTensor.stride());
151 const uint32_t dstInnerStrideCol = tla::get<1, 0>(dstTensor.stride());151 const uint32_t dstInnerStrideCol = tla::get<1, 0>(dstTensor.stride());
152 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());152 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
@@ -194,8 +194,8 @@ struct TileCopyTla<
194 "while TensorDst must be L1 and nZ"194 "while TensorDst must be L1 and nZ"
195 );195 );
196 196 
197- const uint32_t blockCount = tla::get<0, 1>(srcTensor.shape());197+ uint32_t blockCount = CeilDiv<ELE_NUM_PER_C0>(tla::get<0>(srcTensor.originShape()));
198- const uint32_t blockLen = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());198+ uint32_t blockLen = RoundUp<C0_NUM_PER_FRACTAL>(tla::get<1>(srcTensor.originShape()));
199 199 
200 AscendC::DataCopyParams repeatParams;200 AscendC::DataCopyParams repeatParams;
201 201 
@@ -238,7 +238,7 @@ struct TileCopyTla<
238 238 
239 AscendC::DataCopyParams intriParams;239 AscendC::DataCopyParams intriParams;
240 intriParams.blockCount = 1;240 intriParams.blockCount = 1;
241- intriParams.blockLen = CeilDiv(tla::get<0>(srcTensor.shape()), ELE_NUM_PER_C0);241+ intriParams.blockLen = CeilDiv(tla::get<0>(srcTensor.originShape()), ELE_NUM_PER_C0);
242 intriParams.srcStride = 0;242 intriParams.srcStride = 0;
243 intriParams.dstStride = 0;243 intriParams.dstStride = 0;
244 244 
Minclude/catlass/gemm/tile/ascend950/copy_l0c_to_gm.hpp+6-6
@@ -44,8 +44,8 @@ struct CopyL0CToGmTla<
44 44 
45 AscendC::DataCopyCO12DstParams intriParams;45 AscendC::DataCopyCO12DstParams intriParams;
46 46 
47- intriParams.nSize = tla::get<1>(dstTensor.shape());47+ intriParams.nSize = tla::get<1>(dstTensor.originShape());
48- intriParams.mSize = tla::get<0>(dstTensor.shape());48+ intriParams.mSize = tla::get<0>(dstTensor.originShape());
49 intriParams.dstStride = tla::get<0>(dstTensor.stride());49 intriParams.dstStride = tla::get<0>(dstTensor.stride());
50 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());50 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
51 intriParams.quantPre = quantPre;51 intriParams.quantPre = quantPre;
@@ -96,8 +96,8 @@ struct CopyL0CToGmTla<
96 96 
97 AscendC::DataCopyCO12DstParams intriParams;97 AscendC::DataCopyCO12DstParams intriParams;
98 98 
99- intriParams.nSize = tla::get<1>(dstTensor.shape());99+ intriParams.nSize = tla::get<1>(dstTensor.originShape());
100- intriParams.mSize = tla::get<0>(dstTensor.shape());100+ intriParams.mSize = tla::get<0>(dstTensor.originShape());
101 intriParams.dstStride = tla::get<0>(dstTensor.stride());101 intriParams.dstStride = tla::get<0>(dstTensor.stride());
102 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());102 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
103 intriParams.quantPre = quantPre;103 intriParams.quantPre = quantPre;
@@ -136,8 +136,8 @@ struct CopyL0CToGmTla<
136 136 
137 AscendC::DataCopyCO12DstParams intriParams;137 AscendC::DataCopyCO12DstParams intriParams;
138 138 
139- intriParams.nSize = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());139+ intriParams.nSize = tla::get<1>(dstTensor.originShape());
140- intriParams.mSize = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());140+ intriParams.mSize = tla::get<0>(dstTensor.originShape());
141 intriParams.dstStride = tla::get<1, 1>(dstTensor.stride()) / (BYTE_PER_C0 / sizeof(ElementDst));141 intriParams.dstStride = tla::get<1, 1>(dstTensor.stride()) / (BYTE_PER_C0 / sizeof(ElementDst));
142 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());142 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
143 intriParams.quantPre = quantPre;143 intriParams.quantPre = quantPre;
Minclude/catlass/gemm/tile/ascend950/copy_l0c_to_ub.hpp+6-6
@@ -50,8 +50,8 @@ struct CopyL0CToUBTla<
50 AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams;50 AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams;
51 51 
52 // Fixpipe layout information52 // Fixpipe layout information
53- intriParams.nSize = tla::get<1>(dstTensor.shape());53+ intriParams.nSize = tla::get<1>(dstTensor.originShape());
54- intriParams.mSize = tla::get<0>(dstTensor.shape());54+ intriParams.mSize = tla::get<0>(dstTensor.originShape());
55 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());55 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
56 intriParams.dstStride = tla::get<0>(dstTensor.stride());56 intriParams.dstStride = tla::get<0>(dstTensor.stride());
57 57 
@@ -97,8 +97,8 @@ struct CopyL0CToUBTla<
97 AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams;97 AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams;
98 98 
99 // Fixpipe layout information99 // Fixpipe layout information
100- intriParams.nSize = tla::get<1>(dstTensor.shape());100+ intriParams.nSize = tla::get<1>(dstTensor.originShape());
101- intriParams.mSize = RoundUp(tla::get<0>(dstTensor.shape()), 2); // m must be even when spilt m101+ intriParams.mSize = RoundUp(tla::get<0>(dstTensor.originShape()), 2); // m must be even when spilt m
102 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());102 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
103 intriParams.dstStride = tla::get<0>(dstTensor.stride());103 intriParams.dstStride = tla::get<0>(dstTensor.stride());
104 104 
@@ -145,8 +145,8 @@ struct CopyL0CToUBTla<
145 AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams;145 AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams;
146 146 
147 // Fixpipe layout information147 // Fixpipe layout information
148- intriParams.nSize = RoundUp(tla::get<1>(dstTensor.shape()), 32);148+ intriParams.nSize = RoundUp(tla::get<1>(dstTensor.originShape()), 32);
149- intriParams.mSize = tla::get<0>(dstTensor.shape()); // m must be even when spilt m149+ intriParams.mSize = tla::get<0>(dstTensor.originShape()); // m must be even when spilt m
150 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());150 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
151 intriParams.dstStride = tla::get<0>(dstTensor.stride());151 intriParams.dstStride = tla::get<0>(dstTensor.stride());
152 152 
Minclude/catlass/gemm/tile/ascend950/copy_l1_to_bt.hpp+1-1
@@ -46,7 +46,7 @@ struct TileCopyTla<
46 46 
47 AscendC::DataCopyParams intriParams;47 AscendC::DataCopyParams intriParams;
48 intriParams.blockCount = 1;48 intriParams.blockCount = 1;
49- intriParams.blockLen = CeilDiv(tla::get<0>(srcTensor.shape()), ELE_NUM_PER_C0);49+ intriParams.blockLen = CeilDiv(tla::get<0>(srcTensor.originShape()), ELE_NUM_PER_C0);
50 if (sizeof(ElementSrc) == 4) {50 if (sizeof(ElementSrc) == 4) {
51 // the burst length should be even when B3251 // the burst length should be even when B32
52 intriParams.blockLen = RoundUp(intriParams.blockLen, 2);52 intriParams.blockLen = RoundUp(intriParams.blockLen, 2);
Minclude/catlass/gemm/tile/ascend950/copy_l1_to_l0a.hpp+37-24
@@ -43,8 +43,8 @@ struct TileCopyTla<
43 "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zN"43 "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zN"
44 );44 );
45 45 
46- const uint32_t dstOuterShapeRow = tla::get<0, 1>(dstTensor.shape());46+ const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape()));
47- const uint32_t dstOuterShapeCol = tla::get<1, 1>(dstTensor.shape());47+ const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape()));
48 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());48 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
49 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());49 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
50 auto srcCoord = srcTensor.coord();50 auto srcCoord = srcTensor.coord();
@@ -72,8 +72,8 @@ struct TileCopyTla<
72 "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zN"72 "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zN"
73 );73 );
74 74 
75- const uint32_t dstOuterShapeRow = tla::get<0, 1>(dstTensor.shape());75+ const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape()));
76- const uint32_t dstOuterShapeCol = tla::get<1, 1>(dstTensor.shape());76+ const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape()));
77 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());77 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
78 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());78 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
79 79 
@@ -120,8 +120,8 @@ struct TileCopyTla<
120 "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"120 "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"
121 );121 );
122 122 
123- const uint32_t L0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());123+ const uint32_t L0M = tla::get<0>(dstTensor.originShape());
124- const uint32_t L0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());124+ const uint32_t L0K = tla::get<1>(dstTensor.originShape());
125 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());125 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
126 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());126 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
127 auto srcCoord = srcTensor.coord();127 auto srcCoord = srcTensor.coord();
@@ -131,6 +131,9 @@ struct TileCopyTla<
131 loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<0>(srcCoord));131 loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<0>(srcCoord));
132 loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);132 loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);
133 loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0M);133 loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0M);
134+ if constexpr (AscendC::Std::is_one_of_v<typename TensorSrc::Element, float, uint32_t, int32_t>) {
135+ loadDataParams.kStep = RoundUp<2>(loadDataParams.kStep); // for b32 data types, ensure kStep is even
136+ }
134 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);137 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);
135 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);138 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);
136 loadDataParams.ifTranspose = true;139 loadDataParams.ifTranspose = true;
@@ -155,14 +158,19 @@ struct TileCopyTla<
155 const uint32_t L1K = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());158 const uint32_t L1K = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());
156 const uint32_t L0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());159 const uint32_t L0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
157 const uint32_t L0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());160 const uint32_t L0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
161+ const uint32_t L0MOrigin = tla::get<0>(dstTensor.originShape());
162+ const uint32_t L0KOrigin = tla::get<1>(dstTensor.originShape());
158 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());163 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
159 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());164 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
160 165 
161 AscendC::LoadData2DParamsV2 loadDataParams;166 AscendC::LoadData2DParamsV2 loadDataParams;
162 loadDataParams.mStartPosition = 0;167 loadDataParams.mStartPosition = 0;
163 loadDataParams.kStartPosition = 0;168 loadDataParams.kStartPosition = 0;
164- loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);169+ loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0KOrigin);
165- loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0M);170+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0MOrigin);
171+ if constexpr (AscendC::Std::is_one_of_v<typename TensorSrc::Element, float, uint32_t, int32_t>) {
172+ loadDataParams.kStep = RoundUp<2>(loadDataParams.kStep); // for b32 data types, ensure kStep is even
173+ }
166 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);174 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);
167 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);175 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);
168 loadDataParams.ifTranspose = true;176 loadDataParams.ifTranspose = true;
@@ -205,18 +213,19 @@ struct TileCopyTla<
205 "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"213 "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"
206 );214 );
207 215 
208- const uint32_t L0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());216+ const uint32_t L0MPadded = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
209- const uint32_t L0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());217+ const uint32_t L0MOrigin = tla::get<0>(dstTensor.originShape());
218+ const uint32_t L0KOrigin = tla::get<1>(dstTensor.originShape());
210 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());219 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
211 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());220 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
212 auto srcCoord = srcTensor.coord();221 auto srcCoord = srcTensor.coord();
213 222 
214 AscendC::LoadData2DParamsV2 loadDataParams;223 AscendC::LoadData2DParamsV2 loadDataParams;
215- if (L0M % ELE_NUM_PER_C0 == 0) {224+ if (RoundUp<C0_NUM_PER_FRACTAL>(L0MOrigin) % ELE_NUM_PER_C0 == 0) {
216 loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<1>(srcCoord));225 loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<1>(srcCoord));
217 loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<0>(srcCoord));226 loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<0>(srcCoord));
218- loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);227+ loadDataParams.mStep = RoundUp<2>(CeilDiv<C0_NUM_PER_FRACTAL>(L0KOrigin));
219- loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0M);228+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0MOrigin);
220 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);229 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);
221 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);230 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);
222 loadDataParams.ifTranspose = true;231 loadDataParams.ifTranspose = true;
@@ -224,18 +233,18 @@ struct TileCopyTla<
224 auto dstOffset = dstTensor.layout()(dstTensor.coord());233 auto dstOffset = dstTensor.layout()(dstTensor.coord());
225 AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams);234 AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams);
226 } else {235 } else {
227- for (uint32_t kIdx = 0; kIdx < L0K / ELE_NUM_PER_C0; kIdx++) {236+ for (uint32_t kIdx = 0; kIdx < CeilDiv<ELE_NUM_PER_C0>(L0KOrigin); kIdx++) {
228 loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<1>(srcCoord)) + kIdx * 2;237 loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<1>(srcCoord)) + kIdx * 2;
229 loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<0>(srcCoord));238 loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<0>(srcCoord));
230 loadDataParams.mStep = 2;239 loadDataParams.mStep = 2;
231- loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0M);240+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0MOrigin);
232 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);241 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);
233 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);242 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);
234 loadDataParams.ifTranspose = true;243 loadDataParams.ifTranspose = true;
235 244 
236 auto dstOffset = dstTensor.layout()(dstTensor.coord());245 auto dstOffset = dstTensor.layout()(dstTensor.coord());
237 AscendC::LoadData(246 AscendC::LoadData(
238- dstTensor.data()[dstOffset + kIdx * L0M * ELE_NUM_PER_C0], srcTensor.data(), loadDataParams247+ dstTensor.data()[dstOffset + kIdx * L0MPadded * ELE_NUM_PER_C0], srcTensor.data(), loadDataParams
239 );248 );
240 }249 }
241 }250 }
@@ -257,36 +266,40 @@ struct TileCopyTla<
257 const uint32_t L1K = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());266 const uint32_t L1K = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());
258 const uint32_t L0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());267 const uint32_t L0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
259 const uint32_t L0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());268 const uint32_t L0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
269+ const uint32_t L0MPadded = L0M;
270+ const uint32_t L0KPadded = L0K;
271+ const uint32_t L0MOrigin = tla::get<0>(dstTensor.originShape());
272+ const uint32_t L0KOrigin = tla::get<1>(dstTensor.originShape());
260 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());273 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
261 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());274 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
262 275 
263 AscendC::LoadData2DParamsV2 loadDataParams;276 AscendC::LoadData2DParamsV2 loadDataParams;
264- if (L0M % ELE_NUM_PER_C0 == 0) {277+ if (RoundUp<C0_NUM_PER_FRACTAL>(L0MOrigin) % ELE_NUM_PER_C0 == 0) {
265 loadDataParams.mStartPosition = 0;278 loadDataParams.mStartPosition = 0;
266 loadDataParams.kStartPosition = 0;279 loadDataParams.kStartPosition = 0;
267- loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);280+ loadDataParams.mStep = RoundUp<2>(CeilDiv<C0_NUM_PER_FRACTAL>(L0KOrigin));
268- loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0M);281+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0MOrigin);
269 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);282 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);
270 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);283 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);
271 loadDataParams.ifTranspose = true;284 loadDataParams.ifTranspose = true;
272 285 
273 for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) {286 for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) {
274 AscendC::LoadData(287 AscendC::LoadData(
275- dstTensor.data()[l0BatchIdx * L0M * L0K], srcTensor.data()[l0BatchIdx * L1M * L1K], loadDataParams288+ dstTensor.data()[l0BatchIdx * L0MPadded * L0KPadded], srcTensor.data()[l0BatchIdx * L1M * L1K], loadDataParams
276 );289 );
277 }290 }
278 } else {291 } else {
279 loadDataParams.mStartPosition = 0;292 loadDataParams.mStartPosition = 0;
280 loadDataParams.kStartPosition = 0;293 loadDataParams.kStartPosition = 0;
281 loadDataParams.mStep = 2;294 loadDataParams.mStep = 2;
282- loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0M);295+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0MOrigin);
283 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);296 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);
284 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);297 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);
285 loadDataParams.ifTranspose = true;298 loadDataParams.ifTranspose = true;
286 for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) {299 for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) {
287- for (uint32_t kIdx = 0; kIdx < L0K / ELE_NUM_PER_C0; kIdx++) {300+ for (uint32_t kIdx = 0; kIdx < CeilDiv<ELE_NUM_PER_C0>(L0KOrigin); kIdx++) {
288 AscendC::LoadData(301 AscendC::LoadData(
289- dstTensor.data()[l0BatchIdx * L0M * L0K + kIdx * L0M * ELE_NUM_PER_C0],302+ dstTensor.data()[l0BatchIdx * L0MPadded * L0KPadded + kIdx * L0MPadded * ELE_NUM_PER_C0],
290 srcTensor.data()[l0BatchIdx * L1M * L1K + kIdx * ELE_NUM_PER_FRACTAL * 2], loadDataParams303 srcTensor.data()[l0BatchIdx * L1M * L1K + kIdx * ELE_NUM_PER_FRACTAL * 2], loadDataParams
291 );304 );
292 }305 }
@@ -312,7 +325,7 @@ struct TileCopyTla<
312 {325 {
313 uint16_t aL1M = tla::get<0, 0>(srcTensor.stride());326 uint16_t aL1M = tla::get<0, 0>(srcTensor.stride());
314 uint16_t madM = tla::get<1, 1>(dstTensor.stride());327 uint16_t madM = tla::get<1, 1>(dstTensor.stride());
315- uint16_t madK = tla::get<1, 1>(dstTensor.shape());328+ uint16_t madK = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape()));
316 329 
317 AscendC::LoadData2DParamsV2 loadDataParams;330 AscendC::LoadData2DParamsV2 loadDataParams;
318 loadDataParams.mStartPosition = 0;331 loadDataParams.mStartPosition = 0;
Minclude/catlass/gemm/tile/ascend950/copy_l1_to_l0b.hpp+34-23
@@ -48,8 +48,8 @@ struct TileCopyTla<
48 "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"48 "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"
49 );49 );
50 50 
51- const uint32_t L0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());51+ const uint32_t L0KOrigin = tla::get<0>(dstTensor.originShape());
52- const uint32_t L0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());52+ const uint32_t L0NOrigin = tla::get<1>(dstTensor.originShape());
53 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());53 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
54 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());54 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
55 auto srcCoord = srcTensor.coord();55 auto srcCoord = srcTensor.coord();
@@ -57,8 +57,11 @@ struct TileCopyTla<
57 AscendC::LoadData2DParamsV2 loadDataParams;57 AscendC::LoadData2DParamsV2 loadDataParams;
58 loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<0>(srcCoord));58 loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<0>(srcCoord));
59 loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<1>(srcCoord));59 loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<1>(srcCoord));
60- loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);60+ loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0KOrigin);
61- loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0N);61+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0NOrigin);
62+ if constexpr (AscendC::Std::is_one_of_v<typename TensorSrc::Element, float, uint32_t, int32_t>) {
63+ loadDataParams.kStep = RoundUp<2>(loadDataParams.kStep); // for b32 data types, ensure kStep is even
64+ }
62 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);65 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);
63 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);66 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);
64 loadDataParams.ifTranspose = true;67 loadDataParams.ifTranspose = true;
@@ -83,14 +86,19 @@ struct TileCopyTla<
83 const uint32_t L1N = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());86 const uint32_t L1N = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());
84 const uint32_t L0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());87 const uint32_t L0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
85 const uint32_t L0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());88 const uint32_t L0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
89+ const uint32_t L0KOrigin = tla::get<0>(dstTensor.originShape());
90+ const uint32_t L0NOrigin = tla::get<1>(dstTensor.originShape());
86 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());91 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
87 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());92 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
88 93 
89 AscendC::LoadData2DParamsV2 loadDataParams;94 AscendC::LoadData2DParamsV2 loadDataParams;
90 loadDataParams.mStartPosition = 0;95 loadDataParams.mStartPosition = 0;
91 loadDataParams.kStartPosition = 0;96 loadDataParams.kStartPosition = 0;
92- loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);97+ loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0KOrigin);
93- loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0N);98+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0NOrigin);
99+ if constexpr (AscendC::Std::is_one_of_v<typename TensorSrc::Element, float, uint32_t, int32_t>) {
100+ loadDataParams.kStep = RoundUp<2>(loadDataParams.kStep); // for b32 data types, ensure kStep is even
101+ }
94 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);102 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);
95 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);103 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);
96 loadDataParams.ifTranspose = true;104 loadDataParams.ifTranspose = true;
@@ -133,18 +141,19 @@ struct TileCopyTla<
133 "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"141 "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"
134 );142 );
135 143 
136- const uint32_t L0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());144+ const uint32_t L0NPadded = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
137- const uint32_t L0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());145+ const uint32_t L0KOrigin = tla::get<0>(dstTensor.originShape());
146+ const uint32_t L0NOrigin = tla::get<1>(dstTensor.originShape());
138 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());147 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
139 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());148 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
140 auto srcCoord = srcTensor.coord();149 auto srcCoord = srcTensor.coord();
141 150 
142 AscendC::LoadData2DParamsV2 loadDataParams;151 AscendC::LoadData2DParamsV2 loadDataParams;
143- if (L0N % ELE_NUM_PER_C0 == 0) {152+ if (RoundUp<C0_NUM_PER_FRACTAL>(L0NOrigin) % ELE_NUM_PER_C0 == 0) {
144 loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<0>(srcCoord));153 loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<0>(srcCoord));
145 loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<1>(srcCoord));154 loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<1>(srcCoord));
146- loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);155+ loadDataParams.mStep = RoundUp<2>(CeilDiv<C0_NUM_PER_FRACTAL>(L0KOrigin));
147- loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0N);156+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0NOrigin);
148 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);157 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);
149 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);158 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);
150 loadDataParams.ifTranspose = true;159 loadDataParams.ifTranspose = true;
@@ -152,18 +161,18 @@ struct TileCopyTla<
152 auto dstOffset = dstTensor.layout()(dstTensor.coord());161 auto dstOffset = dstTensor.layout()(dstTensor.coord());
153 AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams);162 AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams);
154 } else {163 } else {
155- for (uint32_t kIdx = 0; kIdx < L0K / ELE_NUM_PER_C0; kIdx++) {164+ for (uint32_t kIdx = 0; kIdx < CeilDiv<ELE_NUM_PER_C0>(L0KOrigin); kIdx++) {
156 loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<0>(srcCoord)) + kIdx * 2;165 loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<0>(srcCoord)) + kIdx * 2;
157 loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<1>(srcCoord));166 loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<1>(srcCoord));
158 loadDataParams.mStep = 2;167 loadDataParams.mStep = 2;
159- loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0N);168+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0NOrigin);
160 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);169 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);
161 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);170 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);
162 loadDataParams.ifTranspose = true;171 loadDataParams.ifTranspose = true;
163 172 
164 auto dstOffset = dstTensor.layout()(dstTensor.coord());173 auto dstOffset = dstTensor.layout()(dstTensor.coord());
165 AscendC::LoadData(174 AscendC::LoadData(
166- dstTensor.data()[dstOffset + kIdx * L0N * ELE_NUM_PER_C0], srcTensor.data(), loadDataParams175+ dstTensor.data()[dstOffset + kIdx * L0NPadded * ELE_NUM_PER_C0], srcTensor.data(), loadDataParams
167 );176 );
168 }177 }
169 }178 }
@@ -185,15 +194,17 @@ struct TileCopyTla<
185 const uint32_t L1N = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());194 const uint32_t L1N = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());
186 const uint32_t L0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());195 const uint32_t L0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
187 const uint32_t L0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());196 const uint32_t L0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
197+ const uint32_t L0KOrigin = tla::get<0>(dstTensor.originShape());
198+ const uint32_t L0NOrigin = tla::get<1>(dstTensor.originShape());
188 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());199 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
189 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());200 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
190 201 
191 AscendC::LoadData2DParamsV2 loadDataParams;202 AscendC::LoadData2DParamsV2 loadDataParams;
192- if (L0N % ELE_NUM_PER_C0 == 0) {203+ if (RoundUp<C0_NUM_PER_FRACTAL>(L0NOrigin) % ELE_NUM_PER_C0 == 0) {
193 loadDataParams.mStartPosition = 0;204 loadDataParams.mStartPosition = 0;
194 loadDataParams.kStartPosition = 0;205 loadDataParams.kStartPosition = 0;
195- loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);206+ loadDataParams.mStep = RoundUp<2>(CeilDiv<C0_NUM_PER_FRACTAL>(L0KOrigin));
196- loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0N);207+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0NOrigin);
197 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);208 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);
198 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);209 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);
199 loadDataParams.ifTranspose = true;210 loadDataParams.ifTranspose = true;
@@ -207,12 +218,12 @@ struct TileCopyTla<
207 loadDataParams.mStartPosition = 0;218 loadDataParams.mStartPosition = 0;
208 loadDataParams.kStartPosition = 0;219 loadDataParams.kStartPosition = 0;
209 loadDataParams.mStep = 2;220 loadDataParams.mStep = 2;
210- loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0N);221+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0NOrigin);
211 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);222 loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);
212 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);223 loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);
213 loadDataParams.ifTranspose = true;224 loadDataParams.ifTranspose = true;
214 for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) {225 for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) {
215- for (uint32_t kIdx = 0; kIdx < L0K / ELE_NUM_PER_C0; kIdx++) {226+ for (uint32_t kIdx = 0; kIdx < CeilDiv<ELE_NUM_PER_C0>(L0KOrigin); kIdx++) {
216 AscendC::LoadData(227 AscendC::LoadData(
217 dstTensor.data()[l0BatchIdx * L0N * L0K + kIdx * L0N * ELE_NUM_PER_C0],228 dstTensor.data()[l0BatchIdx * L0N * L0K + kIdx * L0N * ELE_NUM_PER_C0],
218 srcTensor.data()[l0BatchIdx * L1N * L1K + kIdx * ELE_NUM_PER_FRACTAL * 2], loadDataParams229 srcTensor.data()[l0BatchIdx * L1N * L1K + kIdx * ELE_NUM_PER_FRACTAL * 2], loadDataParams
@@ -248,8 +259,8 @@ struct TileCopyTla<
248 "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0B and nZ"259 "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0B and nZ"
249 );260 );
250 261 
251- const uint32_t dstOuterShapeRow = tla::get<0, 1>(dstTensor.shape());262+ const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape()));
252- const uint32_t dstOuterShapeCol = tla::get<1, 1>(dstTensor.shape());263+ const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape()));
253 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());264 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
254 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());265 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
255 auto srcCoord = srcTensor.coord();266 auto srcCoord = srcTensor.coord();
@@ -277,8 +288,8 @@ struct TileCopyTla<
277 "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0B and nZ"288 "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0B and nZ"
278 );289 );
279 290 
280- const uint32_t dstOuterShapeRow = tla::get<0, 1>(dstTensor.shape());291+ const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape()));
281- const uint32_t dstOuterShapeCol = tla::get<1, 1>(dstTensor.shape());292+ const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape()));
282 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());293 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
283 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());294 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
284 295 
Minclude/catlass/gemm/tile/atlasa2/copy_gm_to_l1.hpp+8-8
@@ -1460,8 +1460,8 @@ struct TileCopyTla<Arch::AtlasA2,
1460 TensorDst::position == AscendC::TPosition::A1,1460 TensorDst::position == AscendC::TPosition::A1,
1461 "The input parameters do not match. TensorSrc must be GM and RowMajor, while TensorDst must be L1 and zN");1461 "The input parameters do not match. TensorSrc must be GM and RowMajor, while TensorDst must be L1 and zN");
1462 1462 
1463- const uint32_t nValue = tla::get<0>(srcTensor.shape());1463+ const uint32_t nValue = tla::get<0>(srcTensor.originShape());
1464- const uint32_t dValue = tla::get<1>(srcTensor.shape());1464+ const uint32_t dValue = tla::get<1>(srcTensor.originShape());
1465 const uint32_t srcDValue = tla::get<0>(srcTensor.stride());1465 const uint32_t srcDValue = tla::get<0>(srcTensor.stride());
1466 const uint32_t dstInnerStrideRow = tla::get<0, 0>(dstTensor.stride());1466 const uint32_t dstInnerStrideRow = tla::get<0, 0>(dstTensor.stride());
1467 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());1467 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
@@ -1520,8 +1520,8 @@ struct TileCopyTla<Arch::AtlasA2,
1520 "The input parameters do not match. TensorSrc must be GM and ColumnMajor, "1520 "The input parameters do not match. TensorSrc must be GM and ColumnMajor, "
1521 "while TensorDst must be L1 and nZ");1521 "while TensorDst must be L1 and nZ");
1522 1522 
1523- const uint32_t nValue = tla::get<1>(srcTensor.shape());1523+ const uint32_t nValue = tla::get<1>(srcTensor.originShape());
1524- const uint32_t dValue = tla::get<0>(srcTensor.shape());1524+ const uint32_t dValue = tla::get<0>(srcTensor.originShape());
1525 const uint32_t srcDValue = tla::get<1>(srcTensor.stride());1525 const uint32_t srcDValue = tla::get<1>(srcTensor.stride());
1526 const uint32_t dstInnerStrideRow = tla::get<1, 0>(dstTensor.stride());1526 const uint32_t dstInnerStrideRow = tla::get<1, 0>(dstTensor.stride());
1527 const uint32_t dstOuterStrideCol = tla::get<0, 1>(dstTensor.stride());1527 const uint32_t dstOuterStrideCol = tla::get<0, 1>(dstTensor.stride());
@@ -1582,8 +1582,8 @@ struct TileCopyTla<Arch::AtlasA2,
1582 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());1582 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
1583 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());1583 const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
1584 1584 
1585- uint32_t blockCount = tla::get<1, 1>(srcTensor.shape());1585+ uint32_t blockCount = CeilDiv<ELE_NUM_PER_C0>(tla::get<1>(srcTensor.originShape()));
1586- uint32_t blockLen = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape());1586+ uint32_t blockLen = tla::get<0>(srcTensor.originShape());
1587 1587 
1588 auto dstOffset = dstTensor.layout()(dstTensor.coord());1588 auto dstOffset = dstTensor.layout()(dstTensor.coord());
1589 auto srcOffset = srcTensor.layout()(srcTensor.coord());1589 auto srcOffset = srcTensor.layout()(srcTensor.coord());
@@ -1636,8 +1636,8 @@ struct TileCopyTla<Arch::AtlasA2,
1636 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());1636 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
1637 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());1637 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
1638 1638 
1639- uint32_t blockCount = tla::get<0, 1>(srcTensor.shape());1639+ uint32_t blockCount = CeilDiv<ELE_NUM_PER_C0>(tla::get<0>(srcTensor.originShape()));
1640- uint32_t blockLen = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());1640+ uint32_t blockLen = tla::get<1>(srcTensor.originShape());
1641 1641 
1642 auto dstOffset = dstTensor.layout()(dstTensor.coord());1642 auto dstOffset = dstTensor.layout()(dstTensor.coord());
1643 auto srcOffset = srcTensor.layout()(srcTensor.coord());1643 auto srcOffset = srcTensor.layout()(srcTensor.coord());
Minclude/catlass/gemm/tile/atlasa2/copy_gm_to_ub.hpp+7-4
@@ -44,11 +44,14 @@ struct TileCopyTla<Arch::AtlasA2,
44 "The input parameters do not match. TensorSrc must be GM and RowMajor, "44 "The input parameters do not match. TensorSrc must be GM and RowMajor, "
45 "while TensorDst must be UB and RowMajor");45 "while TensorDst must be UB and RowMajor");
46 46 
47+ const uint16_t row = tla::get<0>(srcTensor.originShape());
48+ const uint16_t col = tla::get<1>(srcTensor.originShape());
49+ 
47 AscendC::DataCopyExtParams dataCopyParams(50 AscendC::DataCopyExtParams dataCopyParams(
48- tla::get<0>(srcTensor.shape()),51+ row,
49- tla::get<1>(srcTensor.shape()) * sizeof(ElementSrc),52+ col * sizeof(ElementSrc),
50- (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) * sizeof(ElementSrc),53+ (tla::get<0>(srcTensor.stride()) - col) * sizeof(ElementSrc),
51- (tla::get<0>(dstTensor.stride()) - tla::get<1>(dstTensor.shape())) / ELE_NUM_PER_BLK,54+ (tla::get<0>(dstTensor.stride()) - col) / ELE_NUM_PER_BLK,
52 055 0
53 );56 );
54 AscendC::DataCopyPadExtParams<ElementSrc> padParams(false, 0, 0, 0);57 AscendC::DataCopyPadExtParams<ElementSrc> padParams(false, 0, 0, 0);
Minclude/catlass/gemm/tile/atlasa2/copy_l0c_to_gm.hpp+4-4
@@ -524,8 +524,8 @@ struct CopyL0CToGmTla<Catlass::Arch::AtlasA2,
524 AscendC::FixpipeParamsV220 intriParams;524 AscendC::FixpipeParamsV220 intriParams;
525 525 
526 // Fixpipe layout information526 // Fixpipe layout information
527- intriParams.nSize = tla::get<1>(dstTensor.shape());527+ intriParams.nSize = tla::get<1>(dstTensor.originShape());
528- intriParams.mSize = tla::get<0>(dstTensor.shape());528+ intriParams.mSize = tla::get<0>(dstTensor.originShape());
529 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());529 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
530 intriParams.dstStride = tla::get<0>(dstTensor.stride());530 intriParams.dstStride = tla::get<0>(dstTensor.stride());
531 531 
@@ -576,8 +576,8 @@ struct CopyL0CToGmTla<Catlass::Arch::AtlasA2,
576 AscendC::FixpipeParamsV220 intriParams;576 AscendC::FixpipeParamsV220 intriParams;
577 577 
578 // Fixpipe layout information578 // Fixpipe layout information
579- intriParams.nSize = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());579+ intriParams.nSize = tla::get<1>(dstTensor.originShape());
Z
Zzjw6668883月17日

L0相关的shape修改涉及点有两处： 1、原先直接获取分块大小，现在需使用CeilDiv(orgshape)
2、原先获取完整大小是使用乘法算出，此时是对齐的，现在直接使用orgshape，有可能是不对齐的，不能保证结果正确

需针对以下几个维度设计测试用例，保证所有tile层的修改正确性： 1、数据类型：int8，half， float 2、转置：00,01,10,11 3、构造shape 5组左右：m,n,k非对齐，大小覆盖1~5000 4、测试kernel：basic、streamK、fa、sparse

likedislike
580- intriParams.mSize = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());580+ intriParams.mSize = tla::get<0>(dstTensor.originShape());
581 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape());581 intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape());
582 intriParams.dstStride = tla::get<1, 1>(dstTensor.stride()) / (BYTE_PER_C0 / sizeof(ElementDst));582 intriParams.dstStride = tla::get<1, 1>(dstTensor.stride()) / (BYTE_PER_C0 / sizeof(ElementDst));
583 583 
Minclude/catlass/gemm/tile/atlasa2/copy_l1_to_l0a.hpp+13-13
@@ -441,8 +441,8 @@ struct TileCopyTla<Arch::AtlasA2,
441 441 
442 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());442 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
443 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());443 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
444- const uint32_t dstOuterShapeRow = tla::get<0, 1>(dstTensor.shape());444+ const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape()));
445- const uint32_t dstOuterShapeCol = tla::get<1, 1>(dstTensor.shape());445+ const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape()));
446 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());446 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
447 447 
448 AscendC::LoadData2DParams loadDataParams;448 AscendC::LoadData2DParams loadDataParams;
@@ -496,9 +496,9 @@ struct TileCopyTla<Arch::AtlasA2,
496 496 
497 constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0};497 constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0};
498 uint16_t l1M = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape());498 uint16_t l1M = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape());
499- uint16_t l1K = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());499+ uint16_t l1K = RoundUp<ELE_NUM_PER_C0>(tla::get<1>(srcTensor.originShape()));
500- uint16_t l0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());500+ uint16_t l0M = RoundUp<ELE_NUM_PER_C0>(tla::get<0>(dstTensor.originShape()));
501- uint16_t l0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());501+ uint16_t l0K = RoundUp<ELE_NUM_PER_C0>(tla::get<1>(dstTensor.originShape()));
502 AscendC::SetFmatrix(1, l1M, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT);502 AscendC::SetFmatrix(1, l1M, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT);
503 static constexpr AscendC::IsResetLoad3dConfig config = {false, false};503 static constexpr AscendC::IsResetLoad3dConfig config = {false, false};
504 AscendC::LoadData3DParamsV2<Element> loadDataParams;504 AscendC::LoadData3DParamsV2<Element> loadDataParams;
@@ -539,8 +539,8 @@ struct TileCopyTla<Arch::AtlasA2,
539 "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zZ");539 "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zZ");
540 540 
541 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());541 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
542- const uint32_t dstOuterShapeRow = tla::get<0, 1>(dstTensor.shape());542+ const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape()));
543- const uint32_t dstOuterShapeCol = tla::get<1, 1>(dstTensor.shape());543+ const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape()));
544 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());544 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
545 545 
546 AscendC::LoadData2DParams loadDataParams;546 AscendC::LoadData2DParams loadDataParams;
@@ -592,9 +592,9 @@ struct TileCopyTla<Arch::AtlasA2,
592 "The input parameters do not match. TensorSrc must be int8_t, L1 and nZ, "592 "The input parameters do not match. TensorSrc must be int8_t, L1 and nZ, "
593 "while TensorDst must be int8_t, L0A and zZ");593 "while TensorDst must be int8_t, L0A and zZ");
594 594 
595- const uint32_t srcOuterShapeRow = tla::get<0, 1>(srcTensor.shape());595+ const uint32_t srcOuterShapeRow = CeilDiv(tla::get<0>(srcTensor.originShape()), tla::get<0, 0>(srcTensor.shape()));
596 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());596 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
597- const uint32_t dstOuterShapeCol = tla::get<1, 1>(dstTensor.shape());597+ const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape()));
598 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());598 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
599 599 
600 AscendC::LoadData2dTransposeParams loadDataParams;600 AscendC::LoadData2dTransposeParams loadDataParams;
@@ -645,10 +645,10 @@ struct TileCopyTla<Arch::AtlasA2,
645 "while TensorDst must be float, L0A and zZ");645 "while TensorDst must be float, L0A and zZ");
646 646 
647 constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0};647 constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0};
648- uint16_t l1M = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape());648+ uint16_t l1M = tla::get<0>(srcTensor.originShape());
649 uint16_t l1K = tla::get<0, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.shape());649 uint16_t l1K = tla::get<0, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.shape());
650- uint16_t l0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());650+ uint16_t l0M = tla::get<0>(dstTensor.originShape());
651- uint16_t l0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());651+ uint16_t l0K = tla::get<1>(dstTensor.originShape());
652 // K, M need to be 16 aligned for f32652 // K, M need to be 16 aligned for f32
653 uint16_t l1MAlign = RoundUp<C0_NUM_PER_FRACTAL>(l1M);653 uint16_t l1MAlign = RoundUp<C0_NUM_PER_FRACTAL>(l1M);
654 uint16_t l1KAlign = RoundUp<C0_NUM_PER_FRACTAL>(l1K);654 uint16_t l1KAlign = RoundUp<C0_NUM_PER_FRACTAL>(l1K);
@@ -699,7 +699,7 @@ struct TileCopySparseTla<Arch::AtlasA2,
699 TensorSrc::position == AscendC::TPosition::A1 &&699 TensorSrc::position == AscendC::TPosition::A1 &&
700 TensorDst::position == AscendC::TPosition::A2,700 TensorDst::position == AscendC::TPosition::A2,
701 "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zZ");701 "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zZ");
702- 702+
703 auto srcShape = srcTensor.shape();703 auto srcShape = srcTensor.shape();
704 auto dstShape = dstTensor.shape();704 auto dstShape = dstTensor.shape();
705 auto coord = srcTensor.coord();705 auto coord = srcTensor.coord();
Minclude/catlass/gemm/tile/atlasa2/copy_l1_to_l0b.hpp+19-19
@@ -538,8 +538,8 @@ struct TileCopyTla<Arch::AtlasA2,
538 538 
539 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());539 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
540 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());540 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
541- const uint32_t dstOuterShapeRow = tla::get<0, 1>(dstTensor.shape());541+ const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape()));
542- const uint32_t dstOuterShapeCol = tla::get<1, 1>(dstTensor.shape());542+ const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape()));
543 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());543 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
544 544 
545 AscendC::LoadData2DParams loadDataParams;545 AscendC::LoadData2DParams loadDataParams;
@@ -590,8 +590,8 @@ struct TileCopyTla<Arch::AtlasA2,
590 590 
591 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());591 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
592 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());592 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
593- const uint32_t dstOuterShapeRow = tla::get<0, 1>(dstTensor.shape());593+ const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape()));
594- const uint32_t dstOuterShapeCol = tla::get<1, 1>(dstTensor.shape());594+ const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape()));
595 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());595 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
596 596 
597 auto dstOffset = dstTensor.layout()(dstTensor.coord());597 auto dstOffset = dstTensor.layout()(dstTensor.coord());
@@ -609,18 +609,18 @@ struct TileCopyTla<Arch::AtlasA2,
609 609 
610 AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], loadDataParams);610 AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], loadDataParams);
611 } else {611 } else {
612- loadDataParams.startIndex = 0;612+ loadDataParams.startIndex = 0;
613- loadDataParams.repeatTimes = dstOuterShapeCol;613+ loadDataParams.repeatTimes = dstOuterShapeCol;
614 loadDataParams.srcStride = 1;614 loadDataParams.srcStride = 1;
615- loadDataParams.sid = 0;615+ loadDataParams.sid = 0;
616- loadDataParams.dstGap = 0;616+ loadDataParams.dstGap = 0;
617- loadDataParams.ifTranspose = false;617+ loadDataParams.ifTranspose = false;
618- loadDataParams.addrMode = 0;618+ loadDataParams.addrMode = 0;
619 619 
620- for (uint32_t i = 0; i < dstOuterShapeRow; i++) {620+ for (uint32_t i = 0; i < dstOuterShapeRow; i++) {
621- AscendC::LoadData(dstTensor.data()[dstOffset + i * dstOuterStrideRow],621+ AscendC::LoadData(dstTensor.data()[dstOffset + i * dstOuterStrideRow],
622- srcTensor.data()[srcOffset + i * srcOuterStrideRow],622+ srcTensor.data()[srcOffset + i * srcOuterStrideRow],
623- loadDataParams);623+ loadDataParams);
624 }624 }
625 }625 }
626 }626 }
@@ -654,10 +654,10 @@ struct TileCopyTla<Arch::AtlasA2,
654 "The input parameters do not match. TensorSrc must be int8_t, L1 and zN, "654 "The input parameters do not match. TensorSrc must be int8_t, L1 and zN, "
655 "while TensorDst must be int8_t, L0B and nZ");655 "while TensorDst must be int8_t, L0B and nZ");
656 656 
657- const uint32_t srcOuterShapeCol = tla::get<1, 1>(srcTensor.shape());657+ const uint32_t srcOuterShapeCol = CeilDiv(tla::get<1>(srcTensor.originShape()), tla::get<1, 0>(srcTensor.shape()));
658 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());658 const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
659 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());659 const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
660- const uint32_t dstOuterShapeRow = tla::get<0, 1>(dstTensor.shape());660+ const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape()));
661 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());661 const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
662 662 
663 AscendC::LoadData2dTransposeParams loadDataParams;663 AscendC::LoadData2dTransposeParams loadDataParams;
@@ -709,9 +709,9 @@ struct TileCopyTla<Arch::AtlasA2,
709 709 
710 constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0};710 constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0};
711 uint16_t l1K = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape());711 uint16_t l1K = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape());
712- uint16_t l1N = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());712+ uint16_t l1N = tla::get<1>(srcTensor.originShape());
713- uint16_t l0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());713+ uint16_t l0K = tla::get<0>(dstTensor.originShape());
714- uint16_t l0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());714+ uint16_t l0N = tla::get<1>(dstTensor.originShape());
715 // K, N need to be 16 aligned for f32715 // K, N need to be 16 aligned for f32
716 uint16_t l1KAlign = RoundUp<C0_NUM_PER_FRACTAL>(l1K);716 uint16_t l1KAlign = RoundUp<C0_NUM_PER_FRACTAL>(l1K);
717 uint16_t l1NAlign = RoundUp<C0_NUM_PER_FRACTAL>(l1N);717 uint16_t l1NAlign = RoundUp<C0_NUM_PER_FRACTAL>(l1N);
Minclude/catlass/gemm/tile/atlasa2/copy_ub_to_gm.hpp+7-4
@@ -44,11 +44,14 @@ struct TileCopyTla<Arch::AtlasA2,
44 "The input parameters do not match. TensorSrc must be GM and RowMajor, "44 "The input parameters do not match. TensorSrc must be GM and RowMajor, "
45 "while TensorDst must be UB and RowMajor");45 "while TensorDst must be UB and RowMajor");
46 46 
47+ const uint16_t row = tla::get<0>(dstTensor.originShape());
48+ const uint16_t col = tla::get<1>(dstTensor.originShape());
49+ 
47 AscendC::DataCopyExtParams dataCopyParams(50 AscendC::DataCopyExtParams dataCopyParams(
48- tla::get<0>(dstTensor.shape()),51+ row,
49- tla::get<1>(dstTensor.shape()) * sizeof(ElementSrc),52+ col * sizeof(ElementSrc),
50- (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) / ELE_NUM_PER_C0,53+ (tla::get<0>(srcTensor.stride()) - col) / ELE_NUM_PER_C0,
51- (tla::get<0>(dstTensor.stride()) - tla::get<1>(dstTensor.shape())) * sizeof(ElementSrc),54+ (tla::get<0>(dstTensor.stride()) - col) * sizeof(ElementSrc),
52 055 0
53 );56 );
54 auto dstOffset = dstTensor.layout()(dstTensor.coord());57 auto dstOffset = dstTensor.layout()(dstTensor.coord());
Minclude/catlass/gemm/tile/tile_copy.hpp+1-0
@@ -328,6 +328,7 @@ struct PackedTileCopyTla {
328 using LayoutTagL1B = typename helper::L1BTypeSelector<Gemm::GemmType<ElementB, LayoutTagB>>::L1BType::Layout;328 using LayoutTagL1B = typename helper::L1BTypeSelector<Gemm::GemmType<ElementB, LayoutTagB>>::L1BType::Layout;
329 using LayoutTagL0A = typename helper::L0ALayoutSelector<ArchTag>::Layout;329 using LayoutTagL0A = typename helper::L0ALayoutSelector<ArchTag>::Layout;
330 using LayoutTagL0B = layout::nZ;330 using LayoutTagL0B = layout::nZ;
331+ using LayoutTagL0C = layout::L0C;
331 332 
332 using LayoutA = detail::TagToLayout_t<ElementA, LayoutTagA>;333 using LayoutA = detail::TagToLayout_t<ElementA, LayoutTagA>;
333 using LayoutB = detail::TagToLayout_t<ElementB, LayoutTagB>;334 using LayoutB = detail::TagToLayout_t<ElementB, LayoutTagB>;
Minclude/catlass/gemm/tile/tile_mmad.hpp+45-0
@@ -227,6 +227,51 @@ struct TileMmadTla {
227 mmadParams);227 mmadParams);
228 }228 }
229 }229 }
230+ 
231+ // automatically extracts actual sizes from tensor originShape
232+ template <class TensorC, class TensorA, class TensorB>
233+ CATLASS_DEVICE
234+ void operator()(TensorC const &l0CTensor,
235+ TensorA const &l0ATensor,
236+ TensorB const &l0BTensor,
237+ bool initC = true, uint8_t unitFlag = 0)
238+ {
239+ // Get actual sizes from tensor originShape
240+ uint32_t m = tla::get<0>(l0CTensor.layout().originShape());
241+ uint32_t n = tla::get<1>(l0CTensor.layout().originShape());
242+ uint32_t k = tla::get<1>(l0ATensor.layout().originShape());
243+ AscendC::MmadParams mmadParams;
244+#if (defined (__NPU_ARCH__) && __NPU_ARCH__ == 2201)
245+ if constexpr (std::is_same_v<ElementA, float> && std::is_same_v<LayoutTagL1A, layout::nZ>) {
246+ mmadParams.kDirectionAlign = true;
247+ }
248+ if constexpr (!std::is_same_v<LayoutTagL1A, layout::VectorLayout>) {
249+ if (m == 1) m = 16; // avoid gemv mode
250+ }
251+#endif
252+#if (defined (__NPU_ARCH__) && __NPU_ARCH__ == 3510)
253+ if constexpr(std::is_same_v<LayoutTagL1A, layout::VectorLayout>) {
254+ mmadParams.disableGemv = false;
255+ } else {
256+ mmadParams.disableGemv = true;
257+ }
258+#endif
259+ mmadParams.m = m;
260+ mmadParams.n = n;
261+ mmadParams.k = k;
262+ mmadParams.unitFlag = unitFlag;
263+ mmadParams.cmatrixInitVal = initC;
264+ 
265+ AscendC::Mmad(l0CTensor.data(),
266+ l0ATensor.data(),
267+ l0BTensor.data(),
268+ mmadParams);
269+ 
270+ const uint32_t PIPE_M_BARRIER_THRESHOLD = 10;
271+ if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) {
272+ AscendC::PipeBarrier<PIPE_M>();
273+ }
274+ }
230};275};
231 276 
232/////////////////////////////////////////////////////////////////////////////////////////////////////////////277/////////////////////////////////////////////////////////////////////////////////////////////////////////////
Minclude/catlass/layout/matrix.hpp+178-1文件内容审核中，请稍后刷新重试
Minclude/tla/layout.hpp+239-74
@@ -13,6 +13,7 @@
13 13 
14#include "catlass/catlass.hpp"14#include "catlass/catlass.hpp"
15#include "tla/numeric/integral_constant.hpp"15#include "tla/numeric/integral_constant.hpp"
16+#include "tla/numeric/math.hpp"
16#include "tla/tuple.hpp"17#include "tla/tuple.hpp"
17#include "tla/int_tuple.hpp"18#include "tla/int_tuple.hpp"
18#include "catlass/layout/layout.hpp"19#include "catlass/layout/layout.hpp"
@@ -46,16 +47,42 @@ Coord<Ts...> MakeCoord(Ts const&... t) {
46 return {t...};47 return {t...};
47}48}
48 49 
50+namespace detail {
51+ 
52+// Type trait to generate OriginShape type with rank=R, depth=1, element type=uint32_t
53+template <int Rank, class Sequence = void>
54+struct MakeOriginShapeTypeImpl;
55+ 
56+template <int Rank, size_t... Is>
57+struct MakeOriginShapeTypeImpl<Rank, tla::index_sequence<Is...>> {
58+ template <size_t>
59+ using repeat_type = uint32_t;
60+ using type = Shape<repeat_type<Is>...>;
61+};
62+ 
63+template <class Stride>
64+using MakeOriginShapeType = typename MakeOriginShapeTypeImpl<rank_v<Stride>, tla::make_index_sequence<rank_v<Stride>>>::type;
65+ 
66+struct UnpackedMakeOriginShapeU32 {
67+ template <class... T>
68+ CATLASS_HOST_DEVICE constexpr
69+ auto operator()(T const&... a) const {
70+ return MakeShape(static_cast<uint32_t>(a)...);
71+ }
72+};
73+ 
74+} // namespace detail
75+ 
49//76//
50// Layout77// Layout
51//78//
52- 79+//自动推导 OriginShape 类型为 Shape<uint32_t...>，rank=rank_v<Stride>
53-template <class Shape, class Stride>80+template <class Shape, class Stride, class OriginShape = detail::MakeOriginShapeType<Stride>>
54-struct Layout : private tla::tuple<Shape, Stride> {81+struct Layout : private tla::tuple<Shape, Stride, OriginShape> {
55 // NOTE: This defaults static Shapes/Strides correctly, but not dynamic82 // NOTE: This defaults static Shapes/Strides correctly, but not dynamic
56 CATLASS_HOST_DEVICE constexpr83 CATLASS_HOST_DEVICE constexpr
57- Layout(Shape const& shape = {}, Stride const& stride = {})84+ Layout(Shape const& shape = {}, Stride const& stride = {}, OriginShape const& originShape = {})
58- : tla::tuple<Shape, Stride>(shape, stride) {}85+ : tla::tuple<Shape, Stride, OriginShape>(shape, stride, originShape) {}
59 86 
60 //87 //
61 // Accessors88 // Accessors
@@ -68,28 +95,42 @@ struct Layout : private tla::tuple<Shape, Stride> {
68 CATLASS_HOST_DEVICE constexpr95 CATLASS_HOST_DEVICE constexpr
69 decltype(auto) shape()96 decltype(auto) shape()
70 {97 {
71- return get<0, I...>(static_cast<tla::tuple<Shape, Stride>&>(*this));98+ return get<0, I...>(static_cast<tla::tuple<Shape, Stride, OriginShape>&>(*this));
72 }99 }
73 100 
74 template <int... I>101 template <int... I>
75 CATLASS_HOST_DEVICE constexpr102 CATLASS_HOST_DEVICE constexpr
76 decltype(auto) shape() const103 decltype(auto) shape() const
77 {104 {
78- return get<0, I...>(static_cast<tla::tuple<Shape, Stride> const&>(*this));105+ return get<0, I...>(static_cast<tla::tuple<Shape, Stride, OriginShape> const&>(*this));
79 }106 }
80 107 
81 template <int... I>108 template <int... I>
82 CATLASS_HOST_DEVICE constexpr109 CATLASS_HOST_DEVICE constexpr
83 decltype(auto) stride()110 decltype(auto) stride()
84 {111 {
85- return get<1, I...>(static_cast<tla::tuple<Shape, Stride>&>(*this));112+ return get<1, I...>(static_cast<tla::tuple<Shape, Stride, OriginShape>&>(*this));
86 }113 }
87 114 
88 template <int... I>115 template <int... I>
89 CATLASS_HOST_DEVICE constexpr116 CATLASS_HOST_DEVICE constexpr
90 decltype(auto) stride() const117 decltype(auto) stride() const
91 {118 {
92- return get<1, I...>(static_cast<tla::tuple<Shape, Stride> const&>(*this));119+ return get<1, I...>(static_cast<tla::tuple<Shape, Stride, OriginShape> const&>(*this));
120+ }
121+ 
122+ template <int... I>
123+ CATLASS_HOST_DEVICE constexpr
124+ decltype(auto) originShape()
125+ {
126+ return get<2, I...>(static_cast<tla::tuple<Shape, Stride, OriginShape>&>(*this));
127+ }
128+ 
129+ template <int... I>
130+ CATLASS_HOST_DEVICE constexpr
131+ decltype(auto) originShape() const
132+ {
133+ return get<2, I...>(static_cast<tla::tuple<Shape, Stride, OriginShape> const&>(*this));
93 }134 }
94 135 
95 template <class Coord>136 template <class Coord>
@@ -102,13 +143,24 @@ struct Layout : private tla::tuple<Shape, Stride> {
102 143 
103// Layout construction144// Layout construction
104 145 
146+template <class Shape, class Stride, class OriginShape>
147+CATLASS_HOST_DEVICE constexpr
148+auto MakeLayout(Shape const& shape, Stride const& stride, OriginShape const& originShape)
149+{
150+ static_assert(is_tuple<Shape>::value || is_integral<Shape>::value);
151+ static_assert(is_tuple<Stride>::value || is_integral<Stride>::value);
152+ static_assert(is_tuple<OriginShape>::value || is_integral<OriginShape>::value);
153+ return Layout<Shape, Stride, OriginShape>(shape, stride, originShape);
154+}
155+ 
105template <class Shape, class Stride>156template <class Shape, class Stride>
106CATLASS_HOST_DEVICE constexpr157CATLASS_HOST_DEVICE constexpr
107auto MakeLayout(Shape const& shape, Stride const& stride)158auto MakeLayout(Shape const& shape, Stride const& stride)
108{159{
109 static_assert(is_tuple<Shape>::value || is_integral<Shape>::value);160 static_assert(is_tuple<Shape>::value || is_integral<Shape>::value);
110 static_assert(is_tuple<Stride>::value || is_integral<Stride>::value);161 static_assert(is_tuple<Stride>::value || is_integral<Stride>::value);
111- return Layout<Shape, Stride>(shape, stride);162+ // 计算默认的 originShape：将 shape 扁平化为 depth=1，并将每个维度归一化为 uint32_t
163+ return MakeLayout(shape, stride, tla::transform_apply(shape, Product{}, detail::UnpackedMakeOriginShapeU32{}));
112}164}
113 165 
114// Convenience tags for common layouts166// Convenience tags for common layouts
@@ -121,65 +173,86 @@ auto MakeLayoutFromTag(LayoutTag const& tag)
121 std::is_same_v<LayoutTag, Catlass::layout::ColumnMajor> ||173 std::is_same_v<LayoutTag, Catlass::layout::ColumnMajor> ||
122 std::is_same_v<LayoutTag, Catlass::layout::VectorLayout> ||174 std::is_same_v<LayoutTag, Catlass::layout::VectorLayout> ||
123 std::is_same_v<LayoutTag, Catlass::layout::zN> ||175 std::is_same_v<LayoutTag, Catlass::layout::zN> ||
124- std::is_same_v<LayoutTag, Catlass::layout::nZ>,176+ std::is_same_v<LayoutTag, Catlass::layout::nZ> ||
177+ std::is_same_v<LayoutTag, Catlass::layout::L0C>,
125 "Unsupported LayoutTag for MakeLayoutFromTag, only support Catlass::layout::RowMajor or"178 "Unsupported LayoutTag for MakeLayoutFromTag, only support Catlass::layout::RowMajor or"
126- "Catlass::layout::ColumnMajor or Catlass::layout::zN or Catlass::layout::nZ or"179+ "Catlass::layout::ColumnMajor or Catlass::layout::VectorLayout or Catlass::layout::zN or Catlass::layout::nZ or Catlass::layout::L0C");
127- "Catlass::layout::VectorLayout");
128 180 
129- if constexpr (std::is_same_v<LayoutTag, Catlass::layout::RowMajor>) {181+ if constexpr (std::is_same_v<LayoutTag, Catlass::layout::VectorLayout>) {
130- return MakeLayout(MakeShape(tag.shape(0), tag.shape(1)), MakeStride(tag.stride(0), Int<1>{}));182+ return MakeLayout(MakeShape(tag.shape(0)),
183+ MakeStride(tag.stride(0)),
184+ MakeShape(tag.shape(0)));
185+ } else if constexpr (std::is_same_v<LayoutTag, Catlass::layout::RowMajor>) {
186+ return MakeLayout(MakeShape(tag.shape(0), tag.shape(1)),
187+ MakeStride(tag.stride(0), Int<1>{}),
188+ MakeShape(tag.shape(0), tag.shape(1)));
131 } else if constexpr (std::is_same_v<LayoutTag, Catlass::layout::ColumnMajor>) {189 } else if constexpr (std::is_same_v<LayoutTag, Catlass::layout::ColumnMajor>) {
132- return MakeLayout(MakeShape(tag.shape(0), tag.shape(1)), MakeStride(Int<1>{}, tag.stride(1)));190+ return MakeLayout(MakeShape(tag.shape(0), tag.shape(1)),
133- } else if constexpr (std::is_same_v<LayoutTag, Catlass::layout::VectorLayout>) {191+ MakeStride(Int<1>{}, tag.stride(1)),
134- return MakeLayout(MakeShape(tag.shape(0)), MakeStride(Int<1>{}));192+ MakeShape(tag.shape(0), tag.shape(1)));
135- } else { // zN or nZ193+ } else { // zN or nZ or L0C
136 return MakeLayout(MakeShape(MakeShape(tag.shape(0), tag.shape(1)), MakeShape(tag.shape(2), tag.shape(3))),194 return MakeLayout(MakeShape(MakeShape(tag.shape(0), tag.shape(1)), MakeShape(tag.shape(2), tag.shape(3))),
137- MakeStride(MakeStride(tag.stride(0), tag.stride(1)), MakeStride(tag.stride(2), tag.stride(3))));195+ MakeStride(MakeStride(tag.stride(0), tag.stride(1)), MakeStride(tag.stride(2), tag.stride(3))),
196+ MakeShape(tag.orgShape(0), tag.orgShape(1)));
138 }197 }
139}198}
140 199 
141// Return the shape of a mode200// Return the shape of a mode
142-template <int... Is, class Shape, class Stride>201+template <int... Is, class Shape, class Stride, class OriginShape>
143CATLASS_HOST_DEVICE constexpr202CATLASS_HOST_DEVICE constexpr
144-decltype(auto) shape(Layout<Shape, Stride>& layout)203+decltype(auto) shape(Layout<Shape, Stride, OriginShape>& layout)
145{204{
146 return layout.template shape<Is...>();205 return layout.template shape<Is...>();
147}206}
148 207 
149-template <int... Is, class Shape, class Stride>208+template <int... Is, class Shape, class Stride, class OriginShape>
150CATLASS_HOST_DEVICE constexpr209CATLASS_HOST_DEVICE constexpr
151-decltype(auto) shape(Layout<Shape, Stride> const& layout)210+decltype(auto) shape(Layout<Shape, Stride, OriginShape> const& layout)
152{211{
153 return layout.template shape<Is...>();212 return layout.template shape<Is...>();
154}213}
155 214 
156// Return the stride of a mode215// Return the stride of a mode
157-template <int... Is, class Shape, class Stride>216+template <int... Is, class Shape, class Stride, class OriginShape>
158CATLASS_HOST_DEVICE constexpr217CATLASS_HOST_DEVICE constexpr
159-decltype(auto) stride(Layout<Shape, Stride>& layout)218+decltype(auto) stride(Layout<Shape, Stride, OriginShape>& layout)
160{219{
161 return layout.template stride<Is...>();220 return layout.template stride<Is...>();
162}221}
163 222 
164-template <int... Is, class Shape, class Stride>223+template <int... Is, class Shape, class Stride, class OriginShape>
165CATLASS_HOST_DEVICE constexpr224CATLASS_HOST_DEVICE constexpr
166-decltype(auto) stride(Layout<Shape, Stride> const& layout)225+decltype(auto) stride(Layout<Shape, Stride, OriginShape> const& layout)
167{226{
168 return layout.template stride<Is...>();227 return layout.template stride<Is...>();
169}228}
170 229 
230+template <int... Is, class Shape, class Stride, class OriginShape>
231+CATLASS_HOST_DEVICE constexpr
232+decltype(auto) originShape(Layout<Shape, Stride, OriginShape>& layout)
233+{
234+ return layout.template originShape<Is...>();
235+}
236+ 
237+template <int... Is, class Shape, class Stride, class OriginShape>
238+CATLASS_HOST_DEVICE constexpr
239+decltype(auto) originShape(Layout<Shape, Stride, OriginShape> const& layout)
240+{
241+ return layout.template originShape<Is...>();
242+}
243+ 
171// Return the rank of layout244// Return the rank of layout
172-template <int... Is, class Shape, class Stride>245+template <int... Is, class Shape, class Stride, class OriginShape>
173CATLASS_HOST_DEVICE constexpr246CATLASS_HOST_DEVICE constexpr
174-auto rank(Layout<Shape, Stride> const& layout)247+auto rank(Layout<Shape, Stride, OriginShape> const& layout)
175{248{
176 return rank(shape<Is...>(layout));249 return rank(shape<Is...>(layout));
177}250}
178 251 
179// Return the depth of the layout252// Return the depth of the layout
180-template <int... Is, class Shape, class Stride>253+template <int... Is, class Shape, class Stride, class OriginShape>
181CATLASS_HOST_DEVICE constexpr254CATLASS_HOST_DEVICE constexpr
182-auto depth(Layout<Shape, Stride> const& layout)255+auto depth(Layout<Shape, Stride, OriginShape> const& layout)
183{256{
184 return depth(shape<Is...>(layout));257 return depth(shape<Is...>(layout));
185}258}
@@ -239,8 +312,8 @@ auto crd2offset(Coord const& coord, Shape const& shape, Stride const& stride)
239 312 
240template <class Layout>313template <class Layout>
241struct is_layout : false_type {};314struct is_layout : false_type {};
242-template <class Shape, class Stride>315+template <class Shape, class Stride, class OriginShape>
243-struct is_layout<Layout<Shape, Stride>> : true_type {};316+struct is_layout<Layout<Shape, Stride, OriginShape>> : true_type {};
244 317 
245// Layout Check318// Layout Check
246namespace detail {319namespace detail {
@@ -323,12 +396,13 @@ struct isnZ<Element, Layout, std::enable_if_t<Layout::depth == 2 && Layout::rank
323} // end namespace detail396} // end namespace detail
324 397 
325// Advanced Layout constructions398// Advanced Layout constructions
326-// Make a vector layout.399+ 
400+// Make a vector layout.
327template <class T>401template <class T>
328CATLASS_HOST_DEVICE constexpr402CATLASS_HOST_DEVICE constexpr
329auto MakeLayout(T const& len)403auto MakeLayout(T const& len)
330{404{
331- return MakeLayout(MakeShape(len), MakeStride(Int<1>{}));405+ return MakeLayout(MakeShape(len), MakeStride(Int<1>{}), MakeShape(len));
332}406}
333 407 
334// Make a inner layout with Rows and Cols.408// Make a inner layout with Rows and Cols.
@@ -341,71 +415,127 @@ auto MakeLayout(T const& rows, U const& cols)
341 std::is_same_v<LayoutTag, Catlass::layout::VectorLayout> ||415 std::is_same_v<LayoutTag, Catlass::layout::VectorLayout> ||
342 std::is_same_v<LayoutTag, Catlass::layout::zN> ||416 std::is_same_v<LayoutTag, Catlass::layout::zN> ||
343 std::is_same_v<LayoutTag, Catlass::layout::nZ> ||417 std::is_same_v<LayoutTag, Catlass::layout::nZ> ||
344- std::is_same_v<LayoutTag, Catlass::layout::zZ>,418+ std::is_same_v<LayoutTag, Catlass::layout::zZ> ||
419+ std::is_same_v<LayoutTag, Catlass::layout::L0C>,
345 "Unsupported LayoutTag for MakeLayoutFromTag, only support Catlass::layout::RowMajor or"420 "Unsupported LayoutTag for MakeLayoutFromTag, only support Catlass::layout::RowMajor or"
346- "Catlass::layout::ColumnMajor or Catlass::layout::zN or Catlass::layout::nZ or Catlass::layout::zZ");421+ "Catlass::layout::ColumnMajor or Catlass::layout::zN or Catlass::layout::nZ or Catlass::layout::zZ or Catlass::layout::L0C");
347 422 
348 constexpr uint32_t ELE_NUM_PER_C0 = Catlass::BYTE_PER_C0 / sizeof(Element);423 constexpr uint32_t ELE_NUM_PER_C0 = Catlass::BYTE_PER_C0 / sizeof(Element);
349 constexpr uint32_t ELE_NUM_PER_FRACTAL = Catlass::BYTE_PER_FRACTAL / sizeof(Element);424 constexpr uint32_t ELE_NUM_PER_FRACTAL = Catlass::BYTE_PER_FRACTAL / sizeof(Element);
350 425 
351 if constexpr (std::is_same_v<LayoutTag, Catlass::layout::VectorLayout>) {426 if constexpr (std::is_same_v<LayoutTag, Catlass::layout::VectorLayout>) {
352- return MakeLayout(MakeShape(cols), MakeStride(Int<1>{}));427+ return MakeLayout(MakeShape(cols), MakeStride(Int<1>{}), MakeShape(cols));
353 } else if constexpr (std::is_same_v<LayoutTag, Catlass::layout::RowMajor>) {428 } else if constexpr (std::is_same_v<LayoutTag, Catlass::layout::RowMajor>) {
354- return MakeLayout(MakeShape(rows, cols), MakeStride((int64_t)cols, Int<1>{}));429+ return MakeLayout(MakeShape(rows, cols),
430+ MakeStride((int64_t)cols, Int<1>{}),
431+ MakeShape(rows, cols));
355 } else if constexpr (std::is_same_v<LayoutTag, Catlass::layout::ColumnMajor>) {432 } else if constexpr (std::is_same_v<LayoutTag, Catlass::layout::ColumnMajor>) {
356- return MakeLayout(MakeShape(rows, cols), MakeStride(Int<1>{}, (int64_t)rows));433+ return MakeLayout(MakeShape(rows, cols),
434+ MakeStride(Int<1>{}, (int64_t)rows),
435+ MakeShape(rows, cols));
357 } else if constexpr (std::is_same_v<LayoutTag, Catlass::layout::zN>) {436 } else if constexpr (std::is_same_v<LayoutTag, Catlass::layout::zN>) {
358 return MakeLayout(437 return MakeLayout(
359 MakeShape(MakeShape(Int<Catlass::C0_NUM_PER_FRACTAL>{}, CeilDiv(rows, Int<Catlass::C0_NUM_PER_FRACTAL>{})),438 MakeShape(MakeShape(Int<Catlass::C0_NUM_PER_FRACTAL>{}, CeilDiv(rows, Int<Catlass::C0_NUM_PER_FRACTAL>{})),
360 MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(cols, Int<ELE_NUM_PER_C0>{}))),439 MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(cols, Int<ELE_NUM_PER_C0>{}))),
361 MakeStride(MakeStride(Int<ELE_NUM_PER_C0>{}, Int<ELE_NUM_PER_FRACTAL>{}),440 MakeStride(MakeStride(Int<ELE_NUM_PER_C0>{}, Int<ELE_NUM_PER_FRACTAL>{}),
362- MakeStride(Int<1>{}, RoundUp((int64_t)rows, Int<Catlass::C0_NUM_PER_FRACTAL>{}) * ELE_NUM_PER_C0)));441+ MakeStride(Int<1>{}, RoundUp((int64_t)rows, Int<Catlass::C0_NUM_PER_FRACTAL>{}) * ELE_NUM_PER_C0)),
442+ MakeShape(rows, cols));
363 } else if constexpr (std::is_same_v<LayoutTag, Catlass::layout::zZ>) {443 } else if constexpr (std::is_same_v<LayoutTag, Catlass::layout::zZ>) {
364 return MakeLayout(444 return MakeLayout(
365 MakeShape(MakeShape(Int<Catlass::C0_NUM_PER_FRACTAL>{}, CeilDiv(rows, Int<Catlass::C0_NUM_PER_FRACTAL>{})),445 MakeShape(MakeShape(Int<Catlass::C0_NUM_PER_FRACTAL>{}, CeilDiv(rows, Int<Catlass::C0_NUM_PER_FRACTAL>{})),
366 MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(cols, Int<ELE_NUM_PER_C0>{}))),446 MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(cols, Int<ELE_NUM_PER_C0>{}))),
367 MakeStride(MakeStride(Int<ELE_NUM_PER_C0>{},447 MakeStride(MakeStride(Int<ELE_NUM_PER_C0>{},
368 RoundUp((int64_t)cols, Int<ELE_NUM_PER_C0>{}) * Catlass::C0_NUM_PER_FRACTAL),448 RoundUp((int64_t)cols, Int<ELE_NUM_PER_C0>{}) * Catlass::C0_NUM_PER_FRACTAL),
369- MakeStride(Int<1>{}, Int<ELE_NUM_PER_FRACTAL>{})));449+ MakeStride(Int<1>{}, Int<ELE_NUM_PER_FRACTAL>{})),
450+ MakeShape(rows, cols));
451+ } else if constexpr (std::is_same_v<LayoutTag, Catlass::layout::L0C>) {
452+ constexpr uint32_t ELE_NUM_PER_FRACTAL = 256;
453+ return MakeLayout(
454+ MakeShape(MakeShape(Int<Catlass::C0_NUM_PER_FRACTAL>{}, CeilDiv(rows, Int<Catlass::C0_NUM_PER_FRACTAL>{})),
455+ MakeShape(Int<Catlass::C0_NUM_PER_FRACTAL>{}, CeilDiv(cols, Int<Catlass::C0_NUM_PER_FRACTAL>{}))),
456+ MakeStride(MakeStride(Int<Catlass::C0_NUM_PER_FRACTAL>{}, Int<ELE_NUM_PER_FRACTAL>{}),
457+ MakeStride(Int<1>{}, RoundUp((int64_t)rows, Int<Catlass::C0_NUM_PER_FRACTAL>{}) * Catlass::C0_NUM_PER_FRACTAL)),
458+ MakeShape(rows, cols));
370 } else {459 } else {
371 return MakeLayout(460 return MakeLayout(
372 MakeShape(MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(rows, Int<ELE_NUM_PER_C0>{})),461 MakeShape(MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(rows, Int<ELE_NUM_PER_C0>{})),
373 MakeShape(Int<Catlass::C0_NUM_PER_FRACTAL>{}, CeilDiv(cols, Int<Catlass::C0_NUM_PER_FRACTAL>{}))),462 MakeShape(Int<Catlass::C0_NUM_PER_FRACTAL>{}, CeilDiv(cols, Int<Catlass::C0_NUM_PER_FRACTAL>{}))),
374 MakeStride(463 MakeStride(
375 MakeStride(Int<1>{}, RoundUp((int64_t)cols, Int<Catlass::C0_NUM_PER_FRACTAL>{}) * ELE_NUM_PER_C0),464 MakeStride(Int<1>{}, RoundUp((int64_t)cols, Int<Catlass::C0_NUM_PER_FRACTAL>{}) * ELE_NUM_PER_C0),
376- MakeStride(Int<ELE_NUM_PER_C0>{}, Int<ELE_NUM_PER_FRACTAL>{})));465+ MakeStride(Int<ELE_NUM_PER_C0>{}, Int<ELE_NUM_PER_FRACTAL>{})),
466+ MakeShape(rows, cols));
377 }467 }
378}468}
379 469 
380-template <class Layout, class ShapeNew>
381-CATLASS_HOST_DEVICE constexpr
382-auto MakeLayoutTile(Layout const& layout, ShapeNew const& shapeNew)
383-{
384- static_assert(
385- is_tuple<ShapeNew>::value && depth_v<ShapeNew> == 1 && (rank_v<ShapeNew> == 1 || rank_v<ShapeNew> == 2)
386- );
387 470 
388- if constexpr (Layout::depth == 1 && (Layout::rank == 1 || Layout::rank == 2)) {471+namespace detail {
389- return MakeLayout(shapeNew, layout.stride());472+ 
390- } else if constexpr (is_static<decltype(shape<0, 0>(layout))>::value &&473+template <class OriginBase, class TileShape, class Coord, int... Is>
391- is_static<decltype(shape<1, 0>(layout))>::value) {474+CATLASS_HOST_DEVICE constexpr
392- const uint32_t rows = get<0>(shapeNew);475+auto CropOriginShape(OriginBase const& originBase, TileShape const& tileShape, Coord const& coord, seq<Is...>)
393- const uint32_t cols = get<1>(shapeNew);476+{
394- constexpr uint32_t dstInnerShapeRow = decltype(shape<0, 0>(layout))::value;477+ return MakeShape(
395- constexpr uint32_t dstInnerShapeCol = decltype(shape<1, 0>(layout))::value;478+ tla::min(
396- return MakeLayout(479+ static_cast<uint32_t>(get<Is>(tileShape)),
397- MakeShape(MakeShape(Int<dstInnerShapeRow>{}, CeilDiv<dstInnerShapeRow>(rows)),480+ (static_cast<uint32_t>(get<Is>(coord)) < static_cast<uint32_t>(get<Is>(originBase)))
398- MakeShape(Int<dstInnerShapeCol>{}, CeilDiv<dstInnerShapeCol>(cols))),481+ ? (static_cast<uint32_t>(get<Is>(originBase)) - static_cast<uint32_t>(get<Is>(coord)))
399- layout.stride());482+ : 0u
483+ )...
484+ );
485+}
486+ 
487+} // namespace detail
488+ 
489+/// 创建 tile layout：使用指定的 tile 尺寸用于内存布局计算，同时携带实际逻辑尺寸（origin_shape）。
490+/// coord 是元素坐标，用于计算实际的 originShape（处理边界情况）。
491+/// Supports layouts of any rank (rank >= 1) for depth==1 layouts.
492+/// For depth>1 (fractal) layouts, currently only rank-2 is supported.
493+template <class Layout, class TileShape, class Coord>
494+CATLASS_HOST_DEVICE constexpr
495+auto GetTileLayout(Layout const& layout, TileShape const& tileShape, Coord const& coord)
496+{
497+ static_assert(is_tuple<TileShape>::value && depth_v<TileShape> == 1 && rank_v<TileShape> >= 1,
498+ "GetTileLayout: TileShape must be a flat tuple with rank >= 1.");
499+ static_assert(is_tuple<Coord>::value && depth_v<Coord> == 1 && rank_v<Coord> == rank_v<TileShape>,
500+ "GetTileLayout: Coord must have the same rank as TileShape.");
501+ 
502+ // 统一计算 tail tile 的逻辑尺寸（originShape 裁剪）
503+ auto tileOriginShape = detail::CropOriginShape(layout.originShape(), tileShape, coord, tuple_seq<TileShape>{});
504+ 
505+ // depth==1 的布局（vector/matrix/tensor）：tile shape 直接作为 memory-layout shape
506+ // 支持任意 rank >= 1（但必须与 layout.rank 匹配）
507+ if constexpr (Layout::depth == 1) {
508+ static_assert(Layout::rank == rank_v<TileShape>,
509+ "GetTileLayout: for depth==1 layouts, TileShape rank must match layout rank.");
510+ return MakeLayout(tileShape, layout.stride(), tileOriginShape);
400 } else {511 } else {
401- const uint32_t rows = get<0>(shapeNew);512+ // depth>1 的布局（fractal layout）：目前只支持 rank=2
402- const uint32_t cols = get<1>(shapeNew);513+ // 因为 fractal layout 通常用于矩阵（rank-2），需要把 (rows, cols) 转为同结构嵌套 shape
403- const uint32_t dstInnerShapeRow = shape<0, 0>(layout);514+ static_assert(rank_v<TileShape> == 2,
404- const uint32_t dstInnerShapeCol = shape<1, 0>(layout);515+ "GetTileLayout: for depth>1 (fractal) layouts, TileShape must be rank-2 (rows, cols).");
405- return MakeLayout(516+ 
406- MakeShape(MakeShape(dstInnerShapeRow, CeilDiv(rows, dstInnerShapeRow)),517+ if constexpr (is_static<decltype(shape<0, 0>(layout))>::value &&
407- MakeShape(dstInnerShapeCol, CeilDiv(cols, dstInnerShapeCol))),518+ is_static<decltype(shape<1, 0>(layout))>::value) {
408- layout.stride());519+ const uint32_t rows = get<0>(tileShape);
520+ const uint32_t cols = get<1>(tileShape);
521+ constexpr uint32_t dstInnerShapeRow = decltype(shape<0, 0>(layout))::value;
522+ constexpr uint32_t dstInnerShapeCol = decltype(shape<1, 0>(layout))::value;
523+ return MakeLayout(
524+ MakeShape(MakeShape(Int<dstInnerShapeRow>{}, CeilDiv<dstInnerShapeRow>(rows)),
525+ MakeShape(Int<dstInnerShapeCol>{}, CeilDiv<dstInnerShapeCol>(cols))),
526+ layout.stride(),
527+ tileOriginShape);
528+ } else {
529+ const uint32_t rows = get<0>(tileShape);
530+ const uint32_t cols = get<1>(tileShape);
531+ const uint32_t dstInnerShapeRow = shape<0, 0>(layout);
532+ const uint32_t dstInnerShapeCol = shape<1, 0>(layout);
533+ return MakeLayout(
534+ MakeShape(MakeShape(dstInnerShapeRow, CeilDiv(rows, dstInnerShapeRow)),
535+ MakeShape(dstInnerShapeCol, CeilDiv(cols, dstInnerShapeCol))),
536+ layout.stride(),
537+ tileOriginShape);
538+ }
409 }539 }
410}540}
411 541 
@@ -416,12 +546,47 @@ auto MakeLayoutL0C(T const& rows, U const& cols)
416 constexpr uint32_t ELE_NUM_PER_FRACTAL = 256;546 constexpr uint32_t ELE_NUM_PER_FRACTAL = 256;
417 return MakeLayout(547 return MakeLayout(
418 MakeShape(MakeShape(Int<Catlass::C0_NUM_PER_FRACTAL>{}, CeilDiv(rows, Int<Catlass::C0_NUM_PER_FRACTAL>{})),548 MakeShape(MakeShape(Int<Catlass::C0_NUM_PER_FRACTAL>{}, CeilDiv(rows, Int<Catlass::C0_NUM_PER_FRACTAL>{})),
419- MakeShape(Int<Catlass::C0_NUM_PER_FRACTAL>{}, CeilDiv(cols, Int<Catlass::C0_NUM_PER_FRACTAL>{}))),549+ MakeShape(Int<Catlass::C0_NUM_PER_FRACTAL>{}, CeilDiv(cols, Int<Catlass::C0_NUM_PER_FRACTAL>{}))),
420 MakeStride(MakeStride(Int<Catlass::C0_NUM_PER_FRACTAL>{}, Int<ELE_NUM_PER_FRACTAL>{}),550 MakeStride(MakeStride(Int<Catlass::C0_NUM_PER_FRACTAL>{}, Int<ELE_NUM_PER_FRACTAL>{}),
421- MakeStride(551+ MakeStride(Int<1>{}, RoundUp((int64_t)rows, Int<Catlass::C0_NUM_PER_FRACTAL>{}) * Catlass::C0_NUM_PER_FRACTAL)),
422- Int<1>{}, RoundUp((int64_t)rows, Int<Catlass::C0_NUM_PER_FRACTAL>{}) * Catlass::C0_NUM_PER_FRACTAL)));552+ MakeShape(rows, cols));
423}553}
424 554 
555+//
556+// Layout transforms
557+//
558+namespace detail {
559+ 
560+// Prepend one leading dimension to a layout type
561+// This is the general form of "make batched layout":
562+// - Given a base layout type of rank R, create a new layout of rank R+1 by
563+// prefixing (batchShape, batchStride, batchOrigin) to the existing shape/stride/originShape.
564+// intentionally preserve the *types* of each stride element from the original layout
565+template <class Layout, class NewShapeT = uint32_t, class NewStrideT = int64_t, class NewOriginT = uint32_t,
566+ class Seq = tla::make_seq<Layout::rank>>
567+struct PrependDimLayout;
568+ 
569+template <class Layout, class NewShapeT, class NewStrideT, class NewOriginT, int... Is>
570+struct PrependDimLayout<Layout, NewShapeT, NewStrideT, NewOriginT, tla::seq<Is...>> {
571+ using ShapeOld = tla::remove_cvref_t<decltype(std::declval<Layout const&>().shape())>;
572+ using StrideOld = tla::remove_cvref_t<decltype(std::declval<Layout const&>().stride())>;
573+ using OriginOld = tla::remove_cvref_t<decltype(std::declval<Layout const&>().originShape())>;
574+ 
575+ using ShapeNew = tla::Shape<NewShapeT, tla::remove_cvref_t<decltype(tla::get<Is>(std::declval<ShapeOld>()))>...>;
576+ using StrideNew = tla::Stride<NewStrideT, tla::remove_cvref_t<decltype(tla::get<Is>(std::declval<StrideOld>()))>...>;
577+ using OriginNew = tla::Shape<NewOriginT, tla::remove_cvref_t<decltype(tla::get<Is>(std::declval<OriginOld>()))>...>;
578+ 
579+ using type = tla::Layout<ShapeNew, StrideNew, OriginNew>;
580+};
581+ 
582+} // namespace detail
583+ 
584+template <class Layout, class NewShapeT = uint32_t, class NewStrideT = int64_t, class NewOriginT = uint32_t>
585+using PrependDimLayout_t = typename detail::PrependDimLayout<Layout, NewShapeT, NewStrideT, NewOriginT>::type;
586+ 
587+template <class Layout>
588+using MakeBatchedLayout_t = PrependDimLayout_t<Layout>;
589+ 
425} // end namespace tla590} // end namespace tla
426 591 
427# endif // TLA_LAYOUT_HPP592# endif // TLA_LAYOUT_HPP
Minclude/tla/numeric/integral_constant.hpp+26-1
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4 * CANN Open Software License Agreement Version 2.0 (the "License").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -92,6 +92,31 @@ using _128 = Int<128>;
92using _256 = Int<256>;92using _256 = Int<256>;
93using _512 = Int<512>;93using _512 = Int<512>;
94 94 
95+//
96+// Underscore placeholder (for slicing semantics)
97+//
98+// Usage:
99+// - `tla::_` is an empty tag value that can be used inside `tla::Coord` / tensor indexing
100+// to indicate "take the whole dimension" (full slice).
101+struct Underscore {
102+ using type = Underscore;
103+};
104+ 
105+CATLASS_HOST_DEVICE constexpr Underscore _{};
106+ 
107+template <class T>
108+struct is_underscore : false_type {};
109+template <>
110+struct is_underscore<Underscore> : true_type {};
111+template <class T>
112+struct is_underscore<T const> : is_underscore<T> {};
113+template <class T>
114+struct is_underscore<T const&> : is_underscore<T> {};
115+template <class T>
116+struct is_underscore<T&> : is_underscore<T> {};
117+template <class T>
118+struct is_underscore<T&&> : is_underscore<T> {};
119+ 
95/***************/120/***************/
96/** Operators **/121/** Operators **/
97/***************/122/***************/
Minclude/tla/tensor.hpp+322-6
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4 * CANN Open Software License Agreement Version 2.0 (the "License").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -17,10 +17,178 @@
17#include "tla/int_tuple.hpp"17#include "tla/int_tuple.hpp"
18 18 
19namespace tla {19namespace tla {
20+ 
21+//
22+// Underscore slicing utilities
23+//
24+namespace detail {
25+ 
26+// Safe "element type at index I" helper that avoids instantiating get<I> for invalid indices.
27+template <class Coord, int I, class Enable = void>
28+struct coord_elem_type {
29+ using type = void;
30+};
31+template <class Coord, int I>
32+struct coord_elem_type<Coord, I,
33+ std::enable_if_t<(I >= 0) && (I < (int)tla::tuple_size<tla::remove_cvref_t<Coord>>::value)>> {
34+ using type = tla::remove_cvref_t<decltype(tla::get<I>(std::declval<Coord>()))>;
35+};
36+ 
37+template <class Coord, int I>
38+struct coord_elem_is_underscore : tla::is_underscore<typename coord_elem_type<Coord, I>::type> {};
39+ 
40+// Count underscores.
41+template <class Coord, int I, class Enable = void>
42+struct underscore_count_from : tla::integral_constant<int, 0> {};
43+template <class Coord, int I>
44+struct underscore_count_from<Coord, I, std::enable_if_t<(I >= 0)>> {
45+ static constexpr int value =
46+ (coord_elem_is_underscore<Coord, I>::value ? 1 : 0) + underscore_count_from<Coord, I - 1>::value;
47+};
48+template <class Coord>
49+struct underscore_count
50+ : tla::integral_constant<int,
51+ underscore_count_from<Coord, (int)tla::tuple_size<tla::remove_cvref_t<Coord>>::value - 1>::value> {};
52+ 
53+// Build index sequences for underscore dims (stable 0..R-1 recursion).
54+template <class Coord, int I, int R, int... Is>
55+struct underscore_indices_impl;
56+template <class Coord, int R, int... Is>
57+struct underscore_indices_impl<Coord, R, R, Is...> {
58+ using type = seq<Is...>;
59+};
60+template <class Coord, int I, int R, int... Is>
61+struct underscore_indices_impl
62+ : std::conditional_t<coord_elem_is_underscore<Coord, I>::value,
63+ underscore_indices_impl<Coord, I + 1, R, Is..., I>,
64+ underscore_indices_impl<Coord, I + 1, R, Is...>> {};
65+ 
66+template <class Coord>
67+using underscore_indices =
68+ typename underscore_indices_impl<Coord, 0, (int)tla::tuple_size<tla::remove_cvref_t<Coord>>::value>::type;
69+ 
70+// Replace every tla::_ with 0 for offset computation.
71+template <class T>
72+CATLASS_HOST_DEVICE constexpr decltype(auto) underscore_to_zero(T const& x) {
73+ if constexpr (tla::is_underscore<T>::value) {
74+ return tla::_0{};
75+ } else {
76+ return x;
77+ }
78+}
79+ 
80+template <class Coord, int... I>
81+CATLASS_HOST_DEVICE constexpr auto replace_underscore_with_zero_impl(Coord const& c, seq<I...>) {
82+ return tla::MakeCoord(underscore_to_zero(tla::get<I>(c))...);
83+}
84+ 
85+template <class Coord>
86+CATLASS_HOST_DEVICE constexpr auto replace_underscore_with_zero(Coord const& c) {
87+ static_assert(tla::is_tuple<tla::remove_cvref_t<Coord>>::value,
88+ "Coord must be tla::tuple for underscore slicing.");
89+ return replace_underscore_with_zero_impl(c, tuple_seq<Coord>{});
90+}
91+ 
92+// Build a layout from selected top-level indices of `layout`.
93+template <class Layout, int... Is>
94+CATLASS_HOST_DEVICE constexpr auto select_layout(Layout const& layout, seq<Is...>) {
95+ auto shape_new = tla::MakeTuple(tla::get<Is>(layout.shape())...);
96+ auto stride_new = tla::MakeTuple(tla::get<Is>(layout.stride())...);
97+ auto origin_new = tla::MakeTuple(tla::get<Is>(layout.originShape())...);
98+ return tla::MakeLayout(shape_new, stride_new, origin_new);
99+}
100+ 
101+} // namespace detail
102+ 
103+//
104+// slice_and_offset
105+//
106+// A lightweight helper that factors underscore slicing into:
107+// - a projected layout (keeping only underscored dimensions, in-order), and
108+// - a base offset computed at the fixed indices (underscored dims treated as 0).
109+//
110+// Notes:
111+// - `coord_arg` must be a one-level `tla::tuple` and contain at least one `tla::_`.
112+// - This function does not perform runtime bounds checks against originShape(); out-of-bounds is undefined behavior.
113+// - Returned `offset` is an element offset intended for `BuiltinTensor::operator[](offset)` view creation.
114+template <class CoordArg, class Layout, class BaseCoord>
115+CATLASS_HOST_DEVICE constexpr auto slice_and_offset(CoordArg const& coord_arg,
116+ Layout const& layout,
117+ BaseCoord const& base_coord)
118+{
119+ static_assert(tla::is_tuple<tla::remove_cvref_t<CoordArg>>::value, "slice_and_offset expects a tuple CoordArg.");
120+ static_assert(depth_v<CoordArg> == 1, "slice_and_offset only supports one-level CoordArg (no nested tuples).");
121+ static_assert((int)tla::tuple_size<tla::remove_cvref_t<CoordArg>>::value == (int)Layout::rank,
122+ "slice_and_offset requires CoordArg rank == Layout::rank.");
123+ static_assert(tla::is_tuple<tla::remove_cvref_t<BaseCoord>>::value, "slice_and_offset expects a tuple BaseCoord.");
124+ static_assert((int)tla::tuple_size<tla::remove_cvref_t<BaseCoord>>::value == (int)Layout::rank,
125+ "slice_and_offset requires BaseCoord rank == Layout::rank.");
126+ 
127+ constexpr int k = detail::underscore_count<CoordArg>::value;
128+ static_assert(k > 0, "slice_and_offset requires at least one underscore.");
129+ static_assert(k <= Layout::rank, "Invalid underscore count.");
130+ 
131+ // Compute base offset using zeros for underscores
132+ auto coord0 = detail::replace_underscore_with_zero(coord_arg);
133+ auto full0 = Add(base_coord, coord0);
134+ auto offset = (int64_t)layout(full0);
135+ 
136+ // Determine output dims (underscored dims) and build projected layout.
137+ using Us = detail::underscore_indices<CoordArg>;
138+ auto layout_proj = detail::select_layout(layout, Us{});
139+ 
140+ return tla::MakeTuple(layout_proj, offset);
141+}
142+ 
143+// Convenience overload (no base_coord): assume base_coord == 0.
144+template <class CoordArg, class Layout>
145+CATLASS_HOST_DEVICE constexpr auto slice_and_offset(CoordArg const& coord_arg, Layout const& layout)
146+{
147+ using Z = detail::MakeZeroTuple<Layout::rank>;
148+ return slice_and_offset(coord_arg, layout, Z{});
149+}
150+ 
151+ 
20//152//
21// Tensor153// Tensor
22//154//
23 155 
156+namespace detail {
157+ 
158+template <class A, class B, int... Is>
159+CATLASS_DEVICE constexpr
160+auto HadamardU32(A const& a, B const& b, seq<Is...>)
161+{
162+ return MakeCoord((static_cast<uint32_t>(get<Is>(a)) * static_cast<uint32_t>(get<Is>(b)))...);
163+}
164+ 
165+template <class TensorT, class CoordT, class ShapeT, int R>
166+CATLASS_DEVICE constexpr
167+auto GetTileImpl(TensorT const& tensor, CoordT const& coord, ShapeT const& shape, Int<R>)
168+{
169+ static_assert(is_tuple<CoordT>::value && depth_v<CoordT> == 1 && rank_v<CoordT> == R, "Coord rank mismatch.");
170+ static_assert(is_tuple<ShapeT>::value && depth_v<ShapeT> == 1 && rank_v<ShapeT> == R, "Shape rank mismatch.");
171+ 
172+ auto layoutNew = GetTileLayout(tensor.layout(), shape, coord);
173+ auto coordNew = Add(tensor.coord(), coord);
174+ return MakeTensor(tensor.data(), layoutNew, coordNew, Catlass::Arch::PositionType<TensorT::position>{});
175+}
176+ 
177+template <class TensorT, class TileCoord, class TileShape, int R>
178+CATLASS_DEVICE constexpr
179+auto TileViewImpl(TensorT const& tensor, TileCoord const& tileCoord, TileShape const& tileShape, Int<R>)
180+{
181+ static_assert(is_tuple<TileCoord>::value && depth_v<TileCoord> == 1 && rank_v<TileCoord> == R, "TileCoord rank mismatch.");
182+ static_assert(is_tuple<TileShape>::value && depth_v<TileShape> == 1 && rank_v<TileShape> == R, "TileShape rank mismatch.");
183+ 
184+ auto elementOffset = HadamardU32(tileCoord, tileShape, tuple_seq<TileCoord>{});
185+ auto layoutNew = GetTileLayout(tensor.layout(), tileShape, elementOffset);
186+ auto coordNew = Add(tensor.coord(), elementOffset);
187+ return MakeTensor(tensor.data(), layoutNew, coordNew, Catlass::Arch::PositionType<TensorT::position>{});
188+}
189+ 
190+} // namespace detail
191+ 
24template <class BuiltinTensor, class Layout_, class Coord_, AscendC::TPosition Position>192template <class BuiltinTensor, class Layout_, class Coord_, AscendC::TPosition Position>
25struct Tensor {193struct Tensor {
26 using Element = typename BuiltinTensor::PrimType;194 using Element = typename BuiltinTensor::PrimType;
@@ -83,6 +251,61 @@ struct Tensor {
83 return layout().stride();251 return layout().stride();
84 }252 }
85 253 
254+ CATLASS_HOST_DEVICE constexpr
255+ decltype(auto) originShape() const
256+ {
257+ return layout().originShape();
258+ }
259+ 
260+ //
261+ // Indexing / slicing
262+ //
263+ // - No underscore: returns `data()[layout()(coord()+coord_arg)]`
264+ // - Underscores (0..rank, one-level coord): returns a subtensor view over the underscored dimensions (kept in-order)
265+ // Notes:
266+ // - Coord must be one-level (no nested tuples in coord elements).
267+ // - Fixed (non-underscore) indices are expected to be within originShape(). This implementation does not
268+ // generally perform runtime bounds checks/cropping; out-of-bounds indices result in undefined behavior.
269+ template <class CoordArg>
270+ CATLASS_HOST_DEVICE constexpr
271+ decltype(auto) operator()(CoordArg const& coord_arg) const
272+ {
273+ if constexpr (tla::is_tuple<tla::remove_cvref_t<CoordArg>>::value) {
274+ static_assert(depth_v<CoordArg> == 1, "Underscore slicing only supports one-level Coord (no nested tuples).");
275+ static_assert(tla::tuple_size<tla::remove_cvref_t<CoordArg>>::value == Layout::rank,
276+ "Tensor::operator()(coord): Coord rank must equal tensor rank (Layout::rank).");
277+ 
278+ constexpr int k = detail::underscore_count<CoordArg>::value;
279+ if constexpr (k > 0) {
280+ static_assert(k <= Layout::rank, "Invalid underscore count.");
281+ 
282+ auto sliced = tla::slice_and_offset(coord_arg, layout(), coord());
283+ auto layout_proj = tla::get<0>(sliced);
284+ auto offset = (int64_t)tla::get<1>(sliced);
285+ 
286+ using CoordZ = detail::MakeZeroTuple<(size_t)k>;
287+ auto data_new = data()[static_cast<uint64_t>(offset)];
288+ return Tensor<decltype(data_new), decltype(layout_proj), CoordZ, position>(data_new, layout_proj, CoordZ{});
289+ } else {
290+ // No underscore: point view at coord() + coord_arg
291+ auto full = Add(coord(), coord_arg);
292+ return data()[layout()(full)];
293+ }
294+ } else {
295+ // Scalar coordinate convenience (rank-1): treat it as a 1D coord tuple.
296+ static_assert(Layout::rank == 1, "Tensor::operator()(scalar) is only supported for rank-1 tensors.");
297+ auto full = Add(coord(), MakeCoord(coord_arg));
298+ return data()[layout()(full)];
299+ }
300+ }
301+ 
302+ template <class Coord0, class Coord1, class... Coords>
303+ CATLASS_HOST_DEVICE constexpr
304+ decltype(auto) operator()(Coord0 const& c0, Coord1 const& c1, Coords const&... cs) const
305+ {
306+ return operator()(MakeCoord(c0, c1, cs...));
307+ }
308+ 
86 tla::tuple<BuiltinTensor, Layout, Coord> rep_;309 tla::tuple<BuiltinTensor, Layout, Coord> rep_;
87};310};
88 311 
@@ -101,15 +324,108 @@ auto MakeTensor(BuiltinTensor const& builtinTensor, Layout const& layout, Coord
101 return Tensor<BuiltinTensor, Layout, Coord, PositionType::value>(builtinTensor, layout, coord);324 return Tensor<BuiltinTensor, Layout, Coord, PositionType::value>(builtinTensor, layout, coord);
102}325}
103 326 
327+// Get a tile from tensor, automatically handling boundary cases (tail tile).
328+// coord is element coordinate; shape is tile size for memory layout.
329+// The returned Tensor's layout.shape() is used for memory layout calculation,
330+// and layout.originShape() is the actual logical size (may be smaller than shape).
331+// Supports tensors of any rank (rank >= 1).
104template <class Tensor, class Coord, class Shape>332template <class Tensor, class Coord, class Shape>
105CATLASS_DEVICE constexpr333CATLASS_DEVICE constexpr
106auto GetTile(Tensor const& tensor, Coord const& coord, Shape const& shape)334auto GetTile(Tensor const& tensor, Coord const& coord, Shape const& shape)
107{335{
108- auto layout = tensor.layout();336+ static_assert(Tensor::rank >= 1, "GetTile requires tensor rank >= 1.");
109- auto builtinTensor = tensor.data();337+ static_assert(Tensor::rank == rank_v<Coord> && Tensor::rank == rank_v<Shape>,
110- auto layoutNew = MakeLayoutTile(layout, shape);338+ "GetTile: coord and shape must have the same rank as the tensor.");
111- auto coordNew = Add(tensor.coord(), coord);339+ return detail::GetTileImpl(tensor, coord, shape, Int<Tensor::rank>{});
112- return MakeTensor(builtinTensor, layoutNew, coordNew, Catlass::Arch::PositionType<Tensor::position>{});340+}
341+ 
342+ 
343+// 从 tensor 中获取一个 tile，自动处理边界情况（tail tile）。
344+// tileCoord 是 tile 单位坐标；tileShape 是用于内存布局的 tile 尺寸。
345+// 返回的 Tensor 的 layout.shape() 用于内存布局计算，layout.originShape() 是实际逻辑尺寸（可能小于 tileShape）。
346+// TileView(tensor, tileCoord, tileShape) = GetTile(tensor, tileCoord ⊙ tileShape, tileShape)
347+// Supports tensors of any rank (rank >= 1).
348+template <class TensorT, class TileCoord, class TileShape>
349+CATLASS_DEVICE constexpr
350+auto TileView(TensorT const& tensor, TileCoord const& tileCoord, TileShape const& tileShape)
351+{
352+ static_assert(TensorT::rank >= 1, "TileView requires tensor rank >= 1.");
353+ static_assert(TensorT::rank == rank_v<TileCoord> && TensorT::rank == rank_v<TileShape>,
354+ "TileView: tileCoord and tileShape must have the same rank as the tensor.");
355+ return detail::TileViewImpl(tensor, tileCoord, tileShape, Int<TensorT::rank>{});
356+}
357+ 
358+// 创建一个与另一个 Tensor 类似的 Tensor：
359+// 目标 layout 根据 LayoutTagDst 构造，从 LikeTensor::Element 推断 ElementDst，从 likeTensor 的 originShape 提取尺寸。
360+template <class LayoutTagDst, class BuiltinTensor, class LikeTensor, class PositionType>
361+CATLASS_HOST_DEVICE constexpr
362+auto MakeTensorLike(BuiltinTensor const& builtinTensor,
363+ LikeTensor const& likeTensor,
364+ PositionType)
365+{
366+ using ElementDst = typename LikeTensor::Element;
367+ static_assert(std::is_same_v<typename BuiltinTensor::PrimType, ElementDst>,
368+ "BuiltinTensor element type must match LikeTensor element type");
369+ return MakeTensorLike<LayoutTagDst, ElementDst>(builtinTensor, likeTensor, PositionType{});
370+}
371+ 
372+// 创建一个与另一个 Tensor 类似的 Tensor：
373+// 使用 layoutBase 的 shape/stride，但继承 likeTensor 的 originShape，从 LikeTensor::Element 推断 ElementDst。
374+template <class LayoutTagDst, class BuiltinTensor, class LikeTensor, class PositionType, class LayoutBase>
375+CATLASS_HOST_DEVICE constexpr
376+auto MakeTensorLike(BuiltinTensor const& builtinTensor,
377+ LikeTensor const& likeTensor,
378+ PositionType,
379+ LayoutBase const& layoutBase)
380+{
381+ using ElementDst = typename LikeTensor::Element;
382+ static_assert(std::is_same_v<typename BuiltinTensor::PrimType, ElementDst>,
383+ "BuiltinTensor element type must match LikeTensor element type");
384+ return MakeTensorLike<LayoutTagDst, ElementDst>(builtinTensor, likeTensor, PositionType{}, layoutBase);
385+}
386+ 
387+// 创建一个与另一个 Tensor 类似的 Tensor：
388+// 目标 layout 根据 LayoutTagDst 构造，从 LikeTensor::Element 推断 ElementDst，从 likeTensor 的 originShape 提取尺寸。(调用MakeLayout，可能会因分型布局合法要求对shape进行以分型为粒度的向上取整)
389+// 允许 BuiltinTensor 的元素类型与 LikeTensor 的元素类型不同（例如 L0C 使用 ElementAccumulator 而不是 ElementC）。
390+template <class LayoutTagDst, class ElementDst, class BuiltinTensor, class LikeTensor, class PositionType>
391+CATLASS_HOST_DEVICE constexpr
392+auto MakeTensorLike(BuiltinTensor const& builtinTensor,
393+ LikeTensor const& likeTensor,
394+ PositionType)
395+{
396+ static_assert(LikeTensor::rank == 1 || LikeTensor::rank == 2,
397+ "MakeTensorLike<LayoutTag, Element>(..., likeTensor) expects rank-1 or rank-2 likeTensor.");
398+ static_assert(std::is_same_v<typename BuiltinTensor::PrimType, ElementDst>,
399+ "BuiltinTensor element type must match specified ElementDst type");
400+ // 根据目标布局格式（LayoutTagDst）和指定的元素类型构造 layout
401+ if constexpr (LikeTensor::rank == 1) {
402+ auto layoutNominal = MakeLayout<ElementDst, LayoutTagDst>(get<0>(likeTensor.layout().originShape()));
403+ using Coord0 = detail::MakeZeroTuple<decltype(layoutNominal)::rank>;
404+ return Tensor<BuiltinTensor, decltype(layoutNominal), Coord0, PositionType::value>(builtinTensor, layoutNominal);
405+ } else {
406+ static_assert(LikeTensor::rank == 2, "MakeTensorLike<LayoutTag, Element>(..., likeTensor) expects rank-1 or rank-2 likeTensor.");
407+ auto layoutNominal = MakeLayout<ElementDst, LayoutTagDst>(get<0>(likeTensor.layout().originShape()), get<1>(likeTensor.layout().originShape()));
408+ using Coord0 = detail::MakeZeroTuple<decltype(layoutNominal)::rank>;
409+ return Tensor<BuiltinTensor, decltype(layoutNominal), Coord0, PositionType::value>(builtinTensor, layoutNominal);
410+ }
411+}
412+ 
413+// 创建一个与另一个 Tensor 类似的 Tensor：
414+// 使用 layoutBase 的 shape/stride，但继承 likeTensor 的 originShape。允许 BuiltinTensor 的元素类型与 LikeTensor 的元素类型不同。
415+template <class LayoutTagDst, class ElementDst, class BuiltinTensor, class LikeTensor, class PositionType, class LayoutBase>
416+CATLASS_HOST_DEVICE constexpr
417+auto MakeTensorLike(BuiltinTensor const& builtinTensor,
418+ LikeTensor const& likeTensor,
419+ PositionType,
420+ LayoutBase const& layoutBase)
421+{
422+ static_assert(LikeTensor::rank == 1 || LikeTensor::rank == 2, "MakeTensorLike<LayoutTag, Element>(..., likeTensor, layoutBase) expects rank-1 or rank-2 likeTensor.");
423+ static_assert(std::is_same_v<typename BuiltinTensor::PrimType, ElementDst>,
424+ "BuiltinTensor element type must match specified ElementDst type");
425+ 
426+ auto layoutFixedStride = MakeLayout(layoutBase.shape(), layoutBase.stride(), likeTensor.originShape());
427+ using Coord0 = detail::MakeZeroTuple<decltype(layoutFixedStride)::rank>;
428+ return Tensor<BuiltinTensor, decltype(layoutFixedStride), Coord0, PositionType::value>(builtinTensor, layoutFixedStride);
113}429}
114 430 
115} // end namespace tla431} // end namespace tla
Mtests/test_example.py+1-0
@@ -206,6 +206,7 @@ normal_cases_2201 = [
206 "34_single_core_splitk_matmul 256 512 1024 0",206 "34_single_core_splitk_matmul 256 512 1024 0",
207 "42_quant_optimized_matmul_tla 256 512 1024 0",207 "42_quant_optimized_matmul_tla 256 512 1024 0",
208 "44_quant_matmul_full_loadA_tla 256 512 1024 0",208 "44_quant_matmul_full_loadA_tla 256 512 1024 0",
209+ "45_strided_batched_matmul_tla 5 256 512 1024 0",
209 "102_dynamic_optimized_matmul 256 512 1024 0 0 0"210 "102_dynamic_optimized_matmul 256 512 1024 0 0 0"
210 "103_dynamic_optimized_quant_matmul_per_token_basic 256 512 1024 0 0 0",211 "103_dynamic_optimized_quant_matmul_per_token_basic 256 512 1024 0 0 0",
211]212]