1617 lines · plain
1// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -canonicalize="test-convergence" %s | FileCheck %s2// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -canonicalize="test-convergence top-down=0" %s | FileCheck %s3 4// CHECK-LABEL: func @f5func.func @f(%arg0: tensor<2x3x4xf32>) -> tensor<3xindex> {6 // CHECK: shape.const_shape [2, 3, 4] : tensor<3xindex>7 %0 = shape.shape_of %arg0 : tensor<2x3x4xf32> -> tensor<3xindex>8 return %0 : tensor<3xindex>9}10 11// -----12 13// Basic case.14// CHECK-LABEL: func @f15func.func @f() -> (!shape.shape, !shape.shape) {16 // CHECK-DAG: shape.const_shape [2, 3] : !shape.shape17 // CHECK-DAG: shape.const_shape [4, 5] : !shape.shape18 %c2 = arith.constant 2 : index19 %0 = shape.const_shape [2, 3, 4, 5] : !shape.shape20 %head, %tail = "shape.split_at"(%0, %c2) : (!shape.shape, index) -> (!shape.shape, !shape.shape)21 return %head, %tail : !shape.shape, !shape.shape22 23}24 25// -----26 27// Negative split point.28// CHECK-LABEL: func @f29func.func @f() -> (!shape.shape, !shape.shape) {30 // CHECK-DAG: shape.const_shape [2, 3, 4] : !shape.shape31 // CHECK-DAG: shape.const_shape [5] : !shape.shape32 %c-1 = arith.constant -1 : index33 %0 = shape.const_shape [2, 3, 4, 5] : !shape.shape34 %head, %tail = "shape.split_at"(%0, %c-1) : (!shape.shape, index) -> (!shape.shape, !shape.shape)35 return %head, %tail : !shape.shape, !shape.shape36}37 38// -----39 40// Out of range split point. No folding.41// CHECK-LABEL: func @f42func.func @f() -> (!shape.shape, !shape.shape) {43 // CHECK: shape.split_at44 %c5 = arith.constant 5 : index45 %0 = shape.const_shape [2, 3, 4, 5] : !shape.shape46 %head, %tail = "shape.split_at"(%0, %c5) : (!shape.shape, index) -> (!shape.shape, !shape.shape)47 return %head, %tail : !shape.shape, !shape.shape48}49 50// -----51 52// Basic case.53// CHECK-LABEL: func @f54func.func @f() -> !shape.shape {55 // CHECK: shape.const_shape [7, 2] : !shape.shape56 %0 = shape.const_shape [1, 2] : !shape.shape57 %1 = shape.const_shape [7, 1] : !shape.shape58 %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape59 return %2 : !shape.shape60}61 62// -----63 64// Basic case including extent tensors.65// CHECK-LABEL: @broadcast66func.func @broadcast() -> tensor<2xindex> {67 // CHECK: shape.const_shape [7, 2] : tensor<2xindex>68 %0 = shape.const_shape [1, 2] : tensor<2xindex>69 %1 = shape.const_shape [7, 1] : tensor<2xindex>70 %2 = shape.broadcast %0, %171 : tensor<2xindex>, tensor<2xindex> -> tensor<2xindex>72 return %2 : tensor<2xindex>73}74 75// -----76 77// Basic case including extent tensors.78// CHECK-LABEL: @broadcast79func.func @broadcast() -> !shape.shape {80 // CHECK: shape.const_shape [7, 2] : !shape.shape81 %0 = shape.const_shape [1, 2] : tensor<2xindex>82 %1 = shape.const_shape [7, 1] : tensor<2xindex>83 %2 = shape.broadcast %0, %1 : tensor<2xindex>, tensor<2xindex> -> !shape.shape84 return %2 : !shape.shape85}86 87// -----88 89// Variadic case including extent tensors.90// CHECK-LABEL: @broadcast_variadic91func.func @broadcast_variadic() -> !shape.shape {92 // CHECK: shape.const_shape [7, 2, 10] : !shape.shape93 %0 = shape.const_shape [2, 1] : tensor<2xindex>94 %1 = shape.const_shape [7, 2, 1] : tensor<3xindex>95 %2 = shape.const_shape [1, 10] : tensor<2xindex>96 %3 = shape.broadcast %0, %1, %2 : tensor<2xindex>, tensor<3xindex>, tensor<2xindex> -> !shape.shape97 return %3 : !shape.shape98}99 100// -----101 102// Rhs is a scalar.103// CHECK-LABEL: func @f104func.func @f(%arg0 : !shape.shape) -> !shape.shape {105 // CHECK: return %arg0106 %0 = shape.const_shape [] : !shape.shape107 %1 = shape.broadcast %arg0, %0 : !shape.shape, !shape.shape -> !shape.shape108 return %1 : !shape.shape109}110 111// -----112 113// Lhs is a scalar.114// CHECK-LABEL: func @f115func.func @f(%arg0 : !shape.shape) -> !shape.shape {116 // CHECK: return %arg0117 %0 = shape.const_shape [] : !shape.shape118 %1 = shape.broadcast %0, %arg0 : !shape.shape, !shape.shape -> !shape.shape119 return %1 : !shape.shape120}121 122// -----123 124// Lhs is a scalar and rhs is constant.125// CHECK-LABEL: func @f126func.func @f() -> !shape.shape {127 // CHECK: %[[CST:.*]] = shape.const_shape [1, 2, 3] : !shape.shape128 // CHECK: return %[[CST]]129 %0 = shape.const_shape [] : !shape.shape130 %1 = shape.const_shape [1, 2, 3] : !shape.shape131 %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape132 return %2 : !shape.shape133}134 135// -----136 137// All but one operands are known empty shapes.138// CHECK-LABEL: @all_but_one_empty139// CHECK-SAME: (%[[ARG:.*]]: !shape.shape)140func.func @all_but_one_empty(%arg0 : !shape.shape) -> !shape.shape {141 // CHECK: return %[[ARG]]142 %0 = shape.const_shape [] : !shape.shape143 %1 = shape.const_shape [] : tensor<0xindex>144 %2 = shape.broadcast %0, %arg0, %1, %0 : !shape.shape, !shape.shape,145 tensor<0xindex>, !shape.shape -> !shape.shape146 return %2 : !shape.shape147}148 149// -----150 151// All operands are known empty shapes.152// CHECK-LABEL: @all_empty153// CHECK-SAME: (%[[ARG_0:.*]]: tensor<f32>, %[[ARG_1:.*]]: tensor<i1>)154func.func @all_empty(%arg0: tensor<f32>, %arg1: tensor<i1>) -> tensor<0xindex> {155 // CHECK: %[[CST:.*]] = shape.const_shape [] : tensor<0xindex>156 // CHECK: return %[[CST]] : tensor<0xindex>157 %1 = shape.shape_of %arg0 : tensor<f32> -> tensor<0xindex>158 %2 = shape.shape_of %arg1 : tensor<i1> -> tensor<0xindex>159 %3 = shape.const_shape [] : tensor<0xindex>160 %4 = shape.broadcast %1, %2, %3 : tensor<0xindex>, tensor<0xindex>, tensor<0xindex> -> tensor<0xindex>161 return %4 : tensor<0xindex>162}163 164// -----165 166// Partial folding.167// CHECK-LABEL: @partial_folding168// CHECK-SAME: (%[[ARG:.*]]: !shape.shape)169func.func @partial_folding(%arg0 : !shape.shape) -> !shape.shape {170 // CHECK: %[[CST_SHAPE:.*]] = shape.const_shape [1, 2, 3] : tensor<3xindex>171 // CHECK: %[[RESULT:.*]] = shape.broadcast %[[ARG]], %[[CST_SHAPE]] : !shape.shape, tensor<3xindex> -> !shape.shape172 // CHECK: return %[[RESULT]]173 %0 = shape.const_shape [2, 1] : !shape.shape174 %1 = shape.const_shape [1, 2, 3] : tensor<3xindex>175 %2 = shape.broadcast %0, %arg0, %1, %0 : !shape.shape, !shape.shape,176 tensor<3xindex>, !shape.shape -> !shape.shape177 return %2 : !shape.shape178}179 180// -----181 182// Incompatible shapes. No folding.183// CHECK-LABEL: func @f184func.func @f() -> !shape.shape {185 // CHECK: shape.broadcast186 %0 = shape.const_shape [2] : !shape.shape187 %1 = shape.const_shape [7] : !shape.shape188 %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape189 return %2 : !shape.shape190}191 192// -----193 194// Dead code195// CHECK-LABEL: @broadcast196func.func @broadcast(%arg0 : !shape.shape, %arg1 : !shape.shape) {197 // CHECK-NEXT: return198 %0 = shape.broadcast %arg0, %arg1199 : !shape.shape, !shape.shape -> !shape.shape200 return201}202 203// -----204 205// Basic case.206// CHECK-LABEL: func @f207func.func @f() -> !shape.shape {208 // CHECK: shape.const_shape [0, 1, 2, 3] : !shape.shape209 %lhs = shape.const_shape [0, 1] : !shape.shape210 %rhs = shape.const_shape [2, 3] : !shape.shape211 %0 = shape.concat %lhs, %rhs : !shape.shape , !shape.shape -> !shape.shape212 return %0 : !shape.shape213}214 215// -----216 217// Basic case.218// CHECK-LABEL: func @f219func.func @f() -> tensor<4xindex> {220 // CHECK: shape.const_shape [0, 1, 2, 3] : tensor<4xindex>221 %lhs = shape.const_shape [0, 1] : tensor<2xindex>222 %rhs = shape.const_shape [2, 3] : tensor<2xindex>223 %0 = shape.concat %lhs, %rhs : tensor<2xindex>, tensor<2xindex> -> tensor<4xindex>224 return %0 : tensor<4xindex>225}226 227// -----228 229// Basic case.230// CHECK-LABEL: func @f231func.func @f() -> tensor<2xindex> {232 // CHECK: shape.const_shape [0, 1] : tensor<2xindex>233 %cs = shape.const_shape [0, 1] : !shape.shape234 %0 = shape.to_extent_tensor %cs : !shape.shape -> tensor<2xindex>235 return %0 : tensor<2xindex>236}237 238// -----239 240// Basic case.241// CHECK-LABEL: func @f()242func.func @f() -> !shape.shape {243 // CHECK: shape.const_shape [3, 5, 11] : !shape.shape244 %e0 = arith.constant 3 : index245 %e1 = arith.constant 5 : index246 %e2 = arith.constant 11 : index247 %ret = shape.from_extents %e0, %e1, %e2 : index, index, index248 return %ret : !shape.shape249}250 251// -----252 253// fold_const_size254// CHECK-LABEL: func @fold_const_size()255func.func @fold_const_size() -> !shape.shape {256 // CHECK: shape.const_shape [3, 5] : !shape.shape257 %e0 = shape.const_size 3258 %e1 = shape.const_size 5259 %ret = shape.from_extents %e0, %e1 : !shape.size, !shape.size260 return %ret : !shape.shape261}262 263// -----264 265// CHECK-LABEL: func @no_fold266func.func @no_fold(%arg0: index) -> !shape.shape {267 // CHECK-NOT: shape.const_shape268 %e0 = arith.constant 3 : index269 %ret = shape.from_extents %e0, %arg0 : index, index270 return %ret : !shape.shape271}272 273// -----274 275// Cast constant size to index and fold it away.276// CHECK-LABEL: func @const_size_to_index277func.func @const_size_to_index() -> index {278 // CHECK-NOT: shape.index_cast279 %cs = shape.const_size 123280 // CHECK: arith.constant 123 : index281 %ci = shape.size_to_index %cs : !shape.size282 return %ci : index283}284 285// -----286 287// Cast constant index to size and fold it away.288// CHECK-LABEL: func @const_index_to_size289func.func @const_index_to_size() -> !shape.size {290 // CHECK-NOT: arith.index_cast291 %ci = arith.constant 123 : index292 // CHECK: shape.const_size 123293 %cs = shape.index_to_size %ci294 return %cs : !shape.size295}296 297// -----298 299// Cast constant index to size, then back, and fold it away.300// CHECK-LABEL: func @const_index_to_size_to_index301func.func @const_index_to_size_to_index() -> index {302 // CHECK-NOT: shape.index_cast303 %ci0 = arith.constant 123 : index304 %cs0 = shape.index_to_size %ci0305 // CHECK: %[[CI:.*]] = arith.constant 123 : index306 // CHECK-NEXT: return %[[CI]] : index307 %ci1 = shape.size_to_index %cs0 : !shape.size308 return %ci1 : index309}310 311// -----312 313// No folding.314// CHECK-LABEL: func @nonfoldable_size_to_index315func.func @nonfoldable_size_to_index(%cs : !shape.size) -> index {316 // CHECK: shape.size_to_index317 %ci = shape.size_to_index %cs : !shape.size318 return %ci : index319}320 321// -----322 323// No folding.324// CHECK-LABEL: func @nonfoldable_index_to_size325func.func @nonfoldable_index_to_size(%ci : index) -> !shape.size {326 // CHECK: shape.index_to_size327 %cs = shape.index_to_size %ci328 return %cs : !shape.size329}330 331// -----332 333// Fold number of elements computation.334// CHECK-LABEL: func @num_elements335func.func @num_elements() -> !shape.size {336 // CHECK-NOT: shape.const_shape337 %shape = shape.const_shape [4, 5, 6] : !shape.shape338 // CHECK-NOT: shape.num_elements339 %num_elements = shape.num_elements %shape : !shape.shape -> !shape.size340 // CHECK: %[[NUM:.*]] = shape.const_size 120341 // CHECK-NEXT: return %[[NUM]] : !shape.size342 return %num_elements : !shape.size343}344 345// -----346 347// No folding.348// CHECK-LABEL: func @nonfoldable_num_elements349func.func @nonfoldable_num_elements(%shape : !shape.shape) -> !shape.size {350 // CHECK-NOT: shape.const_{{.*}}351 %num_elements = shape.num_elements %shape : !shape.shape -> !shape.size352 return %num_elements : !shape.size353}354 355// -----356 357// Basic folding.358// CHECK-LABEL: func @basic359func.func @basic() -> index {360 // CHECK: constant 2 : index361 %0 = shape.const_shape [0, 1, 2] : tensor<3xindex>362 %c2 = arith.constant 2 : index363 %1 = shape.get_extent %0, %c2 : tensor<3xindex>, index -> index364 return %1 : index365}366 367// -----368 369// Should not fold.370// CHECK-LABEL: func @out_of_bounds371func.func @out_of_bounds() -> index {372 // CHECK: shape.const_shape373 // CHECK: shape.get_extent374 %0 = shape.const_shape [0, 1, 2] : tensor<3xindex>375 %c3 = arith.constant 3 : index376 %1 = shape.get_extent %0, %c3 : tensor<3xindex>, index -> index377 return %1 : index378}379 380// -----381 382// Should not fold.383// CHECK-LABEL: func @not_const384func.func @not_const(%arg0: tensor<?xindex>) -> index {385 // CHECK: shape.get_extent386 %c3 = arith.constant 3 : index387 %0 = shape.get_extent %arg0, %c3 : tensor<?xindex>, index -> index388 return %0 : index389}390 391// -----392 393// Basic folding.394// CHECK-LABEL: func @basic395func.func @basic() -> !shape.size {396 // CHECK: shape.const_size 2397 %0 = shape.const_shape [0, 1, 2] : !shape.shape398 %c2 = shape.const_size 2399 %1 = shape.get_extent %0, %c2 : !shape.shape, !shape.size -> !shape.size400 return %1 : !shape.size401}402 403// -----404 405// Should not fold.406// CHECK-LABEL: func @out_of_bounds407func.func @out_of_bounds() -> !shape.size {408 // CHECK: shape.const_shape409 // CHECK: shape.get_extent410 %0 = shape.const_shape [0, 1, 2] : !shape.shape411 %c3 = shape.const_size 3412 %1 = shape.get_extent %0, %c3 : !shape.shape, !shape.size -> !shape.size413 return %1 : !shape.size414}415 416// -----417 418// Should not fold.419// CHECK-LABEL: func @not_const420func.func @not_const(%arg0 : !shape.shape) -> !shape.size {421 // CHECK: shape.get_extent422 %c3 = shape.const_size 3423 %0 = shape.get_extent %arg0, %c3 : !shape.shape, !shape.size -> !shape.size424 return %0 : !shape.size425}426 427// -----428// cstr_eq with non-constant but known equal shapes can be removed.429// CHECK-LABEL: func @f430func.func @f(%arg0 : !shape.shape) {431 // CHECK-NEXT: shape.const_witness true432 // CHECK-NEXT: consume.witness433 // CHECK-NEXT: return434 %0 = shape.cstr_eq %arg0, %arg0, %arg0 : !shape.shape, !shape.shape, !shape.shape435 "consume.witness"(%0) : (!shape.witness) -> ()436 return437}438 439// -----440// cstr_eq with equal const_shapes can be folded441// CHECK-LABEL: func @f442func.func @f() {443 // CHECK-NEXT: shape.const_witness true444 // CHECK-NEXT: consume.witness445 // CHECK-NEXT: return446 %cs0 = shape.const_shape [0, 1] : !shape.shape447 %cs1 = shape.const_shape [0, 1] : !shape.shape448 %cs2 = shape.const_shape [0, 1] : !shape.shape449 %0 = shape.cstr_eq %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape450 "consume.witness"(%0) : (!shape.witness) -> ()451 return452}453 454// -----455// cstr_eq with unequal const_shapes cannot be folded456// CHECK-LABEL: func @f457func.func @f() {458 // CHECK-NEXT: shape.const_shape459 // CHECK-NEXT: shape.const_shape460 // CHECK-NEXT: shape.cstr_eq461 // CHECK-NEXT: consume.witness462 // CHECK-NEXT: return463 %cs0 = shape.const_shape [0, 1] : !shape.shape464 %cs1 = shape.const_shape [3, 1] : !shape.shape465 %0 = shape.cstr_eq %cs0, %cs1 : !shape.shape, !shape.shape466 "consume.witness"(%0) : (!shape.witness) -> ()467 return468}469 470// -----471// cstr_eq without const_shapes cannot be folded472// CHECK-LABEL: func @f473func.func @f(%arg0: !shape.shape, %arg1: !shape.shape) {474 // CHECK-NEXT: shape.cstr_eq475 // CHECK-NEXT: consume.witness476 // CHECK-NEXT: return477 %0 = shape.cstr_eq %arg0, %arg1 : !shape.shape, !shape.shape478 "consume.witness"(%0) : (!shape.witness) -> ()479 return480}481 482// -----483// cstr_require with constant can be folded484// CHECK-LABEL: func @cstr_require_fold485func.func @cstr_require_fold() {486 // CHECK-NEXT: shape.const_witness true487 // CHECK-NEXT: consume.witness488 // CHECK-NEXT: return489 %true = arith.constant true490 %0 = shape.cstr_require %true, "msg"491 "consume.witness"(%0) : (!shape.witness) -> ()492 return493}494 495// -----496// cstr_require without constant cannot be folded497// CHECK-LABEL: func @cstr_require_no_fold498func.func @cstr_require_no_fold(%arg0: i1) {499 // CHECK-NEXT: shape.cstr_require500 // CHECK-NEXT: consume.witness501 // CHECK-NEXT: return502 %0 = shape.cstr_require %arg0, "msg"503 "consume.witness"(%0) : (!shape.witness) -> ()504 return505}506 507// -----508 509// merge assuming_all operations510// CHECK-LABEL: func @f511func.func @f() {512 // CHECK-NEXT: %[[W0:.*]] = "test.source"513 // CHECK-NEXT: %[[W1:.*]] = "test.source"514 // CHECK-NEXT: %[[W2:.*]] = "test.source"515 // CHECK-NEXT: shape.assuming_all %[[W0]], %[[W1]], %[[W2]]516 // CHECK-NEXT: consume.witness517 // CHECK-NEXT: return518 %0 = "test.source"() : () -> !shape.witness519 %1 = "test.source"() : () -> !shape.witness520 %2 = "test.source"() : () -> !shape.witness521 %3 = shape.assuming_all %0, %1522 %4 = shape.assuming_all %3, %2523 "consume.witness"(%4) : (!shape.witness) -> ()524 return525}526 527// -----528// `assuming_all` with all `cstr_eq` and shared operands can be collapsed.529// CHECK-LABEL: func @assuming_all_to_cstr_eq530// CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: tensor<?xindex>, %[[C:.*]]: tensor<3xindex>)531func.func @assuming_all_to_cstr_eq(%a : !shape.shape, %b : tensor<?xindex>,532 %c : tensor<3xindex>) -> !shape.witness {533 // CHECK: %[[RESULT:.*]] = shape.cstr_eq %[[A]], %[[B]], %[[B]], %[[C]]534 // CHECK: return %[[RESULT]]535 %0 = shape.cstr_eq %a, %b : !shape.shape, tensor<?xindex>536 %1 = shape.cstr_eq %b, %c : tensor<?xindex>, tensor<3xindex>537 %2 = shape.assuming_all %0, %1538 return %2 : !shape.witness539}540 541// -----542// `assuming_all` with duplicate operands.543// CHECK-LABEL: func @assuming_all_duplicate_operands544// CHECK-SAME: (%[[ARG0:.*]]: tensor<?xindex>, %[[ARG1:.*]]: tensor<?xindex>)545func.func @assuming_all_duplicate_operands(%arg0 : tensor<?xindex>,546 %arg1 : tensor<?xindex>) -> !shape.witness {547 // CHECK: %[[RES:.*]] = shape.cstr_broadcastable %[[ARG0]], %[[ARG1]]548 // CHECK: return %[[RES]]549 %0 = shape.cstr_broadcastable %arg0, %arg1 : tensor<?xindex>, tensor<?xindex>550 %1 = shape.assuming_all %0, %0, %0551 return %1 : !shape.witness552}553 554// -----555// `assuming_all` with all `cstr_eq` but disjoint operands cannot be collapsed.556// CHECK-LABEL: func @assuming_all_to_cstr_eq557// CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: tensor<?xindex>, %[[C:.*]]: tensor<3xindex>, %[[D:.*]]: tensor<3xindex>)558func.func @assuming_all_to_cstr_eq(%a : !shape.shape, %b : tensor<?xindex>,559 %c : tensor<3xindex>, %d : tensor<3xindex>) -> !shape.witness {560 // CHECK: %[[EQ0:.*]] = shape.cstr_eq %[[A]], %[[B]]561 // CHECK: %[[EQ1:.*]] = shape.cstr_eq %[[C]], %[[D]]562 // CHECK: %[[RESULT:.*]] = shape.assuming_all %[[EQ0]], %[[EQ1]]563 // CHECK: return %[[RESULT]]564 %0 = shape.cstr_eq %a, %b : !shape.shape, tensor<?xindex>565 %1 = shape.cstr_eq %c, %d : tensor<3xindex>, tensor<3xindex>566 %2 = shape.assuming_all %0, %1567 return %2 : !shape.witness568}569 570// -----571// assuming_all with known passing witnesses can be folded572// CHECK-LABEL: func @f573func.func @f() {574 // CHECK-NEXT: shape.const_witness true575 // CHECK-NEXT: consume.witness576 // CHECK-NEXT: return577 %0 = shape.const_witness true578 %1 = shape.const_witness true579 %2 = shape.const_witness true580 %3 = shape.assuming_all %0, %1, %2581 "consume.witness"(%3) : (!shape.witness) -> ()582 return583}584 585// -----586 587// assuming_all should not be removed if more than one witness is not588// statically passing589//590// Additionally check that the attribute is moved to the end as this op is591// commutative.592// CHECK-LABEL: func @f593func.func @f() {594 // CHECK-NEXT: %[[UNKNOWN1:.*]] = "test.source"595 // CHECK-NEXT: %[[UNKNOWN2:.*]] = "test.source"596 // CHECK-NEXT: shape.assuming_all %[[UNKNOWN1]], %[[UNKNOWN2]]597 // CHECK-NEXT: consume.witness598 // CHECK-NEXT: return599 %0 = shape.const_witness true600 %1 = "test.source"() : () -> !shape.witness601 %2 = "test.source"() : () -> !shape.witness602 %3 = shape.assuming_all %0, %1, %2603 "consume.witness"(%3) : (!shape.witness) -> ()604 return605}606 607// -----608 609// merge cstr_broadcastable operations610//611// CHECK-LABEL: func @f612// CHECK: %[[ARG0:[a-z0-9]*]]: !shape.shape613// CHECK-SAME: %[[ARG1:[a-z0-9]*]]: !shape.shape614// CHECK-SAME: %[[ARG2:[a-z0-9]*]]: !shape.shape615func.func @f(%arg0 : !shape.shape, %arg1 : !shape.shape, %arg2 : !shape.shape) {616 // CHECK-NEXT: %[[W:.*]] = shape.cstr_broadcastable %[[ARG0]], %[[ARG1]], %[[ARG2]]617 // CHECK-NEXT: "consume.witness"(%[[W]])618 // CHECK-NEXT: return619 %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape620 %1 = shape.cstr_broadcastable %arg0, %arg1, %arg2 : !shape.shape, !shape.shape, !shape.shape621 %2 = shape.assuming_all %0, %1622 "consume.witness"(%2) : (!shape.witness) -> ()623 return624}625 626// -----627 628// do not merge cstr_broadcastable operations629//630// CHECK-LABEL: func @f631// CHECK: %[[ARG0:[a-z0-9]*]]: !shape.shape632// CHECK-SAME: %[[ARG1:[a-z0-9]*]]: !shape.shape633// CHECK-SAME: %[[ARG2:[a-z0-9]*]]: !shape.shape634func.func @f(%arg0 : !shape.shape, %arg1 : !shape.shape, %arg2 : !shape.shape) {635 // CHECK-NEXT: %[[W0:.*]] = shape.cstr_broadcastable %[[ARG0]], %[[ARG1]]636 // CHECK-NEXT: %[[W1:.*]] = shape.cstr_broadcastable %[[ARG1]], %[[ARG2]]637 // CHECK-NEXT: %[[W2:.*]] = shape.assuming_all %[[W0]], %[[W1]]638 // CHECK-NEXT: "consume.witness"(%[[W2]])639 // CHECK-NEXT: return640 %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape641 %1 = shape.cstr_broadcastable %arg1, %arg2 : !shape.shape, !shape.shape642 %2 = shape.assuming_all %0, %1643 "consume.witness"(%2) : (!shape.witness) -> ()644 return645}646 647// -----648 649// any can be replaced with a constant input if it has one.650// CHECK-LABEL: func @f651func.func @f(%arg : !shape.shape) -> !shape.shape {652 // CHECK-NEXT: %[[CS:.*]] = shape.const_shape653 // CHECK-NEXT: return %[[CS]]654 %0 = shape.const_shape [2, 3, 4] : !shape.shape655 %1 = shape.any %0, %arg : !shape.shape, !shape.shape -> !shape.shape656 return %1 : !shape.shape657}658 659// -----660 661// any can be replaced with a constant input if it has one.662// CHECK-LABEL: func @f663func.func @f(%arg : tensor<?xindex>) -> tensor<3xindex> {664 // CHECK-NEXT: %[[CS:.*]] = shape.const_shape [2, 3, 4] : tensor<3xindex>665 // CHECK-NEXT: return %[[CS]] : tensor<3xindex>666 %0 = shape.const_shape [2, 3, 4] : tensor<3xindex>667 %1 = shape.any %0, %arg : tensor<3xindex>, tensor<?xindex> -> tensor<3xindex>668 return %1 : tensor<3xindex>669}670 671// -----672 673// Folding of any with partially constant operands is not yet implemented.674// CHECK-LABEL: func @f675func.func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> !shape.shape {676 // CHECK-NEXT: %[[CS:.*]] = shape.any677 // CHECK-NEXT: return %[[CS]]678 %1 = shape.any %arg0, %arg1 : !shape.shape, !shape.shape -> !shape.shape679 return %1 : !shape.shape680}681 682// -----683 684// assuming with a known passing witness can be removed685// CHECK-LABEL: func @f686func.func @f() {687 // CHECK-NEXT: source688 // CHECK-NEXT: sink689 // CHECK-NEXT: return690 %0 = shape.const_witness true691 %1 = shape.assuming %0 -> index {692 %2 = "test.source"() : () -> (index)693 shape.assuming_yield %2 : index694 }695 "test.sink"(%1) : (index) -> ()696 return697}698 699// -----700 701// assuming without a known passing passing witness cannot be removed702// CHECK-LABEL: func @f703func.func @f() {704 // CHECK-NEXT: test.source705 // CHECK-NEXT: shape.assuming706 // CHECK-NEXT: test.source707 // CHECK-NEXT: shape.assuming_yield708 // CHECK-NEXT: }709 // CHECK-NEXT: test.sink710 // CHECK-NEXT: return711 %0 = "test.source"() : () -> (!shape.witness)712 %1 = shape.assuming %0 -> index {713 %2 = "test.source"() : () -> (index)714 shape.assuming_yield %2 : index715 }716 "test.sink"(%1) : (index) -> ()717 return718}719 720// -----721 722// Remove unused results from assuming ops.723// CHECK-LABEL: func @unused_assuming_results724func.func @unused_assuming_results() {725 // CHECK: %[[ASSUMING_RESULT:.*]] = shape.assuming %0 -> (f32) {726 // CHECK: %{{.*}} = "produce.redundant"727 // CHECK: %[[MEANINGFUL:.*]] = "produce.meaningful"728 // CHECK: shape.assuming_yield %[[MEANINGFUL]] : f32729 // CHECK: }730 // CHECK: "use"(%[[ASSUMING_RESULT]])731 %0 = "test.source"() : () -> (!shape.witness)732 %1:2 = shape.assuming %0 -> (f32, f32) {733 %2 = "produce.redundant"() : () -> (f32)734 %3 = "produce.meaningful"() : () -> (f32)735 shape.assuming_yield %2, %3 : f32, f32736 }737 "use"(%1#1) : (f32) -> ()738 return739}740 741// -----742// Broadcastable with broadcastable constant shapes can be removed.743// CHECK-LABEL: func @f744func.func @f() {745 // CHECK-NEXT: shape.const_witness true746 // CHECK-NEXT: consume.witness747 // CHECK-NEXT: return748 %cs0 = shape.const_shape [3, 1] : !shape.shape749 %cs1 = shape.const_shape [1, 5] : !shape.shape750 %0 = shape.cstr_broadcastable %cs0, %cs1 : !shape.shape, !shape.shape751 "consume.witness"(%0) : (!shape.witness) -> ()752 return753}754 755// -----756// Empty shape arguments can be removed from broadcastable ops.757// CHECK-LABEL: func @f758// CHECK-SAME: (%[[ARG0:.*]]: tensor<?xindex>, %[[ARG1:.*]]: tensor<?xindex>, %{{.*}}: tensor<0xindex>)759func.func @f(%arg0 : tensor<?xindex>, %arg1 : tensor<?xindex>, %arg2 : tensor<0xindex>) {760 // CHECK-NOT: const_shape761 // CHECK: cstr_broadcastable %[[ARG0]], %[[ARG1]] : tensor<?xindex>, tensor<?xindex>762 %0 = shape.const_shape [] : !shape.shape763 %1 = shape.cstr_broadcastable %arg0, %arg1, %0, %arg2764 : tensor<?xindex>, tensor<?xindex>, !shape.shape, tensor<0xindex>765 "consume.witness"(%1) : (!shape.witness) -> ()766 return767}768 769// -----770// Broadcastable with non-broadcastable constant shapes is always false771// CHECK-LABEL: func @static_non_broadcastable772func.func @static_non_broadcastable() {773 // CHECK-NEXT: shape.const_shape774 // CHECK-NEXT: shape.const_shape775 // CHECK-NEXT: shape.cstr_broadcastable776 // CHECK-NEXT: consume.witness777 // CHECK-NEXT: return778 %cs0 = shape.const_shape [1, 3] : !shape.shape779 %cs1 = shape.const_shape [1, 5] : !shape.shape780 %0 = shape.cstr_broadcastable %cs0, %cs1 : !shape.shape, !shape.shape781 "consume.witness"(%0) : (!shape.witness) -> ()782 return783}784 785// -----786// Broadcastable without guaranteed broadcastable shapes cannot be removed.787// CHECK-LABEL: func @f788func.func @f(%arg0 : !shape.shape) {789 // CHECK-NEXT: shape.const_shape790 // CHECK-NEXT: shape.cstr_broadcastable791 // CHECK-NEXT: consume.witness792 // CHECK-NEXT: return793 %cs0 = shape.const_shape [1, 3] : !shape.shape794 %0 = shape.cstr_broadcastable %arg0, %cs0 : !shape.shape, !shape.shape795 "consume.witness"(%0) : (!shape.witness) -> ()796 return797}798 799// -----800// Broadcastable with non-constant but known equal shapes can be removed.801// CHECK-LABEL: func @f802func.func @f(%arg0 : !shape.shape) {803 // CHECK-NEXT: shape.const_witness true804 // CHECK-NEXT: consume.witness805 // CHECK-NEXT: return806 %0 = shape.cstr_broadcastable %arg0, %arg0 : !shape.shape, !shape.shape807 "consume.witness"(%0) : (!shape.witness) -> ()808 return809}810 811// -----812 813// Broadcastable canonicalization also works on extent tensors.814// CHECK-LABEL: func @broadcastable_on_extent_tensors815func.func @broadcastable_on_extent_tensors(%arg : tensor<?xindex>) {816 // CHECK-NEXT: shape.const_witness true817 // CHECK-NEXT: consume.witness818 // CHECK-NEXT: return819 %0 = shape.cstr_broadcastable %arg, %arg : tensor<?xindex>, tensor<?xindex>820 "consume.witness"(%0) : (!shape.witness) -> ()821 return822}823 824// -----825// Fold ternary broadcastable826// CHECK-LABEL: func @f827func.func @f() {828 // CHECK-NEXT: shape.const_witness true829 // CHECK-NEXT: consume.witness830 // CHECK-NEXT: return831 %cs0 = shape.const_shape [8, 1] : !shape.shape832 %cs1 = shape.const_shape [1, 8] : !shape.shape833 %cs2 = shape.const_shape [1, 1] : !shape.shape834 %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape835 "consume.witness"(%0) : (!shape.witness) -> ()836 return837}838 839// -----840// Fold ternary broadcastable with dynamic ranks841// CHECK-LABEL: func @f842func.func @f() {843 // CHECK-NEXT: shape.const_witness true844 // CHECK-NEXT: consume.witness845 // CHECK-NEXT: return846 %cs0 = shape.const_shape [8, 1] : !shape.shape847 %cs1 = shape.const_shape [1, -9223372036854775808] : !shape.shape848 %0 = shape.cstr_broadcastable %cs0, %cs0, %cs1 : !shape.shape, !shape.shape, !shape.shape849 "consume.witness"(%0) : (!shape.witness) -> ()850 return851}852 853// -----854// One scalar and one non-scalar and one unknown cannot be broadcasted at compile time855// CHECK-LABEL: func @f856func.func @f() {857 // CHECK: shape.cstr_broadcastable858 // CHECK-NEXT: consume.witness859 // CHECK-NEXT: return860 %cs0 = shape.const_shape [8, 1] : !shape.shape861 %cs1 = shape.const_shape [1, 8] : !shape.shape862 %cs2 = shape.const_shape [1, -9223372036854775808] : !shape.shape863 %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape864 "consume.witness"(%0) : (!shape.witness) -> ()865 return866}867 868// -----869// One scalar and two unknowns cannot be broadcasted at compile time870// CHECK-LABEL: func @f871func.func @f() {872 // CHECK: shape.cstr_broadcastable873 // CHECK-NEXT: consume.witness874 // CHECK-NEXT: return875 %cs0 = shape.const_shape [8, 1] : !shape.shape876 %cs1 = shape.const_shape [1, -9223372036854775808] : !shape.shape877 %cs2 = shape.const_shape [8, -9223372036854775808] : !shape.shape878 %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape879 "consume.witness"(%0) : (!shape.witness) -> ()880 return881}882 883// -----884// Broadcastable with scalars and a non-scalar can be constant folded885// CHECK-LABEL: func @f886func.func @f(%arg0 : !shape.shape) {887 // CHECK-NEXT: shape.const_witness true888 // CHECK-NEXT: consume.witness889 // CHECK-NEXT: return890 %cs0 = shape.const_shape [] : !shape.shape891 %0 = shape.cstr_broadcastable %cs0, %cs0, %arg0 : !shape.shape, !shape.shape, !shape.shape892 "consume.witness"(%0) : (!shape.witness) -> ()893 return894}895 896// -----897// One scalar and one non-scalar and one unknown cannot be folded.898// CHECK-LABEL: func @f899func.func @f(%arg0 : !shape.shape) {900 // CHECK: shape.cstr_broadcastable901 // CHECK-NEXT: consume.witness902 // CHECK-NEXT: return903 %cs0 = shape.const_shape [] : !shape.shape904 %cs1 = shape.const_shape [2] : !shape.shape905 %0 = shape.cstr_broadcastable %cs0, %cs1, %arg0 : !shape.shape, !shape.shape, !shape.shape906 "consume.witness"(%0) : (!shape.witness) -> ()907 return908}909 910// -----911 912// Fold `rank` based on constant shape.913// CHECK-LABEL: @fold_rank914func.func @fold_rank() -> !shape.size {915 // CHECK: %[[RESULT:.*]] = shape.const_size 5916 // CHECK: return %[[RESULT]] : !shape.size917 %shape = shape.const_shape [3, 4, 5, 6, 7] : !shape.shape918 %rank = shape.rank %shape : !shape.shape -> !shape.size919 return %rank : !shape.size920}921 922// -----923 924// Do not fold `rank` if shape is dynamic.925// CHECK-LABEL: @dont_fold_rank926// CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape) -> !shape.size927func.func @dont_fold_rank(%shape : !shape.shape) -> !shape.size {928 // CHECK: %[[RESULT:.*]] = shape.rank %[[SHAPE]]929 // CHECK: return %[[RESULT]] : !shape.size930 %rank = shape.rank %shape : !shape.shape -> !shape.size931 return %rank : !shape.size932}933 934// -----935 936// Fold `rank` based on constant extent tensor.937// CHECK-LABEL: @fold_rank938func.func @fold_rank() -> index {939 // CHECK: %[[RESULT:.*]] = arith.constant 5 : index940 // CHECK: return %[[RESULT]] : index941 %shape = shape.const_shape [3, 4, 5, 6, 7] : tensor<5xindex>942 %rank = shape.rank %shape : tensor<5xindex> -> index943 return %rank : index944}945 946// -----947 948// Do not fold `rank` for non-constant extent tensors.949// CHECK-LABEL: @dont_fold_rank950// CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>) -> index951func.func @dont_fold_rank(%shape : tensor<?xindex>) -> index {952 // CHECK: %[[RESULT:.*]] = shape.rank %[[SHAPE]] : tensor<?xindex> -> index953 // CHECK: return %[[RESULT]] : index954 %rank = shape.rank %shape : tensor<?xindex> -> index955 return %rank : index956}957 958// -----959 960// Canonicalize `rank` when shape is derived from ranked tensor.961// CHECK-LABEL: @canonicalize_rank962func.func @canonicalize_rank(%arg : tensor<1x2x?xf32>) -> index {963 // CHECK: %[[RESULT:.*]] = arith.constant 3 : index964 // CHECK: return %[[RESULT]] : index965 %shape = shape.shape_of %arg : tensor<1x2x?xf32> -> tensor<?xindex>966 %rank = shape.rank %shape : tensor<?xindex> -> index967 return %rank : index968}969 970// -----971 972// Canonicalize `rank` when shape is derived from ranked tensor.973// CHECK-LABEL: @canonicalize_rank974func.func @canonicalize_rank_size(%arg : tensor<1x2x?xf32>) -> !shape.size {975 // CHECK: %[[RESULT:.*]] = shape.const_size 3976 // CHECK: return %[[RESULT]] : !shape.size977 %shape = shape.shape_of %arg : tensor<1x2x?xf32> -> !shape.shape978 %rank = shape.rank %shape : !shape.shape -> !shape.size979 return %rank : !shape.size980}981 982// -----983 984// Do not canonicalize `rank` when shape is derived from unranked tensor.985// CHECK-LABEL: @dont_canonicalize_rank986// CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>) -> index987func.func @dont_canonicalize_rank(%arg : tensor<*xf32>) -> index {988 // CHECK: %[[SHAPE:.*]] = shape.shape_of %[[ARG]] : tensor<*xf32> -> tensor<?xindex>989 // CHECK: %[[SIZE:.*]] = shape.rank %[[SHAPE]]990 // CHECK: return %[[SIZE]] : index991 %shape = shape.shape_of %arg : tensor<*xf32> -> tensor<?xindex>992 %rank = shape.rank %shape : tensor<?xindex> -> index993 return %rank : index994}995 996// -----997 998// Canonicalize redundant conversion from `index` to `size` and back.999// CHECK-LABEL: @index_to_size_to_index1000// CHECK-SAME: (%[[IDX:.*]]: index) -> index1001func.func @index_to_size_to_index(%index : index) -> index {1002 // CHECK: return %[[IDX]] : index1003 %size = shape.index_to_size %index1004 %result = shape.size_to_index %size : !shape.size1005 return %result : index1006}1007 1008// -----1009 1010// Canonicalize redundant conversion from `size` to `index` and back.1011// CHECK-LABEL: @size_to_index_to_size1012// CHECK-SAME: (%[[SIZE:.*]]: !shape.size) -> !shape.size1013func.func @size_to_index_to_size(%size : !shape.size) -> !shape.size {1014 // CHECK: return %[[SIZE]] : !shape.size1015 %idx = shape.size_to_index %size : !shape.size1016 %result = shape.index_to_size %idx1017 return %result : !shape.size1018}1019 1020// -----1021 1022// Canonicalize scalar cstr_broadcastable checks1023// CHECK-LABEL: @cstr_broadcastable_scalar1024func.func @cstr_broadcastable_scalar(%arg0 : tensor<?xf32>) {1025 // CHECK-NEXT: shape.const_witness true1026 // CHECK-NEXT: consume.witness1027 // CHECK-NEXT: return1028 %0 = shape.const_shape [] : !shape.shape1029 %1 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>1030 %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, tensor<?xindex>1031 "consume.witness"(%2) : (!shape.witness) -> ()1032 return1033}1034 1035// -----1036 1037// Do not canonicalize cstr_broadcastable checks with 2 unknowns1038// CHECK-LABEL: @cstr_broadcastable_unknown1039func.func @cstr_broadcastable_unknown(%arg0 : tensor<?xf32>, %arg1 : tensor<?xf32>) {1040 // CHECK-NEXT: shape.shape_of %arg01041 // CHECK-NEXT: shape.shape_of %arg11042 // CHECK-NEXT: shape.cstr_broadcastable1043 // CHECK-NEXT: consume.witness1044 // CHECK-NEXT: return1045 %0 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>1046 %1 = shape.shape_of %arg1 : tensor<?xf32> -> tensor<?xindex>1047 %2 = shape.cstr_broadcastable %0, %1 : tensor<?xindex>, tensor<?xindex>1048 "consume.witness"(%2) : (!shape.witness) -> ()1049 return1050}1051 1052// -----1053 1054// Scalars are safe to broadcast to unranked sizes.1055// CHECK-LABEL: @cstr_broadcastable_scalar_unranked1056func.func @cstr_broadcastable_scalar_unranked(%arg0 : tensor<*xf32>, %arg1 : tensor<index>) {1057 // CHECK-NEXT: shape.const_witness true1058 // CHECK-NEXT: consume.witness1059 // CHECK-NEXT: return1060 %0 = shape.shape_of %arg1 : tensor<index> -> tensor<?xindex>1061 %1 = shape.shape_of %arg0 : tensor<*xf32> -> tensor<?xindex>1062 %2 = shape.cstr_broadcastable %0, %1 : tensor<?xindex>, tensor<?xindex>1063 "consume.witness"(%2) : (!shape.witness) -> ()1064 return1065}1066 1067// -----1068 1069// Fold `shape_eq` for equal and constant shapes.1070// CHECK-LABEL: @shape_eq_fold_11071func.func @shape_eq_fold_1() -> i1 {1072 // CHECK: %[[RESULT:.*]] = arith.constant true1073 // CHECK: return %[[RESULT]] : i11074 %a = shape.const_shape [1, 2, 3] : !shape.shape1075 %b = shape.const_shape [1, 2, 3] : tensor<3xindex>1076 %c = shape.const_shape [1, 2, 3] : tensor<3xindex>1077 %result = shape.shape_eq %a, %b, %c : !shape.shape, tensor<3xindex>, tensor<3xindex>1078 return %result : i11079}1080 1081// -----1082 1083// Fold `shape_eq` for different but constant shapes of same length.1084// CHECK-LABEL: @shape_eq_fold_01085func.func @shape_eq_fold_0() -> i1 {1086 // CHECK: %[[RESULT:.*]] = arith.constant false1087 // CHECK: return %[[RESULT]] : i11088 %a = shape.const_shape [1, 2, 3] : tensor<3xindex>1089 %b = shape.const_shape [4, 5, 6] : tensor<3xindex>1090 %c = shape.const_shape [4, 5, 6] : tensor<3xindex>1091 %result = shape.shape_eq %a, %b, %c : tensor<3xindex>, tensor<3xindex>, tensor<3xindex>1092 return %result : i11093}1094 1095// -----1096 1097// Fold `shape_eq` for different but constant shapes of different length.1098// CHECK-LABEL: @shape_eq_fold_01099func.func @shape_eq_fold_0() -> i1 {1100 // CHECK: %[[RESULT:.*]] = arith.constant false1101 // CHECK: return %[[RESULT]] : i11102 %a = shape.const_shape [1, 2, 3, 4, 5, 6] : !shape.shape1103 %b = shape.const_shape [1, 2, 3] : !shape.shape1104 %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape1105 return %result : i11106}1107 1108// -----1109 1110// Do not fold `shape_eq` for non-constant different shapes.1111// CHECK-LABEL: @shape_eq_do_not_fold1112// CHECK-SAME: (%[[A:.*]]: !shape.shape) -> i11113func.func @shape_eq_do_not_fold(%a : !shape.shape) -> i1 {1114 // CHECK: %[[B:.*]] = shape.const_shape [4, 5, 6]1115 // CHECK: %[[RESULT:.*]] = shape.shape_eq %[[A]], %[[B]] : !shape.shape, !shape.shape1116 // CHECK: return %[[RESULT]] : i11117 %b = shape.const_shape [4, 5, 6] : !shape.shape1118 %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape1119 return %result : i11120}1121 1122// -----1123 1124// Fold `add` for constant sizes.1125// CHECK-LABEL: @fold_add_size1126func.func @fold_add_size() -> !shape.size {1127 // CHECK: %[[RESULT:.*]] = shape.const_size 51128 // CHECK: return %[[RESULT]] : !shape.size1129 %c2 = shape.const_size 21130 %c3 = shape.const_size 31131 %result = shape.add %c2, %c3 : !shape.size, !shape.size -> !shape.size1132 return %result : !shape.size1133}1134 1135// -----1136 1137// Fold `mul` for constant sizes.1138// CHECK-LABEL: @fold_mul_size1139func.func @fold_mul_size() -> !shape.size {1140 // CHECK: %[[RESULT:.*]] = shape.const_size 61141 // CHECK: return %[[RESULT]] : !shape.size1142 %c2 = shape.const_size 21143 %c3 = shape.const_size 31144 %result = shape.mul %c2, %c3 : !shape.size, !shape.size -> !shape.size1145 return %result : !shape.size1146}1147 1148// -----1149 1150// Fold `mul` for constant indices.1151// CHECK-LABEL: @fold_mul_index1152func.func @fold_mul_index() -> index {1153 // CHECK: %[[RESULT:.*]] = arith.constant 6 : index1154 // CHECK: return %[[RESULT]] : index1155 %c2 = arith.constant 2 : index1156 %c3 = arith.constant 3 : index1157 %result = shape.mul %c2, %c3 : index, index -> index1158 return %result : index1159}1160 1161// -----1162 1163// Fold `mul` for mixed constants.1164// CHECK-LABEL: @fold_mul_mixed1165func.func @fold_mul_mixed() -> !shape.size {1166 // CHECK: %[[RESULT:.*]] = shape.const_size 61167 // CHECK: return %[[RESULT]] : !shape.size1168 %c2 = shape.const_size 21169 %c3 = arith.constant 3 : index1170 %result = shape.mul %c2, %c3 : !shape.size, index -> !shape.size1171 return %result : !shape.size1172}1173 1174// -----1175 1176// Fold `div` for constant sizes.1177// CHECK-LABEL: @fold_div_size1178func.func @fold_div_size() -> !shape.size {1179 // CHECK: %[[RESULT:.*]] = shape.const_size 31180 // CHECK: return %[[RESULT]] : !shape.size1181 %c2 = shape.const_size 101182 %c3 = shape.const_size 31183 %result = shape.div %c2, %c3 : !shape.size, !shape.size -> !shape.size1184 return %result : !shape.size1185}1186 1187// -----1188 1189// Fold `div` for constant indices.1190// CHECK-LABEL: @fold_div_index1191func.func @fold_div_index() -> index {1192 // CHECK: %[[RESULT:.*]] = arith.constant 2 : index1193 // CHECK: return %[[RESULT]] : index1194 %c2 = arith.constant 10 : index1195 %c3 = arith.constant 4 : index1196 %result = shape.div %c2, %c3 : index, index -> index1197 return %result : index1198}1199 1200// -----1201 1202// Fold `div` for constant indices and lhs is negative.1203// CHECK-LABEL: @fold_div_index_neg_lhs1204func.func @fold_div_index_neg_lhs() -> index {1205 // CHECK: %[[RESULT:.*]] = arith.constant -3 : index1206 // CHECK: return %[[RESULT]] : index1207 %c2 = arith.constant -10 : index1208 %c3 = arith.constant 4 : index1209 %result = shape.div %c2, %c3 : index, index -> index1210 return %result : index1211}1212 1213// -----1214 1215// Fold `div` for constant indices and rhs is negative.1216// CHECK-LABEL: @fold_div_index_neg_rhs1217func.func @fold_div_index_neg_rhs() -> index {1218 // CHECK: %[[RESULT:.*]] = arith.constant -3 : index1219 // CHECK: return %[[RESULT]] : index1220 %c2 = arith.constant 10 : index1221 %c3 = arith.constant -4 : index1222 %result = shape.div %c2, %c3 : index, index -> index1223 return %result : index1224}1225 1226// -----1227 1228// Fold `div` for mixed constants.1229// CHECK-LABEL: @fold_div_mixed1230func.func @fold_div_mixed() -> !shape.size {1231 // CHECK: %[[RESULT:.*]] = shape.const_size 41232 // CHECK: return %[[RESULT]] : !shape.size1233 %c2 = shape.const_size 121234 %c3 = arith.constant 3 : index1235 %result = shape.div %c2, %c3 : !shape.size, index -> !shape.size1236 return %result : !shape.size1237}1238 1239// -----1240 1241// Fold index_cast when already on index.1242// CHECK-LABEL: @fold_index_cast_on_index1243func.func @fold_index_cast_on_index(%arg: index) -> index {1244 // CHECK-NOT: size_to_index1245 %0 = shape.size_to_index %arg : index1246 return %0 : index1247}1248 1249// -----1250 1251// Fold to_extent_tensor when already on tensor.1252// CHECK-LABEL: @fold_to_extent_tensor_on_tensor1253func.func @fold_to_extent_tensor_on_tensor(%arg: tensor<?xindex>) -> tensor<?xindex> {1254 // CHECK-NOT: to_extent_tensor1255 %0 = shape.to_extent_tensor %arg : tensor<?xindex> -> tensor<?xindex>1256 return %0 : tensor<?xindex>1257}1258 1259// -----1260 1261// Fold assuming_all with a single input1262// CHECK-LABEL: @fold_assuming_all_single_element1263func.func @fold_assuming_all_single_element(%arg: tensor<?xindex>) {1264 // CHECK-NOT: assuming_all1265 %0 = "test.source"() : () -> (!shape.witness)1266 %1 = shape.assuming_all %01267 "consume.witness"(%1) : (!shape.witness) -> ()1268 return1269}1270 1271// -----1272 1273// Verify that tensor.cast folding uses the correct type1274// CHECK-LABEL: @fold_tensor.cast_of_const_shape_returned1275func.func @fold_tensor.cast_of_const_shape_returned(%arg: i1) -> tensor<1xindex> {1276 // CHECK: shape.const_shape [2] : tensor<1xindex>1277 // CHECK-NOT: tensor.cast1278 %0 = shape.const_shape [2] : tensor<1xindex>1279 %1 = tensor.cast %0 : tensor<1xindex> to tensor<1xindex>1280 return %1 : tensor<1xindex>1281}1282 1283// -----1284 1285// CHECK-LABEL: @dont_fold_tensor.cast_of_const_shape_returned_dynamic1286func.func @dont_fold_tensor.cast_of_const_shape_returned_dynamic(%arg: i1) -> tensor<?xindex> {1287 // CHECK: %[[CONST_SHAPE:.*]] = shape.const_shape [2] : tensor<1xindex>1288 // CHECK: tensor.cast %[[CONST_SHAPE]] : tensor<1xindex> to tensor<?xindex>1289 %0 = shape.const_shape [2] : tensor<1xindex>1290 %1 = tensor.cast %0 : tensor<1xindex> to tensor<?xindex>1291 return %1 : tensor<?xindex>1292}1293 1294// -----1295 1296// CHECK-LABEL: @is_broadcastable_on_same_shape1297func.func @is_broadcastable_on_same_shape(%shape : !shape.shape) -> i1 {1298 // CHECK-NOT: is_broadcastable1299 // CHECK: %[[RES:.*]] = arith.constant true1300 // CHECK: return %[[RES]]1301 %0 = shape.is_broadcastable %shape, %shape, %shape1302 : !shape.shape, !shape.shape, !shape.shape1303 return %0 : i11304}1305 1306// -----1307 1308// CHECK-LABEL: @is_broadcastable_on_duplicate_shapes1309// CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: !shape.shape)1310func.func @is_broadcastable_on_duplicate_shapes(%a : !shape.shape, %b : !shape.shape)1311 -> i1 {1312 // CHECK: %[[RES:.*]] = shape.is_broadcastable %[[A]], %[[B]] :1313 // CHECK: return %[[RES]]1314 %0 = shape.is_broadcastable %a, %b, %a, %a, %a, %b : !shape.shape,1315 !shape.shape, !shape.shape, !shape.shape, !shape.shape, !shape.shape1316 return %0 : i11317}1318 1319// -----1320 1321// CHECK-LABEL: @cstr_broadcastable_on_duplicate_shapes1322// CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: !shape.shape)1323func.func @cstr_broadcastable_on_duplicate_shapes(%a : !shape.shape,1324 %b : !shape.shape) -> !shape.witness {1325 // CHECK: %[[RES:.*]] = shape.cstr_broadcastable %[[A]], %[[B]] :1326 // CHECK: return %[[RES]]1327 %0 = shape.cstr_broadcastable %a, %b, %a, %a, %a, %b : !shape.shape,1328 !shape.shape, !shape.shape, !shape.shape, !shape.shape, !shape.shape1329 return %0 : !shape.witness1330}1331 1332// -----1333 1334// CHECK-LABEL: @broadcast_on_same_shape1335// CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape)1336func.func @broadcast_on_same_shape(%shape : !shape.shape) -> !shape.shape {1337 // CHECK-NOT: broadcast1338 // CHECK: return %[[SHAPE]]1339 %0 = shape.broadcast %shape, %shape, %shape : !shape.shape, !shape.shape,1340 !shape.shape -> !shape.shape1341 return %0 : !shape.shape1342}1343 1344// -----1345 1346// CHECK-LABEL: @broadcast_on_duplicate_shapes1347// CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: !shape.shape)1348func.func @broadcast_on_duplicate_shapes(%a : !shape.shape, %b : !shape.shape)1349 -> !shape.shape {1350 // CHECK: %[[RES:.*]] = shape.broadcast %[[A]], %[[B]] :1351 // CHECK: return %[[RES]]1352 %0 = shape.broadcast %a, %b, %a, %a, %a, %b : !shape.shape, !shape.shape,1353 !shape.shape, !shape.shape, !shape.shape, !shape.shape -> !shape.shape1354 return %0 : !shape.shape1355}1356 1357// -----1358 1359// CHECK-LABEL: @broadcast_on_single_operand1360// CHECK-SAME: (%[[A:.*]]: tensor<?xindex>)1361func.func @broadcast_on_single_operand(%a : tensor<?xindex>) {1362 // CHECK-NOT: broadcast1363 // CHECK: "use"(%[[A]])1364 %0 = shape.broadcast %a : tensor<?xindex> -> tensor<?xindex>1365 "use"(%0) : (tensor<?xindex>) -> ()1366 return1367}1368 1369// -----1370 1371// CHECK-LABEL: @broadcast_as_tensor_cast1372// CHECK-SAME: (%[[A:.*]]: tensor<3xindex>)1373func.func @broadcast_as_tensor_cast(%a : tensor<3xindex>) -> tensor<?xindex> {1374 // CHECK: %[[RESULT:.*]] = tensor.cast %[[A]] : tensor<3xindex> to tensor<?xindex>1375 // CHECK: return %[[RESULT]] : tensor<?xindex>1376 %0 = shape.broadcast %a : tensor<3xindex> -> tensor<?xindex>1377 return %0 : tensor<?xindex>1378}1379 1380// -----1381 1382// CHECK-LABEL: @broadcast_as_from_extent_tensor1383// CHECK-SAME: (%[[A:.*]]: tensor<?xindex>)1384func.func @broadcast_as_from_extent_tensor(%a : tensor<?xindex>) -> !shape.shape {1385 // CHECK: %[[RESULT:.*]] = shape.from_extent_tensor %[[A]] : tensor<?xindex>1386 // CHECK: return %[[RESULT]] : !shape.shape1387 %0 = shape.broadcast %a : tensor<?xindex> -> !shape.shape1388 return %0 : !shape.shape1389}1390 1391// -----1392 1393// CHECK-LABEL: func @shape_of_from_reshape1394// CHECK-SAME: %[[INPUT:.*]]: tensor<*xf32>1395// CHECK-SAME: %[[SHAPE:.*]]: tensor<?xindex>1396func.func @shape_of_from_reshape(%arg0: tensor<*xf32>, %arg1: tensor<?xindex>) -> tensor<?xindex> {1397 // CHECK: return %[[SHAPE]] : tensor<?xindex>1398 %0 = tensor.reshape %arg0(%arg1) : (tensor<*xf32>, tensor<?xindex>) -> tensor<*xf32>1399 %1 = shape.shape_of %0 : tensor<*xf32> -> tensor<?xindex>1400 return %1 : tensor<?xindex>1401}1402 1403// -----1404 1405// Check statically shaped types, with element types i32 to index.1406// CHECK-LABEL: func @shape_of_from_reshape_int_to_index1407// CHECK-SAME: %[[INPUT:.*]]: tensor<?x1xf32>1408// CHECK-SAME: %[[SHAPE:.*]]: tensor<3xi32>1409func.func @shape_of_from_reshape_int_to_index(%arg0: tensor<?x1xf32>, %arg1: tensor<3xi32>) -> tensor<3xindex> {1410 // CHECK: %[[CAST_SHAPE:.*]] = arith.index_cast %[[SHAPE]] : tensor<3xi32> to tensor<3xindex>1411 // CHECK: return %[[CAST_SHAPE]] : tensor<3xindex>1412 %0 = tensor.reshape %arg0(%arg1) : (tensor<?x1xf32>, tensor<3xi32>) -> tensor<?x1x1xf32>1413 %1 = shape.shape_of %0 : tensor<?x1x1xf32> -> tensor<3xindex>1414 return %1 : tensor<3xindex>1415}1416 1417// -----1418 1419// Check similar element types, with statically shaped to dynamically shaped.1420// CHECK-LABEL: func @shape_of_from_reshape_static_to_dynamic1421// CHECK-SAME: %[[INPUT:.*]]: tensor<*xf32>1422// CHECK-SAME: %[[SHAPE:.*]]: tensor<5xindex>1423func.func @shape_of_from_reshape_static_to_dynamic(%arg0: tensor<*xf32>, %arg1: tensor<5xindex>) -> tensor<?xindex> {1424 // CHECK: %[[CAST_SHAPE:.*]] = tensor.cast %[[SHAPE]] : tensor<5xindex> to tensor<?xindex>1425 // CHECK: return %[[CAST_SHAPE]] : tensor<?xindex>1426 %0 = tensor.reshape %arg0(%arg1) : (tensor<*xf32>, tensor<5xindex>) -> tensor<*xf32>1427 %1 = shape.shape_of %0 : tensor<*xf32> -> tensor<?xindex>1428 return %1 : tensor<?xindex>1429}1430 1431// -----1432 1433// Check similar element types, with dynamically shaped to statically shaped.1434// CHECK-LABEL: func @shape_of_from_reshape_dynamic_to_static1435// CHECK-SAME: %[[INPUT:.*]]: tensor<*xf32>1436// CHECK-SAME: %[[SHAPE:.*]]: tensor<?xindex>1437func.func @shape_of_from_reshape_dynamic_to_static(%arg0: tensor<*xf32>, %arg1: tensor<?xindex>) -> tensor<5xindex> {1438 // CHECK: %[[CAST_SHAPE:.*]] = tensor.cast %[[SHAPE]] : tensor<?xindex> to tensor<5xindex>1439 // CHECK: return %[[CAST_SHAPE]] : tensor<5xindex>1440 %0 = tensor.reshape %arg0(%arg1) : (tensor<*xf32>, tensor<?xindex>) -> tensor<*xf32>1441 %1 = shape.shape_of %0 : tensor<*xf32> -> tensor<5xindex>1442 return %1 : tensor<5xindex>1443}1444 1445// -----1446 1447// Check similar element types and similar static shape.1448// CHECK-LABEL: func @shape_of_from_reshape_identical_types1449// CHECK-SAME: %[[INPUT:.*]]: tensor<*xf32>1450// CHECK-SAME: %[[SHAPE:.*]]: tensor<5xindex>1451func.func @shape_of_from_reshape_identical_types(%arg0: tensor<*xf32>, %arg1: tensor<5xindex>) -> tensor<5xindex> {1452 // CHECK: return %[[SHAPE]] : tensor<5xindex>1453 %0 = tensor.reshape %arg0(%arg1) : (tensor<*xf32>, tensor<5xindex>) -> tensor<*xf32>1454 %1 = shape.shape_of %0 : tensor<*xf32> -> tensor<5xindex>1455 return %1 : tensor<5xindex>1456}1457 1458// -----1459 1460// CHECK-LABEL: func @shape_of_from_reshape_nofold1461// CHECK-SAME: %[[INPUT:.*]]: tensor<*xf32>1462// CHECK-SAME: %[[SHAPE:.*]]: tensor<?xindex>1463func.func @shape_of_from_reshape_nofold(%arg0: tensor<*xf32>, %arg1: tensor<?xindex>) -> !shape.shape {1464 // CHECK: %[[RESHAPED:.*]] = tensor.reshape %[[INPUT]](%[[SHAPE]]) : (tensor<*xf32>, tensor<?xindex>) -> tensor<*xf32>1465 // CHECK: %[[SHAPE_OF:.*]] = shape.shape_of %[[RESHAPED]] : tensor<*xf32> -> !shape.shape1466 // CHECK: return %[[SHAPE_OF]] : !shape.shape1467 %0 = tensor.reshape %arg0(%arg1) : (tensor<*xf32>, tensor<?xindex>) -> tensor<*xf32>1468 %1 = shape.shape_of %0 : tensor<*xf32> -> !shape.shape1469 return %1 : !shape.shape1470}1471 1472// -----1473 1474// CHECK-LABEL: @cast_extent_tensor1475// CHECK-SAME: (%[[ARG:.*]]: tensor<?x?x?xf32>) -> tensor<?xindex>1476func.func @cast_extent_tensor(%arg : tensor<?x?x?xf32>) -> tensor<?xindex> {1477 // CHECK: %[[RESULT:.*]] = shape.shape_of %[[ARG]] : tensor<?x?x?xf32> -> tensor<?xindex>1478 // CHECK: return %[[RESULT]] : tensor<?xindex>1479 %0 = shape.shape_of %arg : tensor<?x?x?xf32> -> tensor<3xindex>1480 %1 = tensor.cast %0 : tensor<3xindex> to tensor<?xindex>1481 return %1 : tensor<?xindex>1482}1483 1484// -----1485 1486// CHECK-LABEL: @cast_extent_tensor1487// CHECK-SAME: (%[[ARG:.*]]: tensor<?x?x?xf32>) -> tensor<3xindex>1488func.func @cast_extent_tensor(%arg : tensor<?x?x?xf32>) -> tensor<3xindex> {1489 // CHECK: %[[RESULT:.*]] = shape.shape_of %[[ARG]] : tensor<?x?x?xf32> -> tensor<3xindex>1490 // CHECK: return %[[RESULT]] : tensor<3xindex>1491 %0 = shape.shape_of %arg : tensor<?x?x?xf32> -> tensor<?xindex>1492 %1 = tensor.cast %0 : tensor<?xindex> to tensor<3xindex>1493 return %1 : tensor<3xindex>1494}1495 1496// -----1497 1498// CHECK-LABEL: @cast_extent_tensor1499func.func @cast_extent_tensor(%arg : tensor<?x?x?x?xf32>) -> tensor<3xindex> {1500 // CHECK: tensor.cast %{{.*}} : tensor<?xindex> to tensor<3xindex>1501 %0 = shape.shape_of %arg : tensor<?x?x?x?xf32> -> tensor<?xindex>1502 %1 = tensor.cast %0 : tensor<?xindex> to tensor<3xindex>1503 return %1 : tensor<3xindex>1504}1505 1506// -----1507 1508// CHECK-LABEL: @cast_extent_tensor1509func.func @cast_extent_tensor(%arg : tensor<*xf32>) -> tensor<3xindex> {1510 // CHECK: tensor.cast %{{.*}} : tensor<?xindex> to tensor<3xindex>1511 %0 = shape.shape_of %arg : tensor<*xf32> -> tensor<?xindex>1512 %1 = tensor.cast %0 : tensor<?xindex> to tensor<3xindex>1513 return %1 : tensor<3xindex>1514}1515 1516// -----1517 1518// CHECK-LABEL: max_same_arg1519// CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape)1520func.func @max_same_arg(%a: !shape.shape) -> !shape.shape {1521 %1 = shape.max %a, %a : !shape.shape, !shape.shape -> !shape.shape1522 // CHECK: return %[[SHAPE]]1523 return %1 : !shape.shape1524}1525 1526// -----1527 1528// CHECK-LABEL: min_same_arg1529// CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape)1530func.func @min_same_arg(%a: !shape.shape) -> !shape.shape {1531 %1 = shape.min %a, %a : !shape.shape, !shape.shape -> !shape.shape1532 // CHECK: return %[[SHAPE]]1533 return %1 : !shape.shape1534}1535// -----1536 1537// CHECK-LABEL: @cstr_broadcastable_folding1538func.func @cstr_broadcastable_folding(%arg : tensor<?x4xf32>) {1539 // CHECK: const_witness true1540 %0 = shape.shape_of %arg : tensor<?x4xf32> -> tensor<2xindex>1541 %1 = shape.const_shape [4] : tensor<1xindex>1542 %2 = shape.cstr_broadcastable %0, %1: tensor<2xindex>, tensor<1xindex>1543 "use"(%2) : (!shape.witness) -> ()1544}1545 1546// -----1547 1548// CHECK-LABEL: @cast_extent_tensor_operands1549// CHECK-SAME: (%[[ARG0:.*]]: tensor<?xindex>, %[[ARG1:.*]]: tensor<3xindex>)1550func.func @cast_extent_tensor_operands(%arg0 : tensor<?xindex>,1551 %arg1 : tensor<3xindex>) -> (!shape.witness, tensor<?xindex>) {1552 // CHECK: %[[CAST_ARG0:.*]] = tensor.cast %[[ARG0]] : tensor<?xindex> to tensor<3xindex>1553 // CHECK: %[[WIT:.*]] = shape.cstr_broadcastable %[[CAST_ARG0]], %[[ARG1]] : tensor<3xindex>, tensor<3xindex>1554 // CHECK: %[[UNCAST_RES:.*]] = shape.broadcast %[[CAST_ARG0]], %[[ARG1]] : tensor<3xindex>, tensor<3xindex> -> tensor<3xindex>1555 // CHECK: %[[RES:.*]] = tensor.cast %[[UNCAST_RES]] : tensor<3xindex> to tensor<?xindex>1556 // CHECK: return %[[WIT]], %[[RES]]1557 %0 = tensor.cast %arg0 : tensor<?xindex> to tensor<3xindex>1558 %1 = tensor.cast %arg1 : tensor<3xindex> to tensor<?xindex>1559 %2 = shape.cstr_broadcastable %0, %1 : tensor<3xindex>, tensor<?xindex>1560 %3 = shape.broadcast %0, %1 :tensor<3xindex>, tensor<?xindex>1561 -> tensor<?xindex>1562 return %2, %3 : !shape.witness, tensor<?xindex>1563}1564 1565// -----1566 1567// CHECK-LABEL: @concretize_broadcast_result_type1568// CHECK-SAME: (%[[ARG0:.*]]: tensor<2xindex>, %[[ARG1:.*]]: tensor<3xindex>)1569func.func @concretize_broadcast_result_type(%arg0 : tensor<2xindex>,1570 %arg1 : tensor<3xindex>) -> tensor<?xindex> {1571 // CHECK: %[[CONCR:.*]] = shape.broadcast %[[ARG0]], %[[ARG1]] : tensor<2xindex>, tensor<3xindex> -> tensor<3xindex>1572 // CHECK: %[[RES:.*]] = tensor.cast %[[CONCR]] : tensor<3xindex> to tensor<?xindex>1573 // CHECK: return %[[RES]]1574 %0 = shape.broadcast %arg0, %arg1 : tensor<2xindex>, tensor<3xindex>1575 -> tensor<?xindex>1576 return %0 : tensor<?xindex>1577}1578 1579// -----1580 1581// CHECK-LABEL: func @extract_shapeof1582// CHECK-SAME: %[[ARG0:.*]]: tensor<?x?xf64>1583func.func @extract_shapeof(%arg0 : tensor<?x?xf64>) -> index {1584 %c1 = arith.constant 1 : index1585// CHECK: %[[C1:.*]] = arith.constant 11586 %shape = shape.shape_of %arg0 : tensor<?x?xf64> -> tensor<2xindex>1587// CHECK: %[[DIM:.*]] = tensor.dim %[[ARG0]], %[[C1]]1588 %result = tensor.extract %shape[%c1] : tensor<2xindex>1589// CHECK: return %[[DIM]]1590 return %result : index1591}1592 1593 1594// -----1595 1596// CHECK-LABEL: @add_poison1597// CHECK: %[[P:.*]] = ub.poison : !shape.siz1598// CHECK: return %[[P]]1599func.func @add_poison() -> !shape.size {1600 %1 = shape.const_size 21601 %2 = ub.poison : !shape.size1602 %result = shape.add %1, %2 : !shape.size, !shape.size -> !shape.size1603 return %result : !shape.size1604}1605 1606// -----1607 1608// CHECK-LABEL: func @shape_of_0d(1609// CHECK-SAME: %[[arg0:.*]]: tensor<f32>1610// CHECK: %[[const:.*]] = shape.const_shape [] : tensor<0xindex>1611// CHECK: %[[cast:.*]] = tensor.cast %[[const]] : tensor<0xindex> to tensor<?xindex>1612// CHECK: return %[[cast]]1613func.func @shape_of_0d(%arg0: tensor<f32>) -> tensor<?xindex> {1614 %0 = shape.shape_of %arg0 : tensor<f32> -> tensor<?xindex>1615 return %0 : tensor<?xindex>1616}1617