[rv32] Expand `arith.mulsi_extended` before going to LLVM (#12241)

This enables `tosa.apply_scale` to be vectorized, and thus fixes a code
size regression.

Fixes: https://github.com/iree-org/iree/issues/12239

compiler/src/iree/compiler/Codegen/LLVMCPU/ConvertToLLVM.cpp

diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/ConvertToLLVM.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/ConvertToLLVM.cpp
index 1c7aae3..b19f1ca 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/ConvertToLLVM.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/ConvertToLLVM.cpp
@@ -48,6 +48,9 @@
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
@@ -358,6 +361,49 @@
   LLVMTypeConverter &typeConverter;
 };
 
+/// The 32-bit RISC-V backend is very sensitive to how extended multiplication
+/// is lowered. This pattern lowers `arith.mulsi_extended` before going to the
+/// LLVM dialect, in a way compatible with that backend, so that we break down
+/// any 64-bit constants that would otherwise prevent the code from being
+/// vectorized.
+class ExpandMulSIExtended : public OpRewritePattern<arith::MulSIExtendedOp> {
+ public:
+  using OpRewritePattern<arith::MulSIExtendedOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(arith::MulSIExtendedOp op,
+                                PatternRewriter &rewriter) const override {
+    Type resultType = op.getLhs().getType();
+    if (getElementTypeOrSelf(resultType).getIntOrFloatBitWidth() != 32) {
+      return failure();
+    }
+
+    Location loc = op.getLoc();
+
+    Type wideType = rewriter.getIntegerType(64);
+    // Shift amount necessary to extract the high bits from widened result.
+    Attribute shiftValAttr = rewriter.getI64IntegerAttr(32);
+    if (auto vecTy = resultType.dyn_cast<VectorType>()) {
+      wideType = VectorType::get(vecTy.getShape(), wideType);
+      shiftValAttr = SplatElementsAttr::get(wideType, shiftValAttr);
+    }
+    Value shiftVal = rewriter.create<arith::ConstantOp>(loc, shiftValAttr);
+
+    Value lhsExt = rewriter.create<arith::ExtSIOp>(loc, wideType, op.getLhs());
+    Value rhsExt = rewriter.create<arith::ExtSIOp>(loc, wideType, op.getRhs());
+    Value mulExt =
+        rewriter.create<arith::MulIOp>(loc, wideType, lhsExt, rhsExt);
+    Value low = rewriter.create<arith::MulIOp>(loc, resultType, op.getLhs(),
+                                               op.getRhs());
+
+    // Produce two 32-bit results.
+    Value highExt = rewriter.create<arith::ShRUIOp>(loc, mulExt, shiftVal);
+    Value high = rewriter.create<arith::TruncIOp>(loc, resultType, highExt);
+
+    rewriter.replaceOp(op, {low, high});
+    return success();
+  }
+};
+
 class ConvertToLLVMPass : public ConvertToLLVMBase<ConvertToLLVMPass> {
  public:
   ConvertToLLVMPass(bool reassociateFpReductions) {
@@ -465,6 +511,12 @@
   }
   tosa::populateTosaRescaleToArithConversionPatterns(&patterns, use32BitImpl);
 
+  // Make sure we expand any `arith.mulsi_extended` before going to the LLVM
+  // dialect.
+  if (use32BitImpl) {
+    patterns.add<ExpandMulSIExtended>(patterns.getContext(), /*benefit=*/1024);
+  }
+
   populateAffineToStdConversionPatterns(patterns);
   populateSCFToControlFlowConversionPatterns(patterns);
   cf::populateControlFlowToLLVMConversionPatterns(typeConverter, patterns);

compiler/src/iree/compiler/Codegen/LLVMCPU/test/apply_scale_lowering.mlir

diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/apply_scale_lowering.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/apply_scale_lowering.mlir
index c5602fd..a86a5af 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/apply_scale_lowering.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/apply_scale_lowering.mlir
@@ -180,8 +180,13 @@
 
 // 32-bit lowering is used with '+zve32x'. Note that the 32-bit lowering
 // generates 64-bit mul operations that are decomposed into 32-bit operations by
-// the LLVM backend.
+// the LLVM backend. The backend expects both the low half to be an `llvm.mul` op.
 // CHECK-LABEL: llvm.func @apply_scale_zve32x
-//       CHECK:   %[[MUL:.*]] = llvm.mul %{{.*}}, %{{.*}} : vector<2xi64>
-//       CHECK:   %[[SHR:.*]] = llvm.lshr %{{.*}}, %{{.*}} : vector<2xi64>
+//   CHECK-DAG:   %[[RHS:.+]]    = llvm.mlir.constant(dense<19689> : vector<2xi32>) : vector<2xi32>
+//   CHECK-DAG:   %[[RHSEXT:.+]] = llvm.mlir.constant(dense<19689> : vector<2xi64>) : vector<2xi64>
+//       CHECK:   %[[LHS:.+]]    = llvm.load %{{.+}} {alignment = 4 : i64} : !llvm.ptr<vector<2xi32>>
+//       CHECK:   %[[LHSEXT:.+]] = llvm.sext %[[LHS]] : vector<2xi32> to vector<2xi64>
+//       CHECK:   %[[MULEXT:.*]] = llvm.mul %[[LHSEXT]], %[[RHSEXT]] : vector<2xi64>
+//       CHECK:   %[[MULLOW:.*]] = llvm.mul %[[LHS]], %[[RHS]] : vector<2xi32>
+//       CHECK:   %[[SHR:.*]]    = llvm.lshr %[[MULEXT]], %{{.*}} : vector<2xi64>
 //  CHECK-NEXT:   llvm.trunc %[[SHR]] : vector<2xi64> to vector<2xi32>