1// RUN: mlir-opt <%s | mlir-opt | FileCheck %s 2 3// CHECK-LABEL: func @cast( 4func.func @cast(%arg0: tensor<*xf32>, %arg1 : tensor<4x4xf32>, %arg2: tensor<?x?xf32>) { 5 // CHECK: tensor.cast %arg0 : tensor<*xf32> to tensor<?x?xf32> 6 %0 = tensor.cast %arg0 : tensor<*xf32> to tensor<?x?xf32> 7 // CHECK: tensor.cast %arg1 : tensor<4x4xf32> to tensor<*xf32> 8 %1 = tensor.cast %arg1 : tensor<4x4xf32> to tensor<*xf32> 9 // CHECK: tensor.cast %arg2 : tensor<?x?xf32> to tensor<4x?xf32> 10 %2 = tensor.cast %arg2 : tensor<?x?xf32> to tensor<4x?xf32> 11 // CHECK: tensor.cast %2 : tensor<4x?xf32> to tensor<?x?xf32> 12 %3 = tensor.cast %2 : tensor<4x?xf32> to tensor<?x?xf32> 13 return 14} 15 16// CHECK-LABEL: func @extract( 17// CHECK-SAME: %[[TENSOR:.*]]: tensor<?x?x?xf32>, 18// CHECK-SAME: %[[INDEX:.*]]: index) { 19func.func @extract(%arg0: tensor<?x?x?xf32>, %arg1: index) { 20 // CHECK: tensor.extract %[[TENSOR]][%[[INDEX]], %[[INDEX]], %[[INDEX]]] : tensor<?x?x?xf32> 21 %0 = tensor.extract %arg0[%arg1, %arg1, %arg1] : tensor<?x?x?xf32> 22 return 23} 24 25// CHECK-LABEL: func @insert( 26// CHECK-SAME: %[[SCALAR:.*]]: f32 27// CHECK-SAME: %[[INDEX:.*]]: index 28// CHECK-SAME: %[[DEST1:.*]]: tensor<?x?x?xf32> 29// CHECK-SAME: %[[DEST2:.*]]: tensor<*xf32> 30func.func @insert(%arg0: f32, %arg1: index, %arg2: tensor<?x?x?xf32>, %arg3: tensor<*xf32>) { 31 // CHECK: tensor.insert %[[SCALAR]] into %[[DEST1]][%[[INDEX]], %[[INDEX]], %[[INDEX]]] : tensor<?x?x?xf32> 32 %0 = tensor.insert %arg0 into %arg2[%arg1, %arg1, %arg1] : tensor<?x?x?xf32> 33 // CHECK: tensor.insert %[[SCALAR]] into %[[DEST2]][%[[INDEX]], %[[INDEX]], %[[INDEX]]] : tensor<*xf32> 34 %1 = tensor.insert %arg0 into %arg3[%arg1, %arg1, %arg1] : tensor<*xf32> 35 return 36} 37 38// CHECK-LABEL: func @tensor.from_elements() { 39func.func @tensor.from_elements() { 40 %c0 = "arith.constant"() {value = 0: index} : () -> index 41 // CHECK: tensor.from_elements %c0 : tensor<1xindex> 42 %0 = tensor.from_elements %c0 : tensor<1xindex> 43 44 %c1 = "arith.constant"() {value = 1: index} : () -> index 45 // CHECK: tensor.from_elements %c0, %c1 : tensor<2xindex> 46 %1 = tensor.from_elements %c0, %c1 : tensor<2xindex> 47 48 %c0_f32 = "arith.constant"() {value = 0.0: f32} : () -> f32 49 // CHECK: [[C0_F32:%.*]] = arith.constant 50 // CHECK: tensor.from_elements [[C0_F32]] : tensor<1xf32> 51 %2 = tensor.from_elements %c0_f32 : tensor<1xf32> 52 53 // CHECK: tensor.from_elements : tensor<0xindex> 54 %3 = tensor.from_elements : tensor<0xindex> 55 56 // CHECK: tensor.from_elements %c0, %c1, %c0, %c1, %c0, %c1 : tensor<2x3xindex> 57 %4 = tensor.from_elements %c0, %c1, %c0, %c1, %c0, %c1 : tensor<2x3xindex> 58 59 // CHECK: tensor.from_elements %c0 : tensor<index> 60 %5 = tensor.from_elements %c0 : tensor<index> 61 return 62} 63 64// CHECK-LABEL: @tensor.generate 65func.func @tensor.generate(%m : index, %n : index) 66 -> tensor<?x3x?xf32> { 67 %tnsr = tensor.generate %m, %n { 68 ^bb0(%i : index, %j : index, %k : index): 69 %elem = arith.constant 8.0 : f32 70 tensor.yield %elem : f32 71 } : tensor<?x3x?xf32> 72 return %tnsr : tensor<?x3x?xf32> 73} 74 75// CHECK-LABEL: func @tensor_reshape 76func.func @tensor_reshape(%unranked: tensor<*xf32>, %shape1: tensor<1xi32>, 77 %shape2: tensor<2xi32>, %shape3: tensor<?xi32>) -> tensor<*xf32> { 78 %dyn_vec = tensor.reshape %unranked(%shape1) 79 : (tensor<*xf32>, tensor<1xi32>) -> tensor<?xf32> 80 %dyn_mat = tensor.reshape %dyn_vec(%shape2) 81 : (tensor<?xf32>, tensor<2xi32>) -> tensor<?x?xf32> 82 %new_unranked = tensor.reshape %dyn_mat(%shape3) 83 : (tensor<?x?xf32>, tensor<?xi32>) -> tensor<*xf32> 84 return %new_unranked : tensor<*xf32> 85} 86 87// CHECK-LABEL: func @slice({{.*}}) { 88func.func @slice(%t: tensor<8x16x4xf32>, %idx : index) { 89 %c0 = arith.constant 0 : index 90 %c1 = arith.constant 1 : index 91 92 // CHECK: tensor.extract_slice 93 // CHECK-SAME: tensor<8x16x4xf32> to tensor<?x?x?xf32> 94 %1 = tensor.extract_slice %t[%c0, %c0, %c0][%idx, %idx, %idx][%c1, %c1, %c1] 95 : tensor<8x16x4xf32> to tensor<?x?x?xf32> 96 97 // CHECK: tensor.extract_slice 98 // CHECK-SAME: tensor<8x16x4xf32> to tensor<4x4x4xf32> 99 %2 = tensor.extract_slice %t[0, 2, 0][4, 4, 4][1, 1, 1] 100 : tensor<8x16x4xf32> to tensor<4x4x4xf32> 101 102 // CHECK: tensor.extract_slice 103 // CHECK-SAME: tensor<8x16x4xf32> to tensor<4x4xf32> 104 %3 = tensor.extract_slice %t[0, 2, 0][4, 1, 4][1, 1, 1] 105 : tensor<8x16x4xf32> to tensor<4x4xf32> 106 107 return 108} 109 110// ----- 111 112// CHECK-LABEL: func @insert_slice({{.*}}) { 113func.func @insert_slice( 114 %t: tensor<8x16x4xf32>, 115 %td: tensor<8x?x4xf32>, 116 %t2: tensor<16x32x8xf32>, 117 %t3: tensor<4x4xf32>, 118 %idx : index, 119 %sz : index) { 120 %c0 = arith.constant 0 : index 121 %c1 = arith.constant 1 : index 122 123 // CHECK: tensor.insert_slice 124 // CHECK-SAME: tensor<8x16x4xf32> into tensor<16x32x8xf32> 125 %1 = tensor.insert_slice %t into %t2[%c0, %c0, %c0][8, 16, 4][%c1, %c1, %c1] 126 : tensor<8x16x4xf32> into tensor<16x32x8xf32> 127 128 // CHECK: tensor.insert_slice 129 // CHECK-SAME: tensor<8x16x4xf32> into tensor<16x32x8xf32> 130 %2 = tensor.insert_slice %t into %t2[%c0, %idx, %c0][8, 16, 4][%c1, 1, %c1] 131 : tensor<8x16x4xf32> into tensor<16x32x8xf32> 132 133 // CHECK: tensor.insert_slice 134 // CHECK-SAME: tensor<4x4xf32> into tensor<8x16x4xf32> 135 %3 = tensor.insert_slice %t3 into %t[0, 2, 0][4, 1, 4][1, 1, 1] 136 : tensor<4x4xf32> into tensor<8x16x4xf32> 137 138 // CHECK: tensor.insert_slice 139 // CHECK-SAME: tensor<8x?x4xf32> into tensor<8x16x4xf32> 140 %4 = tensor.insert_slice %td into %t[0, %idx, 0][8, %sz, 4][1, 1, 1] 141 : tensor<8x?x4xf32> into tensor<8x16x4xf32> 142 143 return 144} 145 146// ----- 147 148func.func @tensor_reshape_zero_dim(%arg0 : tensor<1x1xf32>, %arg1 : tensor<f32>) 149 -> (tensor<f32>, tensor<1x1xf32>) { 150 %0 = tensor.collapse_shape %arg0 [] : tensor<1x1xf32> into tensor<f32> 151 %1 = tensor.expand_shape %0 [] : tensor<f32> into tensor<1x1xf32> 152 return %0, %1 : tensor<f32>, tensor<1x1xf32> 153} 154// CHECK-LABEL: func @tensor_reshape_zero_dim 155// CHECK: tensor.collapse_shape %{{.*}} [] : tensor<1x1xf32> into tensor<f32> 156// CHECK: tensor.expand_shape %{{.*}} [] : tensor<f32> into tensor<1x1xf32> 157 158func.func @legal_collapsing_reshape_dynamic_tensor 159 (%arg0: tensor<?x?x?x4x?xf32>) -> tensor<?x?x?xf32> 160{ 161 %0 = tensor.collapse_shape %arg0 [[0], [1], [2, 3, 4]] : 162 tensor<?x?x?x4x?xf32> into tensor<?x?x?xf32> 163 return %0 : tensor<?x?x?xf32> 164} 165// CHECK: func @legal_collapsing_reshape_dynamic_tensor 166// CHECK: tensor.collapse_shape 167// CHECK-SAME: [0], [1], [2, 3, 4] 168 169// ----- 170 171func.func @rank(%t : tensor<4x4x?xf32>) { 172 // CHECK: %{{.*}} = tensor.rank %{{.*}} : tensor<4x4x?xf32> 173 %0 = "tensor.rank"(%t) : (tensor<4x4x?xf32>) -> index 174 175 // CHECK: %{{.*}} = tensor.rank %{{.*}} : tensor<4x4x?xf32> 176 %1 = tensor.rank %t : tensor<4x4x?xf32> 177 return 178} 179 180// ----- 181 182func.func @pad_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index, 183 %pad_value: f32) -> tensor<6x?x?x?xf32> { 184 %0 = tensor.pad %arg0 low[2, %low, 3, 3] high[3, 3, %high, 2] { 185 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index): 186 tensor.yield %pad_value : f32 187 } : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32> 188 return %0 : tensor<6x?x?x?xf32> 189} 190// CHECK-LABEL: func @pad_dynamic 191// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]] 192// CHECK-SAME: %[[LOW:[a-zA-Z0-9_]*]] 193// CHECK-SAME: %[[HIGH:[a-zA-Z0-9_]*]] 194// CHECK: tensor.pad %[[ARG0]] 195// CHECK-SAME: low[2, %[[LOW]], 3, 3] 196// CHECK-SAME: high[3, 3, %[[HIGH]], 2] 197// CHECK: : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32> 198 199// ----- 200 201func.func @pad_static(%arg0: tensor<3x4xf32>, %pad_value: f32) -> tensor<6x9xf32> { 202 %0 = tensor.pad %arg0 low[1, 2] high[2, 3] { 203 ^bb0(%arg1 : index, %arg2 : index): 204 tensor.yield %pad_value : f32 205 } : tensor<3x4xf32> to tensor<6x9xf32> 206 return %0 : tensor<6x9xf32> 207} 208// CHECK-LABEL: func @pad_static 209// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]] 210// CHECK: tensor.pad %[[ARG0]] low[1, 2] high[2, 3] 211// CHECK: : tensor<3x4xf32> to tensor<6x9xf32> 212 213// ----- 214 215func.func @pad_asymmetrical(%arg0: tensor<2x3xf32>, %ub0: index, %ub1: index, 216 %pad_value: f32) -> tensor<?x?xf32> { 217 %0 = tensor.pad %arg0 low[0, 0] high[%ub0, %ub1] { 218 ^bb0(%arg1: index, %arg2: index): 219 tensor.yield %pad_value : f32 220 } : tensor<2x3xf32> to tensor<?x?xf32> 221 return %0 : tensor<?x?xf32> 222} 223// CHECK-LABEL: func @pad_asymmetrical 224// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]] 225// CHECK-SAME: %[[UB0:[a-zA-Z0-9_]*]] 226// CHECK-SAME: %[[UB1:[a-zA-Z0-9_]*]] 227// CHECK: tensor.pad %[[ARG0]] 228// CHECK-SAME: low[0, 0] 229// CHECK-SAME: high[%[[UB0]], %[[UB1]]] 230// CHECK: : tensor<2x3xf32> to tensor<?x?xf32> 231 232// ----- 233 234func.func @pad_to_static_size(%arg0: tensor<?x?xf32>, %ub0: index, %ub1: index, 235 %pad_value: f32) -> tensor<2x3xf32> { 236 %0 = tensor.pad %arg0 low[0, 0] high[%ub0, %ub1] { 237 ^bb0(%arg1: index, %arg2: index): 238 tensor.yield %pad_value : f32 239 } : tensor<?x?xf32> to tensor<2x3xf32> 240 return %0 : tensor<2x3xf32> 241} 242// CHECK-LABEL: func @pad_to_static_size 243// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]] 244// CHECK-SAME: %[[UB0:[a-zA-Z0-9_]*]] 245// CHECK-SAME: %[[UB1:[a-zA-Z0-9_]*]] 246// CHECK: tensor.pad %[[ARG0]] 247// CHECK-SAME: low[0, 0] 248// CHECK-SAME: high[%[[UB0]], %[[UB1]]] 249// CHECK: : tensor<?x?xf32> to tensor<2x3xf32> 250 251// ----- 252 253// CHECK-LABEL: func @test_splat_op 254// CHECK-SAME: [[S:%arg[0-9]+]]: f32 255func.func @test_splat_op(%s : f32) { 256 // CHECK: tensor.splat [[S]] : tensor<8xf32> 257 %v = tensor.splat %s : tensor<8xf32> 258 259 // CHECK: tensor.splat [[S]] : tensor<4xf32> 260 %u = "tensor.splat"(%s) : (f32) -> tensor<4xf32> 261 return 262} 263