86 lines · plain
1// RUN: mlir-opt --split-input-file --tosa-to-scf %s -verify-diagnostics -o -| FileCheck %s2 3// CHECK-LABEL: func @while_test4// CHECK-SAME: ([[ARG0:%.+]]: tensor<i32>)5func.func @while_test(%arg0 : tensor<i32>) -> (tensor<i32>) {6 // CHECK: [[WHILE:%.+]] = scf.while ([[ARG1:%.+]] = [[ARG0]])7 %0 = tosa.while_loop (%arg1 = %arg0) : (tensor<i32>) -> tensor<i32> {8 // CHECK: tosa.const9 %1 = "tosa.const"() {values = dense<3> : tensor<i32>} : () -> tensor<i32>10 11 // CHECK: [[COMPARE:%.+]] = tosa.greater_equal12 %2 = tosa.greater_equal %1, %arg1 : (tensor<i32>, tensor<i32>) -> tensor<i1>13 14 // CHECK: [[EX:%.+]] = tensor.extract [[COMPARE]]15 // CHECK: scf.condition([[EX]]) [[ARG1]]16 tosa.yield %2 : tensor<i1>17 } do {18 // CHECK: ^bb0([[ARG1:%.+]]: tensor<i32>)19 ^bb0(%arg1: tensor<i32>):20 // CHECK: tosa.const21 %1 = "tosa.const"() {values = dense<1> : tensor<i32>} : () -> tensor<i32>22 23 // CHECK: [[ADD:%.+]] = tosa.add24 %2 = tosa.add %arg1, %1 : (tensor<i32>, tensor<i32>) -> tensor<i32>25 26 // CHECK: scf.yield [[ADD]]27 tosa.yield %2 : tensor<i32>28 }29 return %0 : tensor<i32>30}31 32// -----33 34// CHECK-LABEL: func @if_test35// CHECK-SAME: ([[ARG0:%.+]]: tensor<f32>, [[ARG1:%.+]]: tensor<f32>, [[ARG2:%.+]]: tensor<i1>)36func.func @if_test(%arg0 : tensor<f32>, %arg1 : tensor<f32>, %arg2 : tensor<i1>) -> (tensor<f32>) {37 // CHECK: [[EX:%.+]] = tensor.extract [[ARG2]]38 // CHECK: [[IF:%.+]] = scf.if [[EX]] -> (tensor<f32>) {39 %0 = tosa.cond_if %arg2 : tensor<i1> -> tensor<f32> {40 41 // CHECK: scf.yield [[ARG0]]42 tosa.yield %arg0 : tensor<f32>43 44 // CHECK: } else {45 } else {46 47 // CHECK: scf.yield [[ARG1]]48 tosa.yield %arg1 : tensor<f32>49 50 // CHECK: }51 // CHECK: return [[IF]]52 }53 54 return %0 : tensor<f32>55}56 57// -----58 59// CHECK-LABEL: func @scatter_test60// CHECK-SAME: ([[VALUES_IN:%.+]]: tensor<3x7x5xi32>, [[INDICES:%.+]]: tensor<3x6xi32>, [[INPUT:%.+]]: tensor<3x6x5xi32>)61func.func @scatter_test(%values_in: tensor<3x7x5xi32>, %indices : tensor<3x6xi32>, %input: tensor<3x6x5xi32>) -> tensor<3x7x5xi32> {62 63 // CHECK-DAG: [[C_0:%.+]] = arith.constant 0 : index64 // CHECK-DAG: [[C_1:%.+]] = arith.constant 1 : index65 // CHECK-DAG: [[C_2:%.+]] = arith.constant 2 : index66 // CHECK-DAG: [[C_3:%.+]] = arith.constant 3 : index67 // CHECK-DAG: [[C_5:%.+]] = arith.constant 5 : index68 // CHECK-DAG: [[C_6:%.+]] = arith.constant 6 : index69 // CHECK-DAG: [[C_0_0:%.+]] = arith.constant 0 : index70 // CHECK-DAG: [[C_1_0:%.+]] = arith.constant 1 : index71 // CHECK: [[RESULT_0:%.+]] = scf.for [[ITER_VAR_0:%.+]] = [[C_0_0]] to [[C_3]] step [[C_1_0]] iter_args([[ITER_ARG_0:%.+]] = [[VALUES_IN]]) -> (tensor<3x7x5xi32>) {72 // CHECK: [[RESULT_1:%.+]] = scf.for [[ITER_VAR_1:%.+]] = [[C_0_0]] to [[C_6]] step [[C_1_0]] iter_args([[ITER_ARG_1:%.+]] = [[ITER_ARG_0]]) -> (tensor<3x7x5xi32>) {73 // CHECK-DAG: [[EXTRACTED:%.+]] = tensor.extract [[INDICES]][[[ITER_VAR_0]], [[ITER_VAR_1]]] : tensor<3x6xi32>74 // CHECK-DAG: [[EXTRACTED_CAST:%.+]] = arith.index_cast [[EXTRACTED]] : i32 to index75 // CHECK-DAG: [[EXTRACTED_SLICE:%.+]] = tensor.extract_slice [[INPUT]][[[ITER_VAR_0]], [[ITER_VAR_1]], [[C_0_0]]] [[[C_1_0]], [[C_1_0]], [[C_5]]] [[[C_1_0]], [[C_1_0]], [[C_1_0]]] : tensor<3x6x5xi32> to tensor<?x?x?xi32>76 // CHECK-DAG: [[INSERTED_SLICE:%.+]] = tensor.insert_slice [[EXTRACTED_SLICE]] into [[ITER_ARG_1]][[[ITER_VAR_0]], [[EXTRACTED_CAST]], [[C_0_0]]] [[[C_1_0]], [[C_1_0]], [[C_5]]] [[[C_1_0]], [[C_1_0]], [[C_1_0]]] : tensor<?x?x?xi32> into tensor<3x7x5xi32>77 // CHECK: scf.yield [[INSERTED_SLICE]] : tensor<3x7x5xi32>78 // CHECK: }79 // CHECK: scf.yield [[RESULT_1]] : tensor<3x7x5xi32>80 // CHECK: }81 %0 = "tosa.scatter"(%values_in, %indices, %input) : (tensor<3x7x5xi32>, tensor<3x6xi32>, tensor<3x6x5xi32>) -> (tensor<3x7x5xi32>)82 83 // CHECK: return [[RESULT_0]] : tensor<3x7x5xi32>84 return %0 : tensor<3x7x5xi32>85}86