<commit_before>{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
-- | We generate code for non-segmented/single-segment SegRed using
-- the basic approach outlined in the paper "Design and GPGPU
-- Performance of Futhark’s Redomap Construct" (ARRAY '16).  The main
-- deviations are:
--
-- * While we still use two-phase reduction, we use only a single
--   kernel, with the final workgroup to write a result (tracked via
--   an atomic counter) performing the final reduction as well.
--
-- * Instead of depending on storage layout transformations to handle
--   non-commutative reductions efficiently, we slide a
--   'groupsize'-sized window over the input, and perform a parallel
--   reduction for each window.  This sacrifices the notion of
--   efficient sequentialisation, but is sometimes faster and
--   definitely simpler and more predictable (and uses less auxiliary
--   storage).
--
-- For segmented reductions we use the approach from "Strategies for
-- Regular Segmented Reductions on GPU" (FHPC '17).  This involves
-- having two different strategies, and dynamically deciding which one
-- to use based on the number of segments and segment size. We use the
-- (static) @group_size@ to decide which of the following two
-- strategies to choose:
--
-- * Large: uses one or more groups to process a single segment. If
--   multiple groups are used per segment, the intermediate reduction
--   results must be recursively reduced, until there is only a single
--   value per segment.
--
--   Each thread /can/ read multiple elements, which will greatly
--   increase performance; however, if the reduction is
--   non-commutative we will have to use a less efficient traversal
--   (with interim group-wide reductions) to enable coalesced memory
--   accesses, just as in the non-segmented case.
--
-- * Small: is used to let each group process *multiple* segments
--   within a group. We will only use this approach when we can
--   process at least two segments within a single group. In those
--   cases, we would allocate a /whole/ group per segment with the
--   large strategy, but at most 50% of the threads in the group would
--   have any element to read, which becomes highly inefficient.
module Futhark.CodeGen.ImpGen.Kernels.SegRed
  ( compileSegRed
  )
  where

import Control.Monad.Except
import Data.Maybe
import qualified Data.Set as S
import Data.List

import Prelude hiding (quot, rem)

import Futhark.MonadFreshNames
import Futhark.Transform.Rename
import Futhark.Representation.ExplicitMemory
import qualified Futhark.CodeGen.ImpCode.Kernels as Imp
import qualified Futhark.CodeGen.ImpGen as ImpGen
import Futhark.CodeGen.ImpGen ((<--),
                               sFor, sComment, sIf, sWhen,
                               sOp,
                               dPrim, dPrimV)
import Futhark.CodeGen.ImpGen.Kernels.Base
import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun
import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem)

-- | For many kernels, we may not have enough physical groups to cover
-- the logical iteration space.  Some groups thus have to perform
-- double duty; we put an outer loop to accomplish this.  The
-- advantage over just launching a bazillion threads is that the cost
-- of memory expansion should be proportional to the number of
-- *physical* threads (hardware parallelism), not the amount of
-- application parallelism.
virtualiseGroups :: KernelConstants
                 -> Imp.Exp
                 -> (Imp.Exp -> ImpGen.ImpM lore op ())
                 -> ImpGen.ImpM lore op ()
virtualiseGroups constants required_groups m = do
  let group_id = kernelGroupId constants
      iterations = (required_groups - group_id) `quotRoundingUp` kernelNumGroups constants
  i <- newVName "i"
  sFor i Int32 iterations $ m $ group_id + Imp.var i int32 * kernelNumGroups constants

-- Compile 'SegRed' instance to host-level code with calls to various
-- kernels.
compileSegRed :: Pattern ExplicitMemory
              -> KernelSpace
              -> Commutativity -> Lambda InKernel -> [SubExp]
              -> Body InKernel
              -> CallKernelGen ()
compileSegRed pat space comm red_op nes body
  | [(_, Constant (IntValue (Int32Value 1))), _] <- spaceDimensions space =
      nonsegmentedReduction pat space comm red_op nes body
  | otherwise = do
      segment_size <-
        ImpGen.compileSubExp $ last $ map snd $ spaceDimensions space
      group_size <- ImpGen.compileSubExp $ spaceGroupSize space
      let use_small_segments = segment_size * 2 .<. group_size
      sIf (segment_size .==. 1)
        (unitSegmentsReduction pat space nes body) $
        sIf use_small_segments
        (smallSegmentsReduction pat space red_op nes body)
        (largeSegmentsReduction pat space comm red_op nes body)

-- Handle degenerate case where segments are of size 1, meaning
-- that it is really just a 'map' in disguise.
unitSegmentsReduction :: Pattern ExplicitMemory
                      -> KernelSpace
                      -> [SubExp]
                      -> Body InKernel
                      -> CallKernelGen ()
unitSegmentsReduction (Pattern _ segred_pes) space nes body = do
  (constants, init_constants) <- kernelInitialisationSetSpace space $ return ()

  let (gtids, dims) = unzip $ spaceDimensions space
      (redout_pes, mapout_pes) = splitAt (length nes) segred_pes

  dims' <- mapM ImpGen.compileSubExp dims

  let num_segments = product $ init dims'
      required_groups = num_segments `quotRoundingUp` kernelGroupSize constants

  ImpGen.emit $ Imp.DebugPrint "num_segments" int32 num_segments
  ImpGen.emit $ Imp.DebugPrint "required_groups" int32 required_groups

  sKernel constants "segred_mapseg" $ do
    init_constants
    virtualiseGroups constants required_groups $ \group_id -> do
      setSpaceIndices (group_id * kernelGroupSize constants + kernelLocalThreadId constants) space
      ImpGen.compileStms mempty (stmsToList $ bodyStms body) $
        sWhen (kernelThreadActive constants) $ do
        let (redout_ses, mapout_ses) = splitAt (length nes) $ bodyResult body
        forM_ (zip redout_pes redout_ses) $ \(pe, se) ->
          ImpGen.copyDWIM (patElemName pe)
          (map (`Imp.var` int32) (init gtids)) se []

        forM_ (zip mapout_pes mapout_ses) $ \(pe, se) ->
          ImpGen.copyDWIM (patElemName pe)
          (map (`Imp.var` int32) gtids) se []

nonsegmentedReduction :: Pattern ExplicitMemory
                      -> KernelSpace
                      -> Commutativity -> Lambda InKernel -> [SubExp]
                      -> Body InKernel
                      -> CallKernelGen ()
nonsegmentedReduction segred_pat space comm red_op nes body = do
  (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()
  let constants = base_constants { kernelThreadActive = true }
      global_tid = kernelGlobalThreadId constants
      (_, w) = last $ spaceDimensions space

  let red_op_params = lambdaParams red_op
      (red_acc_params, _) = splitAt (length nes) red_op_params
  red_arrs <- forM red_acc_params $ \p ->
    case paramAttr p of<commit_msg>Might as well move this up.<commit_after>{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
-- | We generate code for non-segmented/single-segment SegRed using
-- the basic approach outlined in the paper "Design and GPGPU
-- Performance of Futhark’s Redomap Construct" (ARRAY '16).  The main
-- deviations are:
--
-- * While we still use two-phase reduction, we use only a single
--   kernel, with the final workgroup to write a result (tracked via
--   an atomic counter) performing the final reduction as well.
--
-- * Instead of depending on storage layout transformations to handle
--   non-commutative reductions efficiently, we slide a
--   'groupsize'-sized window over the input, and perform a parallel
--   reduction for each window.  This sacrifices the notion of
--   efficient sequentialisation, but is sometimes faster and
--   definitely simpler and more predictable (and uses less auxiliary
--   storage).
--
-- For segmented reductions we use the approach from "Strategies for
-- Regular Segmented Reductions on GPU" (FHPC '17).  This involves
-- having two different strategies, and dynamically deciding which one
-- to use based on the number of segments and segment size. We use the
-- (static) @group_size@ to decide which of the following two
-- strategies to choose:
--
-- * Large: uses one or more groups to process a single segment. If
--   multiple groups are used per segment, the intermediate reduction
--   results must be recursively reduced, until there is only a single
--   value per segment.
--
--   Each thread /can/ read multiple elements, which will greatly
--   increase performance; however, if the reduction is
--   non-commutative we will have to use a less efficient traversal
--   (with interim group-wide reductions) to enable coalesced memory
--   accesses, just as in the non-segmented case.
--
-- * Small: is used to let each group process *multiple* segments
--   within a group. We will only use this approach when we can
--   process at least two segments within a single group. In those
--   cases, we would allocate a /whole/ group per segment with the
--   large strategy, but at most 50% of the threads in the group would
--   have any element to read, which becomes highly inefficient.
module Futhark.CodeGen.ImpGen.Kernels.SegRed
  ( compileSegRed
  )
  where

import Control.Monad.Except
import Data.Maybe
import qualified Data.Set as S
import Data.List

import Prelude hiding (quot, rem)

import Futhark.MonadFreshNames
import Futhark.Transform.Rename
import Futhark.Representation.ExplicitMemory
import qualified Futhark.CodeGen.ImpCode.Kernels as Imp
import qualified Futhark.CodeGen.ImpGen as ImpGen
import Futhark.CodeGen.ImpGen ((<--),
                               sFor, sComment, sIf, sWhen,
                               sOp,
                               dPrim, dPrimV)
import Futhark.CodeGen.ImpGen.Kernels.Base
import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun
import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem)

-- | For many kernels, we may not have enough physical groups to cover
-- the logical iteration space.  Some groups thus have to perform
-- double duty; we put an outer loop to accomplish this.  The
-- advantage over just launching a bazillion threads is that the cost
-- of memory expansion should be proportional to the number of
-- *physical* threads (hardware parallelism), not the amount of
-- application parallelism.
virtualiseGroups :: KernelConstants
                 -> Imp.Exp
                 -> (Imp.Exp -> ImpGen.ImpM lore op ())
                 -> ImpGen.ImpM lore op ()
virtualiseGroups constants required_groups m = do
  let group_id = kernelGroupId constants
      iterations = (required_groups - group_id) `quotRoundingUp` kernelNumGroups constants
  i <- newVName "i"
  sFor i Int32 iterations $ m $ group_id + Imp.var i int32 * kernelNumGroups constants

-- Compile 'SegRed' instance to host-level code with calls to various
-- kernels.
compileSegRed :: Pattern ExplicitMemory
              -> KernelSpace
              -> Commutativity -> Lambda InKernel -> [SubExp]
              -> Body InKernel
              -> CallKernelGen ()
compileSegRed pat space comm red_op nes body
  | [(_, Constant (IntValue (Int32Value 1))), _] <- spaceDimensions space =
      nonsegmentedReduction pat space comm red_op nes body
  | otherwise = do
      segment_size <-
        ImpGen.compileSubExp $ last $ map snd $ spaceDimensions space
      group_size <- ImpGen.compileSubExp $ spaceGroupSize space
      let use_small_segments = segment_size * 2 .<. group_size
      sIf (segment_size .==. 1)
        (unitSegmentsReduction pat space nes body) $
        sIf use_small_segments
        (smallSegmentsReduction pat space red_op nes body)
        (largeSegmentsReduction pat space comm red_op nes body)

-- Handle degenerate case where segments are of size 1, meaning
-- that it is really just a 'map' in disguise.
unitSegmentsReduction :: Pattern ExplicitMemory
                      -> KernelSpace
                      -> [SubExp]
                      -> Body InKernel
                      -> CallKernelGen ()
unitSegmentsReduction (Pattern _ segred_pes) space nes body = do
  (constants, init_constants) <- kernelInitialisationSetSpace space $ return ()

  let (gtids, dims) = unzip $ spaceDimensions space
      (redout_pes, mapout_pes) = splitAt (length nes) segred_pes

  dims' <- mapM ImpGen.compileSubExp dims

  let num_segments = product $ init dims'
      required_groups = num_segments `quotRoundingUp` kernelGroupSize constants

  ImpGen.emit $ Imp.DebugPrint "num_segments" int32 num_segments
  ImpGen.emit $ Imp.DebugPrint "required_groups" int32 required_groups

  sKernel constants "segred_mapseg" $ do
    init_constants
    virtualiseGroups constants required_groups $ \group_id -> do
      setSpaceIndices (group_id * kernelGroupSize constants + kernelLocalThreadId constants) space
      sWhen (kernelThreadActive constants) $
        ImpGen.compileStms mempty (stmsToList $ bodyStms body) $ do
        let (redout_ses, mapout_ses) = splitAt (length nes) $ bodyResult body
        forM_ (zip redout_pes redout_ses) $ \(pe, se) ->
          ImpGen.copyDWIM (patElemName pe)
          (map (`Imp.var` int32) (init gtids)) se []

        forM_ (zip mapout_pes mapout_ses) $ \(pe, se) ->
          ImpGen.copyDWIM (patElemName pe)
          (map (`Imp.var` int32) gtids) se []

nonsegmentedReduction :: Pattern ExplicitMemory
                      -> KernelSpace
                      -> Commutativity -> Lambda InKernel -> [SubExp]
                      -> Body InKernel
                      -> CallKernelGen ()
nonsegmentedReduction segred_pat space comm red_op nes body = do
  (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()
  let constants = base_constants { kernelThreadActive = true }
      global_tid = kernelGlobalThreadId constants
      (_, w) = last $ spaceDimensions space

  let red_op_params = lambdaParams red_op
      (red_acc_params, _) = splitAt (length nes) red_op_params
  red_arrs <- forM red_acc_params $ \p ->
    case paramAttr p of
{'fraction_non_alphanumeric': 0.07575545607162842, 'fraction_numerical': 0.003357582540570789, 'mean_word_length': 2.9266684976654767, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 51, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2002351', 'n_tokens_mistral': 4033, 'n_tokens_neox': 3898, 'n_words': 1492}