158 lines · plain
1// RUN: mlir-opt --split-input-file --tosa-to-arith="include-apply-rescale=true use-32-bit=true" %s -verify-diagnostics -o -| FileCheck %s2// RUN: mlir-opt --split-input-file --tosa-to-arith="include-apply-rescale=false" %s -verify-diagnostics -o -| FileCheck --check-prefix="SCALE" %s3 4// CHECK-LABEL: func @const_test5func.func @const_test() -> (tensor<i32>) {6 // CHECK: [[C3:%.+]] = arith.constant dense<3> : tensor<i32>7 %result = "tosa.const"() {values = dense<3> : tensor<i32>} : () -> tensor<i32>8 9 // CHECK: return [[C3]]10 return %result : tensor<i32>11}12 13// -----14 15// CHECK-LABEL: @apply_scale_test_i3216// SCALE: tosa.apply_scale17func.func @apply_scale_test_i32(%arg0 : i32, %arg1 : i32, %arg2 : i8) -> (i32) {18 // CHECK-DAG: %[[S32:.+]] = arith.extui %arg2 : i8 to i3219 // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : i3220 // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : i3221 // CHECK-DAG: %[[C2:.+]] = arith.constant 2 : i3222 // CHECK-DAG: %[[C30:.+]] = arith.constant 30 : i3223 // CHECK-DAG: %[[C32:.+]] = arith.constant 32 : i3224 25 // Compute the high-low values of the matmul in 64-bits.26 // CHECK-DAG: %[[LOW:.+]], %[[HI:.+]] = arith.mulsi_extended %arg0, %arg127 28 // Determine whether the high bits need to shift left or right and by how much.29 // CHECK-DAG: %[[OVER31:.+]] = arith.cmpi sge, %[[S32]], %[[C32]]30 // CHECK-DAG: %[[OVER32:.+]] = arith.cmpi sgt, %[[S32]], %[[C32]]31 // CHECK-DAG: %[[HISHLN:.+]] = arith.subi %[[C32]], %[[S32]]32 // CHECK-DAG: %[[HISHRN:.+]] = arith.subi %[[S32]], %[[C32]]33 // CHECK-DAG: %[[HISHL:.+]] = arith.select %[[OVER31]], %[[C0]], %[[HISHLN]]34 // CHECK-DAG: %[[HISHR:.+]] = arith.select %[[OVER31]], %[[HISHRN]], %[[C0]]35 36 // Apply double rounding.37 // CHECK-DAG: %[[CN1:.+]] = arith.constant -138 // CHECK-DAG: %[[POS:.+]] = arith.cmpi sge, %arg0, %[[C0]]39 // CHECK-DAG: %[[DIR:.+]] = arith.select %[[POS]], %[[C1]], %[[CN1]]40 // CHECK-DAG: %[[DRND:.+]] = arith.select %[[OVER31]], %[[DIR]], %[[C0]]41 // CHECK-DAG: %[[DSHFTR:.+]] = arith.shrui %[[LOW]], %[[C30]]42 // CHECK-DAG: %[[DRNDED:.+]] = arith.addi %[[DSHFTR]], %[[DRND]]43 // CHECK-DAG: %[[DCARRY:.+]] = arith.shrsi %[[DRNDED]], %[[C2:.+]]44 // CHECK-DAG: %[[DBIT:.+]] = arith.shli %[[DRND]], %[[C30]]45 // CHECK-DAG: %[[DLOW:.+]] = arith.addi %[[LOW]], %[[DBIT]]46 // CHECK-DAG: %[[DHI:.+]] = arith.addi %[[HI]], %[[DCARRY]]47 48 // Apply low-bit rounding.49 // CHECK-DAG: %[[SHFTM1:.+]] = arith.subi %[[S32]], %[[C1]]50 // CHECK-DAG: %[[LBIT:.+]] = arith.shli %[[C1]], %[[SHFTM1]]51 // CHECK-DAG: %[[HALF:.+]] = arith.select %[[OVER32]], %[[C0]], %[[LBIT]]52 // CHECK-DAG: %[[LADD:.+]] = arith.addi %[[DLOW]], %[[HALF]]53 // CHECK-DAG: %[[LLO:.+]] = arith.cmpi ugt, %[[DLOW]], %[[LADD]]54 // CHECK-DAG: %[[LCARRY:.+]] = arith.extui %[[LLO]] : i1 to i3255 // CHECK-DAG: %[[LRNDED:.+]] = arith.addi %[[DHI]], %[[LCARRY]]56 57 // Apply high-bit rounding.58 // CHECK-DAG: %[[HISHRM1:.+]] = arith.subi %[[HISHR]], %[[C1]]59 // CHECK-DAG: %[[LHISHFT:.+]] = arith.shli %[[C1]], %[[HISHRM1]]60 // CHECK-DAG: %[[LHI:.+]] = arith.select %[[OVER32]], %[[LHISHFT]], %[[C0]]61 // CHECK-DAG: %[[FHI:.+]] = arith.addi %[[LRNDED]], %[[LHI]]62 63 // Combine hi-low into the final result.64 // CHECK-DAG: %[[HIL:.+]] = arith.shli %[[FHI]], %[[HISHL]]65 // CHECK-DAG: %[[HIALIGN:.+]] = arith.shrsi %[[HIL:.+]], %[[HISHR]]66 // CHECK-DAG: %[[LOR:.+]] = arith.shrui %[[LADD]], %[[S32]]67 // CHECK-DAG: %[[LOWALIGN:.+]] = arith.select %[[OVER31]], %[[C0]], %[[LOR]]68 // CHECK-DAG: %[[RESULT:.+]] = arith.addi %[[LOWALIGN]], %[[HIALIGN]]69 // CHECK: return %[[RESULT]]70 %res = tosa.apply_scale %arg0, %arg1, %arg2 {rounding_mode = DOUBLE_ROUND} : (i32, i32, i8) -> i3271 return %res : i3272}73 74// -----75 76// CHECK-LABEL: @apply_scale_test_vector77// SCALE: tosa.apply_scale78func.func @apply_scale_test_vector(%arg0 : vector<4xi32>, %arg1 : vector<4xi32>, %arg2 : vector<4xi8>) -> (vector<4xi32>) {79 // CHECK-NOT: "tosa.apply_scale"80 %res = tosa.apply_scale %arg0, %arg1, %arg2 {rounding_mode = DOUBLE_ROUND} : (vector<4xi32>, vector<4xi32>, vector<4xi8>) -> vector<4xi32>81 return %res : vector<4xi32>82}83 84// -----85 86// CHECK-LABEL: @apply_scale_test_i4887// SCALE: tosa.apply_scale88func.func @apply_scale_test_i48(%arg0 : i48, %arg1 : i32, %arg2 : i8) -> (i32) {89 // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : i4890 // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : i6491 // CHECK-DAG: %[[C31:.+]] = arith.constant 31 : i3292 93 // Multiply in 64 bits.94 // CHECK-DAG: %[[V64:.+]] = arith.extsi %arg0 : i48 to i6495 // CHECK-DAG: %[[M64:.+]] = arith.extsi %arg1 : i32 to i6496 // CHECK-DAG: %[[MUL:.+]] = arith.muli %[[V64]], %[[M64]]97 98 // Round normally.99 // CHECK-DAG: %[[S32:.+]] = arith.extui %arg2 : i8 to i32100 // CHECK-DAG: %[[S64:.+]] = arith.extui %[[S32]] : i32 to i64101 // CHECK-DAG: %[[ONEL:.+]] = arith.shli %[[C1]], %[[S64]] : i64102 // CHECK-DAG: %[[ONER:.+]] = arith.shrui %[[ONEL]], %[[C1]]103 // CHECK-DAG: %[[ROUND:.+]] = arith.addi %[[MUL]], %[[ONER]]104 105 // Apply double rounding.106 // CHECK-DAG: %[[DUP:.+]] = arith.constant 1073741824 : i64107 // CHECK-DAG: %[[DDOWN:.+]] = arith.constant -1073741824 : i64108 // CHECK-DAG: %[[POS:.+]] = arith.cmpi sge, %arg0, %[[C0]]109 // CHECK-DAG: %[[DBIT:.+]] = arith.select %[[POS]], %[[DUP]], %[[DDOWN]]110 // CHECK-DAG: %[[DRND:.+]] = arith.addi %[[DBIT]], %[[ROUND]]111 // CHECK-DAG: %[[USED:.+]] = arith.cmpi sgt, %[[S32]], %[[C31]] : i32112 // CHECK-DAG: %[[RES64:.+]] = arith.select %[[USED]], %[[DRND]], %[[ROUND]] : i64113 114 // Shift and truncate final answer.115 // CHECK-DAG: %[[SHR:.+]] = arith.shrsi %[[RES64]], %[[S64]]116 // CHECK-DAG: %[[TRUNC:.+]] = arith.trunci %[[SHR]] : i64 to i32117 // CHECK: return %[[TRUNC]]118 %res = tosa.apply_scale %arg0, %arg1, %arg2 {rounding_mode = DOUBLE_ROUND} : (i48, i32, i8) -> i32119 return %res : i32120}121 122// -----123 124// CHECK-LABEL: @apply_scale_test_i64125// SCALE: tosa.apply_scale126func.func @apply_scale_test_i64(%arg0 : i64, %arg1 : i32, %arg2 : i8) -> (i32) {127 // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : i64128 // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : i64129 // CHECK-DAG: %[[C31:.+]] = arith.constant 31 : i32130 131 // Multiply in 64 bits.132 // CHECK-DAG: %[[M64:.+]] = arith.extsi %arg1 : i32 to i64133 // CHECK-DAG: %[[MUL:.+]] = arith.muli %arg0, %[[M64]]134 135 // Round normally.136 // CHECK-DAG: %[[S32:.+]] = arith.extui %arg2 : i8 to i32137 // CHECK-DAG: %[[S64:.+]] = arith.extui %[[S32]] : i32 to i64138 // CHECK-DAG: %[[ONEL:.+]] = arith.shli %[[C1]], %[[S64]] : i64139 // CHECK-DAG: %[[ONER:.+]] = arith.shrui %[[ONEL]], %[[C1]]140 // CHECK-DAG: %[[ROUND:.+]] = arith.addi %[[MUL]], %[[ONER]]141 142 // Apply double rounding.143 // CHECK-DAG: %[[DUP:.+]] = arith.constant 1073741824 : i64144 // CHECK-DAG: %[[DDOWN:.+]] = arith.constant -1073741824 : i64145 // CHECK-DAG: %[[POS:.+]] = arith.cmpi sge, %arg0, %[[C0]]146 // CHECK-DAG: %[[DBIT:.+]] = arith.select %[[POS]], %[[DUP]], %[[DDOWN]]147 // CHECK-DAG: %[[DRND:.+]] = arith.addi %[[DBIT]], %[[ROUND]]148 // CHECK-DAG: %[[USED:.+]] = arith.cmpi sgt, %[[S32]], %[[C31]] : i32149 // CHECK-DAG: %[[RES64:.+]] = arith.select %[[USED]], %[[DRND]], %[[ROUND]] : i64150 151 // Shift and truncate final answer.152 // CHECK-DAG: %[[SHR:.+]] = arith.shrsi %[[RES64]], %[[S64]]153 // CHECK-DAG: %[[TRUNC:.+]] = arith.trunci %[[SHR]] : i64 to i32154 // CHECK: return %[[TRUNC]]155 %res = tosa.apply_scale %arg0, %arg1, %arg2 {rounding_mode = DOUBLE_ROUND} : (i64, i32, i8) -> i32156 return %res : i32157}158