Lines Matching full:tensor
4 func.func @cast(%arg0: tensor<*xf32>, %arg1 : tensor<4x4xf32>, %arg2: tensor<?x?xf32>) {
5 // CHECK: tensor.cast %{{.*}} : tensor<*xf32> to tensor<?x?xf32>
6 %0 = tensor.cast %arg0 : tensor<*xf32> to tensor<?x?xf32>
7 // CHECK: tensor.cast %{{.*}} : tensor<4x4xf32> to tensor<*xf32>
8 %1 = tensor.cast %arg1 : tensor<4x4xf32> to tensor<*xf32>
9 // CHECK: tensor.cast %{{.*}} : tensor<?x?xf32> to tensor<4x?xf32>
10 %2 = tensor.cast %arg2 : tensor<?x?xf32> to tensor<4x?xf32>
11 // CHECK: tensor.cast %{{.*}} : tensor<4x?xf32> to tensor<?x?xf32>
12 %3 = tensor.cast %2 : tensor<4x?xf32> to tensor<?x?xf32>
19 func.func @concat(%arg0: tensor<4x7x3xf32>, %arg1 : tensor<4x4x3xf32>, %arg2: tensor<?x?x?xf32>) {
20 // CHECK: tensor.concat dim(0) %{{.*}} : (tensor<4x7x3xf32>) -> tensor<4x7x3xf32>
21 %0 = tensor.concat dim(0) %arg0 : (tensor<4x7x3xf32>) -> tensor<4x7x3xf32>
22 …// CHECK: tensor.concat dim(1) %{{.*}} : (tensor<4x7x3xf32>, tensor<4x4x3xf32>) -> tensor<4x11x3xf…
23 …%1 = tensor.concat dim(1) %arg0, %arg1 : (tensor<4x7x3xf32>, tensor<4x4x3xf32>) -> tensor<4x11x3xf…
24 …// CHECK: tensor.concat dim(2) %{{.*}} : (tensor<4x7x3xf32>, tensor<?x?x?xf32>) -> tensor<?x?x?xf3…
25 …%2 = tensor.concat dim(2) %arg0, %arg2 : (tensor<4x7x3xf32>, tensor<?x?x?xf32>) -> tensor<?x?x?xf3…
26 …// CHECK: tensor.concat dim(1) %{{.*}} : (tensor<?x?x?xf32>, tensor<?x?x?xf32>) -> tensor<?x10x?xf…
27 …%3 = tensor.concat dim(1) %arg2, %arg2 : (tensor<?x?x?xf32>, tensor<?x?x?xf32>) -> tensor<?x10x?xf…
28 …// CHECK: tensor.concat dim(1) %{{.*}} : (tensor<?x?x?xf32>, tensor<4x4x3xf32>, tensor<4x7x3xf32>)…
29 …%4 = tensor.concat dim(1) %arg2, %arg1, %arg0 : (tensor<?x?x?xf32>, tensor<4x4x3xf32>, tensor<4x7x…
37 func.func @empty(%sz: index) -> tensor<5x?x6xf32> {
38 // CHECK: tensor.empty(%[[sz]]) : tensor<5x?x6xf32>
39 %0 = tensor.empty(%sz) : tensor<5x?x6xf32>
40 return %0 : tensor<5x?x6xf32>
47 func.func @empty_with_encoding(%sz: index) -> tensor<5x?x6xf32, "foo"> {
48 // CHECK: tensor.empty(%[[sz]]) : tensor<5x?x6xf32, "foo">
49 %0 = tensor.empty(%sz) : tensor<5x?x6xf32, "foo">
50 return %0 : tensor<5x?x6xf32, "foo">
56 // CHECK-SAME: %[[TENSOR:.*]]: tensor<?x?x?xf32>,
58 func.func @extract(%arg0: tensor<?x?x?xf32>, %arg1: index) {
59 // CHECK: tensor.extract %[[TENSOR]][%[[INDEX]], %[[INDEX]], %[[INDEX]]] : tensor<?x?x?xf32>
60 %0 = tensor.extract %arg0[%arg1, %arg1, %arg1] : tensor<?x?x?xf32>
69 // CHECK-SAME: %[[DEST1:.*]]: tensor<?x?x?xf32>
70 func.func @insert(%arg0: f32, %arg1: index, %arg2: tensor<?x?x?xf32>) {
71 …// CHECK: tensor.insert %[[SCALAR]] into %[[DEST1]][%[[INDEX]], %[[INDEX]], %[[INDEX]]] : tensor<?…
72 %0 = tensor.insert %arg0 into %arg2[%arg1, %arg1, %arg1] : tensor<?x?x?xf32>
78 // CHECK-LABEL: func @tensor.from_elements() {
79 func.func @tensor.from_elements() {
81 // CHECK: tensor.from_elements %c0 : tensor<1xindex>
82 %0 = tensor.from_elements %c0 : tensor<1xindex>
85 // CHECK: tensor.from_elements %c0, %c1 : tensor<2xindex>
86 %1 = tensor.from_elements %c0, %c1 : tensor<2xindex>
90 // CHECK: tensor.from_elements [[C0_F32]] : tensor<1xf32>
91 %2 = tensor.from_elements %c0_f32 : tensor<1xf32>
93 // CHECK: tensor.from_elements : tensor<0xindex>
94 %3 = tensor.from_elements : tensor<0xindex>
96 // CHECK: tensor.from_elements %c0, %c1, %c0, %c1, %c0, %c1 : tensor<2x3xindex>
97 %4 = tensor.from_elements %c0, %c1, %c0, %c1, %c0, %c1 : tensor<2x3xindex>
99 // CHECK: tensor.from_elements %c0 : tensor<index>
100 %5 = tensor.from_elements %c0 : tensor<index>
106 // CHECK-LABEL: @tensor.generate
107 func.func @tensor.generate(%m : index, %n : index)
108 -> tensor<?x3x?xf32> {
109 %tnsr = tensor.generate %m, %n {
112 tensor.yield %elem : f32
113 } : tensor<?x3x?xf32>
114 return %tnsr : tensor<?x3x?xf32>
120 func.func @tensor_reshape(%unranked: tensor<*xf32>, %shape1: tensor<1xi32>,
121 %shape2: tensor<2xi32>, %shape3: tensor<?xi32>) -> tensor<*xf32> {
122 %dyn_vec = tensor.reshape %unranked(%shape1)
123 : (tensor<*xf32>, tensor<1xi32>) -> tensor<?xf32>
124 %dyn_mat = tensor.reshape %dyn_vec(%shape2)
125 : (tensor<?xf32>, tensor<2xi32>) -> tensor<?x?xf32>
126 %new_unranked = tensor.reshape %dyn_mat(%shape3)
127 : (tensor<?x?xf32>, tensor<?xi32>) -> tensor<*xf32>
128 return %new_unranked : tensor<*xf32>
134 func.func @slice(%t: tensor<8x16x4xf32>, %idx : index) {
138 // CHECK: tensor.extract_slice
139 // CHECK-SAME: tensor<8x16x4xf32> to tensor<?x?x?xf32>
140 %1 = tensor.extract_slice %t[%c0, %c0, %c0][%idx, %idx, %idx][%c1, %c1, %c1]
141 : tensor<8x16x4xf32> to tensor<?x?x?xf32>
143 // CHECK: tensor.extract_slice
144 // CHECK-SAME: tensor<8x16x4xf32> to tensor<4x4x4xf32>
145 %2 = tensor.extract_slice %t[0, 2, 0][4, 4, 4][1, 1, 1]
146 : tensor<8x16x4xf32> to tensor<4x4x4xf32>
148 // CHECK: tensor.extract_slice
149 // CHECK-SAME: tensor<8x16x4xf32> to tensor<4x4xf32>
150 %3 = tensor.extract_slice %t[0, 2, 0][4, 1, 4][1, 1, 1]
151 : tensor<8x16x4xf32> to tensor<4x4xf32>
160 %t: tensor<8x16x4xf32>,
161 %td: tensor<8x?x4xf32>,
162 %t2: tensor<16x32x8xf32>,
163 %t3: tensor<4x4xf32>,
169 // CHECK: tensor.insert_slice
170 // CHECK-SAME: tensor<8x16x4xf32> into tensor<16x32x8xf32>
171 %1 = tensor.insert_slice %t into %t2[%c0, %c0, %c0][8, 16, 4][%c1, %c1, %c1]
172 : tensor<8x16x4xf32> into tensor<16x32x8xf32>
174 // CHECK: tensor.insert_slice
175 // CHECK-SAME: tensor<8x16x4xf32> into tensor<16x32x8xf32>
176 %2 = tensor.insert_slice %t into %t2[%c0, %idx, %c0][8, 16, 4][%c1, 1, %c1]
177 : tensor<8x16x4xf32> into tensor<16x32x8xf32>
179 // CHECK: tensor.insert_slice
180 // CHECK-SAME: tensor<4x4xf32> into tensor<8x16x4xf32>
181 %3 = tensor.insert_slice %t3 into %t[0, 2, 0][4, 1, 4][1, 1, 1]
182 : tensor<4x4xf32> into tensor<8x16x4xf32>
184 // CHECK: tensor.insert_slice
185 // CHECK-SAME: tensor<8x?x4xf32> into tensor<8x16x4xf32>
186 %4 = tensor.insert_slice %td into %t[0, %idx, 0][8, %sz, 4][1, 1, 1]
187 : tensor<8x?x4xf32> into tensor<8x16x4xf32>
194 func.func @tensor_reshape_zero_dim(%arg0 : tensor<1x1xf32>, %arg1 : tensor<f32>)
195 -> (tensor<f32>, tensor<1x1xf32>) {
196 %0 = tensor.collapse_shape %arg0 [] : tensor<1x1xf32> into tensor<f32>
197 %1 = tensor.expand_shape %0 [] output_shape [1, 1] : tensor<f32> into tensor<1x1xf32>
198 return %0, %1 : tensor<f32>, tensor<1x1xf32>
201 // CHECK: tensor.collapse_shape %{{.*}} [] : tensor<1x1xf32> into tensor<f32>
202 // CHECK: tensor.expand_shape %{{.*}} [] output_shape [1, 1] : tensor<f32> into tensor<1x1x…
206 func.func @tensor_expand_shape_dynamic_dim(%arg0 : tensor<?x?xf32>, %sz0 : index, %sz1 : index, %sz…
207 -> (tensor<5x?x?x?xf32>) {
208 …%1 = tensor.expand_shape %arg0 [[0, 1], [2, 3]] output_shape [5, %sz0, %sz1, %sz2] : tensor<?x?xf3…
209 return %1 : tensor<5x?x?x?xf32>
212 …c @tensor_expand_shape_dynamic_dim(%arg0: tensor<?x?xf32>, %arg1: index, %arg2: index, %arg3: inde…
213 … %expanded = tensor.expand_shape %arg0 {{\[\[}}0, 1], [2, 3{{\]\]}} output_shape [5, %arg1, %arg…
214 // CHECK: return %expanded : tensor<5x?x?x?xf32>
221 (%arg0: tensor<?x?x?x4x?xf32>) -> tensor<?x?x?xf32>
223 %0 = tensor.collapse_shape %arg0 [[0], [1], [2, 3, 4]] :
224 tensor<?x?x?x4x?xf32> into tensor<?x?x?xf32>
225 return %0 : tensor<?x?x?xf32>
228 // CHECK: tensor.collapse_shape
233 func.func @rank(%t : tensor<4x4x?xf32>) {
234 // CHECK: %{{.*}} = tensor.rank %{{.*}} : tensor<4x4x?xf32>
235 %0 = "tensor.rank"(%t) : (tensor<4x4x?xf32>) -> index
237 // CHECK: %{{.*}} = tensor.rank %{{.*}} : tensor<4x4x?xf32>
238 %1 = tensor.rank %t : tensor<4x4x?xf32>
244 func.func @pad_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index,
245 %pad_value: f32) -> tensor<6x?x?x?xf32> {
246 %0 = tensor.pad %arg0 low[2, %low, 3, 3] high[3, 3, %high, 2] {
248 tensor.yield %pad_value : f32
249 } : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32>
250 return %0 : tensor<6x?x?x?xf32>
256 // CHECK: tensor.pad %[[ARG0]]
259 // CHECK: : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32>
263 func.func @pad_static(%arg0: tensor<3x4xf32>, %pad_value: f32) -> tensor<6x9xf32> {
264 %0 = tensor.pad %arg0 low[1, 2] high[2, 3] {
266 tensor.yield %pad_value : f32
267 } : tensor<3x4xf32> to tensor<6x9xf32>
268 return %0 : tensor<6x9xf32>
272 // CHECK: tensor.pad %[[ARG0]] low[1, 2] high[2, 3]
273 // CHECK: : tensor<3x4xf32> to tensor<6x9xf32>
277 func.func @pad_asymmetrical(%arg0: tensor<2x3xf32>, %ub0: index, %ub1: index,
278 %pad_value: f32) -> tensor<?x?xf32> {
279 %0 = tensor.pad %arg0 low[0, 0] high[%ub0, %ub1] {
281 tensor.yield %pad_value : f32
282 } : tensor<2x3xf32> to tensor<?x?xf32>
283 return %0 : tensor<?x?xf32>
289 // CHECK: tensor.pad %[[ARG0]]
292 // CHECK: : tensor<2x3xf32> to tensor<?x?xf32>
296 func.func @pad_to_static_size(%arg0: tensor<?x?xf32>, %ub0: index, %ub1: index,
297 %pad_value: f32) -> tensor<2x3xf32> {
298 %0 = tensor.pad %arg0 low[0, 0] high[%ub0, %ub1] {
300 tensor.yield %pad_value : f32
301 } : tensor<?x?xf32> to tensor<2x3xf32>
302 return %0 : tensor<2x3xf32>
308 // CHECK: tensor.pad %[[ARG0]]
311 // CHECK: : tensor<?x?xf32> to tensor<2x3xf32>
318 // CHECK: tensor.splat [[S]] : tensor<8xf32>
319 %v = tensor.splat %s : tensor<8xf32>
321 // CHECK: tensor.splat [[S]] : tensor<4xf32>
322 %u = "tensor.splat"(%s) : (f32) -> tensor<4xf32>
331 // CHECK: tensor.splat %[[S]][%[[M]], %[[N]]] : tensor<?x8x?xf32>
332 %v = tensor.splat %s[%m, %n] : tensor<?x8x?xf32>
339 // CHECK-SAME: %[[ARG0:.*]]: tensor<4x5x6xf32>,
340 // CHECK-SAME: %[[ARG1:.*]]: tensor<1x3x2xindex>,
341 // CHECK-SAME: %[[ARG2:.*]]: tensor<1x3x2xi32>) {
343 %dest : tensor<4x5x6xf32>, %indices: tensor<1x3x2xindex>, %indices_i32: tensor<1x3x2xi32>) {
344 …K: %[[GATHER:.*]] = tensor.gather %[[ARG0]][%[[ARG2]]] gather_dims([1, 2]) unique : (tensor<4x5x6x…
345 %gathered = tensor.gather %dest[%indices_i32] gather_dims([1, 2]) unique:
346 (tensor<4x5x6xf32>, tensor<1x3x2xi32>) -> tensor<1x3x4x1x1xf32>
347 … %[[GATHER0:.*]] = tensor.gather %[[ARG0]][%[[ARG1]]] gather_dims([1, 2]) unique : (tensor<4x5x6xf…
348 %rank_reduced_gathered = tensor.gather %dest[%indices] gather_dims([1, 2]) unique:
349 (tensor<4x5x6xf32>, tensor<1x3x2xindex>) -> tensor<1x3x4xf32>
351 …tensor.scatter %[[GATHER]] into %[[ARG0]][%[[ARG1]]] scatter_dims([1, 2]) unique : (tensor<1x3x4x1…
352 %scattered = tensor.scatter %gathered into %dest[%indices]
354 (tensor<1x3x4x1x1xf32>, tensor<4x5x6xf32>, tensor<1x3x2xindex>) -> tensor<4x5x6xf32>
355 … tensor.scatter %[[GATHER0]] into %[[ARG0]][%[[ARG2]]] scatter_dims([1, 2]) unique : (tensor<1x3x4…
356 %rank_reduced_scattered = tensor.scatter %rank_reduced_gathered into %dest[%indices_i32]
358 (tensor<1x3x4xf32>, tensor<4x5x6xf32>, tensor<1x3x2xi32>) -> tensor<4x5x6xf32>
364 func.func @pack_nc_to_ncnc(%source: tensor<128x256xf32>, %dest: tensor<4x16x32x16xf32>) -> tensor<1…
365 …%0 = tensor.pack %source inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %dest : tensor<128x25…
366 %1 = tensor.empty() : tensor<128x256xf32>
367 …%2 = tensor.unpack %0 inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %1 : tensor<4x16x32x16xf…
368 return %2 : tensor<128x256xf32>
372 // CHECK-SAME: %[[SOURCE:.*]]: tensor<128x256xf32>,
373 // CHECK-SAME: %[[DEST:.*]]: tensor<4x16x32x16xf32>)
374 … %[[PACKED:.*]] = tensor.pack %[[SOURCE]] inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %[[D…
375 // CHECK: %[[BUFF:.*]] = tensor.empty() : tensor<128x256xf32>
376 …ECK: %{{.*}} = tensor.unpack %[[PACKED]] inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %[[BU…
380 …c @pack_nc_to_ncnc_with_padding(%source: tensor<13x15xf32>, %dest: tensor<2x8x8x2xf32>, %padding: …
381 …%0 = tensor.pack %source padding_value(%padding : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 2…
382 %1 = tensor.empty() : tensor<13x15xf32>
383 …%2 = tensor.unpack %0 inner_dims_pos = [0, 1] inner_tiles = [8, 2] into %1 : tensor<2x8x8x2xf32> -…
384 return %2 : tensor<13x15xf32>
388 // CHECK-SAME: %[[SOURCE:.*]]: tensor<13x15xf32>,
389 // CHECK-SAME: %[[DEST:.*]]: tensor<2x8x8x2xf32>,
391 … = tensor.pack %[[SOURCE]] padding_value(%[[PADDING]] : f32) inner_dims_pos = [0, 1] inner_tiles =…
392 // CHECK: %[[BUFF:.*]] = tensor.empty() : tensor<13x15xf32>
393 … CHECK: %{{.*}} = tensor.unpack %[[PACKED]] inner_dims_pos = [0, 1] inner_tiles = [8, 2] into %[[B…
397 func.func @pack_ck_to_kcck(%source: tensor<128x256xf32>, %dest: tensor<16x4x32x16xf32>) -> tensor<1…
398 …%0 = tensor.pack %source outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 16] i…
399 %1 = tensor.empty() : tensor<128x256xf32>
400 …%2 = tensor.unpack %0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 16] into…
401 return %2 : tensor<128x256xf32>
405 // CHECK-SAME: %[[SOURCE:.*]]: tensor<128x256xf32>,
406 // CHECK-SAME: %[[DEST:.*]]: tensor<16x4x32x16xf32>)
407 …]] = tensor.pack %[[SOURCE]] outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 1…
408 // CHECK: %[[BUFF:.*]] = tensor.empty() : tensor<128x256xf32>
409 … = tensor.unpack %[[PACKED]] outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 1…
413 …ack_fully_dynamic(%source: tensor<?x?xf32>, %dest: tensor<?x?x?x?xf32>, %pad: f32, %tile_n : index…
414 …%0 = tensor.pack %source padding_value(%pad : f32) inner_dims_pos = [0, 1] inner_tiles = [%tile_n,…
415 return %0 : tensor<?x?x?x?xf32>
419 // CHECK-SAME: %[[SOURCE:.*]]: tensor<?x?xf32>,
420 // CHECK-SAME: %[[DEST:.*]]: tensor<?x?x?x?xf32>,
424 …tensor.pack %[[SOURCE]] padding_value(%[[PAD]] : f32) inner_dims_pos = [0, 1] inner_tiles = [%[[TI…
428 …nc @pad_and_pack_partially_dynamic(%source: tensor<?x?xf32>, %dest: tensor<?x?x8x2xf32>, %pad: f32…
429 …%0 = tensor.pack %source padding_value(%pad : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 2] in…
430 return %0 : tensor<?x?x8x2xf32>
434 // CHECK-SAME: %[[SOURCE:.*]]: tensor<?x?xf32>,
435 // CHECK-SAME: %[[DEST:.*]]: tensor<?x?x8x2xf32>,
437 …*}} = tensor.pack %[[SOURCE]] padding_value(%[[PAD]] : f32) inner_dims_pos = [0, 1] inner_tiles = …
441 … @unpack_fully_dynamic(%source: tensor<?x?x?x?xf32>, %dest: tensor<?x?xf32>, %tile_n : index, %til…
442 …%0 = tensor.unpack %source inner_dims_pos = [0, 1] inner_tiles = [%tile_n, %tile_m] into %dest : t…
443 return %0 : tensor<?x?xf32>
447 // CHECK-SAME: %[[SOURCE:.*]]: tensor<?x?x?x?xf32>,
448 // CHECK-SAME: %[[DEST:.*]]: tensor<?x?xf32>,
451 …{.*}} = tensor.unpack %[[SOURCE]] inner_dims_pos = [0, 1] inner_tiles = [%[[TILE_N]], %[[TILE_M]]]…
455 func.func @unpack_partially_dynamic(%source: tensor<?x?x8x2xf32>, %dest: tensor<?x?xf32>) -> tensor…
456 …%0 = tensor.unpack %source inner_dims_pos = [0, 1] inner_tiles = [8, 2] into %dest : tensor<?x?x8x…
457 return %0: tensor<?x?xf32>
461 // CHECK-SAME: %[[SOURCE:.*]]: tensor<?x?x8x2xf32>,
462 // CHECK-SAME: %[[DEST:.*]]: tensor<?x?xf32>)
463 … CHECK: %{{.*}} = tensor.unpack %[[SOURCE]] inner_dims_pos = [0, 1] inner_tiles = [8, 2] into %[[D…