@@ -580,133 +580,6 @@ module attributes {transform.with_named_sequence} {
580580 }
581581}
582582
583- // -----
584-
585- // CHECK-LABEL: func @test_masked_vectorize_pad
586- func.func @test_masked_vectorize_pad (
587- %0 : tensor <?x?xf32 >, %h0 : index , %h1 : index )
588- -> tensor <2 x4 xf32 >
589- {
590- // CHECK-DAG: %[[c42:.*]] = arith.constant 4.243000e+01 : f32
591- // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
592- // CHECK-DAG: %[[c0_0:.*]] = arith.constant 0 : index
593- // CHECK: %[[d0:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
594- // CHECK: %[[d1:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
595- // CHECK: %[[mask:.*]] = vector.create_mask %[[d0]], %[[d1]] : vector<2x4xi1>
596- // CHECK: %[[masked_read:.*]] = vector.mask %[[mask]] {
597- // CHECK-SAME: vector.transfer_read %{{.*}}[%[[c0_0]], %[[c0_0]]], %[[c42]]
598- // CHECK-SAME: {in_bounds = [true, true]} : tensor<?x?xf32>, vector<2x4xf32>
599- // CHECK-SAME: } : vector<2x4xi1> -> vector<2x4xf32>
600- // CHECK-DAG: %[[c0_1:.*]] = arith.constant 0 : index
601- // CHECK-DAG: %[[empty:.*]] = tensor.empty() : tensor<2x4xf32>
602- // CHECK: vector.transfer_write %[[masked_read]], %[[empty]][%[[c0_1]], %[[c0_1]]]
603- // CHECK-SAME: {in_bounds = [true, true]} : vector<2x4xf32>, tensor<2x4xf32>
604- %cst = arith.constant 42.43 : f32
605- %c0 = arith.constant 0 : index
606- %1 = tensor.pad %0 low [0 , %c0 ] high [%h0 , %h1 ] {
607- ^bb0 (%hh1: index , %hh2: index ):
608- tensor.yield %cst : f32
609- } : tensor <?x?xf32 > to tensor <2 x4 xf32 >
610- return %1: tensor <2 x4 xf32 >
611- }
612-
613- module attributes {transform.with_named_sequence } {
614- transform.named_sequence @__transform_main (%arg1: !transform.any_op {transform.readonly }) {
615- %0 = transform.structured.match ops {[" tensor.pad" ]} in %arg1
616- : (!transform.any_op ) -> !transform.any_op
617- transform.structured.vectorize %0 vector_sizes [2 , 4 ] : !transform.any_op
618- transform.yield
619- }
620- }
621-
622- // -----
623-
624- // CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>
625- // CHECK: func @test_masked_vectorize_dynamic_pad
626- func.func @test_masked_vectorize_dynamic_pad (
627- %0 : tensor <?x?xf32 >, %h0 : index , %h1 : index )
628- -> tensor <?x?xf32 >
629- {
630- // CHECK-DAG: %[[c42:.*]] = arith.constant 4.243000e+01 : f32
631- // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
632- // CHECK-DAG: %[[res_d0:.+]] = affine.apply #[[MAP]]()
633- // CHECK-DAG: %[[res_d1:.+]] = affine.apply #[[MAP]]()
634- // CHECK: %[[c0_2:.*]] = arith.constant 0 : index
635- // CHECK: %[[d0:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
636- // CHECK: %[[d1:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
637- // CHECK: %[[mask:.*]] = vector.create_mask %[[d0]], %[[d1]] : vector<2x4xi1>
638- // CHECK: %[[masked_read:.*]] = vector.mask %[[mask]] {
639- // CHECK-SAME: vector.transfer_read %{{.*}}[%[[c0_2]], %[[c0_2]]], %[[c42]]
640- // CHECK-SAME: {in_bounds = [true, true]} : tensor<?x?xf32>, vector<2x4xf32>
641- // CHECK-SAME: } : vector<2x4xi1> -> vector<2x4xf32>
642- // CHECK-DAG: %[[empty:.*]] = tensor.empty(%[[res_d0]], %[[res_d1]]) : tensor<?x?xf32>
643- // CHECK-DAG: %[[c0_3:.*]] = arith.constant 0 : index
644- // CHECK-DAG: %[[d2:.*]] = tensor.dim %[[empty]], {{.*}} : tensor<?x?xf32>
645- // CHECK-DAG: %[[d3:.*]] = tensor.dim %[[empty]], {{.*}} : tensor<?x?xf32>
646- // CHECK: %[[mask_2:.*]] = vector.create_mask %[[d2]], %[[d3]] : vector<2x4xi1>
647- // CHECK: %[[masked_write:.*]] = vector.mask %[[mask_2]] {
648- // CHECK-SAME: vector.transfer_write %[[masked_read]], %[[empty]][%[[c0_3]], %[[c0_3]]]
649- // CHECK-SAME: {in_bounds = [true, true]} : vector<2x4xf32>, tensor<?x?xf32>
650- // CHECK: return %[[masked_write]] : tensor<?x?xf32>
651- %cst = arith.constant 42.43 : f32
652- %c0 = arith.constant 0 : index
653- %1 = tensor.pad %0 low [0 , %c0 ] high [%h0 , %h1 ] {
654- ^bb0 (%hh1: index , %hh2: index ):
655- tensor.yield %cst : f32
656- } : tensor <?x?xf32 > to tensor <?x?xf32 >
657- return %1: tensor <?x?xf32 >
658- }
659-
660- module attributes {transform.with_named_sequence } {
661- transform.named_sequence @__transform_main (%arg1: !transform.any_op {transform.readonly }) {
662- %0 = transform.structured.match ops {[" tensor.pad" ]} in %arg1
663- : (!transform.any_op ) -> !transform.any_op
664- transform.structured.vectorize %0 vector_sizes [2 , 4 ] : !transform.any_op
665- transform.yield
666- }
667- }
668-
669- // -----
670- // This case is supported because low padding `%l0` is applied on
671- // a unit dimension which is supported, non unit result dimension low
672- // padding is currently unsupported.
673- // CHECK-LABEL: func @test_masked_vectorize_non_zero_low_pad_unit_res_dim
674- func.func @test_masked_vectorize_non_zero_low_pad_unit_res_dim (
675- %0 : tensor <?x?xf32 >, %h0 : index , %h1 : index , %l0 : index )
676- -> tensor <1 x4 xf32 >
677- {
678- // CHECK-DAG: %[[C42:.*]] = arith.constant 4.243000e+01 : f32
679- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
680- // CHECK: %[[C0_1:.*]] = arith.constant 0 : index
681- // CHECK-DAG: %[[D0:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
682- // CHECK-DAG: %[[D1:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
683- // CHECK: %[[MASK:.*]] = vector.create_mask %[[D0]], %[[D1]] : vector<1x4xi1>
684- // CHECK: %[[MASKED_READ:.*]] = vector.mask %[[MASK]] {
685- // CHECK-SAME: vector.transfer_read %{{.*}}[%[[C0_1]], %[[C0_1]]], %[[C42]]
686- // CHECK-SAME: {in_bounds = [true, true]} : tensor<?x?xf32>, vector<1x4xf32>
687- // CHECK-SAME: } : vector<1x4xi1> -> vector<1x4xf32>
688- // CHECK-DAG: %[[EMPTY:.*]] = tensor.empty() : tensor<1x4xf32>
689- // CHECK-DAG: %[[C0_2:.*]] = arith.constant 0 : index
690- // CHECK: %[[MASKED_WRITE:.*]] = vector.transfer_write %[[MASKED_READ]], %[[EMPTY]][%[[C0_2]], %[[C0_2]]]
691- // CHECK-SAME: {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>
692- // CHECK: return %[[MASKED_WRITE]] : tensor<1x4xf32>
693- %cst = arith.constant 42.43 : f32
694- %c0 = arith.constant 0 : index
695- %1 = tensor.pad %0 low [%l0 , %c0 ] high [%h0 , %h1 ] {
696- ^bb0 (%hh1: index , %hh2: index ):
697- tensor.yield %cst : f32
698- } : tensor <?x?xf32 > to tensor <1 x4 xf32 >
699- return %1: tensor <1 x4 xf32 >
700- }
701-
702- module attributes {transform.with_named_sequence } {
703- transform.named_sequence @__transform_main (%arg1: !transform.any_op {transform.readonly }) {
704- %0 = transform.structured.match ops {[" tensor.pad" ]} in %arg1
705- : (!transform.any_op ) -> !transform.any_op
706- transform.structured.vectorize %0 vector_sizes [1 , 4 ] : !transform.any_op
707- transform.yield
708- }
709- }
710583
711584// -----
712585
@@ -1155,153 +1028,3 @@ func.func @test_vectorize_unpack_no_vector_sizes_permute(%source: tensor<4x7x4xf
11551028 }
11561029 }
11571030
1158- // -----
1159-
1160- ///----------------------------------------------------------------------------------------
1161- /// tensor.insert_slice
1162- ///----------------------------------------------------------------------------------------
1163-
1164- func.func private @insert_slice_static_sizes (%source: tensor <?x3 x?x1 xi32 >) -> tensor <5 x3 xi32 > {
1165- %c2 = arith.constant 2 : index
1166- %init = tensor.empty () : tensor <5 x3 xi32 >
1167-
1168- %source_slice = tensor.extract_slice %source [0 , %c2 , 0 , 0 ] [1 , 1 , 5 , 1 ] [1 , 1 , 1 , 1 ] : tensor <?x3 x?x1 xi32 > to tensor <5 x1 xi32 >
1169- %res = tensor.insert_slice %source_slice into %init [0 , %c2 ] [5 , 1 ] [1 , 1 ] : tensor <5 x1 xi32 > into tensor <5 x3 xi32 >
1170-
1171- return %res : tensor <5 x3 xi32 >
1172- }
1173-
1174- // CHECK-LABEL: func.func private @insert_slice_static_sizes(
1175- // CHECK-SAME: %[[SEC:.*]]: tensor<?x3x?x1xi32>) -> tensor<5x3xi32> {
1176- // CHECK: %[[C_2:.*]] = arith.constant 2 : index
1177- // CHECK: %[[INIT:.*]] = tensor.empty() : tensor<5x3xi32>
1178- // CHECK: %[[SRC_SLICE:.*]] = tensor.extract_slice %[[SEC]][0, %[[C_2]], 0, 0] [1, 1, 5, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<5x1xi32>
1179- // CHECK-DAG: %[[PAD:.*]] = arith.constant 0 : i32
1180- // CHECK-DAG: %[[C_5:.*]] = arith.constant 5 : index
1181- // CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : index
1182- // CHECK: %[[MASK:.*]] = vector.create_mask %[[C_5]], %[[C_1]] : vector<8x1xi1>
1183- // CHECK: %[[C0:.*]] = arith.constant 0 : index
1184- // CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C0]], %[[C0]]], %[[PAD]] : tensor<5x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
1185- // CHECK: %[[C_0:.*]] = arith.constant 0 : index
1186- // CHECK: %[[RES:.*]] = vector.mask %[[MASK]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0]], %[[C_2]]] : vector<8x1xi32>, tensor<5x3xi32> } : vector<8x1xi1> -> tensor<5x3xi32>
1187- // CHECK: return %[[RES]] : tensor<5x3xi32>
1188-
1189- module attributes {transform.with_named_sequence } {
1190- transform.named_sequence @__transform_main (%arg0: !transform.any_op {transform.readonly }) {
1191- %0 = transform.structured.match ops {[" tensor.insert_slice" ]} in %arg0 : (!transform.any_op ) -> !transform.any_op
1192- transform.structured.vectorize %0 vector_sizes [8 , 1 ] : !transform.any_op
1193- transform.yield
1194- }
1195- }
1196-
1197- // -----
1198-
1199- // One of the _source_ dimensions is dynamic (but _destination_ dimensions are static).
1200-
1201- func.func private @insert_slice_dynamic_src_dim (%source: tensor <?x3 x?x1 xi32 >, %size: index ) -> tensor <5 x3 xi32 > {
1202- %c2 = arith.constant 2 : index
1203- %init = tensor.empty () : tensor <5 x3 xi32 >
1204-
1205- %source_slice = tensor.extract_slice %source [0 , %c2 , 0 , 0 ] [1 , 1 , %size , 1 ] [1 , 1 , 1 , 1 ] : tensor <?x3 x?x1 xi32 > to tensor <?x1 xi32 >
1206- %res = tensor.insert_slice %source_slice into %init [0 , %c2 ] [%size , 1 ] [1 , 1 ] : tensor <?x1 xi32 > into tensor <5 x3 xi32 >
1207-
1208- return %res : tensor <5 x3 xi32 >
1209- }
1210-
1211- // CHECK-LABEL: func.func private @insert_slice_dynamic_src_dim(
1212- // CHECK-SAME: %[[SRC:.*]]: tensor<?x3x?x1xi32>,
1213- // CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<5x3xi32> {
1214- // CHECK: %[[C_2:.*]] = arith.constant 2 : index
1215- // CHECK: %[[INIT:.*]] = tensor.empty() : tensor<5x3xi32>
1216- // CHECK: %[[SRC_SLICE:.*]] = tensor.extract_slice %[[SRC]][0, %[[C_2]], 0, 0] [1, 1, %[[SIZE]], 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<?x1xi32>
1217- // CHECK-DAG: %[[PAD:.*]] = arith.constant 0 : i32
1218- // CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : index
1219- // CHECK: %[[MASK:.*]] = vector.create_mask %[[SIZE]], %[[C_1]] : vector<8x1xi1>
1220- // CHECK: %[[C_0:.*]] = arith.constant 0 : index
1221- // CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C_0]], %[[C_0]]], %[[PAD]] : tensor<?x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
1222- // CHECK: %[[C_0_1:.*]] = arith.constant 0 : index
1223- // CHECK: %[[RES:.*]] = vector.mask %[[MASK]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_1]], %[[C_2]]] : vector<8x1xi32>, tensor<5x3xi32> } : vector<8x1xi1> -> tensor<5x3xi32>
1224- // CHECK: return %[[RES]] : tensor<5x3xi32>
1225-
1226- module attributes {transform.with_named_sequence } {
1227- transform.named_sequence @__transform_main (%arg0: !transform.any_op {transform.readonly }) {
1228- %0 = transform.structured.match ops {[" tensor.insert_slice" ]} in %arg0 : (!transform.any_op ) -> !transform.any_op
1229- transform.structured.vectorize %0 vector_sizes [8 , 1 ] : !transform.any_op
1230- transform.yield
1231- }
1232- }
1233-
1234- // -----
1235-
1236- // One of the _destination_ dimensions is dynamic (but _source_ dimensions are static).
1237-
1238- func.func private @insert_slice_dynamic_dest_dim (%source: tensor <?x3 x?x1 xi32 >, %size: index ) -> tensor <?x3 xi32 > {
1239- %c2 = arith.constant 2 : index
1240- %init = tensor.empty (%size ) : tensor <?x3 xi32 >
1241-
1242- %source_slice = tensor.extract_slice %source [0 , %c2 , 0 , 0 ] [1 , 1 , 5 , 1 ] [1 , 1 , 1 , 1 ] : tensor <?x3 x?x1 xi32 > to tensor <5 x1 xi32 >
1243- %res = tensor.insert_slice %source_slice into %init [0 , %c2 ] [5 , 1 ] [1 , 1 ] : tensor <5 x1 xi32 > into tensor <?x3 xi32 >
1244-
1245- return %res : tensor <?x3 xi32 >
1246- }
1247-
1248- // CHECK-LABEL: func.func private @insert_slice_dynamic_dest_dim(
1249- // CHECK-SAME: %[[SRC:.*]]: tensor<?x3x?x1xi32>,
1250- // CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<?x3xi32> {
1251- // CHECK: %[[C_2:.*]] = arith.constant 2 : index
1252- // CHECK: %[[INIT:.*]] = tensor.empty(%[[SIZE]]) : tensor<?x3xi32>
1253- // CHECK: %[[SRC_SLICE:.*]] = tensor.extract_slice %[[SRC]][0, %[[C_2]], 0, 0] [1, 1, 5, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<5x1xi32>
1254- // CHECK: %[[PAD:.*]] = arith.constant 0 : i32
1255- // CHECK: %[[C_5:.*]] = arith.constant 5 : index
1256- // CHECK: %[[C_1:.*]] = arith.constant 1 : index
1257- // CHECK: %[[MASK:.*]] = vector.create_mask %[[C_5]], %[[C_1]] : vector<8x1xi1>
1258- // CHECK: %[[C_0:.*]] = arith.constant 0 : index
1259- // CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C_0]], %[[C_0]]], %[[PAD]] : tensor<5x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
1260- // CHECK: %[[C_0_1:.*]] = arith.constant 0 : index
1261- // CHECK: %[[WRITE:.*]] = vector.mask %[[MASK]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_1]], %[[C_2]]] : vector<8x1xi32>, tensor<?x3xi32> } : vector<8x1xi1> -> tensor<?x3xi32>
1262- // CHECK: return %[[WRITE]] : tensor<?x3xi32>
1263-
1264- module attributes {transform.with_named_sequence } {
1265- transform.named_sequence @__transform_main (%arg0: !transform.any_op {transform.readonly }) {
1266- %0 = transform.structured.match ops {[" tensor.insert_slice" ]} in %arg0 : (!transform.any_op ) -> !transform.any_op
1267- transform.structured.vectorize %0 vector_sizes [8 , 1 ] : !transform.any_op
1268- transform.yield
1269- }
1270- }
1271-
1272- // -----
1273-
1274- // At least one _source_ and one _destination_ dimensions are dynamic.
1275-
1276- func.func private @insert_slice_dynamic_source_and_dest_dim (%source: tensor <?x3 x?x1 xi32 >, %size: index ) -> tensor <?x3 xi32 > {
1277- %c2 = arith.constant 2 : index
1278- %init = tensor.empty (%size ) : tensor <?x3 xi32 >
1279-
1280- %source_slice = tensor.extract_slice %source [0 , %c2 , 0 , 0 ] [1 , 1 , %size , 1 ] [1 , 1 , 1 , 1 ] : tensor <?x3 x?x1 xi32 > to tensor <?x1 xi32 >
1281- %res = tensor.insert_slice %source_slice into %init [0 , %c2 ] [%size , 1 ] [1 , 1 ] : tensor <?x1 xi32 > into tensor <?x3 xi32 >
1282-
1283- return %res : tensor <?x3 xi32 >
1284- }
1285-
1286- // CHECK-LABEL: func.func private @insert_slice_dynamic_source_and_dest_dim(
1287- // CHECK-SAME: %[[SRC:.*]]: tensor<?x3x?x1xi32>,
1288- // CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<?x3xi32> {
1289- // CHECK: %[[C_2:.*]] = arith.constant 2 : index
1290- // CHECK: %[[INIT:.*]] = tensor.empty(%[[SIZE]]) : tensor<?x3xi32>
1291- // CHECK: %[[SRC_SIZE:.*]] = tensor.extract_slice %[[SRC]][0, %[[C_2]], 0, 0] [1, 1, %[[SIZE]], 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<?x1xi32>
1292- // CHECK: %[[PAD:.*]] = arith.constant 0 : i32
1293- // CHECK: %[[C1:.*]] = arith.constant 1 : index
1294- // CHECK: %[[MASK:.*]] = vector.create_mask %[[SIZE]], %[[C1]] : vector<8x1xi1>
1295- // CHECK: %[[C0:.*]] = arith.constant 0 : index
1296- // CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SIZE]]{{\[}}%[[C0]], %[[C0]]], %[[PAD]] : tensor<?x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
1297- // CHECK: %[[C_0_1:.*]] = arith.constant 0 : index
1298- // CHECK: %[[WRITE:.*]] = vector.mask %[[MASK]] { vector.transfer_write %[[READ]], %[[INIT]]{{\[}}%[[C_0_1]], %[[C_2]]] : vector<8x1xi32>, tensor<?x3xi32> } : vector<8x1xi1> -> tensor<?x3xi32>
1299- // CHECK: return %[[WRITE]] : tensor<?x3xi32>
1300-
1301- module attributes {transform.with_named_sequence } {
1302- transform.named_sequence @__transform_main (%arg0: !transform.any_op {transform.readonly }) {
1303- %0 = transform.structured.match ops {[" tensor.insert_slice" ]} in %arg0 : (!transform.any_op ) -> !transform.any_op
1304- transform.structured.vectorize %0 vector_sizes [8 , 1 ] : !transform.any_op
1305- transform.yield
1306- }
1307- }
0 commit comments