Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Release rai 1 2 #201

Merged
merged 9 commits into from
Jun 20, 2024
9 changes: 6 additions & 3 deletions mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -99,9 +99,12 @@ struct SqrtReciprocalOptimization : public OpRewritePattern<tosa::PowOp> {
// An improvement for the future would be to generate a tile operator here instead
if (inputType != outputType)
return rewriter.notifyMatchFailure(op, "input type and output type are different, tiling is not supported for this canonicalization");

rewriter.replaceOpWithNewOp<tosa::RsqrtOp>(user, outputType, op.getInput1());

auto rsqrtOp = rewriter.create<tosa::RsqrtOp>(
rewriter.getFusedLoc({op.getLoc(), user->getLoc()}), outputType,
op.getInput1());
rewriter.replaceOp(user, rsqrtOp);

return success();
}
};
Expand Down Expand Up @@ -888,7 +891,7 @@ OpFoldResult CastOp::fold(FoldAdaptor adaptor) {
llvm::cast<IntegerType>(outETy).getIntOrFloatBitWidth(), unsign);
auto floatVal = operand.getSplatValue<APFloat>();
bool exact;
floatVal.convertToInteger(intVal, llvm::RoundingMode::TowardZero, &exact);
floatVal.convertToInteger(intVal, llvm::RoundingMode::NearestTiesToEven, &exact);
return SplatElementsAttr::get(outTy, intVal);
}

Expand Down
36 changes: 34 additions & 2 deletions mlir/test/Dialect/Tosa/canonicalize_with_debuginfo.mlir
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// RUN: mlir-opt -mlir-print-debuginfo -canonicalize="test-convergence" %s | FileCheck %s
// RUN: mlir-opt -split-input-file -mlir-print-debuginfo -canonicalize="test-convergence" %s | FileCheck %s

// CHECK-LABEL: @clamp_twice_is_single_clamp
func.func @clamp_twice_is_single_clamp(%arg0: tensor<4xi8>) -> tensor<4xi8> {
Expand All @@ -11,4 +11,36 @@ func.func @clamp_twice_is_single_clamp(%arg0: tensor<4xi8>) -> tensor<4xi8> {
return %1 : tensor<4xi8>
}
#loc0 = loc("Clamp_A")
#loc1 = loc("Clamp_B")
#loc1 = loc("Clamp_B")

// -----

// CHECK-LABEL: @canonicalize_optimize_sqrt_reciprocal
func.func @canonicalize_optimize_sqrt_reciprocal_with_debinfo(%arg0: tensor<1x5x1x1xf32>) -> tensor<1x5x1x1xf32> {
// CHECK: %[[RSQRT:.*]] = tosa.rsqrt %arg{{.*}} : (tensor<1x5x1x1xf32>) -> tensor<1x5x1x1xf32> loc([[LOC:.*]])
// CHECK-DAG: #[[A:.*]] = loc("Pow_A")
// CHECK-DAG: #[[B:.*]] = loc("Reciprocal_B")
// CHECK-DAG: [[LOC]] = loc(fused[#[[A]], #[[B]]])
%0 = "tosa.const"() <{value = dense<5.000000e-01> : tensor<1x1x1x1xf32>}> : () -> tensor<1x1x1x1xf32>
%1 = tosa.pow %arg0, %0 : (tensor<1x5x1x1xf32>, tensor<1x1x1x1xf32>) -> tensor<1x5x1x1xf32> loc(#loc0)
%2 = tosa.reciprocal %1 : (tensor<1x5x1x1xf32>) -> tensor<1x5x1x1xf32> loc(#loc1)
return %2 : tensor<1x5x1x1xf32>
}
#loc0 = loc("Pow_A")
#loc1 = loc("Reciprocal_B")

// -----

// CHECK-LABEL: @canonicalize_optimize_sqrt_reciprocal
func.func @canonicalize_optimize_sqrt_reciprocal_bf16(%arg0: tensor<1x5x1x1xbf16>) -> tensor<1x5x1x1xbf16> {
// CHECK: %[[RSQRT:.*]] = tosa.rsqrt %arg{{.*}} : (tensor<1x5x1x1xbf16>) -> tensor<1x5x1x1xbf16> loc([[LOC:.*]])
// CHECK-DAG: #[[A:.*]] = loc("Pow_B")
// CHECK-DAG: #[[B:.*]] = loc("Reciprocal_C")
// CHECK-DAG: [[LOC]] = loc(fused[#[[A]], #[[B]]])
%0 = "tosa.const"() <{value = dense<5.000000e-01> : tensor<1x1x1x1xbf16>}> : () -> tensor<1x1x1x1xbf16>
%1 = tosa.pow %arg0, %0 : (tensor<1x5x1x1xbf16>, tensor<1x1x1x1xbf16>) -> tensor<1x5x1x1xbf16> loc(#loc0)
%2 = tosa.reciprocal %1 : (tensor<1x5x1x1xbf16>) -> tensor<1x5x1x1xbf16> loc(#loc1)
return %2 : tensor<1x5x1x1xbf16>
}
#loc0 = loc("Pow_B")
#loc1 = loc("Reciprocal_C")
11 changes: 11 additions & 0 deletions mlir/test/Dialect/Tosa/constant-op-fold.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -602,6 +602,17 @@ func.func @cast_float_to_int() -> tensor<i16> {

// -----

// CHECK: func.func @cast_float_to_int_round
func.func @cast_float_to_int_round() -> tensor<i16> {
%splat = "tosa.const"() {value = dense<-3.5> : tensor<f32>} : () -> tensor<f32>
// CHECK: %[[SPLAT:.+]] = "tosa.const"() <{value = dense<-4> : tensor<i16>}
%cast = tosa.cast %splat : (tensor<f32>) -> tensor<i16>
// CHECK: return %[[SPLAT]]
return %cast : tensor<i16>
}

// -----

// CHECK: func.func @cast_int_to_int_trunc
func.func @cast_int_to_int_trunc() -> tensor<i16> {
%splat = "tosa.const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
Expand Down