// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -canonicalize="test-convergence" %s | FileCheck %s // CHECK-LABEL: func @f func.func @f(%arg0: tensor<2x3x4xf32>) -> tensor<3xindex> { // CHECK: shape.const_shape [2, 3, 4] : tensor<3xindex> %0 = shape.shape_of %arg0 : tensor<2x3x4xf32> -> tensor<3xindex> return %0 : tensor<3xindex> } // ----- // Basic case. // CHECK-LABEL: func @f func.func @f() -> (!shape.shape, !shape.shape) { // CHECK-DAG: shape.const_shape [2, 3] : !shape.shape // CHECK-DAG: shape.const_shape [4, 5] : !shape.shape %c2 = arith.constant 2 : index %0 = shape.const_shape [2, 3, 4, 5] : !shape.shape %head, %tail = "shape.split_at"(%0, %c2) : (!shape.shape, index) -> (!shape.shape, !shape.shape) return %head, %tail : !shape.shape, !shape.shape } // ----- // Negative split point. // CHECK-LABEL: func @f func.func @f() -> (!shape.shape, !shape.shape) { // CHECK-DAG: shape.const_shape [2, 3, 4] : !shape.shape // CHECK-DAG: shape.const_shape [5] : !shape.shape %c-1 = arith.constant -1 : index %0 = shape.const_shape [2, 3, 4, 5] : !shape.shape %head, %tail = "shape.split_at"(%0, %c-1) : (!shape.shape, index) -> (!shape.shape, !shape.shape) return %head, %tail : !shape.shape, !shape.shape } // ----- // Out of range split point. No folding. // CHECK-LABEL: func @f func.func @f() -> (!shape.shape, !shape.shape) { // CHECK: shape.split_at %c5 = arith.constant 5 : index %0 = shape.const_shape [2, 3, 4, 5] : !shape.shape %head, %tail = "shape.split_at"(%0, %c5) : (!shape.shape, index) -> (!shape.shape, !shape.shape) return %head, %tail : !shape.shape, !shape.shape } // ----- // Basic case. // CHECK-LABEL: func @f func.func @f() -> !shape.shape { // CHECK: shape.const_shape [7, 2] : !shape.shape %0 = shape.const_shape [1, 2] : !shape.shape %1 = shape.const_shape [7, 1] : !shape.shape %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape return %2 : !shape.shape } // ----- // Basic case including extent tensors. // CHECK-LABEL: @broadcast func.func @broadcast() -> tensor<2xindex> { // CHECK: shape.const_shape [7, 2] : tensor<2xindex> %0 = shape.const_shape [1, 2] : tensor<2xindex> %1 = shape.const_shape [7, 1] : tensor<2xindex> %2 = shape.broadcast %0, %1 : tensor<2xindex>, tensor<2xindex> -> tensor<2xindex> return %2 : tensor<2xindex> } // ----- // Basic case including extent tensors. // CHECK-LABEL: @broadcast func.func @broadcast() -> !shape.shape { // CHECK: shape.const_shape [7, 2] : !shape.shape %0 = shape.const_shape [1, 2] : tensor<2xindex> %1 = shape.const_shape [7, 1] : tensor<2xindex> %2 = shape.broadcast %0, %1 : tensor<2xindex>, tensor<2xindex> -> !shape.shape return %2 : !shape.shape } // ----- // Rhs is a scalar. // CHECK-LABEL: func @f func.func @f(%arg0 : !shape.shape) -> !shape.shape { // CHECK: return %arg0 %0 = shape.const_shape [] : !shape.shape %1 = shape.broadcast %arg0, %0 : !shape.shape, !shape.shape -> !shape.shape return %1 : !shape.shape } // ----- // Lhs is a scalar. // CHECK-LABEL: func @f func.func @f(%arg0 : !shape.shape) -> !shape.shape { // CHECK: return %arg0 %0 = shape.const_shape [] : !shape.shape %1 = shape.broadcast %0, %arg0 : !shape.shape, !shape.shape -> !shape.shape return %1 : !shape.shape } // ----- // Lhs is a scalar and rhs is constant. // CHECK-LABEL: func @f func.func @f() -> !shape.shape { // CHECK: %[[CST:.*]] = shape.const_shape [1, 2, 3] : !shape.shape // CHECK: return %[[CST]] %0 = shape.const_shape [] : !shape.shape %1 = shape.const_shape [1, 2, 3] : !shape.shape %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape return %2 : !shape.shape } // ----- // All but one operands are known empty shapes. // CHECK-LABEL: @all_but_one_empty // CHECK-SAME: (%[[ARG:.*]]: !shape.shape) func.func @all_but_one_empty(%arg0 : !shape.shape) -> !shape.shape { // CHECK: return %[[ARG]] %0 = shape.const_shape [] : !shape.shape %1 = shape.const_shape [] : tensor<0xindex> %2 = shape.broadcast %0, %arg0, %1, %0 : !shape.shape, !shape.shape, tensor<0xindex>, !shape.shape -> !shape.shape return %2 : !shape.shape } // ----- // Partial folding. // CHECK-LABEL: @partial_folding // CHECK-SAME: (%[[ARG:.*]]: !shape.shape) func.func @partial_folding(%arg0 : !shape.shape) -> !shape.shape { // CHECK: %[[CST_SHAPE:.*]] = shape.const_shape [1, 2, 3] : tensor<3xindex> // CHECK: %[[RESULT:.*]] = shape.broadcast %[[ARG]], %[[CST_SHAPE]] : !shape.shape, tensor<3xindex> -> !shape.shape // CHECK: return %[[RESULT]] %0 = shape.const_shape [2, 1] : !shape.shape %1 = shape.const_shape [1, 2, 3] : tensor<3xindex> %2 = shape.broadcast %0, %arg0, %1, %0 : !shape.shape, !shape.shape, tensor<3xindex>, !shape.shape -> !shape.shape return %2 : !shape.shape } // ----- // Incompatible shapes. No folding. // CHECK-LABEL: func @f func.func @f() -> !shape.shape { // CHECK: shape.broadcast %0 = shape.const_shape [2] : !shape.shape %1 = shape.const_shape [7] : !shape.shape %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape return %2 : !shape.shape } // ----- // Dead code // CHECK-LABEL: @broadcast func.func @broadcast(%arg0 : !shape.shape, %arg1 : !shape.shape) { // CHECK-NEXT: return %0 = shape.broadcast %arg0, %arg1 : !shape.shape, !shape.shape -> !shape.shape return } // ----- // Basic case. // CHECK-LABEL: func @f func.func @f() -> !shape.shape { // CHECK: shape.const_shape [0, 1, 2, 3] : !shape.shape %lhs = shape.const_shape [0, 1] : !shape.shape %rhs = shape.const_shape [2, 3] : !shape.shape %0 = shape.concat %lhs, %rhs : !shape.shape , !shape.shape -> !shape.shape return %0 : !shape.shape } // ----- // Basic case. // CHECK-LABEL: func @f func.func @f() -> tensor<4xindex> { // CHECK: shape.const_shape [0, 1, 2, 3] : tensor<4xindex> %lhs = shape.const_shape [0, 1] : tensor<2xindex> %rhs = shape.const_shape [2, 3] : tensor<2xindex> %0 = shape.concat %lhs, %rhs : tensor<2xindex>, tensor<2xindex> -> tensor<4xindex> return %0 : tensor<4xindex> } // ----- // Basic case. // CHECK-LABEL: func @f func.func @f() -> tensor<2xindex> { // CHECK: shape.const_shape [0, 1] : tensor<2xindex> %cs = shape.const_shape [0, 1] : !shape.shape %0 = shape.to_extent_tensor %cs : !shape.shape -> tensor<2xindex> return %0 : tensor<2xindex> } // ----- // Basic case. // CHECK-LABEL: func @f() func.func @f() -> !shape.shape { // CHECK: shape.const_shape [3, 5, 11] : !shape.shape %e0 = arith.constant 3 : index %e1 = arith.constant 5 : index %e2 = arith.constant 11 : index %ret = shape.from_extents %e0, %e1, %e2 : index, index, index return %ret : !shape.shape } // ----- // fold_const_size // CHECK-LABEL: func @fold_const_size() func.func @fold_const_size() -> !shape.shape { // CHECK: shape.const_shape [3, 5] : !shape.shape %e0 = shape.const_size 3 %e1 = shape.const_size 5 %ret = shape.from_extents %e0, %e1 : !shape.size, !shape.size return %ret : !shape.shape } // ----- // CHECK-LABEL: func @no_fold func.func @no_fold(%arg0: index) -> !shape.shape { // CHECK-NOT: shape.const_shape %e0 = arith.constant 3 : index %ret = shape.from_extents %e0, %arg0 : index, index return %ret : !shape.shape } // ----- // Cast constant size to index and fold it away. // CHECK-LABEL: func @const_size_to_index func.func @const_size_to_index() -> index { // CHECK-NOT: shape.index_cast %cs = shape.const_size 123 // CHECK: arith.constant 123 : index %ci = shape.size_to_index %cs : !shape.size return %ci : index } // ----- // Cast constant index to size and fold it away. // CHECK-LABEL: func @const_index_to_size func.func @const_index_to_size() -> !shape.size { // CHECK-NOT: arith.index_cast %ci = arith.constant 123 : index // CHECK: shape.const_size 123 %cs = shape.index_to_size %ci return %cs : !shape.size } // ----- // Cast constant index to size, then back, and fold it away. // CHECK-LABEL: func @const_index_to_size_to_index func.func @const_index_to_size_to_index() -> index { // CHECK-NOT: shape.index_cast %ci0 = arith.constant 123 : index %cs0 = shape.index_to_size %ci0 // CHECK: %[[CI:.*]] = arith.constant 123 : index // CHECK-NEXT: return %[[CI]] : index %ci1 = shape.size_to_index %cs0 : !shape.size return %ci1 : index } // ----- // No folding. // CHECK-LABEL: func @nonfoldable_size_to_index func.func @nonfoldable_size_to_index(%cs : !shape.size) -> index { // CHECK: shape.size_to_index %ci = shape.size_to_index %cs : !shape.size return %ci : index } // ----- // No folding. // CHECK-LABEL: func @nonfoldable_index_to_size func.func @nonfoldable_index_to_size(%ci : index) -> !shape.size { // CHECK: shape.index_to_size %cs = shape.index_to_size %ci return %cs : !shape.size } // ----- // Fold number of elements computation. // CHECK-LABEL: func @num_elements func.func @num_elements() -> !shape.size { // CHECK-NOT: shape.const_shape %shape = shape.const_shape [4, 5, 6] : !shape.shape // CHECK-NOT: shape.num_elements %num_elements = shape.num_elements %shape : !shape.shape -> !shape.size // CHECK: %[[NUM:.*]] = shape.const_size 120 // CHECK-NEXT: return %[[NUM]] : !shape.size return %num_elements : !shape.size } // ----- // No folding. // CHECK-LABEL: func @nonfoldable_num_elements func.func @nonfoldable_num_elements(%shape : !shape.shape) -> !shape.size { // CHECK-NOT: shape.const_{{.*}} %num_elements = shape.num_elements %shape : !shape.shape -> !shape.size return %num_elements : !shape.size } // ----- // Basic folding. // CHECK-LABEL: func @basic func.func @basic() -> index { // CHECK: constant 2 : index %0 = shape.const_shape [0, 1, 2] : tensor<3xindex> %c2 = arith.constant 2 : index %1 = shape.get_extent %0, %c2 : tensor<3xindex>, index -> index return %1 : index } // ----- // Should not fold. // CHECK-LABEL: func @out_of_bounds func.func @out_of_bounds() -> index { // CHECK: shape.const_shape // CHECK: shape.get_extent %0 = shape.const_shape [0, 1, 2] : tensor<3xindex> %c3 = arith.constant 3 : index %1 = shape.get_extent %0, %c3 : tensor<3xindex>, index -> index return %1 : index } // ----- // Should not fold. // CHECK-LABEL: func @not_const func.func @not_const(%arg0: tensor) -> index { // CHECK: shape.get_extent %c3 = arith.constant 3 : index %0 = shape.get_extent %arg0, %c3 : tensor, index -> index return %0 : index } // ----- // Basic folding. // CHECK-LABEL: func @basic func.func @basic() -> !shape.size { // CHECK: shape.const_size 2 %0 = shape.const_shape [0, 1, 2] : !shape.shape %c2 = shape.const_size 2 %1 = shape.get_extent %0, %c2 : !shape.shape, !shape.size -> !shape.size return %1 : !shape.size } // ----- // Should not fold. // CHECK-LABEL: func @out_of_bounds func.func @out_of_bounds() -> !shape.size { // CHECK: shape.const_shape // CHECK: shape.get_extent %0 = shape.const_shape [0, 1, 2] : !shape.shape %c3 = shape.const_size 3 %1 = shape.get_extent %0, %c3 : !shape.shape, !shape.size -> !shape.size return %1 : !shape.size } // ----- // Should not fold. // CHECK-LABEL: func @not_const func.func @not_const(%arg0 : !shape.shape) -> !shape.size { // CHECK: shape.get_extent %c3 = shape.const_size 3 %0 = shape.get_extent %arg0, %c3 : !shape.shape, !shape.size -> !shape.size return %0 : !shape.size } // ----- // cstr_eq with non-constant but known equal shapes can be removed. // CHECK-LABEL: func @f func.func @f(%arg0 : !shape.shape) { // CHECK-NEXT: shape.const_witness true // CHECK-NEXT: consume.witness // CHECK-NEXT: return %0 = shape.cstr_eq %arg0, %arg0, %arg0 : !shape.shape, !shape.shape, !shape.shape "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // cstr_eq with equal const_shapes can be folded // CHECK-LABEL: func @f func.func @f() { // CHECK-NEXT: shape.const_witness true // CHECK-NEXT: consume.witness // CHECK-NEXT: return %cs0 = shape.const_shape [0, 1] : !shape.shape %cs1 = shape.const_shape [0, 1] : !shape.shape %cs2 = shape.const_shape [0, 1] : !shape.shape %0 = shape.cstr_eq %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // cstr_eq with unequal const_shapes cannot be folded // CHECK-LABEL: func @f func.func @f() { // CHECK-NEXT: shape.const_shape // CHECK-NEXT: shape.const_shape // CHECK-NEXT: shape.cstr_eq // CHECK-NEXT: consume.witness // CHECK-NEXT: return %cs0 = shape.const_shape [0, 1] : !shape.shape %cs1 = shape.const_shape [3, 1] : !shape.shape %0 = shape.cstr_eq %cs0, %cs1 : !shape.shape, !shape.shape "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // cstr_eq without const_shapes cannot be folded // CHECK-LABEL: func @f func.func @f(%arg0: !shape.shape, %arg1: !shape.shape) { // CHECK-NEXT: shape.cstr_eq // CHECK-NEXT: consume.witness // CHECK-NEXT: return %0 = shape.cstr_eq %arg0, %arg1 : !shape.shape, !shape.shape "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // cstr_require with constant can be folded // CHECK-LABEL: func @cstr_require_fold func.func @cstr_require_fold() { // CHECK-NEXT: shape.const_witness true // CHECK-NEXT: consume.witness // CHECK-NEXT: return %true = arith.constant true %0 = shape.cstr_require %true, "msg" "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // cstr_require without constant cannot be folded // CHECK-LABEL: func @cstr_require_no_fold func.func @cstr_require_no_fold(%arg0: i1) { // CHECK-NEXT: shape.cstr_require // CHECK-NEXT: consume.witness // CHECK-NEXT: return %0 = shape.cstr_require %arg0, "msg" "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // merge assuming_all operations // CHECK-LABEL: func @f func.func @f() { // CHECK-NEXT: %[[W0:.*]] = "test.source" // CHECK-NEXT: %[[W1:.*]] = "test.source" // CHECK-NEXT: %[[W2:.*]] = "test.source" // CHECK-NEXT: shape.assuming_all %[[W0]], %[[W1]], %[[W2]] // CHECK-NEXT: consume.witness // CHECK-NEXT: return %0 = "test.source"() : () -> !shape.witness %1 = "test.source"() : () -> !shape.witness %2 = "test.source"() : () -> !shape.witness %3 = shape.assuming_all %0, %1 %4 = shape.assuming_all %3, %2 "consume.witness"(%4) : (!shape.witness) -> () return } // ----- // `assuming_all` with all `cstr_eq` and shared operands can be collapsed. // CHECK-LABEL: func @assuming_all_to_cstr_eq // CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: tensor, %[[C:.*]]: tensor<3xindex>) func.func @assuming_all_to_cstr_eq(%a : !shape.shape, %b : tensor, %c : tensor<3xindex>) -> !shape.witness { // CHECK: %[[RESULT:.*]] = shape.cstr_eq %[[A]], %[[B]], %[[B]], %[[C]] // CHECK: return %[[RESULT]] %0 = shape.cstr_eq %a, %b : !shape.shape, tensor %1 = shape.cstr_eq %b, %c : tensor, tensor<3xindex> %2 = shape.assuming_all %0, %1 return %2 : !shape.witness } // ----- // `assuming_all` with duplicate operands. // CHECK-LABEL: func @assuming_all_duplicate_operands // CHECK-SAME: (%[[ARG0:.*]]: tensor, %[[ARG1:.*]]: tensor) func.func @assuming_all_duplicate_operands(%arg0 : tensor, %arg1 : tensor) -> !shape.witness { // CHECK: %[[RES:.*]] = shape.cstr_broadcastable %[[ARG0]], %[[ARG1]] // CHECK: return %[[RES]] %0 = shape.cstr_broadcastable %arg0, %arg1 : tensor, tensor %1 = shape.assuming_all %0, %0, %0 return %1 : !shape.witness } // ----- // `assuming_all` with all `cstr_eq` but disjoint operands cannot be collapsed. // CHECK-LABEL: func @assuming_all_to_cstr_eq // CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: tensor, %[[C:.*]]: tensor<3xindex>, %[[D:.*]]: tensor<3xindex>) func.func @assuming_all_to_cstr_eq(%a : !shape.shape, %b : tensor, %c : tensor<3xindex>, %d : tensor<3xindex>) -> !shape.witness { // CHECK: %[[EQ0:.*]] = shape.cstr_eq %[[A]], %[[B]] // CHECK: %[[EQ1:.*]] = shape.cstr_eq %[[C]], %[[D]] // CHECK: %[[RESULT:.*]] = shape.assuming_all %[[EQ0]], %[[EQ1]] // CHECK: return %[[RESULT]] %0 = shape.cstr_eq %a, %b : !shape.shape, tensor %1 = shape.cstr_eq %c, %d : tensor<3xindex>, tensor<3xindex> %2 = shape.assuming_all %0, %1 return %2 : !shape.witness } // ----- // assuming_all with known passing witnesses can be folded // CHECK-LABEL: func @f func.func @f() { // CHECK-NEXT: shape.const_witness true // CHECK-NEXT: consume.witness // CHECK-NEXT: return %0 = shape.const_witness true %1 = shape.const_witness true %2 = shape.const_witness true %3 = shape.assuming_all %0, %1, %2 "consume.witness"(%3) : (!shape.witness) -> () return } // ----- // assuming_all should not be removed if more than one witness is not // statically passing // // Additionally check that the attribute is moved to the end as this op is // commutative. // CHECK-LABEL: func @f func.func @f() { // CHECK-NEXT: %[[UNKNOWN1:.*]] = "test.source" // CHECK-NEXT: %[[UNKNOWN2:.*]] = "test.source" // CHECK-NEXT: shape.assuming_all %[[UNKNOWN1]], %[[UNKNOWN2]] // CHECK-NEXT: consume.witness // CHECK-NEXT: return %0 = shape.const_witness true %1 = "test.source"() : () -> !shape.witness %2 = "test.source"() : () -> !shape.witness %3 = shape.assuming_all %0, %1, %2 "consume.witness"(%3) : (!shape.witness) -> () return } // ----- // merge cstr_broadcastable operations // // CHECK-LABEL: func @f // CHECK: %[[ARG0:[a-z0-9]*]]: !shape.shape // CHECK-SAME: %[[ARG1:[a-z0-9]*]]: !shape.shape // CHECK-SAME: %[[ARG2:[a-z0-9]*]]: !shape.shape func.func @f(%arg0 : !shape.shape, %arg1 : !shape.shape, %arg2 : !shape.shape) { // CHECK-NEXT: %[[W:.*]] = shape.cstr_broadcastable %[[ARG0]], %[[ARG1]], %[[ARG2]] // CHECK-NEXT: "consume.witness"(%[[W]]) // CHECK-NEXT: return %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape %1 = shape.cstr_broadcastable %arg0, %arg1, %arg2 : !shape.shape, !shape.shape, !shape.shape %2 = shape.assuming_all %0, %1 "consume.witness"(%2) : (!shape.witness) -> () return } // ----- // do not merge cstr_broadcastable operations // // CHECK-LABEL: func @f // CHECK: %[[ARG0:[a-z0-9]*]]: !shape.shape // CHECK-SAME: %[[ARG1:[a-z0-9]*]]: !shape.shape // CHECK-SAME: %[[ARG2:[a-z0-9]*]]: !shape.shape func.func @f(%arg0 : !shape.shape, %arg1 : !shape.shape, %arg2 : !shape.shape) { // CHECK-NEXT: %[[W0:.*]] = shape.cstr_broadcastable %[[ARG0]], %[[ARG1]] // CHECK-NEXT: %[[W1:.*]] = shape.cstr_broadcastable %[[ARG1]], %[[ARG2]] // CHECK-NEXT: %[[W2:.*]] = shape.assuming_all %[[W0]], %[[W1]] // CHECK-NEXT: "consume.witness"(%[[W2]]) // CHECK-NEXT: return %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape %1 = shape.cstr_broadcastable %arg1, %arg2 : !shape.shape, !shape.shape %2 = shape.assuming_all %0, %1 "consume.witness"(%2) : (!shape.witness) -> () return } // ----- // any can be replaced with a constant input if it has one. // CHECK-LABEL: func @f func.func @f(%arg : !shape.shape) -> !shape.shape { // CHECK-NEXT: %[[CS:.*]] = shape.const_shape // CHECK-NEXT: return %[[CS]] %0 = shape.const_shape [2, 3, 4] : !shape.shape %1 = shape.any %0, %arg : !shape.shape, !shape.shape -> !shape.shape return %1 : !shape.shape } // ----- // any can be replaced with a constant input if it has one. // CHECK-LABEL: func @f func.func @f(%arg : tensor) -> tensor<3xindex> { // CHECK-NEXT: %[[CS:.*]] = shape.const_shape [2, 3, 4] : tensor<3xindex> // CHECK-NEXT: return %[[CS]] : tensor<3xindex> %0 = shape.const_shape [2, 3, 4] : tensor<3xindex> %1 = shape.any %0, %arg : tensor<3xindex>, tensor -> tensor<3xindex> return %1 : tensor<3xindex> } // ----- // Folding of any with partially constant operands is not yet implemented. // CHECK-LABEL: func @f func.func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> !shape.shape { // CHECK-NEXT: %[[CS:.*]] = shape.any // CHECK-NEXT: return %[[CS]] %1 = shape.any %arg0, %arg1 : !shape.shape, !shape.shape -> !shape.shape return %1 : !shape.shape } // ----- // assuming with a known passing witness can be removed // CHECK-LABEL: func @f func.func @f() { // CHECK-NEXT: source // CHECK-NEXT: sink // CHECK-NEXT: return %0 = shape.const_witness true %1 = shape.assuming %0 -> index { %2 = "test.source"() : () -> (index) shape.assuming_yield %2 : index } "test.sink"(%1) : (index) -> () return } // ----- // assuming without a known passing passing witness cannot be removed // CHECK-LABEL: func @f func.func @f() { // CHECK-NEXT: test.source // CHECK-NEXT: shape.assuming // CHECK-NEXT: test.source // CHECK-NEXT: shape.assuming_yield // CHECK-NEXT: } // CHECK-NEXT: test.sink // CHECK-NEXT: return %0 = "test.source"() : () -> (!shape.witness) %1 = shape.assuming %0 -> index { %2 = "test.source"() : () -> (index) shape.assuming_yield %2 : index } "test.sink"(%1) : (index) -> () return } // ----- // Remove unused results from assuming ops. // CHECK-LABEL: func @unused_assuming_results func.func @unused_assuming_results() { // CHECK: %[[ASSUMING_RESULT:.*]] = shape.assuming %0 -> (f32) { // CHECK: %{{.*}} = "produce.redundant" // CHECK: %[[MEANINGFUL:.*]] = "produce.meaningful" // CHECK: shape.assuming_yield %[[MEANINGFUL]] : f32 // CHECK: } // CHECK: "use"(%[[ASSUMING_RESULT]]) %0 = "test.source"() : () -> (!shape.witness) %1:2 = shape.assuming %0 -> (f32, f32) { %2 = "produce.redundant"() : () -> (f32) %3 = "produce.meaningful"() : () -> (f32) shape.assuming_yield %2, %3 : f32, f32 } "use"(%1#1) : (f32) -> () return } // ----- // Broadcastable with broadcastable constant shapes can be removed. // CHECK-LABEL: func @f func.func @f() { // CHECK-NEXT: shape.const_witness true // CHECK-NEXT: consume.witness // CHECK-NEXT: return %cs0 = shape.const_shape [3, 1] : !shape.shape %cs1 = shape.const_shape [1, 5] : !shape.shape %0 = shape.cstr_broadcastable %cs0, %cs1 : !shape.shape, !shape.shape "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // Empty shape arguments can be removed from broadcastable ops. // CHECK-LABEL: func @f // CHECK-SAME: (%[[ARG0:.*]]: tensor, %[[ARG1:.*]]: tensor, %{{.*}}: tensor<0xindex>) func.func @f(%arg0 : tensor, %arg1 : tensor, %arg2 : tensor<0xindex>) { // CHECK-NOT: const_shape // CHECK: cstr_broadcastable %[[ARG0]], %[[ARG1]] : tensor, tensor %0 = shape.const_shape [] : !shape.shape %1 = shape.cstr_broadcastable %arg0, %arg1, %0, %arg2 : tensor, tensor, !shape.shape, tensor<0xindex> "consume.witness"(%1) : (!shape.witness) -> () return } // ----- // Broadcastable with non-broadcastable constant shapes is always false // CHECK-LABEL: func @static_non_broadcastable func.func @static_non_broadcastable() { // CHECK-NEXT: shape.const_shape // CHECK-NEXT: shape.const_shape // CHECK-NEXT: shape.cstr_broadcastable // CHECK-NEXT: consume.witness // CHECK-NEXT: return %cs0 = shape.const_shape [1, 3] : !shape.shape %cs1 = shape.const_shape [1, 5] : !shape.shape %0 = shape.cstr_broadcastable %cs0, %cs1 : !shape.shape, !shape.shape "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // Broadcastable without guaranteed broadcastable shapes cannot be removed. // CHECK-LABEL: func @f func.func @f(%arg0 : !shape.shape) { // CHECK-NEXT: shape.const_shape // CHECK-NEXT: shape.cstr_broadcastable // CHECK-NEXT: consume.witness // CHECK-NEXT: return %cs0 = shape.const_shape [1, 3] : !shape.shape %0 = shape.cstr_broadcastable %arg0, %cs0 : !shape.shape, !shape.shape "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // Broadcastable with non-constant but known equal shapes can be removed. // CHECK-LABEL: func @f func.func @f(%arg0 : !shape.shape) { // CHECK-NEXT: shape.const_witness true // CHECK-NEXT: consume.witness // CHECK-NEXT: return %0 = shape.cstr_broadcastable %arg0, %arg0 : !shape.shape, !shape.shape "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // Broadcastable canonicalization also works on extent tensors. // CHECK-LABEL: func @broadcastable_on_extent_tensors func.func @broadcastable_on_extent_tensors(%arg : tensor) { // CHECK-NEXT: shape.const_witness true // CHECK-NEXT: consume.witness // CHECK-NEXT: return %0 = shape.cstr_broadcastable %arg, %arg : tensor, tensor "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // Fold ternary broadcastable // CHECK-LABEL: func @f func.func @f() { // CHECK-NEXT: shape.const_witness true // CHECK-NEXT: consume.witness // CHECK-NEXT: return %cs0 = shape.const_shape [8, 1] : !shape.shape %cs1 = shape.const_shape [1, 8] : !shape.shape %cs2 = shape.const_shape [1, 1] : !shape.shape %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // Fold ternary broadcastable with dynamic ranks // CHECK-LABEL: func @f func.func @f() { // CHECK-NEXT: shape.const_witness true // CHECK-NEXT: consume.witness // CHECK-NEXT: return %cs0 = shape.const_shape [8, 1] : !shape.shape %cs1 = shape.const_shape [1, -9223372036854775808] : !shape.shape %0 = shape.cstr_broadcastable %cs0, %cs0, %cs1 : !shape.shape, !shape.shape, !shape.shape "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // One scalar and one non-scalar and one unknown cannot be broadcasted at compile time // CHECK-LABEL: func @f func.func @f() { // CHECK: shape.cstr_broadcastable // CHECK-NEXT: consume.witness // CHECK-NEXT: return %cs0 = shape.const_shape [8, 1] : !shape.shape %cs1 = shape.const_shape [1, 8] : !shape.shape %cs2 = shape.const_shape [1, -9223372036854775808] : !shape.shape %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // One scalar and two unknowns cannot be broadcasted at compile time // CHECK-LABEL: func @f func.func @f() { // CHECK: shape.cstr_broadcastable // CHECK-NEXT: consume.witness // CHECK-NEXT: return %cs0 = shape.const_shape [8, 1] : !shape.shape %cs1 = shape.const_shape [1, -9223372036854775808] : !shape.shape %cs2 = shape.const_shape [8, -9223372036854775808] : !shape.shape %0 = shape.cstr_broadcastable %cs0, %cs1, %cs2 : !shape.shape, !shape.shape, !shape.shape "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // Broadcastable with scalars and a non-scalar can be constant folded // CHECK-LABEL: func @f func.func @f(%arg0 : !shape.shape) { // CHECK-NEXT: shape.const_witness true // CHECK-NEXT: consume.witness // CHECK-NEXT: return %cs0 = shape.const_shape [] : !shape.shape %0 = shape.cstr_broadcastable %cs0, %cs0, %arg0 : !shape.shape, !shape.shape, !shape.shape "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // One scalar and one non-scalar and one unknown cannot be folded. // CHECK-LABEL: func @f func.func @f(%arg0 : !shape.shape) { // CHECK: shape.cstr_broadcastable // CHECK-NEXT: consume.witness // CHECK-NEXT: return %cs0 = shape.const_shape [] : !shape.shape %cs1 = shape.const_shape [2] : !shape.shape %0 = shape.cstr_broadcastable %cs0, %cs1, %arg0 : !shape.shape, !shape.shape, !shape.shape "consume.witness"(%0) : (!shape.witness) -> () return } // ----- // Fold `rank` based on constant shape. // CHECK-LABEL: @fold_rank func.func @fold_rank() -> !shape.size { // CHECK: %[[RESULT:.*]] = shape.const_size 5 // CHECK: return %[[RESULT]] : !shape.size %shape = shape.const_shape [3, 4, 5, 6, 7] : !shape.shape %rank = shape.rank %shape : !shape.shape -> !shape.size return %rank : !shape.size } // ----- // Do not fold `rank` if shape is dynamic. // CHECK-LABEL: @dont_fold_rank // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape) -> !shape.size func.func @dont_fold_rank(%shape : !shape.shape) -> !shape.size { // CHECK: %[[RESULT:.*]] = shape.rank %[[SHAPE]] // CHECK: return %[[RESULT]] : !shape.size %rank = shape.rank %shape : !shape.shape -> !shape.size return %rank : !shape.size } // ----- // Fold `rank` based on constant extent tensor. // CHECK-LABEL: @fold_rank func.func @fold_rank() -> index { // CHECK: %[[RESULT:.*]] = arith.constant 5 : index // CHECK: return %[[RESULT]] : index %shape = shape.const_shape [3, 4, 5, 6, 7] : tensor<5xindex> %rank = shape.rank %shape : tensor<5xindex> -> index return %rank : index } // ----- // Do not fold `rank` for non-constant extent tensors. // CHECK-LABEL: @dont_fold_rank // CHECK-SAME: (%[[SHAPE:.*]]: tensor) -> index func.func @dont_fold_rank(%shape : tensor) -> index { // CHECK: %[[RESULT:.*]] = shape.rank %[[SHAPE]] : tensor -> index // CHECK: return %[[RESULT]] : index %rank = shape.rank %shape : tensor -> index return %rank : index } // ----- // Canonicalize `rank` when shape is derived from ranked tensor. // CHECK-LABEL: @canonicalize_rank func.func @canonicalize_rank(%arg : tensor<1x2x?xf32>) -> index { // CHECK: %[[RESULT:.*]] = arith.constant 3 : index // CHECK: return %[[RESULT]] : index %shape = shape.shape_of %arg : tensor<1x2x?xf32> -> tensor %rank = shape.rank %shape : tensor -> index return %rank : index } // ----- // Canonicalize `rank` when shape is derived from ranked tensor. // CHECK-LABEL: @canonicalize_rank func.func @canonicalize_rank_size(%arg : tensor<1x2x?xf32>) -> !shape.size { // CHECK: %[[RESULT:.*]] = shape.const_size 3 // CHECK: return %[[RESULT]] : !shape.size %shape = shape.shape_of %arg : tensor<1x2x?xf32> -> !shape.shape %rank = shape.rank %shape : !shape.shape -> !shape.size return %rank : !shape.size } // ----- // Do not canonicalize `rank` when shape is derived from unranked tensor. // CHECK-LABEL: @dont_canonicalize_rank // CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>) -> index func.func @dont_canonicalize_rank(%arg : tensor<*xf32>) -> index { // CHECK: %[[SHAPE:.*]] = shape.shape_of %[[ARG]] : tensor<*xf32> -> tensor // CHECK: %[[SIZE:.*]] = shape.rank %[[SHAPE]] // CHECK: return %[[SIZE]] : index %shape = shape.shape_of %arg : tensor<*xf32> -> tensor %rank = shape.rank %shape : tensor -> index return %rank : index } // ----- // Canonicalize redundant conversion from `index` to `size` and back. // CHECK-LABEL: @index_to_size_to_index // CHECK-SAME: (%[[IDX:.*]]: index) -> index func.func @index_to_size_to_index(%index : index) -> index { // CHECK: return %[[IDX]] : index %size = shape.index_to_size %index %result = shape.size_to_index %size : !shape.size return %result : index } // ----- // Canonicalize redundant conversion from `size` to `index` and back. // CHECK-LABEL: @size_to_index_to_size // CHECK-SAME: (%[[SIZE:.*]]: !shape.size) -> !shape.size func.func @size_to_index_to_size(%size : !shape.size) -> !shape.size { // CHECK: return %[[SIZE]] : !shape.size %idx = shape.size_to_index %size : !shape.size %result = shape.index_to_size %idx return %result : !shape.size } // ----- // Canonicalize scalar cstr_broadcastable checks // CHECK-LABEL: @cstr_broadcastable_scalar func.func @cstr_broadcastable_scalar(%arg0 : tensor) { // CHECK-NEXT: shape.const_witness true // CHECK-NEXT: consume.witness // CHECK-NEXT: return %0 = shape.const_shape [] : !shape.shape %1 = shape.shape_of %arg0 : tensor -> tensor %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, tensor "consume.witness"(%2) : (!shape.witness) -> () return } // ----- // Do not canonicalize cstr_broadcastable checks with 2 unknowns // CHECK-LABEL: @cstr_broadcastable_unknown func.func @cstr_broadcastable_unknown(%arg0 : tensor, %arg1 : tensor) { // CHECK-NEXT: shape.shape_of %arg0 // CHECK-NEXT: shape.shape_of %arg1 // CHECK-NEXT: shape.cstr_broadcastable // CHECK-NEXT: consume.witness // CHECK-NEXT: return %0 = shape.shape_of %arg0 : tensor -> tensor %1 = shape.shape_of %arg1 : tensor -> tensor %2 = shape.cstr_broadcastable %0, %1 : tensor, tensor "consume.witness"(%2) : (!shape.witness) -> () return } // ----- // Scalars are safe to broadcast to unranked sizes. // CHECK-LABEL: @cstr_broadcastable_scalar_unranked func.func @cstr_broadcastable_scalar_unranked(%arg0 : tensor<*xf32>, %arg1 : tensor) { // CHECK-NEXT: shape.const_witness true // CHECK-NEXT: consume.witness // CHECK-NEXT: return %0 = shape.shape_of %arg1 : tensor -> tensor %1 = shape.shape_of %arg0 : tensor<*xf32> -> tensor %2 = shape.cstr_broadcastable %0, %1 : tensor, tensor "consume.witness"(%2) : (!shape.witness) -> () return } // ----- // Fold `shape_eq` for equal and constant shapes. // CHECK-LABEL: @shape_eq_fold_1 func.func @shape_eq_fold_1() -> i1 { // CHECK: %[[RESULT:.*]] = arith.constant true // CHECK: return %[[RESULT]] : i1 %a = shape.const_shape [1, 2, 3] : !shape.shape %b = shape.const_shape [1, 2, 3] : tensor<3xindex> %c = shape.const_shape [1, 2, 3] : tensor<3xindex> %result = shape.shape_eq %a, %b, %c : !shape.shape, tensor<3xindex>, tensor<3xindex> return %result : i1 } // ----- // Fold `shape_eq` for different but constant shapes of same length. // CHECK-LABEL: @shape_eq_fold_0 func.func @shape_eq_fold_0() -> i1 { // CHECK: %[[RESULT:.*]] = arith.constant false // CHECK: return %[[RESULT]] : i1 %a = shape.const_shape [1, 2, 3] : tensor<3xindex> %b = shape.const_shape [4, 5, 6] : tensor<3xindex> %c = shape.const_shape [4, 5, 6] : tensor<3xindex> %result = shape.shape_eq %a, %b, %c : tensor<3xindex>, tensor<3xindex>, tensor<3xindex> return %result : i1 } // ----- // Fold `shape_eq` for different but constant shapes of different length. // CHECK-LABEL: @shape_eq_fold_0 func.func @shape_eq_fold_0() -> i1 { // CHECK: %[[RESULT:.*]] = arith.constant false // CHECK: return %[[RESULT]] : i1 %a = shape.const_shape [1, 2, 3, 4, 5, 6] : !shape.shape %b = shape.const_shape [1, 2, 3] : !shape.shape %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape return %result : i1 } // ----- // Do not fold `shape_eq` for non-constant different shapes. // CHECK-LABEL: @shape_eq_do_not_fold // CHECK-SAME: (%[[A:.*]]: !shape.shape) -> i1 func.func @shape_eq_do_not_fold(%a : !shape.shape) -> i1 { // CHECK: %[[B:.*]] = shape.const_shape [4, 5, 6] // CHECK: %[[RESULT:.*]] = shape.shape_eq %[[A]], %[[B]] : !shape.shape, !shape.shape // CHECK: return %[[RESULT]] : i1 %b = shape.const_shape [4, 5, 6] : !shape.shape %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape return %result : i1 } // ----- // Fold `add` for constant sizes. // CHECK-LABEL: @fold_add_size func.func @fold_add_size() -> !shape.size { // CHECK: %[[RESULT:.*]] = shape.const_size 5 // CHECK: return %[[RESULT]] : !shape.size %c2 = shape.const_size 2 %c3 = shape.const_size 3 %result = shape.add %c2, %c3 : !shape.size, !shape.size -> !shape.size return %result : !shape.size } // ----- // Fold `mul` for constant sizes. // CHECK-LABEL: @fold_mul_size func.func @fold_mul_size() -> !shape.size { // CHECK: %[[RESULT:.*]] = shape.const_size 6 // CHECK: return %[[RESULT]] : !shape.size %c2 = shape.const_size 2 %c3 = shape.const_size 3 %result = shape.mul %c2, %c3 : !shape.size, !shape.size -> !shape.size return %result : !shape.size } // ----- // Fold `mul` for constant indices. // CHECK-LABEL: @fold_mul_index func.func @fold_mul_index() -> index { // CHECK: %[[RESULT:.*]] = arith.constant 6 : index // CHECK: return %[[RESULT]] : index %c2 = arith.constant 2 : index %c3 = arith.constant 3 : index %result = shape.mul %c2, %c3 : index, index -> index return %result : index } // ----- // Fold `mul` for mixed constants. // CHECK-LABEL: @fold_mul_mixed func.func @fold_mul_mixed() -> !shape.size { // CHECK: %[[RESULT:.*]] = shape.const_size 6 // CHECK: return %[[RESULT]] : !shape.size %c2 = shape.const_size 2 %c3 = arith.constant 3 : index %result = shape.mul %c2, %c3 : !shape.size, index -> !shape.size return %result : !shape.size } // ----- // Fold `div` for constant sizes. // CHECK-LABEL: @fold_div_size func.func @fold_div_size() -> !shape.size { // CHECK: %[[RESULT:.*]] = shape.const_size 3 // CHECK: return %[[RESULT]] : !shape.size %c2 = shape.const_size 10 %c3 = shape.const_size 3 %result = shape.div %c2, %c3 : !shape.size, !shape.size -> !shape.size return %result : !shape.size } // ----- // Fold `div` for constant indices. // CHECK-LABEL: @fold_div_index func.func @fold_div_index() -> index { // CHECK: %[[RESULT:.*]] = arith.constant 2 : index // CHECK: return %[[RESULT]] : index %c2 = arith.constant 10 : index %c3 = arith.constant 4 : index %result = shape.div %c2, %c3 : index, index -> index return %result : index } // ----- // Fold `div` for constant indices and lhs is negative. // CHECK-LABEL: @fold_div_index_neg_lhs func.func @fold_div_index_neg_lhs() -> index { // CHECK: %[[RESULT:.*]] = arith.constant -3 : index // CHECK: return %[[RESULT]] : index %c2 = arith.constant -10 : index %c3 = arith.constant 4 : index %result = shape.div %c2, %c3 : index, index -> index return %result : index } // ----- // Fold `div` for constant indices and rhs is negative. // CHECK-LABEL: @fold_div_index_neg_rhs func.func @fold_div_index_neg_rhs() -> index { // CHECK: %[[RESULT:.*]] = arith.constant -3 : index // CHECK: return %[[RESULT]] : index %c2 = arith.constant 10 : index %c3 = arith.constant -4 : index %result = shape.div %c2, %c3 : index, index -> index return %result : index } // ----- // Fold `div` for mixed constants. // CHECK-LABEL: @fold_div_mixed func.func @fold_div_mixed() -> !shape.size { // CHECK: %[[RESULT:.*]] = shape.const_size 4 // CHECK: return %[[RESULT]] : !shape.size %c2 = shape.const_size 12 %c3 = arith.constant 3 : index %result = shape.div %c2, %c3 : !shape.size, index -> !shape.size return %result : !shape.size } // ----- // Fold index_cast when already on index. // CHECK-LABEL: @fold_index_cast_on_index func.func @fold_index_cast_on_index(%arg: index) -> index { // CHECK-NOT: size_to_index %0 = shape.size_to_index %arg : index return %0 : index } // ----- // Fold to_extent_tensor when already on tensor. // CHECK-LABEL: @fold_to_extent_tensor_on_tensor func.func @fold_to_extent_tensor_on_tensor(%arg: tensor) -> tensor { // CHECK-NOT: to_extent_tensor %0 = shape.to_extent_tensor %arg : tensor -> tensor return %0 : tensor } // ----- // Fold assuming_all with a single input // CHECK-LABEL: @fold_assuming_all_single_element func.func @fold_assuming_all_single_element(%arg: tensor) { // CHECK-NOT: assuming_all %0 = "test.source"() : () -> (!shape.witness) %1 = shape.assuming_all %0 "consume.witness"(%1) : (!shape.witness) -> () return } // ----- // Verify that tensor.cast folding uses the correct type // CHECK-LABEL: @fold_tensor.cast_of_const_shape_returned func.func @fold_tensor.cast_of_const_shape_returned(%arg: i1) -> tensor<1xindex> { // CHECK: shape.const_shape [2] : tensor<1xindex> // CHECK-NOT: tensor.cast %0 = shape.const_shape [2] : tensor<1xindex> %1 = tensor.cast %0 : tensor<1xindex> to tensor<1xindex> return %1 : tensor<1xindex> } // ----- // CHECK-LABEL: @dont_fold_tensor.cast_of_const_shape_returned_dynamic func.func @dont_fold_tensor.cast_of_const_shape_returned_dynamic(%arg: i1) -> tensor { // CHECK: %[[CONST_SHAPE:.*]] = shape.const_shape [2] : tensor<1xindex> // CHECK: tensor.cast %[[CONST_SHAPE]] : tensor<1xindex> to tensor %0 = shape.const_shape [2] : tensor<1xindex> %1 = tensor.cast %0 : tensor<1xindex> to tensor return %1 : tensor } // ----- // CHECK-LABEL: @is_broadcastable_on_same_shape func.func @is_broadcastable_on_same_shape(%shape : !shape.shape) -> i1 { // CHECK-NOT: is_broadcastable // CHECK: %[[RES:.*]] = arith.constant true // CHECK: return %[[RES]] %0 = shape.is_broadcastable %shape, %shape, %shape : !shape.shape, !shape.shape, !shape.shape return %0 : i1 } // ----- // CHECK-LABEL: @is_broadcastable_on_duplicate_shapes // CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: !shape.shape) func.func @is_broadcastable_on_duplicate_shapes(%a : !shape.shape, %b : !shape.shape) -> i1 { // CHECK: %[[RES:.*]] = shape.is_broadcastable %[[A]], %[[B]] : // CHECK: return %[[RES]] %0 = shape.is_broadcastable %a, %b, %a, %a, %a, %b : !shape.shape, !shape.shape, !shape.shape, !shape.shape, !shape.shape, !shape.shape return %0 : i1 } // ----- // CHECK-LABEL: @cstr_broadcastable_on_duplicate_shapes // CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: !shape.shape) func.func @cstr_broadcastable_on_duplicate_shapes(%a : !shape.shape, %b : !shape.shape) -> !shape.witness { // CHECK: %[[RES:.*]] = shape.cstr_broadcastable %[[A]], %[[B]] : // CHECK: return %[[RES]] %0 = shape.cstr_broadcastable %a, %b, %a, %a, %a, %b : !shape.shape, !shape.shape, !shape.shape, !shape.shape, !shape.shape, !shape.shape return %0 : !shape.witness } // ----- // CHECK-LABEL: @broadcast_on_same_shape // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape) func.func @broadcast_on_same_shape(%shape : !shape.shape) -> !shape.shape { // CHECK-NOT: broadcast // CHECK: return %[[SHAPE]] %0 = shape.broadcast %shape, %shape, %shape : !shape.shape, !shape.shape, !shape.shape -> !shape.shape return %0 : !shape.shape } // ----- // CHECK-LABEL: @broadcast_on_duplicate_shapes // CHECK-SAME: (%[[A:.*]]: !shape.shape, %[[B:.*]]: !shape.shape) func.func @broadcast_on_duplicate_shapes(%a : !shape.shape, %b : !shape.shape) -> !shape.shape { // CHECK: %[[RES:.*]] = shape.broadcast %[[A]], %[[B]] : // CHECK: return %[[RES]] %0 = shape.broadcast %a, %b, %a, %a, %a, %b : !shape.shape, !shape.shape, !shape.shape, !shape.shape, !shape.shape, !shape.shape -> !shape.shape return %0 : !shape.shape } // ----- // CHECK-LABEL: @broadcast_on_single_operand // CHECK-SAME: (%[[A:.*]]: tensor) func.func @broadcast_on_single_operand(%a : tensor) { // CHECK-NOT: broadcast // CHECK: "use"(%[[A]]) %0 = shape.broadcast %a : tensor -> tensor "use"(%0) : (tensor) -> () return } // ----- // CHECK-LABEL: @broadcast_as_tensor_cast // CHECK-SAME: (%[[A:.*]]: tensor<3xindex>) func.func @broadcast_as_tensor_cast(%a : tensor<3xindex>) -> tensor { // CHECK: %[[RESULT:.*]] = tensor.cast %[[A]] : tensor<3xindex> to tensor // CHECK: return %[[RESULT]] : tensor %0 = shape.broadcast %a : tensor<3xindex> -> tensor return %0 : tensor } // ----- // CHECK-LABEL: @broadcast_as_from_extent_tensor // CHECK-SAME: (%[[A:.*]]: tensor) func.func @broadcast_as_from_extent_tensor(%a : tensor) -> !shape.shape { // CHECK: %[[RESULT:.*]] = shape.from_extent_tensor %[[A]] : tensor // CHECK: return %[[RESULT]] : !shape.shape %0 = shape.broadcast %a : tensor -> !shape.shape return %0 : !shape.shape } // ----- // CHECK-LABEL: @cast_extent_tensor // CHECK-SAME: (%[[ARG:.*]]: tensor) -> tensor func.func @cast_extent_tensor(%arg : tensor) -> tensor { // CHECK: %[[RESULT:.*]] = shape.shape_of %[[ARG]] : tensor -> tensor // CHECK: return %[[RESULT]] : tensor %0 = shape.shape_of %arg : tensor -> tensor<3xindex> %1 = tensor.cast %0 : tensor<3xindex> to tensor return %1 : tensor } // ----- // CHECK-LABEL: @cast_extent_tensor // CHECK-SAME: (%[[ARG:.*]]: tensor) -> tensor<3xindex> func.func @cast_extent_tensor(%arg : tensor) -> tensor<3xindex> { // CHECK: %[[RESULT:.*]] = shape.shape_of %[[ARG]] : tensor -> tensor<3xindex> // CHECK: return %[[RESULT]] : tensor<3xindex> %0 = shape.shape_of %arg : tensor -> tensor %1 = tensor.cast %0 : tensor to tensor<3xindex> return %1 : tensor<3xindex> } // ----- // CHECK-LABEL: @cast_extent_tensor func.func @cast_extent_tensor(%arg : tensor) -> tensor<3xindex> { // CHECK: tensor.cast %{{.*}} : tensor to tensor<3xindex> %0 = shape.shape_of %arg : tensor -> tensor %1 = tensor.cast %0 : tensor to tensor<3xindex> return %1 : tensor<3xindex> } // ----- // CHECK-LABEL: @cast_extent_tensor func.func @cast_extent_tensor(%arg : tensor<*xf32>) -> tensor<3xindex> { // CHECK: tensor.cast %{{.*}} : tensor to tensor<3xindex> %0 = shape.shape_of %arg : tensor<*xf32> -> tensor %1 = tensor.cast %0 : tensor to tensor<3xindex> return %1 : tensor<3xindex> } // ----- // CHECK-LABEL: max_same_arg // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape) func.func @max_same_arg(%a: !shape.shape) -> !shape.shape { %1 = shape.max %a, %a : !shape.shape, !shape.shape -> !shape.shape // CHECK: return %[[SHAPE]] return %1 : !shape.shape } // ----- // CHECK-LABEL: min_same_arg // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape) func.func @min_same_arg(%a: !shape.shape) -> !shape.shape { %1 = shape.min %a, %a : !shape.shape, !shape.shape -> !shape.shape // CHECK: return %[[SHAPE]] return %1 : !shape.shape } // ----- // CHECK-LABEL: @cstr_broadcastable_folding func.func @cstr_broadcastable_folding(%arg : tensor) { // CHECK: const_witness true %0 = shape.shape_of %arg : tensor -> tensor<2xindex> %1 = shape.const_shape [4] : tensor<1xindex> %2 = shape.cstr_broadcastable %0, %1: tensor<2xindex>, tensor<1xindex> "use"(%2) : (!shape.witness) -> () } // ----- // CHECK-LABEL: @cast_extent_tensor_operands // CHECK-SAME: (%[[ARG0:.*]]: tensor, %[[ARG1:.*]]: tensor<3xindex>) func.func @cast_extent_tensor_operands(%arg0 : tensor, %arg1 : tensor<3xindex>) -> (!shape.witness, tensor) { // CHECK: %[[CAST_ARG0:.*]] = tensor.cast %[[ARG0]] : tensor to tensor<3xindex> // CHECK: %[[WIT:.*]] = shape.cstr_broadcastable %[[CAST_ARG0]], %[[ARG1]] : tensor<3xindex>, tensor<3xindex> // CHECK: %[[UNCAST_RES:.*]] = shape.broadcast %[[CAST_ARG0]], %[[ARG1]] : tensor<3xindex>, tensor<3xindex> -> tensor<3xindex> // CHECK: %[[RES:.*]] = tensor.cast %[[UNCAST_RES]] : tensor<3xindex> to tensor // CHECK: return %[[WIT]], %[[RES]] %0 = tensor.cast %arg0 : tensor to tensor<3xindex> %1 = tensor.cast %arg1 : tensor<3xindex> to tensor %2 = shape.cstr_broadcastable %0, %1 : tensor<3xindex>, tensor %3 = shape.broadcast %0, %1 :tensor<3xindex>, tensor -> tensor return %2, %3 : !shape.witness, tensor } // ----- // CHECK-LABEL: @concretize_broadcast_result_type // CHECK-SAME: (%[[ARG0:.*]]: tensor<2xindex>, %[[ARG1:.*]]: tensor<3xindex>) func.func @concretize_broadcast_result_type(%arg0 : tensor<2xindex>, %arg1 : tensor<3xindex>) -> tensor { // CHECK: %[[CONCR:.*]] = shape.broadcast %[[ARG0]], %[[ARG1]] : tensor<2xindex>, tensor<3xindex> -> tensor<3xindex> // CHECK: %[[RES:.*]] = tensor.cast %[[CONCR]] : tensor<3xindex> to tensor // CHECK: return %[[RES]] %0 = shape.broadcast %arg0, %arg1 : tensor<2xindex>, tensor<3xindex> -> tensor return %0 : tensor } // ----- // CHECK-LABEL: func @extract_shapeof // CHECK-SAME: %[[ARG0:.*]]: tensor func.func @extract_shapeof(%arg0 : tensor) -> index { %c1 = arith.constant 1 : index // CHECK: %[[C1:.*]] = arith.constant 1 %shape = shape.shape_of %arg0 : tensor -> tensor<2xindex> // CHECK: %[[DIM:.*]] = tensor.dim %[[ARG0]], %[[C1]] %result = tensor.extract %shape[%c1] : tensor<2xindex> // CHECK: return %[[DIM]] return %result : index } // ----- // CHECK-LABEL: @add_poison // CHECK: %[[P:.*]] = ub.poison : !shape.siz // CHECK: return %[[P]] func.func @add_poison() -> !shape.size { %1 = shape.const_size 2 %2 = ub.poison : !shape.size %result = shape.add %1, %2 : !shape.size, !shape.size -> !shape.size return %result : !shape.size } // ----- // CHECK-LABEL: func @shape_of_0d( // CHECK-SAME: %[[arg0:.*]]: tensor // CHECK: %[[const:.*]] = shape.const_shape [] : tensor<0xindex> // CHECK: %[[cast:.*]] = tensor.cast %[[const]] : tensor<0xindex> to tensor // CHECK: return %[[cast]] func.func @shape_of_0d(%arg0: tensor) -> tensor { %0 = shape.shape_of %arg0 : tensor -> tensor return %0 : tensor }