[mlir][arith][tosa] Use extended mul in 32-bit `tosa.apply_scale`

[kuhar]

Summary

To not introduce 64-bit types that may be difficult to handle for some
targets.

[antiagainst]

Awesome! Please wait for @rsuderman to take a look too!

mlir/lib/Conversion/TosaToArith/TosaToArith.cpp

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.getValue();
     if (getElementTypeOrSelf(value.getType()).getIntOrFloatBitWidth() > 32) {
       return failure();
     }
 
     Value value32 = op.getValue();
     Value multiplier32 = op.getMultiplier();
     Value shift32 = rewriter.create<arith::ExtUIOp>(loc, i32Ty, op.getShift());
 
     // 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 multiply64 =
-        rewriter.create<arith::MulIOp>(loc, value64, multiplier64);
-
     // 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);
+    auto value64 =
+        rewriter.create<arith::MulSIExtendedOp>(loc, value32, multiplier32);
+    Value low32 = value64.getLow();
+    Value high32 = value64.getHigh();
 
     // 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);
 
     Value shiftHighL =

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

 // SCALE: "tosa.apply_scale"
 func.func @apply_scale_test_i32(%arg0 : i32, %arg1 : i32, %arg2 : i8) -> (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
+  // CHECK-DAG: %[[LOW:.+]], %[[HI:.+]] = arith.mulsi_extended %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]]