<commit_before>{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
-- | Multiversion segmented reduction.
module Futhark.Pass.ExtractKernels.Segmented
       ( regularSegmentedRedomapAsScan
       , regularSegmentedRedomap
       , regularSegmentedScan
       )
       where

import Control.Monad
import qualified Data.HashMap.Lazy as HM
import Data.Monoid

import Prelude

import Futhark.Transform.Rename
import qualified Futhark.Analysis.ScalExp as SE
import Futhark.Representation.Kernels
import Futhark.MonadFreshNames
import Futhark.Tools
import Futhark.Pass.ExtractKernels.BlockedKernel

data SegmentedVersion = OneGroupOneSegment
                      | ManyGroupsOneSegment
                      deriving (Eq, Ord, Show)

regularSegmentedRedomap :: (HasScope Kernels m, MonadBinder m, Lore m ~ Kernels) =>
                           SubExp            -- segment_size
                        -> SubExp            -- num_segments
                        -> [SubExp]          -- nest_sizes = the sizes of the maps on "top" of this redomap
                        -> Pattern Kernels   -- flat_pat ... pat where each type is array with dim [w]
                        -> Pattern Kernels   -- pat
                        -> Certificates      -- cs
                        -> SubExp            -- w = total_num_elements
                        -> Commutativity     -- comm
                        -> Lambda InKernel   -- reduce_lam
                        -> Lambda InKernel   -- fold_lam = this lambda performs both the map-part and
                                             -- reduce-part of a redomap (described in redomap paper)
                        -> [(VName, SubExp)] -- ispace = pair of (gtid, size) for the maps on "top" of this redomap
                        -> [KernelInput]     -- inps = inputs that can be looked up by using the gtids from ispace
                        -> [SubExp]          -- nes
                        -> [VName]           -- arrs_flat
                        -> m ()
regularSegmentedRedomap segment_size num_segments _nest_sizes flat_pat
                        pat cs w comm reduce_lam fold_lam ispace inps nes arrs_flat = do
  unless (null $ patternContextElements pat) $ fail "regularSegmentedRedomap result pattern contains context elements, and Rasmus did not think this would ever happen."

  -- the result of the "map" part of a redomap has to be stored somewhere within
  -- the chunking loop of a kernel. The current way to do this is to make some
  -- scratch space initially, and each thread will get a part of this by
  -- splitting it. Finally it is returned as a result of the kernel (to not
  -- break functional semantics).
  map_out_arrs <- forM (drop num_redres $ patternIdents pat) $ \(Ident name t) -> do
    tmp <- letExp (baseString name <> "_out_in") $
           BasicOp $ Scratch (elemType t) (arrayDims t)
    -- This reshape will not always work.
    -- For example if the "map" part takes an input a 1D array and produces a 2D
    -- array, this is clearly wrong. See ex3.fut
    letExp (baseString name ++ "_out_in") $
            BasicOp $ Reshape cs [DimNew w] tmp

  -- Check that we're only dealing with arrays with dimension [w]
  forM_ arrs_flat $ \arr -> do
    tp <- lookupType arr
    case tp of
      -- FIXME: this won't work if the reduction operator works on lists... but
      -- they seem to be handled in some other way (which makes sense). Talk
      -- with troels if I should worry about this.
      Array _primtp (Shape [flatsize]) _uniqness ->
        when (flatsize /= w) $
          fail$ "regularSegmentedRedomap: flat array, with incorrect size encountered " ++ pretty arr
      _ -> fail $ "regularSegmentedRedomap: non-flat array encountered " ++ pretty arr

  -- The pattern passed to chunkLambda must have exactly *one* array dimension,
  -- to get the correct size of [chunk_size]type.
  --
  -- FIXME: not sure if this will work when result of map is multidimensional,
  -- or if reduction operator uses lists... must check
  chunk_pat <- fmap (Pattern []) $ forM (patternValueElements pat) $ \pat_e ->
    case patElemType pat_e of
      (Array ty (Shape (dim0:_)) u) -> do
          vn' <- newName $ patElemName pat_e
          return $ PatElem vn' BindVar $ Array ty (Shape [dim0]) u
      _ -> fail $ "segmentedRedomap: result pattern is not array " ++ pretty pat_e

  chunk_fold_lam <- chunkLambda chunk_pat nes fold_lam

  -- kernliseLambda intializes the value of the merge pattern for the reduction
  -- to the neutral element.
  kern_chunk_fold_lam <- kerneliseLambda nes chunk_fold_lam

  let chunk_red_pat = Pattern [] $ take num_redres $ patternValueElements chunk_pat
  kern_chunk_reduce_lam <- kerneliseLambda nes =<< chunkLambda chunk_red_pat nes reduce_lam

  -- the lambda for a GroupReduce needs these two extra parameters
  my_index <- newVName "my_index"
  other_offset <- newVName "other_offset"
  let my_index_param = Param my_index (Prim int32)
  let other_offset_param = Param other_offset (Prim int32)
  let reduce_lam' = reduce_lam { lambdaParams = my_index_param :
                                                other_offset_param :
                                                lambdaParams reduce_lam
                               }

  -- FIXME: do we need to copy arrays here? :S
  -- see 'blockedReductionStream' in BlockedKernel.hs

  let all_arrs = arrs_flat ++ map_out_arrs
  (ogps_ses, ogps_stms) <- runBinder $ oneGroupPerSeg all_arrs reduce_lam' kern_chunk_fold_lam
  (mgps_ses, mgps_stms) <- runBinder $ manyGroupPerSeg all_arrs reduce_lam' kern_chunk_fold_lam kern_chunk_reduce_lam

  -- TODO: should unify this calculation somewhere, with the one taking place in
  -- 'groupPerSegmentKernel'
  num_groups_hint <- letSubExp "num_groups_hint" $ Op NumGroups
  num_groups_per_segment <- letSubExp "num_groups_per_segment" =<<
    eDivRoundingUp Int32 (eSubExp num_groups_hint) (eSubExp num_segments)

  e <- eIf (eCmpOp (CmpEq $ IntType Int32) (eSubExp num_groups_per_segment)
                                           (eSubExp one))
           (mkBodyM ogps_stms ogps_ses)
           (mkBodyM mgps_stms mgps_ses)

  redres_pes <- forM (take num_redres (patternValueElements pat)) $ \pe -> do
    vn' <- newName $ patElemName pe
    return $ PatElem vn' BindVar $ patElemType pe `setArrayDims` [num_segments]
  let mapres_pes = drop num_redres $ patternValueElements flat_pat
  let unreshaped_pat = Pattern [] $ redres_pes ++ mapres_pes

  addStm $ Let unreshaped_pat () e

  forM_ (zip (patternValueElements unreshaped_pat)
             (patternValueElements pat)) $ \(kpe, pe) ->
    addStm $ Let (Pattern [] [pe]) () $
             BasicOp $ Reshape cs [DimNew se | se <- arrayDims $ patElemAttr pe]
                               (patElemName kpe)

  where
    one = constant (1 :: Int32)

    -- number of reduction results (tuple size for reduction operator)
    num_redres = length nes

    ----------------------------------------------------------------------------
    -- The functions below generate all the needed code for the two different
    -- version of segmented-redomap.
    --
    -- We rename statements before adding them because the same lambdas
    -- (reduce/fold) are used multiple times, and we do not want to bind the
    -- same VName twice (as this is a type error)
    ----------------------------------------------------------------------------
    oneGroupPerSeg all_arrs reduce_lam' kern_chunk_fold_lam = do
      mapres_pes <- forM (drop num_redres $ patternValueElements flat_pat) $ \pe -> do
        vn' <- newName $ patElemName pe
        return $ PatElem vn' BindVar $ patElemType pe

      (kernel, _, _) <- groupPerSegmentKernel segment_size num_segments cs
        all_arrs comm reduce_lam' kern_chunk_fold_lam
        nes w OneGroupOneSegment
        ispace inps
<commit_msg>segredomap: remember why multiversioning decision is like this<commit_after>{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
-- | Multiversion segmented reduction.
module Futhark.Pass.ExtractKernels.Segmented
       ( regularSegmentedRedomapAsScan
       , regularSegmentedRedomap
       , regularSegmentedScan
       )
       where

import Control.Monad
import qualified Data.HashMap.Lazy as HM
import Data.Monoid

import Prelude

import Futhark.Transform.Rename
import qualified Futhark.Analysis.ScalExp as SE
import Futhark.Representation.Kernels
import Futhark.MonadFreshNames
import Futhark.Tools
import Futhark.Pass.ExtractKernels.BlockedKernel

data SegmentedVersion = OneGroupOneSegment
                      | ManyGroupsOneSegment
                      deriving (Eq, Ord, Show)

regularSegmentedRedomap :: (HasScope Kernels m, MonadBinder m, Lore m ~ Kernels) =>
                           SubExp            -- segment_size
                        -> SubExp            -- num_segments
                        -> [SubExp]          -- nest_sizes = the sizes of the maps on "top" of this redomap
                        -> Pattern Kernels   -- flat_pat ... pat where each type is array with dim [w]
                        -> Pattern Kernels   -- pat
                        -> Certificates      -- cs
                        -> SubExp            -- w = total_num_elements
                        -> Commutativity     -- comm
                        -> Lambda InKernel   -- reduce_lam
                        -> Lambda InKernel   -- fold_lam = this lambda performs both the map-part and
                                             -- reduce-part of a redomap (described in redomap paper)
                        -> [(VName, SubExp)] -- ispace = pair of (gtid, size) for the maps on "top" of this redomap
                        -> [KernelInput]     -- inps = inputs that can be looked up by using the gtids from ispace
                        -> [SubExp]          -- nes
                        -> [VName]           -- arrs_flat
                        -> m ()
regularSegmentedRedomap segment_size num_segments _nest_sizes flat_pat
                        pat cs w comm reduce_lam fold_lam ispace inps nes arrs_flat = do
  unless (null $ patternContextElements pat) $ fail "regularSegmentedRedomap result pattern contains context elements, and Rasmus did not think this would ever happen."

  -- the result of the "map" part of a redomap has to be stored somewhere within
  -- the chunking loop of a kernel. The current way to do this is to make some
  -- scratch space initially, and each thread will get a part of this by
  -- splitting it. Finally it is returned as a result of the kernel (to not
  -- break functional semantics).
  map_out_arrs <- forM (drop num_redres $ patternIdents pat) $ \(Ident name t) -> do
    tmp <- letExp (baseString name <> "_out_in") $
           BasicOp $ Scratch (elemType t) (arrayDims t)
    -- This reshape will not always work.
    -- For example if the "map" part takes an input a 1D array and produces a 2D
    -- array, this is clearly wrong. See ex3.fut
    letExp (baseString name ++ "_out_in") $
            BasicOp $ Reshape cs [DimNew w] tmp

  -- Check that we're only dealing with arrays with dimension [w]
  forM_ arrs_flat $ \arr -> do
    tp <- lookupType arr
    case tp of
      -- FIXME: this won't work if the reduction operator works on lists... but
      -- they seem to be handled in some other way (which makes sense). Talk
      -- with troels if I should worry about this.
      Array _primtp (Shape [flatsize]) _uniqness ->
        when (flatsize /= w) $
          fail$ "regularSegmentedRedomap: flat array, with incorrect size encountered " ++ pretty arr
      _ -> fail $ "regularSegmentedRedomap: non-flat array encountered " ++ pretty arr

  -- The pattern passed to chunkLambda must have exactly *one* array dimension,
  -- to get the correct size of [chunk_size]type.
  --
  -- FIXME: not sure if this will work when result of map is multidimensional,
  -- or if reduction operator uses lists... must check
  chunk_pat <- fmap (Pattern []) $ forM (patternValueElements pat) $ \pat_e ->
    case patElemType pat_e of
      (Array ty (Shape (dim0:_)) u) -> do
          vn' <- newName $ patElemName pat_e
          return $ PatElem vn' BindVar $ Array ty (Shape [dim0]) u
      _ -> fail $ "segmentedRedomap: result pattern is not array " ++ pretty pat_e

  chunk_fold_lam <- chunkLambda chunk_pat nes fold_lam

  -- kernliseLambda intializes the value of the merge pattern for the reduction
  -- to the neutral element.
  kern_chunk_fold_lam <- kerneliseLambda nes chunk_fold_lam

  let chunk_red_pat = Pattern [] $ take num_redres $ patternValueElements chunk_pat
  kern_chunk_reduce_lam <- kerneliseLambda nes =<< chunkLambda chunk_red_pat nes reduce_lam

  -- the lambda for a GroupReduce needs these two extra parameters
  my_index <- newVName "my_index"
  other_offset <- newVName "other_offset"
  let my_index_param = Param my_index (Prim int32)
  let other_offset_param = Param other_offset (Prim int32)
  let reduce_lam' = reduce_lam { lambdaParams = my_index_param :
                                                other_offset_param :
                                                lambdaParams reduce_lam
                               }

  -- FIXME: do we need to copy arrays here? :S
  -- see 'blockedReductionStream' in BlockedKernel.hs

  let all_arrs = arrs_flat ++ map_out_arrs
  (ogps_ses, ogps_stms) <- runBinder $ oneGroupPerSeg all_arrs reduce_lam' kern_chunk_fold_lam
  (mgps_ses, mgps_stms) <- runBinder $ manyGroupPerSeg all_arrs reduce_lam' kern_chunk_fold_lam kern_chunk_reduce_lam

  -- TODO: should unify this calculation somewhere, with the one taking place in
  -- 'groupPerSegmentKernel'
  num_groups_hint <- letSubExp "num_groups_hint" $ Op NumGroups
  num_groups_per_segment <- letSubExp "num_groups_per_segment" =<<
    eDivRoundingUp Int32 (eSubExp num_groups_hint) (eSubExp num_segments)

  -- Instead of making this @if@ here, we could run "multiple groups per
  -- segment" first, and if the number of groups per segment was 1, we could
  -- simply return the result immediately without running the second kernel. I
  -- chose not to do this, as the oneGroupPerSegment kernel will contain
  -- slightly fewer instructions -- and hopefully be a bit faster (at the cost
  -- of large code size). TODO: test how much we win by doing this.
  e <- eIf (eCmpOp (CmpEq $ IntType Int32) (eSubExp num_groups_per_segment)
                                           (eSubExp one))
           (mkBodyM ogps_stms ogps_ses)
           (mkBodyM mgps_stms mgps_ses)

  redres_pes <- forM (take num_redres (patternValueElements pat)) $ \pe -> do
    vn' <- newName $ patElemName pe
    return $ PatElem vn' BindVar $ patElemType pe `setArrayDims` [num_segments]
  let mapres_pes = drop num_redres $ patternValueElements flat_pat
  let unreshaped_pat = Pattern [] $ redres_pes ++ mapres_pes

  addStm $ Let unreshaped_pat () e

  forM_ (zip (patternValueElements unreshaped_pat)
             (patternValueElements pat)) $ \(kpe, pe) ->
    addStm $ Let (Pattern [] [pe]) () $
             BasicOp $ Reshape cs [DimNew se | se <- arrayDims $ patElemAttr pe]
                               (patElemName kpe)

  where
    one = constant (1 :: Int32)

    -- number of reduction results (tuple size for reduction operator)
    num_redres = length nes

    ----------------------------------------------------------------------------
    -- The functions below generate all the needed code for the two different
    -- version of segmented-redomap (one group per segment, and many groups per
    -- segment).
    --
    -- We rename statements before adding them because the same lambdas
    -- (reduce/fold) are used multiple times, and we do not want to bind the
    -- same VName twice (as this is a type error)
    ----------------------------------------------------------------------------
    oneGroupPerSeg all_arrs reduce_lam' kern_chunk_fold_lam = do
      mapres_pes <- forM (drop num_redres $ patternValueElements flat_pat) $ \pe -> do
        vn' <- newName $ patElemName pe
        return $ PatElem vn' BindVar $ patElemType pe

      (kernel, _, _) <- groupPerSegmentKernel segment_size num_segments cs
        all_arrs comm reduce_lam' kern_chunk_fold_lam
        nes w OneGroupOneSegment
        ispace inps

{'fraction_non_alphanumeric': 0.09013980868285504, 'fraction_numerical': 0.0020235467255334807, 'mean_word_length': 2.157599225556631, 'pattern_counts': {'":': 0, '<': 55, '<?xml version=': 0, '>': 55, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2969758', 'n_tokens_mistral': 4676, 'n_tokens_neox': 4506, 'n_words': 1631}