[Mlir-commits] [mlir] 9294a1e - [mlir][tosa] Rework tosa.apply_scale lowering for 32-bit
Rob Suderman
llvmlistbot at llvm.org
Tue May 17 16:13:27 PDT 2022
Author: Robert Suderman
Date: 2022-05-17T16:01:12-07:00
New Revision: 9294a1e9a8ab8fed60ee1b0737e944077b5339cd
URL: https://github.com/llvm/llvm-project/commit/9294a1e9a8ab8fed60ee1b0737e944077b5339cd
DIFF: https://github.com/llvm/llvm-project/commit/9294a1e9a8ab8fed60ee1b0737e944077b5339cd.diff
LOG: [mlir][tosa] Rework tosa.apply_scale lowering for 32-bit
Added handling rounding behavior in 32-bits for when possible. This
avoids kernel compilation generating scalarized code on platforms where
64-bit vectors are not available.
As the 48-bit lowering requires 64-bit anyway, we added a full 64-bit
solution simplifying the old path.
Reviewed By: dcaballe, mravishankar
Differential Revision: https://reviews.llvm.org/D125583
Added:
Modified:
mlir/include/mlir/Conversion/Passes.td
mlir/include/mlir/Conversion/TosaToArith/TosaToArith.h
mlir/lib/Conversion/TosaToArith/TosaToArith.cpp
mlir/lib/Conversion/TosaToArith/TosaToArithPass.cpp
mlir/test/Conversion/TosaToArith/tosa-to-arith.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 0c4aca31968d3..d0122f9051ff2 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -756,7 +756,10 @@ def TosaToArith : Pass<"tosa-to-arith"> {
let options = [
Option<"includeApplyRescale", "include-apply-rescale",
"bool", /*default=*/"false",
- "Whether to include the lowering for tosa.apply_rescale to arith">
+ "Whether to include the lowering for tosa.apply_rescale to arith">,
+ Option<"use32Bit", "use-32-bit",
+ "bool", /*default=*/"false",
+ "Whether to prioritze lowering to 32-bit operations">
];
let constructor = "tosa::createTosaToArith()";
diff --git a/mlir/include/mlir/Conversion/TosaToArith/TosaToArith.h b/mlir/include/mlir/Conversion/TosaToArith/TosaToArith.h
index 91099fbb4b378..62f5a0f6d8712 100644
--- a/mlir/include/mlir/Conversion/TosaToArith/TosaToArith.h
+++ b/mlir/include/mlir/Conversion/TosaToArith/TosaToArith.h
@@ -22,7 +22,8 @@ std::unique_ptr<Pass> createTosaToArith();
void populateTosaToArithConversionPatterns(RewritePatternSet *patterns);
-void populateTosaRescaleToArithConversionPatterns(RewritePatternSet *patterns);
+void populateTosaRescaleToArithConversionPatterns(RewritePatternSet *patterns,
+ bool include32Bit = false);
} // namespace tosa
} // namespace mlir
diff --git a/mlir/lib/Conversion/TosaToArith/TosaToArith.cpp b/mlir/lib/Conversion/TosaToArith/TosaToArith.cpp
index 225cd66158340..fce21e224c414 100644
--- a/mlir/lib/Conversion/TosaToArith/TosaToArith.cpp
+++ b/mlir/lib/Conversion/TosaToArith/TosaToArith.cpp
@@ -14,6 +14,7 @@
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
using namespace mlir;
@@ -49,103 +50,194 @@ Attribute getConstantAttr(Type type, int64_t value, PatternRewriter &rewriter) {
return rewriter.getIntegerAttr(type, value);
}
+Value getConstantValue(Location loc, Type type, int64_t value,
+ PatternRewriter &rewriter) {
+ return rewriter.create<arith::ConstantOp>(
+ loc, getConstantAttr(type, value, rewriter));
+}
+
// This converts the TOSA ApplyScale operator to a set of arithmetic ops,
// using 64-bit operations to perform the necessary multiply, bias, and shift.
-// Multiple types are used to use minimal bit width operations.
-class ApplyScaleOpConverter : public OpRewritePattern<tosa::ApplyScaleOp> {
+class ApplyScaleGenericOpConverter
+ : public OpRewritePattern<tosa::ApplyScaleOp> {
public:
using OpRewritePattern<tosa::ApplyScaleOp>::OpRewritePattern;
LogicalResult matchAndRewrite(tosa::ApplyScaleOp op,
PatternRewriter &rewriter) const final {
Location loc = op.getLoc();
- Value value32 = op.value();
+ Value value = op.value();
Value multiplier32 = op.multiplier();
- Value shift8 = op.shift();
- bool doubleRound = op.double_round();
- Type inType = op.value().getType();
Type resultTy = op.getType();
-
- Type i8Ty = matchContainerType(rewriter.getIntegerType(8), resultTy);
+ Type valueTy = value.getType();
Type i32Ty = matchContainerType(rewriter.getI32Type(), resultTy);
Type i64Ty = matchContainerType(rewriter.getI64Type(), resultTy);
- Value one8 = rewriter.create<arith::ConstantOp>(
- loc, getConstantAttr(i8Ty, 1, rewriter));
- Value one64 = rewriter.create<arith::ConstantOp>(
- loc, getConstantAttr(i64Ty, 1, rewriter));
-
- Value shiftSubOne8 = rewriter.create<arith::SubIOp>(loc, shift8, one8);
-
- // The rounding value semantics below equate to the following code:
- // int64_t round = 1 << (shift - 1);
- // if (double_round) {
- // if (shift > 31 && value >= 0) round += 1<<30;
- // if (shift > 31 && value < 0) round -= 1<<30;
- // }
- //
- // Note that minimal bitwidth operators are used throughout the block.
-
- Value round64 = rewriter.create<arith::ShLIOp>(
- loc, one64, rewriter.create<arith::ExtSIOp>(loc, i64Ty, shiftSubOne8));
-
- // Double rounding is performing a round operation before the shift
- if (doubleRound) {
- Value one32 = rewriter.create<arith::ConstantOp>(
- loc, getConstantAttr(i32Ty, 1, rewriter));
- Value shift32 = rewriter.create<arith::ExtSIOp>(loc, i32Ty, shift8);
- Value thirty32 = rewriter.create<arith::ConstantOp>(
- loc, getConstantAttr(i32Ty, 30, rewriter));
-
- Value shiftThirty32 =
- rewriter.create<arith::ShLIOp>(loc, one32, thirty32);
- Value shiftThirty64 =
- rewriter.create<arith::ExtSIOp>(loc, i64Ty, shiftThirty32);
-
- // Round value needs to with be added or subtracted depending on the sign
- // of the input value.
- Value roundAdd64 =
- rewriter.create<arith::AddIOp>(loc, round64, shiftThirty64);
- Value roundSub64 =
- rewriter.create<arith::SubIOp>(loc, round64, shiftThirty64);
-
- Value zero32 =
- rewriter.create<arith::ConstantOp>(loc, rewriter.getZeroAttr(inType));
- Value valueGreaterThanZero = rewriter.create<arith::CmpIOp>(
- loc, arith::CmpIPredicate::sge, value32, zero32);
+ Value zero = getConstantValue(loc, valueTy, 0, rewriter);
+ Value one64 = getConstantValue(loc, i64Ty, 1, rewriter);
+ Value thirtyOne32 = getConstantValue(loc, i32Ty, 31, rewriter);
+
+ Value shift32 = rewriter.create<arith::ExtUIOp>(loc, i32Ty, op.shift());
+
+ // Compute the multiplication in 64-bits then select the high / low parts.
+ Value value64 = rewriter.create<arith::ExtSIOp>(loc, i64Ty, value);
+ Value multiplier64 =
+ rewriter.create<arith::ExtSIOp>(loc, i64Ty, multiplier32);
+ Value multiply64 =
+ rewriter.create<arith::MulIOp>(loc, value64, multiplier64);
- Value doubleRound64 = rewriter.create<arith::SelectOp>(
- loc, valueGreaterThanZero, roundAdd64, roundSub64);
+ // Apply normal rounding.
+ Value shift64 = rewriter.create<arith::ExtUIOp>(loc, i64Ty, shift32);
+ Value round = rewriter.create<arith::ShLIOp>(loc, one64, shift64);
+ round = rewriter.create<arith::ShRUIOp>(loc, round, one64);
+ multiply64 = rewriter.create<arith::AddIOp>(loc, multiply64, round);
- // We only perform double rounding if the shift value is greater than 32.
- Value thirtyTwo32 = rewriter.create<arith::ConstantOp>(
- loc, getConstantAttr(i32Ty, 32, rewriter));
- Value shiftGreaterThanThirtyTwo = rewriter.create<arith::CmpIOp>(
- loc, arith::CmpIPredicate::sge, shift32, thirtyTwo32);
- round64 = rewriter.create<arith::SelectOp>(loc, shiftGreaterThanThirtyTwo,
- doubleRound64, round64);
+ // Apply double rounding if necessary.
+ if (op.double_round()) {
+ int64_t roundInt = 1 << 30;
+ Value roundUp = getConstantValue(loc, i64Ty, roundInt, rewriter);
+ Value roundDown = getConstantValue(loc, i64Ty, -roundInt, rewriter);
+ Value positive = rewriter.create<arith::CmpIOp>(
+ loc, arith::CmpIPredicate::sge, value, zero);
+ Value dir =
+ rewriter.create<arith::SelectOp>(loc, positive, roundUp, roundDown);
+ Value val = rewriter.create<arith::AddIOp>(loc, dir, multiply64);
+ Value valid = rewriter.create<arith::CmpIOp>(
+ loc, arith::CmpIPredicate::sgt, shift32, thirtyOne32);
+ multiply64 =
+ rewriter.create<arith::SelectOp>(loc, valid, val, multiply64);
}
- // The computation below equates to the following pseudocode:
- // int64_t result = (int64_t)value * multiplier + round;
- // result = result >> shift;
- //
- // Note that multiply and shift need to be perform in i64 to preserve bits.
+ Value result64 = rewriter.create<arith::ShRSIOp>(loc, multiply64, shift64);
+ Value result32 = rewriter.create<arith::TruncIOp>(loc, i32Ty, result64);
+
+ rewriter.replaceOp(op, result32);
+ return success();
+ }
+};
+
+class ApplyScale32BitOpConverter : public OpRewritePattern<tosa::ApplyScaleOp> {
+public:
+ using OpRewritePattern<tosa::ApplyScaleOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tosa::ApplyScaleOp op,
+ PatternRewriter &rewriter) const final {
+ Location loc = op.getLoc();
+
+ Type resultTy = op.getType();
+ Type i32Ty = matchContainerType(rewriter.getI32Type(), resultTy);
+ Type i64Ty = matchContainerType(rewriter.getI64Type(), resultTy);
+
+ Value value = op.value();
+ if (getElementTypeOrSelf(value.getType()).getIntOrFloatBitWidth() > 32) {
+ return failure();
+ }
+ Value value32 = op.value();
+ Value multiplier32 = op.multiplier();
+ Value shift32 = rewriter.create<arith::ExtUIOp>(loc, i32Ty, op.shift());
+
+ // Constants used during the scaling operation.
+ Value zero32 = getConstantValue(loc, i32Ty, 0, rewriter);
+ Value one32 = getConstantValue(loc, i32Ty, 1, rewriter);
+ Value two32 = getConstantValue(loc, i32Ty, 2, rewriter);
+ Value thirty32 = getConstantValue(loc, i32Ty, 30, rewriter);
+ Value thirtyTwo32 = getConstantValue(loc, i32Ty, 32, rewriter);
+ Value thirtyTwo64 = getConstantValue(loc, i64Ty, 32, rewriter);
+
+ // Compute the multiplication in 64-bits then select the high / low parts.
Value value64 = rewriter.create<arith::ExtSIOp>(loc, i64Ty, value32);
Value multiplier64 =
rewriter.create<arith::ExtSIOp>(loc, i64Ty, multiplier32);
- Value shift64 = rewriter.create<arith::ExtSIOp>(loc, i64Ty, shift8);
+ Value multiply64 =
+ rewriter.create<arith::MulIOp>(loc, value64, multiplier64);
- // Multiply as a pair of i64 values to guarantee the end value fits.
- Value result64 = rewriter.create<arith::MulIOp>(loc, value64, multiplier64);
- result64 = rewriter.create<arith::AddIOp>(loc, result64, round64);
- result64 = rewriter.create<arith::ShRSIOp>(loc, result64, shift64);
+ // Grab out the high/low of the computation
+ Value high64 =
+ rewriter.create<arith::ShRUIOp>(loc, multiply64, thirtyTwo64);
+ Value high32 = rewriter.create<arith::TruncIOp>(loc, i32Ty, high64);
+ Value low32 = rewriter.create<arith::MulIOp>(loc, value32, multiplier32);
- Value result32 = rewriter.create<arith::TruncIOp>(loc, resultTy, result64);
+ // Determine the direction and amount to shift the high bits.
+ Value shiftOver32 = rewriter.create<arith::CmpIOp>(
+ loc, arith::CmpIPredicate::sge, shift32, thirtyTwo32);
+ Value roundHighBits = rewriter.create<arith::CmpIOp>(
+ loc, arith::CmpIPredicate::sgt, shift32, thirtyTwo32);
- rewriter.replaceOp(op, result32);
+ Value shiftHighL =
+ rewriter.create<arith::SubIOp>(loc, thirtyTwo32, shift32);
+ Value shiftHighR =
+ rewriter.create<arith::SubIOp>(loc, shift32, thirtyTwo32);
+
+ shiftHighL =
+ rewriter.create<arith::SelectOp>(loc, shiftOver32, zero32, shiftHighL);
+ shiftHighR =
+ rewriter.create<arith::SelectOp>(loc, shiftOver32, shiftHighR, zero32);
+
+ // Conditionally perform our double round.
+ if (op.double_round()) {
+ Value negOne32 = getConstantValue(loc, i32Ty, -1, rewriter);
+ Value valuePositive = rewriter.create<arith::CmpIOp>(
+ loc, arith::CmpIPredicate::sge, value32, zero32);
+
+ Value roundDir =
+ rewriter.create<arith::SelectOp>(loc, valuePositive, one32, negOne32);
+ roundDir =
+ rewriter.create<arith::SelectOp>(loc, shiftOver32, roundDir, zero32);
+
+ Value shiftLow = rewriter.create<arith::ShRUIOp>(loc, low32, thirty32);
+ Value rounded = rewriter.create<arith::AddIOp>(loc, shiftLow, roundDir);
+ Value carry = rewriter.create<arith::ShRSIOp>(loc, rounded, two32);
+
+ Value shiftRound =
+ rewriter.create<arith::ShLIOp>(loc, roundDir, thirty32);
+
+ low32 = rewriter.create<arith::AddIOp>(loc, low32, shiftRound);
+ high32 = rewriter.create<arith::AddIOp>(loc, high32, carry);
+ }
+
+ // Conditionally apply rounding in the low bits.
+ {
+ Value shiftSubOne = rewriter.create<arith::SubIOp>(loc, shift32, one32);
+ Value roundBit = rewriter.create<arith::ShLIOp>(loc, one32, shiftSubOne);
+ roundBit = rewriter.create<arith::SelectOp>(loc, roundHighBits, zero32,
+ roundBit);
+
+ Value newLow32 = rewriter.create<arith::AddIOp>(loc, low32, roundBit);
+ Value wasRounded = rewriter.create<arith::CmpIOp>(
+ loc, arith::CmpIPredicate::ugt, low32, newLow32);
+ low32 = newLow32;
+
+ Value rounded32 = rewriter.create<arith::ExtUIOp>(loc, i32Ty, wasRounded);
+ high32 = rewriter.create<arith::AddIOp>(loc, high32, rounded32);
+ }
+
+ // Conditionally apply rounding in the high bits.
+ {
+ Value shiftSubOne =
+ rewriter.create<arith::SubIOp>(loc, shiftHighR, one32);
+ Value roundBit = rewriter.create<arith::ShLIOp>(loc, one32, shiftSubOne);
+ roundBit = rewriter.create<arith::SelectOp>(loc, roundHighBits, roundBit,
+ zero32);
+ high32 = rewriter.create<arith::AddIOp>(loc, high32, roundBit);
+ }
+
+ // Combine the correct high/low bits into the final rescale result.
+ high32 = rewriter.create<arith::ShLIOp>(loc, high32, shiftHighL);
+ high32 = rewriter.create<arith::ShRSIOp>(loc, high32, shiftHighR);
+ low32 = rewriter.create<arith::ShRUIOp>(loc, low32, shift32);
+ low32 = rewriter.create<arith::SelectOp>(loc, shiftOver32, zero32, low32);
+
+ // Apply the rounding behavior and shift to the final alignment.
+ Value result = rewriter.create<arith::AddIOp>(loc, low32, high32);
+
+ // Truncate if necessary.
+ if (!getElementTypeOrSelf(resultTy).isInteger(32)) {
+ result = rewriter.create<arith::TruncIOp>(loc, resultTy, result);
+ }
+
+ rewriter.replaceOp(op, result);
return success();
}
};
@@ -158,6 +250,9 @@ void mlir::tosa::populateTosaToArithConversionPatterns(
}
void mlir::tosa::populateTosaRescaleToArithConversionPatterns(
- RewritePatternSet *patterns) {
- patterns->add<ApplyScaleOpConverter>(patterns->getContext());
+ RewritePatternSet *patterns, bool include32Bit) {
+ patterns->add<ApplyScaleGenericOpConverter>(patterns->getContext(), 100);
+ if (include32Bit) {
+ patterns->add<ApplyScale32BitOpConverter>(patterns->getContext(), 200);
+ }
}
diff --git a/mlir/lib/Conversion/TosaToArith/TosaToArithPass.cpp b/mlir/lib/Conversion/TosaToArith/TosaToArithPass.cpp
index 03bdd6d2a8053..a4e4de6ef7a80 100644
--- a/mlir/lib/Conversion/TosaToArith/TosaToArithPass.cpp
+++ b/mlir/lib/Conversion/TosaToArith/TosaToArithPass.cpp
@@ -36,7 +36,8 @@ struct TosaToArith : public TosaToArithBase<TosaToArith> {
mlir::tosa::populateTosaToArithConversionPatterns(&patterns);
if (this->includeApplyRescale) {
- mlir::tosa::populateTosaRescaleToArithConversionPatterns(&patterns);
+ mlir::tosa::populateTosaRescaleToArithConversionPatterns(&patterns,
+ this->use32Bit);
target.addIllegalOp<tosa::ApplyScaleOp>();
}
diff --git a/mlir/test/Conversion/TosaToArith/tosa-to-arith.mlir b/mlir/test/Conversion/TosaToArith/tosa-to-arith.mlir
index 9e0c72f1686ae..5fb1a6b4c51b5 100644
--- a/mlir/test/Conversion/TosaToArith/tosa-to-arith.mlir
+++ b/mlir/test/Conversion/TosaToArith/tosa-to-arith.mlir
@@ -1,119 +1,126 @@
-// RUN: mlir-opt --split-input-file --tosa-to-arith="include-apply-rescale=true" %s -verify-diagnostics -o -| FileCheck %s
+// RUN: mlir-opt --split-input-file --tosa-to-arith="include-apply-rescale=true use-32-bit=true" %s -verify-diagnostics -o -| FileCheck %s
// RUN: mlir-opt --split-input-file --tosa-to-arith="include-apply-rescale=false" %s -verify-diagnostics -o -| FileCheck --check-prefix="SCALE" %s
// CHECK-LABEL: func @const_test
func.func @const_test() -> (tensor<i32>) {
// CHECK: [[C3:%.+]] = arith.constant dense<3> : tensor<i32>
- %0 = "tosa.const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>
+ %result = "tosa.const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>
// CHECK: return [[C3]]
- return %0 : tensor<i32>
+ return %result : tensor<i32>
}
// -----
// CHECK-LABEL: @apply_scale_test_i32
+// SCALE: "tosa.apply_scale"
func.func @apply_scale_test_i32(%arg0 : i32, %arg1 : i32, %arg2 : i8) -> (i32) {
- // CHECK-DAG: [[C1_8:%.+]] = arith.constant 1 : i8
- // CHECK-DAG: [[C1_32:%.+]] = arith.constant 1 : i32
- // CHECK-DAG: [[C1_64:%.+]] = arith.constant 1 : i64
- // CHECK-DAG: [[SHIFT_MINUS_ONE_8:%.+]] = arith.subi %arg2, [[C1_8]]
-
- // CHECK-DAG: [[SHIFT_32:%.+]] = arith.extsi %arg2 : i8 to i32
- // CHECK-DAG: [[SHIFT_MINUS_ONE_64:%.+]] = arith.extsi [[SHIFT_MINUS_ONE_8]] : i8 to i64
- // CHECK-DAG: [[SHIFTED_64:%.+]] = arith.shli [[C1_64]], [[SHIFT_MINUS_ONE_64]]
-
- // CHECK-DAG: [[C0_32:%.+]] = arith.constant 0 : i32
- // CHECK-DAG: [[C30_32:%.+]] = arith.constant 30 : i32
- // CHECK-DAG: [[SECOND_BIAS:%.+]] = arith.shli [[C1_32]], [[C30_32]]
- // CHECK-DAG: [[SECOND_BIAS_64:%.+]] = arith.extsi [[SECOND_BIAS]] : i32 to i64
- // CHECK-DAG: [[POSITIVE_ROUND:%.+]] = arith.addi [[SHIFTED_64]], [[SECOND_BIAS_64]]
- // CHECK-DAG: [[NEGATIVE_ROUND:%.+]] = arith.subi [[SHIFTED_64]], [[SECOND_BIAS_64]]
- // CHECK-DAG: [[VALUE_NEGATIVE:%.+]] = arith.cmpi sge, %arg0, [[C0_32]] : i32
- // CHECK-DAG: [[DOUBLE_ROUNDED:%.+]] = arith.select [[VALUE_NEGATIVE]], [[POSITIVE_ROUND]], [[NEGATIVE_ROUND]] : i64
- // CHECK-DAG: [[C32_32:%.+]] = arith.constant 32 : i32
- // CHECK-DAG: [[IS_32BIT_SHIFT:%.+]] = arith.cmpi sge, [[SHIFT_32]], [[C32_32]]
- // CHECK-DAG: [[ROUND:%.+]] = arith.select [[IS_32BIT_SHIFT]], [[DOUBLE_ROUNDED]], [[SHIFTED_64]]
-
- // CHECK-DAG: [[VAL_64:%.+]] = arith.extsi %arg0 : i32 to i64
- // CHECK-DAG: [[MULTIPLY_64:%.+]] = arith.extsi %arg1 : i32 to i64
- // CHECK-DAG: [[SHIFT_64:%.+]] = arith.extsi %arg2 : i8 to i64
- // CHECK-DAG: [[SCALED:%.+]] = arith.muli [[VAL_64]], [[MULTIPLY_64]]
- // CHECK-DAG: [[BIASED:%.+]] = arith.addi [[SCALED]], [[ROUND]]
- // CHECK-DAG: [[DOWNSHIFTED:%.+]] = arith.shrsi [[BIASED]], [[SHIFT_64]]
- // CHECK: [[TRUNCATED:%.+]] = arith.trunci [[DOWNSHIFTED]]
-
- // SCALE: "tosa.apply_scale"
- %0 = "tosa.apply_scale"(%arg0, %arg1, %arg2) {double_round = true} : (i32, i32, i8) -> i32
- return %0 : i32
+ // CHECK-DAG: %[[S32:.+]] = arith.extui %arg2 : i8 to i32
+ // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : i32
+ // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : i32
+ // CHECK-DAG: %[[C2:.+]] = arith.constant 2 : i32
+ // CHECK-DAG: %[[C30:.+]] = arith.constant 30 : i32
+ // CHECK-DAG: %[[C32:.+]] = arith.constant 32 : i32
+ // CHECK-DAG: %[[C32L:.+]] = arith.constant 32 : i64
+
+ // Compute the high-low values of the matmul in 64-bits.
+ // CHECK-DAG: %[[V64:.+]] = arith.extsi %arg0 : i32 to i64
+ // CHECK-DAG: %[[M64:.+]] = arith.extsi %arg1 : i32 to i64
+ // CHECK-DAG: %[[MUL64:.+]] = arith.muli %[[V64]], %[[M64]]
+ // CHECK-DAG: %[[HI64:.+]] = arith.shrui %[[MUL64]], %[[C32L]]
+ // CHECK-DAG: %[[HI:.+]] = arith.trunci %[[HI64]] : i64 to i32
+ // CHECK-DAG: %[[LOW:.+]] = arith.muli %arg0, %arg1
+
+ // Determine whether the high bits need to shift left or right and by how much.
+ // CHECK-DAG: %[[OVER31:.+]] = arith.cmpi sge, %[[S32]], %[[C32]]
+ // CHECK-DAG: %[[OVER32:.+]] = arith.cmpi sgt, %[[S32]], %[[C32]]
+ // CHECK-DAG: %[[HISHLN:.+]] = arith.subi %[[C32]], %[[S32]]
+ // CHECK-DAG: %[[HISHRN:.+]] = arith.subi %[[S32]], %[[C32]]
+ // CHECK-DAG: %[[HISHL:.+]] = arith.select %[[OVER31]], %[[C0]], %[[HISHLN]]
+ // CHECK-DAG: %[[HISHR:.+]] = arith.select %[[OVER31]], %[[HISHRN]], %[[C0]]
+
+ // Apply double rounding.
+ // CHECK-DAG: %[[CN1:.+]] = arith.constant -1
+ // CHECK-DAG: %[[POS:.+]] = arith.cmpi sge, %arg0, %[[C0]]
+ // CHECK-DAG: %[[DIR:.+]] = arith.select %[[POS]], %[[C1]], %[[CN1]]
+ // CHECK-DAG: %[[DRND:.+]] = arith.select %[[OVER31]], %[[DIR]], %[[C0]]
+ // CHECK-DAG: %[[DSHFTR:.+]] = arith.shrui %[[LOW]], %[[C30]]
+ // CHECK-DAG: %[[DRNDED:.+]] = arith.addi %[[DSHFTR]], %[[DRND]]
+ // CHECK-DAG: %[[DCARRY:.+]] = arith.shrsi %[[DRNDED]], %[[C2:.+]]
+ // CHECK-DAG: %[[DBIT:.+]] = arith.shli %[[DRND]], %[[C30]]
+ // CHECK-DAG: %[[DLOW:.+]] = arith.addi %[[LOW]], %[[DBIT]]
+ // CHECK-DAG: %[[DHI:.+]] = arith.addi %[[HI]], %[[DCARRY]]
+
+ // Apply low-bit rounding.
+ // CHECK-DAG: %[[SHFTM1:.+]] = arith.subi %[[S32]], %[[C1]]
+ // CHECK-DAG: %[[LBIT:.+]] = arith.shli %[[C1]], %[[SHFTM1]]
+ // CHECK-DAG: %[[HALF:.+]] = arith.select %[[OVER32]], %[[C0]], %[[LBIT]]
+ // CHECK-DAG: %[[LADD:.+]] = arith.addi %[[DLOW]], %[[HALF]]
+ // CHECK-DAG: %[[LLO:.+]] = arith.cmpi ugt, %[[DLOW]], %[[LADD]]
+ // CHECK-DAG: %[[LCARRY:.+]] = arith.extui %[[LLO]] : i1 to i32
+ // CHECK-DAG: %[[LRNDED:.+]] = arith.addi %[[DHI]], %[[LCARRY]]
+
+ // Apply high-bit rounding.
+ // CHECK-DAG: %[[HISHRM1:.+]] = arith.subi %[[HISHR]], %[[C1]]
+ // CHECK-DAG: %[[LHISHFT:.+]] = arith.shli %[[C1]], %[[HISHRM1]]
+ // CHECK-DAG: %[[LHI:.+]] = arith.select %[[OVER32]], %[[LHISHFT]], %[[C0]]
+ // CHECK-DAG: %[[FHI:.+]] = arith.addi %[[LRNDED]], %[[LHI]]
+
+ // Combine hi-low into the final result.
+ // CHECK-DAG: %[[HIL:.+]] = arith.shli %[[FHI]], %[[HISHL]]
+ // CHECK-DAG: %[[HIALIGN:.+]] = arith.shrsi %[[HIL:.+]], %[[HISHR]]
+ // CHECK-DAG: %[[LOR:.+]] = arith.shrui %[[LADD]], %[[S32]]
+ // CHECK-DAG: %[[LOWALIGN:.+]] = arith.select %[[OVER31]], %[[C0]], %[[LOR]]
+ // CHECK-DAG: %[[RESULT:.+]] = arith.addi %[[LOWALIGN]], %[[HIALIGN]]
+ // CHECK: return %[[RESULT]]
+ %res = "tosa.apply_scale"(%arg0, %arg1, %arg2) {double_round = true} : (i32, i32, i8) -> i32
+ return %res : i32
}
// -----
// CHECK-LABEL: @apply_scale_test_vector
+// SCALE: "tosa.apply_scale"
func.func @apply_scale_test_vector(%arg0 : vector<4xi32>, %arg1 : vector<4xi32>, %arg2 : vector<4xi8>) -> (vector<4xi32>) {
- // CHECK-DAG: [[C1_8:%.+]] = arith.constant dense<1> : vector<4xi8>
- // CHECK-DAG: [[C1_32:%.+]] = arith.constant dense<1> : vector<4xi32>
- // CHECK-DAG: [[C1_64:%.+]] = arith.constant dense<1> : vector<4xi64>
- // CHECK-DAG: [[SHIFT_MINUS_ONE_8:%.+]] = arith.subi %arg2, [[C1_8]]
-
- // CHECK-DAG: [[SHIFT_32:%.+]] = arith.extsi %arg2 : vector<4xi8> to vector<4xi32>
- // CHECK-DAG: [[SHIFT_MINUS_ONE_64:%.+]] = arith.extsi [[SHIFT_MINUS_ONE_8]] : vector<4xi8> to vector<4xi64>
- // CHECK-DAG: [[SHIFTED_64:%.+]] = arith.shli [[C1_64]], [[SHIFT_MINUS_ONE_64]]
-
- // CHECK-DAG: [[C0_32:%.+]] = arith.constant dense<0> : vector<4xi32>
- // CHECK-DAG: [[C30_32:%.+]] = arith.constant dense<30> : vector<4xi32>
- // CHECK-DAG: [[SECOND_BIAS:%.+]] = arith.shli [[C1_32]], [[C30_32]]
- // CHECK-DAG: [[SECOND_BIAS_64:%.+]] = arith.extsi [[SECOND_BIAS]] : vector<4xi32> to vector<4xi64>
- // CHECK-DAG: [[POSITIVE_ROUND:%.+]] = arith.addi [[SHIFTED_64]], [[SECOND_BIAS_64]]
- // CHECK-DAG: [[NEGATIVE_ROUND:%.+]] = arith.subi [[SHIFTED_64]], [[SECOND_BIAS_64]]
- // CHECK-DAG: [[VALUE_NEGATIVE:%.+]] = arith.cmpi sge, %arg0, [[C0_32]] : vector<4xi32>
- // CHECK-DAG: [[DOUBLE_ROUNDED:%.+]] = arith.select [[VALUE_NEGATIVE]], [[POSITIVE_ROUND]], [[NEGATIVE_ROUND]] : vector<4xi1>, vector<4xi64>
- // CHECK-DAG: [[C32_32:%.+]] = arith.constant dense<32> : vector<4xi32>
- // CHECK-DAG: [[IS_32BIT_SHIFT:%.+]] = arith.cmpi sge, [[SHIFT_32]], [[C32_32]]
- // CHECK-DAG: [[ROUND:%.+]] = arith.select [[IS_32BIT_SHIFT]], [[DOUBLE_ROUNDED]], [[SHIFTED_64]]
-
- // CHECK-DAG: [[VAL_64:%.+]] = arith.extsi %arg0 : vector<4xi32> to vector<4xi64>
- // CHECK-DAG: [[MULTIPLY_64:%.+]] = arith.extsi %arg1 : vector<4xi32> to vector<4xi64>
- // CHECK-DAG: [[SHIFT_64:%.+]] = arith.extsi %arg2 : vector<4xi8> to vector<4xi64>
- // CHECK-DAG: [[SCALED:%.+]] = arith.muli [[VAL_64]], [[MULTIPLY_64]]
- // CHECK-DAG: [[BIASED:%.+]] = arith.addi [[SCALED]], [[ROUND]]
- // CHECK-DAG: [[DOWNSHIFTED:%.+]] = arith.shrsi [[BIASED]], [[SHIFT_64]]
- // CHECK: [[TRUNCATED:%.+]] = arith.trunci [[DOWNSHIFTED]]
-
- %0 = "tosa.apply_scale"(%arg0, %arg1, %arg2) {double_round = true} : (vector<4xi32>, vector<4xi32>, vector<4xi8>) -> vector<4xi32>
- return %0 : vector<4xi32>
+ // CHECK-NOT: "tosa.apply_scale"
+ %res = "tosa.apply_scale"(%arg0, %arg1, %arg2) {double_round = true} : (vector<4xi32>, vector<4xi32>, vector<4xi8>) -> vector<4xi32>
+ return %res : vector<4xi32>
}
// -----
// CHECK-LABEL: @apply_scale_test_i48
+// SCALE: "tosa.apply_scale"
func.func @apply_scale_test_i48(%arg0 : i48, %arg1 : i32, %arg2 : i8) -> (i32) {
- // CHECK-DAG: [[C1_8:%.+]] = arith.constant 1 : i8
- // CHECK-DAG: [[C1_32:%.+]] = arith.constant 1 : i32
- // CHECK-DAG: [[C1_64:%.+]] = arith.constant 1 : i64
- // CHECK-DAG: [[C30_32:%.+]] = arith.constant 30 : i32
- // CHECK-DAG: [[C0_32:%.+]] = arith.constant 0 : i48
- // CHECK-DAG: [[C32_32:%.+]] = arith.constant 32 : i32
- // CHECK-DAG: [[SHIFT_MINUS_ONE_8:%.+]] = arith.subi %arg2, [[C1_8]]
- // CHECK-DAG: [[SHIFT_32:%.+]] = arith.extsi %arg2 : i8 to i32
- // CHECK-DAG: [[SHIFT_MINUS_ONE_64:%.+]] = arith.extsi [[SHIFT_MINUS_ONE_8]] : i8 to i64
- // CHECK-DAG: [[SHIFTED_64:%.+]] = arith.shli [[C1_64]], [[SHIFT_MINUS_ONE_64]]
- // CHECK-DAG: [[SECOND_BIAS:%.+]] = arith.shli [[C1_32]], [[C30_32]]
- // CHECK-DAG: [[SECOND_BIAS_64:%.+]] = arith.extsi [[SECOND_BIAS]] : i32 to i64
- // CHECK-DAG: [[POSITIVE_ROUND:%.+]] = arith.addi [[SHIFTED_64]], [[SECOND_BIAS_64]]
- // CHECK-DAG: [[NEGATIVE_ROUND:%.+]] = arith.subi [[SHIFTED_64]], [[SECOND_BIAS_64]]
- // CHECK-DAG: [[VALUE_NEGATIVE:%.+]] = arith.cmpi sge, %arg0, [[C0_32]] : i48
- // CHECK-DAG: [[DOUBLE_ROUNDED:%.+]] = arith.select [[VALUE_NEGATIVE]], [[POSITIVE_ROUND]], [[NEGATIVE_ROUND]] : i64
- // CHECK-DAG: [[IS_32BIT_SHIFT:%.+]] = arith.cmpi sge, [[SHIFT_32]], [[C32_32]]
- // CHECK-DAG: [[ROUND:%.+]] = arith.select [[IS_32BIT_SHIFT]], [[DOUBLE_ROUNDED]], [[SHIFTED_64]]
- // CHECK-DAG: [[VAL_64:%.+]] = arith.extsi %arg0 : i48 to i64
- // CHECK-DAG: [[MULTIPLY_64:%.+]] = arith.extsi %arg1 : i32 to i64
- // CHECK-DAG: [[SHIFT_64:%.+]] = arith.extsi %arg2 : i8 to i64
- // CHECK-DAG: [[SCALED:%.+]] = arith.muli [[VAL_64]], [[MULTIPLY_64]]
- // CHECK-DAG: [[BIASED:%.+]] = arith.addi [[SCALED]], [[ROUND]]
- // CHECK-DAG: [[DOWNSHIFTED:%.+]] = arith.shrsi [[BIASED]], [[SHIFT_64]]
- // CHECK: [[TRUNCATED:%.+]] = arith.trunci [[DOWNSHIFTED]]
- %0 = "tosa.apply_scale"(%arg0, %arg1, %arg2) {double_round = true} : (i48, i32, i8) -> i32
- return %0 : i32
+ // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : i48
+ // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : i64
+ // CHECK-DAG: %[[C31:.+]] = arith.constant 31 : i32
+
+ // Multiply in 64 bits.
+ // CHECK-DAG: %[[V64:.+]] = arith.extsi %arg0 : i48 to i64
+ // CHECK-DAG: %[[M64:.+]] = arith.extsi %arg1 : i32 to i64
+ // CHECK-DAG: %[[MUL:.+]] = arith.muli %[[V64]], %[[M64]]
+
+ // Round normally.
+ // CHECK-DAG: %[[S32:.+]] = arith.extui %arg2 : i8 to i32
+ // CHECK-DAG: %[[S64:.+]] = arith.extui %[[S32]] : i32 to i64
+ // CHECK-DAG: %[[ONEL:.+]] = arith.shli %[[C1]], %[[S64]] : i64
+ // CHECK-DAG: %[[ONER:.+]] = arith.shrui %[[ONEL]], %[[C1]]
+ // CHECK-DAG: %[[ROUND:.+]] = arith.addi %[[MUL]], %[[ONER]]
+
+ // Apply double rounding.
+ // CHECK-DAG: %[[DUP:.+]] = arith.constant 1073741824 : i64
+ // CHECK-DAG: %[[DDOWN:.+]] = arith.constant -1073741824 : i64
+ // CHECK-DAG: %[[POS:.+]] = arith.cmpi sge, %arg0, %[[C0]]
+ // CHECK-DAG: %[[DBIT:.+]] = arith.select %[[POS]], %[[DUP]], %[[DDOWN]]
+ // CHECK-DAG: %[[DRND:.+]] = arith.addi %[[DBIT]], %[[ROUND]]
+ // CHECK-DAG: %[[USED:.+]] = arith.cmpi sgt, %[[S32]], %[[C31]] : i32
+ // CHECK-DAG: %[[RES64:.+]] = arith.select %[[USED]], %[[DRND]], %[[ROUND]] : i64
+
+ // Shift and truncate final answer.
+ // CHECK-DAG: %[[SHR:.+]] = arith.shrsi %[[RES64]], %[[S64]]
+ // CHECK-DAG: %[[TRUNC:.+]] = arith.trunci %[[SHR]] : i64 to i32
+ // CHECK: return %[[TRUNC]]
+ %res = "tosa.apply_scale"(%arg0, %arg1, %arg2) {double_round = true} : (i48, i32, i8) -> i32
+ return %res : i32
}
More information about the Mlir-commits
mailing list