[mlir][tosa-to-linalg] fix arithmetic_right_shift conversion with round (#159930)

Fixed: #154259

According to TOSA spec, `tosa.arithmetic_right_shift` should handle
round.
```
if (round == true && static_cast<int32_t>(value2) > 0 &&
    (apply_arith_rshift<in_out_t>(value1, apply_sub_s<in_out_t>(value2, 1)) & 1 != 0)) {
    result = result + 1;
}
```

The original conversion is the similar as definition, and will convert
to pseudo code
```c++
result = (value1 >> value2) +
              ( (i1)(value2 > 0) & (i1)((value1 >> (value2 - 1)) & 1) )
```

But when value2 is 0,`value1 >> (value2 - 1)` will produce poison value
because performing arithmetic right shift on a negative number. Then the
poison value propagate to the final result.

This PR wants to change the conversion to `arith.select` to stop poison
propagation.

```c++
result = (value1 >> value2) + 
             (value2 > 0) ? (i1)((value1 >> (value2 - 1)) & 1) : (i1)(0)
```

Author: HerrCai0907

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

@@ -305,6 +305,8 @@ static Value createLinalgBodyCalculationForElementwiseOp(
                                          IntegerAttr::get(elementTy, 1));
     auto zero = arith::ConstantOp::create(rewriter, loc,
                                           IntegerAttr::get(elementTy, 0));
+    auto i1zero =
+        arith::ConstantOp::create(rewriter, loc, IntegerAttr::get(i1Ty, 0));
     auto i1one =
         arith::ConstantOp::create(rewriter, loc, IntegerAttr::get(i1Ty, 1));
 
@@ -322,9 +324,9 @@ static Value createLinalgBodyCalculationForElementwiseOp(
                                              ArrayRef<NamedAttribute>());
     auto isInputOdd =
         arith::AndIOp::create(rewriter, loc, i1Ty, truncated, i1one);
-
-    auto shouldRound = arith::AndIOp::create(
-        rewriter, loc, i1Ty, shiftValueGreaterThanZero, isInputOdd);
+    // shifted, truncated, isInputOdd can be poison when input2 is 0.
+    auto shouldRound = arith::SelectOp::create(
+        rewriter, loc, i1Ty, shiftValueGreaterThanZero, isInputOdd, i1zero);
     auto extended =
         arith::ExtUIOp::create(rewriter, loc, resultTypes, shouldRound);
     return arith::AddIOp::create(rewriter, loc, resultTypes, result, extended);

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

@@ -698,13 +698,14 @@ func.func @test_simple_i32(%arg0: tensor<1xi32>, %unsigned: tensor<1xui32>, %uns
   // CHECK: linalg.generic
   // CHECK: arith.constant 1
   // CHECK: arith.constant 0
+  // CHECK: arith.constant false
   // CHECK: arith.constant true
   // CHECK: arith.cmpi
   // CHECK: arith.subi
   // CHECK: arith.shrsi
   // CHECK: arith.trunci
   // CHECK: and
-  // CHECK: and
+  // CHECK: arith.select
   // CHECK: arith.extui
   // CHECK: arith.addi
   %12 = tosa.arithmetic_right_shift %arg0, %arg0 {round = 1 : i1} : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>