Use factorization

Author: el-ev

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -32,6 +32,7 @@
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/SDPatternMatch.h"
 #include "llvm/CodeGen/SelectionDAGAddressAnalysis.h"
+#include "llvm/CodeGen/SelectionDAGNodes.h"
 #include "llvm/CodeGen/TargetLoweringObjectFileImpl.h"
 #include "llvm/CodeGen/ValueTypes.h"
 #include "llvm/IR/DiagnosticInfo.h"
@@ -49,6 +50,7 @@
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
+#include <cstdint>
 #include <optional>
 
 using namespace llvm;
@@ -15437,15 +15439,10 @@ static SDValue performXORCombine(SDNode *N, SelectionDAG &DAG,
   return combineSelectAndUseCommutative(N, DAG, /*AllOnes*/ false, Subtarget);
 }
 
-// Try to expand a multiply to a sequence of shifts and add/subs,
-// for a machine w/o native mul instruction.
-static SDValue expandMulToBasicOps(SDNode *N, SelectionDAG &DAG,
-                                   uint64_t MulAmt) {
-  const uint64_t BitWidth = N->getValueType(0).getFixedSizeInBits();
-  SDLoc DL(N);
-
-  if (MulAmt == 0)
-    return DAG.getConstant(0, DL, N->getValueType(0));
+static SDValue expandMulToNAFSequence(SDNode *N, SelectionDAG &DAG,
+                                      const SDLoc &DL, uint64_t MulAmt) {
+  EVT VT = N->getValueType(0);
+  const uint64_t BitWidth = VT.getFixedSizeInBits();
 
   // Find the Non-adjacent form of the multiplier.
   llvm::SmallVector<std::pair<bool, uint64_t>> Sequence; // {isAdd, shamt}
@@ -15470,17 +15467,89 @@ static SDValue expandMulToBasicOps(SDNode *N, SelectionDAG &DAG,
     SDValue ShiftVal;
     if (Op.second > 0)
       ShiftVal =
-          DAG.getNode(ISD::SHL, DL, N->getValueType(0), N0,
-                      DAG.getConstant(Op.second, DL, N->getValueType(0)));
+          DAG.getNode(ISD::SHL, DL, VT, N0, DAG.getConstant(Op.second, DL, VT));
     else
       ShiftVal = N0;
 
     ISD::NodeType AddSubOp = Op.first ? ISD::ADD : ISD::SUB;
-    Result = DAG.getNode(AddSubOp, DL, N->getValueType(0), Result, ShiftVal);
+    Result = DAG.getNode(AddSubOp, DL, VT, Result, ShiftVal);
   }
-
   return Result;
 }
+// Try to expand a multiply to a sequence of shifts and add/subs,
+// for a machine w/o native mul instruction.
+static SDValue expandMulToBasicOps(SDNode *N, SelectionDAG &DAG,
+                                   uint64_t MulAmt) {
+  EVT VT = N->getValueType(0);
+  const uint64_t BitWidth = VT.getFixedSizeInBits();
+  SDLoc DL(N);
+
+  if (MulAmt == 0)
+    return DAG.getConstant(0, DL, N->getValueType(0));
+
+  // Try to factorize into (2^N) * (2^M_1 +/- 1) + (2^M_2 +/- 1) + ...
+  uint64_t E = MulAmt;
+  uint64_t TrailingZeros = 0;
+
+  while (E > 0 && (E & 1) == 0) {
+    E >>= 1;
+    TrailingZeros++;
+  }
+
+  llvm::SmallVector<std::pair<bool, uint64_t>> Factors; // {is_2^M+1, M}
+  
+  while (E > 1) { 
+      bool Found = false;
+      for (int64_t I = BitWidth - 1; I >= 2; --I) {    
+          uint64_t Factor = 1ULL << I;
+
+          if (E % (Factor - 1) == 0) {
+              Factors.push_back({false, I}); 
+              E /= Factor - 1;
+              Found = true;
+              break;
+          }
+
+          if (E % (Factor + 1) == 0) {
+              Factors.push_back({true, I}); 
+              E /= Factor + 1;
+              Found = true;
+              break; 
+          }
+      }
+      if (!Found)
+        break;
+  }
+
+  SDValue Result;
+  SDValue N0 = N->getOperand(0);
+
+  bool UseFactorization =
+      !Factors.empty() && (E < MulAmt) && (Factors.size() < 5);
+
+  if (UseFactorization) {
+    if (E == 1)
+      Result = N0;
+    else
+      Result = expandMulToNAFSequence(N, DAG, DL, E);
+
+    for (const auto &F : Factors) {
+      SDValue ShiftVal = DAG.getNode(ISD::SHL, DL, VT, Result,
+                                     DAG.getConstant(F.second, DL, VT));
+
+      ISD::NodeType AddSubOp = F.first ? ISD::ADD : ISD::SUB;
+      Result = DAG.getNode(AddSubOp, DL, N->getValueType(0), ShiftVal, Result);
+    }
+
+    if (TrailingZeros > 0)
+      Result = DAG.getNode(ISD::SHL, DL, VT, Result,
+                           DAG.getConstant(TrailingZeros, DL, VT));
+
+    return Result;
+  }
+
+  return expandMulToNAFSequence(N, DAG, DL, MulAmt);
+}
 
 // 2^N +/- 2^M -> (add/sub (shl X, C1), (shl X, C2))
 static SDValue expandMulToAddOrSubOfShl(SDNode *N, SelectionDAG &DAG,