Support MO_U_QuotRem2 in LLVM backend

This adds support for MO_U_QuotRem2 in LLVM backend.  Similarly to
MO_U_Mul2 we use the standard LLVM instructions (in this case 'udiv'
and 'urem') but do the computation on double the word width (e.g., for
64-bit we will do them on 128 registers).

Test Plan: validate

Reviewers: rwbarton, austin, bgamari

Reviewed By: bgamari

Subscribers: thomie

Differential Revision: https://phabricator.haskell.org/D1100

GHC Trac Issues: #9430

2 files changed, 74 insertions, 1 deletions

diff --git a/compiler/codeGen/StgCmmPrim.hs b/compiler/codeGen/StgCmmPrim.hs
index 243e2a32ac..d201eaf20b 100644
--- a/compiler/codeGen/StgCmmPrim.hs
+++ b/compiler/codeGen/StgCmmPrim.hs
@@ -808,7 +808,8 @@ callishPrimOpSupported dflags op
       WordQuotRemOp  | ncg && x86ish  -> Left (MO_U_QuotRem  (wordWidth dflags))
                      | otherwise      -> Right (genericWordQuotRemOp dflags)
 
-      WordQuotRem2Op | ncg && x86ish  -> Left (MO_U_QuotRem2 (wordWidth dflags))
+      WordQuotRem2Op | (ncg && x86ish)
+                          || llvm     -> Left (MO_U_QuotRem2 (wordWidth dflags))
                      | otherwise      -> Right (genericWordQuotRem2Op dflags)
 
       WordAdd2Op     | (ncg && x86ish)
diff --git a/compiler/llvmGen/LlvmCodeGen/CodeGen.hs b/compiler/llvmGen/LlvmCodeGen/CodeGen.hs
index fb02120747..517da53cab 100644
--- a/compiler/llvmGen/LlvmCodeGen/CodeGen.hs
+++ b/compiler/llvmGen/LlvmCodeGen/CodeGen.hs
@@ -32,6 +32,13 @@ import UniqSupply
 import Unique
 import Util
 
+import Control.Monad.Trans.Class
+import Control.Monad.Trans.Writer
+
+#if MIN_VERSION_base(4,8,0)
+#else
+import Data.Monoid ( Monoid, mappend, mempty )
+#endif
 import Data.List ( nub )
 import Data.Maybe ( catMaybes )
 
@@ -288,6 +295,53 @@ genCall (PrimTarget (MO_U_Mul2 w)) [dstH, dstL] [lhs, rhs] = do
            toOL [ stmt3 , stmt4, stmt5, stmt6, stmt7, stmt8, storeL, storeH ]
     return (stmts, decls1 ++ decls2)
 
+-- MO_U_QuotRem2 is another case we handle by widening the registers to double
+-- the width and use normal LLVM instructions (similarly to the MO_U_Mul2). The
+-- main difference here is that we need to conmbine two words into one register
+-- and then use both 'udiv' and 'urem' instructions to compute the result.
+genCall (PrimTarget (MO_U_QuotRem2 w)) [dstQ, dstR] [lhsH, lhsL, rhs] = run $ do
+    let width = widthToLlvmInt w
+        bitWidth = widthInBits w
+        width2x = LMInt (bitWidth * 2)
+    -- First zero-extend all parameters to double width.
+    let zeroExtend expr = do
+            var <- liftExprData $ exprToVar expr
+            doExprW width2x $ Cast LM_Zext var width2x
+    lhsExtH <- zeroExtend lhsH
+    lhsExtL <- zeroExtend lhsL
+    rhsExt <- zeroExtend rhs
+    -- Now we combine the first two parameters (that represent the high and low
+    -- bits of the value). So first left-shift the high bits to their position
+    -- and then bit-or them with the low bits.
+    let widthLlvmLit = LMLitVar $ LMIntLit (fromIntegral bitWidth) width
+    lhsExtHShifted <- doExprW width2x $ LlvmOp LM_MO_Shl lhsExtH widthLlvmLit
+    lhsExt <- doExprW width2x $ LlvmOp LM_MO_Or lhsExtHShifted lhsExtL
+    -- Finally, we can call 'udiv' and 'urem' to compute the results.
+    retExtDiv <- doExprW width2x $ LlvmOp LM_MO_UDiv lhsExt rhsExt
+    retExtRem <- doExprW width2x $ LlvmOp LM_MO_URem lhsExt rhsExt
+    -- And since everything is in 2x width, we need to truncate the results and
+    -- then return them.
+    let narrow var = doExprW width $ Cast LM_Trunc var width
+    retDiv <- narrow retExtDiv
+    retRem <- narrow retExtRem
+    dstRegQ <- lift $ getCmmReg (CmmLocal dstQ)
+    dstRegR <- lift $ getCmmReg (CmmLocal dstR)
+    statement $ Store retDiv dstRegQ
+    statement $ Store retRem dstRegR
+  where
+    -- TODO(michalt): Consider extracting this and using in more places.
+    -- Hopefully this should cut down on the noise of accumulating the
+    -- statements and declarations.
+    doExprW :: LlvmType -> LlvmExpression -> WriterT LlvmAccum LlvmM LlvmVar
+    doExprW a b = do
+        (var, stmt) <- lift $ doExpr a b
+        statement stmt
+        return var
+    run :: WriterT LlvmAccum LlvmM () -> LlvmM (LlvmStatements, [LlvmCmmDecl])
+    run action = do
+        LlvmAccum stmts decls <- execWriterT action
+        return (stmts, decls)
+
 -- Handle the MO_{Add,Sub}IntC separately. LLVM versions return a record from
 -- which we need to extract the actual values.
 genCall t@(PrimTarget (MO_AddIntC w)) [dstV, dstO] [lhs, rhs] =
@@ -1767,3 +1821,21 @@ getTBAAMeta u = do
 -- | Returns TBAA meta data for given register
 getTBAARegMeta :: GlobalReg -> LlvmM [MetaAnnot]
 getTBAARegMeta = getTBAAMeta . getTBAA
+
+
+-- | A more convenient way of accumulating LLVM statements and declarations.
+data LlvmAccum = LlvmAccum LlvmStatements [LlvmCmmDecl]
+
+instance Monoid LlvmAccum where
+    mempty = LlvmAccum nilOL []
+    LlvmAccum stmtsA declsA `mappend` LlvmAccum stmtsB declsB =
+        LlvmAccum (stmtsA `mappend` stmtsB) (declsA `mappend` declsB)
+
+liftExprData :: LlvmM ExprData -> WriterT LlvmAccum LlvmM LlvmVar
+liftExprData action = do
+    (var, stmts, decls) <- lift action
+    tell $ LlvmAccum stmts decls
+    return var
+
+statement :: LlvmStatement -> WriterT LlvmAccum LlvmM ()
+statement stmt = tell $ LlvmAccum (unitOL stmt) []