[mlir][tosa] Rework tosa.apply_scale lowering for 32-bit

[rsuderman]

Summary

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.

[rsuderman]

Tweaked to clean up a bit.

[rsuderman]

Tweaked the 64-bit path to slightly decrease the amount of IR.

Added different populate methods depending on the bitwith supported.

[rsuderman]

Dag-ified the checks

[mravishankar]

Dont have an opinion about the transformation itself, but more about expected usage.

mlir/lib/Conversion/TosaToArith/TosaToArith.cpp

61

Why is this ApplyScale48OpConverter?

263

This probably needs to be added with a higher benefit. If both this pattern and the non-32 bit patterns are added by a caller, then it is undefined which will trigger first. That can lead to strange compilation.

[dcaballe]

Thanks for working on this, Rob! I really can't comment on the actual implementation from the TOSA perspective. However, I would suggest not enabling this by default. Architectures with full 64-bit support should do better with the 64-bit lowering. I can follow up with a patch to enable this for the RISC-V specific cases.

mlir/lib/Conversion/TosaToArith/TosaToArith.cpp

263

+1.
Nit: Alternatively, we could have a single populate function with a flag to enable the 32-bit lowering. That would keep both related patterns and their priorities in a single function and avoid potential misuses.

[rsuderman]

Updated for change how TosaToArith target 32-bit vs 64-bit.

[rsuderman]

mlir/lib/Conversion/TosaToArith/TosaToArith.cpp

61

48 meant that it was meant for 48-bit conversion. Because we want to use this for 32-bit operations in some cases I changed it to Generic.

263

Updated to include benefit values and just a bool to decide whether to include the 32-bit lowering path. Note that on llvm integrate we are going to need an integration fix.

[mravishankar]

Removing my blocker, but I didnt check the math. Probably worth getting someone to review that.

mlir/lib/Conversion/TosaToArith/TosaToArithPass.cpp

40 ↗

(On Diff #429841)

You could set this to default false?

[rsuderman]

Added default 32-bit setting.

[rsuderman]

mlir/lib/Conversion/TosaToArith/TosaToArithPass.cpp

40 ↗

(On Diff #429841)

It is set to default to false in the Passes.td file. Added a default on the function definition as well.

[rsuderman]

Resyncing

[rsuderman]

Update relative benefit

[dcaballe]

Thanks for addressing the changes. LGTM but, as Mahesh mentioned, it's probably worth getting someone with experience in TOSA to take a look.

mlir/lib/Conversion/TosaToArith/TosaToArith.cpp

 //===- TosaToArith.cpp - Lowering Tosa to Arith Dialect -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // These rewriters lower from the Tosa to the Arith dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/TosaToArith/TosaToArith.h"
 #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;
 using namespace tosa;
 
 namespace {
 
 class ConstOpConverter : public OpRewritePattern<tosa::ConstOp> {
if (auto shapedTy = type.dyn_cast<ShapedType>()) {
     Type eTy = shapedTy.getElementType();
     APInt valueInt(eTy.getIntOrFloatBitWidth(), value);
     return DenseIntElementsAttr::get(shapedTy, valueInt);
   }
 
   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 ApplyScale48OpConverter : 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 valueETy = getElementTypeOrSelf(valueTy);
     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.
+    if (valueETy.getIntOrFloatBitWidth() <= 32)
+      return failure();
 
-    Value round64 = rewriter.create<arith::ShLIOp>(
-        loc, one64, rewriter.create<arith::ExtSIOp>(loc, i64Ty, shiftSubOne8));
+    Value shift32 = rewriter.create<arith::ExtUIOp>(loc, i32Ty, op.shift());
 
-    // 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 doubleRound64 = rewriter.create<arith::SelectOp>(
-          loc, valueGreaterThanZero, roundAdd64, roundSub64);
+    Value one64 = getConstantValue(loc, i64Ty, 1, rewriter);
+    Value negOne64 = getConstantValue(loc, i64Ty, -1, rewriter);
+    Value thirty64 = getConstantValue(loc, i64Ty, 30, rewriter);
+    Value thirtyOne32 = getConstantValue(loc, i32Ty, 31, rewriter);
 
-      // 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);
+    // 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);
+
+    // 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);
+
+    // Apply double rounding if necessary.
+    if (op.double_round()) {
+      Value positive = rewriter.create<arith::CmpIOp>(
+          loc, arith::CmpIPredicate::sge, value, zero);
+      Value dir =
+          rewriter.create<arith::SelectOp>(loc, positive, one64, negOne64);
+      Value update = rewriter.create<arith::ShLIOp>(loc, dir, thirty64);
+      Value val = rewriter.create<arith::AddIOp>(loc, update, 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 ApplyScale32OpConverter : 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();
   }
 };
 
 } // namespace
 
 void mlir::tosa::populateTosaToArithConversionPatterns(
     RewritePatternSet *patterns) {
   patterns->add<ConstOpConverter>(patterns->getContext());
 }
 
 void mlir::tosa::populateTosaRescaleToArithConversionPatterns(
     RewritePatternSet *patterns) {
-  patterns->add<ApplyScaleOpConverter>(patterns->getContext());
+  patterns->add<ApplyScale32OpConverter>(patterns->getContext());
+  patterns->add<ApplyScale48OpConverter>(patterns->getContext());
 }

mlir/test/Conversion/TosaToArith/tosa-to-arith.mlir

 // 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=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
-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
+func.func @apply_scale_test_i32(%arg0 : i32, %arg1 : i32, %arg2 : i8) -> (i32) {
+  // CHECK: %[[S32:.+]] = arith.extui %arg2 : i8 to i32
+  // CHECK: %[[C0:.+]] = arith.constant 0 : i32
+  // CHECK: %[[C1:.+]] = arith.constant 1 : i32
+  // CHECK: %[[C2:.+]] = arith.constant 2 : i32
+  // CHECK: %[[C30:.+]] = arith.constant 30 : i32
+  // CHECK: %[[C32:.+]] = arith.constant 32 : i32
+  // CHECK: %[[C32L:.+]] = arith.constant 32 : i64
+
+  // Compute the high-low values of the matmul in 64-bits.
+  // CHECK: %[[V64:.+]] = arith.extsi %arg0 : i32 to i64
+  // CHECK: %[[M64:.+]] = arith.extsi %arg1 : i32 to i64
+  // CHECK: %[[MUL64:.+]] = arith.muli %[[V64]], %[[M64]]
+  // CHECK: %[[HI64:.+]] = arith.shrui %[[MUL64]], %[[C32L]]
+  // CHECK: %[[HI:.+]] = arith.trunci %[[HI64]] : i64 to i32
+  // CHECK: %[[LOW:.+]] = arith.muli %arg0, %arg1
+
+  // Determine whether the high bits need to shift left or right and by how much.
+  // CHECK: %[[OVER31:.+]] = arith.cmpi sge, %[[S32]], %[[C32]]
+  // CHECK: %[[OVER32:.+]] = arith.cmpi sgt, %[[S32]], %[[C32]]
+  // CHECK: %[[HISHLN:.+]] = arith.subi %[[C32]], %[[S32]]
+  // CHECK: %[[HISHRN:.+]] = arith.subi %[[S32]], %[[C32]]
+  // CHECK: %[[HISHL:.+]] = arith.select %[[OVER31]], %[[C0]], %[[HISHLN]]
+  // CHECK: %[[HISHR:.+]] = arith.select %[[OVER31]], %[[HISHRN]], %[[C0]]
+
+  // Apply double rounding.
+  // CHECK: %[[CN1:.+]] = arith.constant -1
+  // CHECK: %[[POS:.+]] = arith.cmpi sge, %arg0, %[[C0]]
+  // CHECK: %[[DIR:.+]] = arith.select %[[POS]], %[[C1]], %[[CN1]]
+  // CHECK: %[[DRND:.+]] = arith.select %[[OVER31]], %[[DIR]], %[[C0]]
+  // CHECK: %[[DSHFTR:.+]] = arith.shrui %[[LOW]], %[[C30]]
+  // CHECK: %[[DRNDED:.+]] = arith.addi %[[DSHFTR]], %[[DRND]]
+  // CHECK: %[[DCARRY:.+]] = arith.shrsi %[[DRNDED]], %[[C2:.+]]
+  // CHECK: %[[DBIT:.+]] = arith.shli %[[DRND]], %[[C30]]
+  // CHECK: %[[DLOW:.+]] = arith.addi %[[LOW]], %[[DBIT]]
+  // CHECK: %[[DHI:.+]] = arith.addi %[[HI]], %[[DCARRY]]
+
+  // Apply low-bit rounding.
+  // CHECK: %[[SHFTM1:.+]] = arith.subi %[[S32]], %[[C1]]
+  // CHECK: %[[LBIT:.+]] = arith.shli %[[C1]], %[[SHFTM1]]
+  // CHECK: %[[HALF:.+]] = arith.select %[[OVER32]], %[[C0]], %[[LBIT]]
+  // CHECK: %[[LADD:.+]] = arith.addi %[[DLOW]], %[[HALF]]
+  // CHECK: %[[LLO:.+]] = arith.cmpi ugt, %[[DLOW]], %[[LADD]]
+  // CHECK: %[[LCARRY:.+]] = arith.extui %[[LLO]] : i1 to i32
+  // CHECK: %[[LRNDED:.+]] = arith.addi %[[DHI]], %[[LCARRY]]
+
+  // Apply high-bit rounding.
+  // CHECK: %[[HISHRM1:.+]] = arith.subi %[[HISHR]], %[[C1]]
+  // CHECK: %[[LHISHFT:.+]] = arith.shli %[[C1]], %[[HISHRM1]]
+  // CHECK: %[[LHI:.+]] = arith.select %[[OVER32]], %[[LHISHFT]], %[[C0]]
+  // CHECK: %[[FHI:.+]] = arith.addi %[[LRNDED]], %[[LHI]]
+
+  // Combine hi-low into the final result.
+  // CHECK: %[[HIL:.+]] = arith.shli %[[FHI]], %[[HISHL]]
+  // CHECK: %[[HIALIGN:.+]] = arith.shrsi %[[HIL:.+]], %[[HISHR]] 
+  // CHECK: %[[LOR:.+]] = arith.shrui %[[LADD]], %[[S32]]
+  // CHECK: %[[LOWALIGN:.+]] = arith.select %[[OVER31]], %[[C0]], %[[LOR]]
+  // CHECK: %[[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: %[[CN1:.+]] = arith.constant -1 : i64
+  // CHECK-DAG: %[[C30:.+]] = arith.constant 30 : i64
+  // CHECK-DAG: %[[C31:.+]] = arith.constant 31 : i32
+  // 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]]
+
+  // 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]]
+  // CHECK-DAG: %[[POS:.+]] = arith.cmpi sge, %arg0, %[[C0]]
+  // CHECK-DAG: %[[DBIT:.+]] = arith.select %[[POS]], %[[C1]], %[[CN1]]
+  // CHECK-DAG: %[[DSHFT:.+]] = arith.shli %[[DBIT]], %[[C30]]
+  // CHECK-DAG: %[[DRND:.+]] = arith.addi %[[DSHFT]], %[[ROUND]]
+  // CHECK-DAG: %[[USED:.+]] = arith.cmpi sgt, %[[S32]], %[[C31]] : i32
+  // CHECK-DAG: %[[RES64:.+]] = arith.select %[[USED]], %[[DRND]], %[[ROUND]] : i64
+  // 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
 }