Properly compute unpacked sizes for -funpack-small-strict-fields.wip/andreask/unpack_unboxed_tuples

Base unpacking under -funpack-small-strict-fields on the rep size
of the unpacked constructor instead of the number of reps it's
represented by.

Fixes #22309

6 files changed, 199 insertions, 8 deletions

diff --git a/compiler/GHC/Core/TyCon.hs b/compiler/GHC/Core/TyCon.hs
index 0311ba32bd..994035b2cf 100644
--- a/compiler/GHC/Core/TyCon.hs
+++ b/compiler/GHC/Core/TyCon.hs
@@ -126,8 +126,8 @@ module GHC.Core.TyCon(
         PrimRep(..), PrimElemRep(..),
         primElemRepToPrimRep,
         isVoidRep, isGcPtrRep,
-        primRepSizeB,
-        primElemRepSizeB,
+        primRepSizeB, primRepSizeW64_B,
+        primElemRepSizeB, primElemRepSizeW64_B,
         primRepIsFloat,
         primRepsCompatible,
         primRepCompatible,
@@ -1659,9 +1659,40 @@ primRepSizeB platform = \case
    VoidRep          -> 0
    (VecRep len rep) -> len * primElemRepSizeB platform rep
 
+-- | Like primRepSizeB but assumes pointers/words are 8 words wide.
+--
+-- This can be useful to compute the size of a rep as if we were compiling
+-- for a 64bit platform.
+primRepSizeW64_B :: PrimRep -> Int
+primRepSizeW64_B = \case
+   IntRep           -> 8
+   WordRep          -> 8
+   Int8Rep          -> 1
+   Int16Rep         -> 2
+   Int32Rep         -> 4
+   Int64Rep         -> 8
+   Word8Rep         -> 1
+   Word16Rep        -> 2
+   Word32Rep        -> 4
+   Word64Rep        -> 8
+   FloatRep         -> fLOAT_SIZE
+   DoubleRep        -> dOUBLE_SIZE
+   AddrRep          -> 8
+   LiftedRep        -> 8
+   UnliftedRep      -> 8
+   VoidRep          -> 0
+   (VecRep len rep) -> len * primElemRepSizeW64_B rep
+
 primElemRepSizeB :: Platform -> PrimElemRep -> Int
 primElemRepSizeB platform = primRepSizeB platform . primElemRepToPrimRep
 
+-- | Like primElemRepSizeB but assumes pointers/words are 8 words wide.
+--
+-- This can be useful to compute the size of a rep as if we were compiling
+-- for a 64bit platform.
+primElemRepSizeW64_B :: PrimElemRep -> Int
+primElemRepSizeW64_B = primRepSizeW64_B . primElemRepToPrimRep
+
 primElemRepToPrimRep :: PrimElemRep -> PrimRep
 primElemRepToPrimRep Int8ElemRep   = Int8Rep
 primElemRepToPrimRep Int16ElemRep  = Int16Rep
diff --git a/compiler/GHC/Types/Id/Make.hs b/compiler/GHC/Types/Id/Make.hs
index 4baa335db1..d549bce19d 100644
--- a/compiler/GHC/Types/Id/Make.hs
+++ b/compiler/GHC/Types/Id/Make.hs
@@ -67,6 +67,7 @@ import GHC.Types.Literal
 import GHC.Types.SourceText
 import GHC.Types.Name.Set
 import GHC.Types.Name
+import GHC.Types.RepType
 import GHC.Types.ForeignCall
 import GHC.Types.Id
 import GHC.Types.Id.Info
@@ -1362,9 +1363,20 @@ shouldUnpackTy bang_opts prag fam_envs ty
           | otherwise
           -> bang_opt_unbox_strict bang_opts
              || (bang_opt_unbox_small bang_opts
-                 && rep_tys `lengthAtMost` 1)  -- See Note [Unpack one-wide fields]
-      where (rep_tys, _) = dataConArgUnpack ty
-
+                 && is_small_rep rep_tys)  -- See Note [Unpack one-wide fields]
+      where
+        (rep_tys, _) = dataConArgUnpack ty
+
+    -- Takes in the list of reps used to represent the dataCon after it's unpacked
+    -- and tells us if they can fit into 8 bytes. See Note [Unpack one-wide fields]
+    is_small_rep rep_tys =
+      let -- Neccesary to look through unboxed tuples.
+          prim_reps = concatMap (typePrimRep . scaledThing . fst) $ rep_tys
+          -- Void types are erased when unpacked so we
+          nv_prim_reps = filter (not . isVoidRep) prim_reps
+          -- And then get the actual size of the unpacked constructor.
+          rep_size = sum $ map primRepSizeW64_B nv_prim_reps
+      in rep_size <= 8
 
 -- Given a type already assumed to have been normalized by topNormaliseType,
 -- unpackable_type_datacons ty = Just datacons
@@ -1424,6 +1436,14 @@ However
 
 Here we can represent T with an Int#.
 
+Special care has to be taken to make sure we don't mistake fields with unboxed
+tuple/sum rep or very large reps. See #22309
+
+For consistency we unpack anything that fits into 8 bytes on a 64-bit platform,
+even when compiling for 32bit platforms. This way unpacking decisions will be the
+same for 32bit and 64bit systems. To do so we use primRepSizeW64_B instead of
+primRepSizeB. See also the tests in test case T22309.
+
 Note [Recursive unboxing]
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 Consider
diff --git a/docs/users_guide/using-optimisation.rst b/docs/users_guide/using-optimisation.rst
index d3ca68a1df..4b523a95cf 100644
--- a/docs/users_guide/using-optimisation.rst
+++ b/docs/users_guide/using-optimisation.rst
@@ -1468,9 +1468,9 @@ by saying ``-fno-wombat``.
     default you can disable it for certain constructor fields using the
     ``NOUNPACK`` pragma (see :ref:`nounpack-pragma`).
 
-    Note that for consistency ``Double``, ``Word64``, and ``Int64``
-    constructor fields are unpacked on 32-bit platforms, even though
-    they are technically larger than a pointer on those platforms.
+    Note that for consistency constructor fields are unpacked on 32-bit platforms
+    as if it we were compiling for a 64-bit target even if fields are larger
+    than a pointer on those platforms.
 
 .. ghc-flag:: -funbox-strict-fields
     :shortdesc: Flatten strict constructor fields
diff --git a/testsuite/tests/simplCore/should_compile/T22309.hs b/testsuite/tests/simplCore/should_compile/T22309.hs
new file mode 100644
index 0000000000..4d2d73d49f
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T22309.hs
@@ -0,0 +1,51 @@
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnboxedTuples #-}
+
+module ShouldCompile where
+
+import GHC.Int
+import GHC.Exts
+
+-- These should unbox into another constructor
+data UA = Mk_A !Int
+data UB = Mk_B !Int64
+data UC = Mk_C !Int32
+data UD = Mk_D !Int32 !Int32
+data UE = Mk_E !(# Int# #)
+data UF = Mk_F !(# Double #)
+
+-- These should not be unpacked into another constructor.
+data NU_A = NU_MkA (# Int, Int #)
+data NU_B = NU_MkB !Int64 !Int64
+
+-- The types we unbox into
+
+-- These should unpack their fields.
+data WU_A = MkW_A !UA
+data WU_B = MkW_B !UB
+data WU_C = MkW_C !UC
+data WU_D = MkW_D !UD
+data WU_E = MkW_E !UE
+data WU_F = MkW_F !UF
+
+-- These should not unpack their fields.
+data WNU_A = MkW_NA !NU_A
+data WNU_B = MkW_NB !NU_B
+
+-- data WrapInt = MkWrapInt (# Int# #)
+-- data WrapInts = MkWrapInts (# Int#,Int#,Int# #)
+
+-- -- We should unbox WrapInt here
+-- data WrapWrapInt = MkWrapWrapInt !WrapInt
+-- -- We should not unbox WrapIntLarge as it's large
+-- data MkWrapWrapInts = MkWrapWrapInts !WrapInts
+
+
+-- -- W will contain the ints directly
+-- data T = T !(# Int, Int, Int, Int #)
+-- data W = W !T
+
+-- -- W2 will unbox T2, since it assumes the unboxed tuple to be a single value and doesn't
+-- -- account for (# #) already being unboxed.
+-- data T2 = T2 !( Int, Int, Int)
+-- data W2 = W2 !T2
diff --git a/testsuite/tests/simplCore/should_compile/T22309.stderr b/testsuite/tests/simplCore/should_compile/T22309.stderr
new file mode 100644
index 0000000000..ac0c768688
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T22309.stderr
@@ -0,0 +1,88 @@
+
+==================== Final STG: ====================
+$WMkW_NB :: NU_B %1 -> WNU_B =
+    \r [conrep]
+        case conrep of conrep1 { __DEFAULT -> MkW_NB [conrep1]; };
+
+$WMkW_NA :: NU_A %1 -> WNU_A =
+    \r [conrep]
+        case conrep of conrep1 { __DEFAULT -> MkW_NA [conrep1]; };
+
+$WMkW_F :: UF %1 -> WU_F =
+    \r [conrep] case conrep of { Mk_F us -> MkW_F [us]; };
+
+$WMkW_E :: UE %1 -> WU_E =
+    \r [conrep] case conrep of { Mk_E us -> MkW_E [us]; };
+
+$WMkW_D :: UD %1 -> WU_D =
+    \r [conrep]
+        case conrep of { Mk_D unbx unbx1 -> MkW_D [unbx unbx1]; };
+
+$WMkW_C :: UC %1 -> WU_C =
+    \r [conrep] case conrep of { Mk_C unbx -> MkW_C [unbx]; };
+
+$WMkW_B :: UB %1 -> WU_B =
+    \r [conrep] case conrep of { Mk_B unbx -> MkW_B [unbx]; };
+
+$WMkW_A :: UA %1 -> WU_A =
+    \r [conrep] case conrep of { Mk_A unbx -> MkW_A [unbx]; };
+
+$WNU_MkB :: Int64 %1 -> Int64 %1 -> NU_B =
+    \r [conrep conrep1]
+        case conrep of {
+        I64# unbx ->
+        case conrep1 of { I64# unbx1 -> NU_MkB [unbx unbx1]; };
+        };
+
+$WMk_D :: Int32 %1 -> Int32 %1 -> UD =
+    \r [conrep conrep1]
+        case conrep of {
+        I32# unbx -> case conrep1 of { I32# unbx1 -> Mk_D [unbx unbx1]; };
+        };
+
+$WMk_C :: Int32 %1 -> UC =
+    \r [conrep] case conrep of { I32# unbx -> Mk_C [unbx]; };
+
+$WMk_B :: Int64 %1 -> UB =
+    \r [conrep] case conrep of { I64# unbx -> Mk_B [unbx]; };
+
+$WMk_A :: Int %1 -> UA =
+    \r [conrep] case conrep of { I# unbx -> Mk_A [unbx]; };
+
+MkW_NB :: NU_B %1 -> WNU_B =
+    \r [eta] case eta of eta { __DEFAULT -> MkW_NB [eta]; };
+
+MkW_NA :: NU_A %1 -> WNU_A =
+    \r [eta] case eta of eta { __DEFAULT -> MkW_NA [eta]; };
+
+MkW_F :: (# Double #) %1 -> WU_F = \r [us] MkW_F [us];
+
+MkW_E :: (# Int# #) %1 -> WU_E = \r [us] MkW_E [us];
+
+MkW_D :: Int32# %1 -> Int32# %1 -> WU_D =
+    \r [eta eta] MkW_D [eta eta];
+
+MkW_C :: Int32# %1 -> WU_C = \r [eta] MkW_C [eta];
+
+MkW_B :: Int64# %1 -> WU_B = \r [eta] MkW_B [eta];
+
+MkW_A :: Int# %1 -> WU_A = \r [eta] MkW_A [eta];
+
+NU_MkB :: Int64# %1 -> Int64# %1 -> NU_B =
+    \r [eta eta] NU_MkB [eta eta];
+
+NU_MkA :: (# Int, Int #) %1 -> NU_A = \r [us us] NU_MkA [us us];
+
+Mk_F :: (# Double #) %1 -> UF = \r [us] Mk_F [us];
+
+Mk_E :: (# Int# #) %1 -> UE = \r [us] Mk_E [us];
+
+Mk_D :: Int32# %1 -> Int32# %1 -> UD = \r [eta eta] Mk_D [eta eta];
+
+Mk_C :: Int32# %1 -> UC = \r [eta] Mk_C [eta];
+
+Mk_B :: Int64# %1 -> UB = \r [eta] Mk_B [eta];
+
+Mk_A :: Int# %1 -> UA = \r [eta] Mk_A [eta];
+
+
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 068b830a51..3d613acbb7 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -469,4 +469,5 @@ test('T22662', normal, compile, [''])
 test('T22725', normal, compile, ['-O'])
 test('T22502', normal, compile, ['-O'])
 test('T22611', [when(wordsize(32), skip), grep_errmsg(r'\$salterF') ], compile, ['-O -ddump-simpl -dsuppress-uniques -dsuppress-all'])
+test('T22309', [grep_errmsg(r'^MkW'), only_ways(['optasm']) ], compile, ['-O -ddump-stg-final -dsuppress-uniques -dsuppress-all -dno-typeable-binds -dno-suppress-type-signatures -dsuppress-module-prefixes'])