<commit_before>{-# LANGUAGE BangPatterns #-}

module Kernel.CPU.GCF (kernel) where

import Foreign.C
import Foreign.Ptr
import Foreign.Storable
import Foreign.Marshal.Array
import Data.Complex

import Kernel.CPU.FFT
import Data
import Vector

foreign import ccall mkGCFLayer ::
     FftPlan
  -> Ptr (Complex Double) -- destination buffer
  -> Ptr (Ptr (Complex Double)) -- table
  -> Ptr (Complex Double) -- work buffer
  -> CInt                 -- support size to extract
  -> CInt                 -- maximum support size
  -> CInt                 -- lead dim padding
  -> Double               -- Theta/2
  -> Double               -- w
  -> IO FftPlan

foreign import ccall calcAccums ::
     Ptr Double -- uvws
  -> Ptr Double -- sums (out)
  -> Ptr CInt   -- npts (out)
  -> Double     -- wstep
  -> CInt       -- # of points baselines * channels * timesteps
  -> CInt       -- # of planes
  -> IO ()

kernel ::
     GridPar
  -> GCFPar
  -> Vis
  -> IO GCFSet
kernel gp gcfp vis =
  allocaArray numOfPlanes $ \wsump ->
  allocaArray numOfPlanes $ \np ->
  allocaArray arenaSize $ \arenap -> do
    fftInitThreading
    gcfDataPtr <- mallocArray gcfDataSize
    gcfTablePtr <- mallocArray gcfTableSize

    calcAccums uvwp wsump np wstep (f $ visTimesteps vis * (length $ visBaselines vis)) (f numOfPlanes)
    -- I'm using an explicit recursion explicitly.
    -- Usually this generates much better code.
    -- OTOH, higher-order functions could potentially be fused.
    let
      mkLayers !p0 !destPtr !dTabPtr (s:ss) wsumpc npc = do
        wsumc <- peek wsumpc
        nc <- peek npc
        !p <- mkGCFLayer p0 destPtr dTabPtr arenap (f s) (f gcfMaxSize) (f pad) t2 (wsumc / fromIntegral nc)
        mkLayers p (advancePtr destPtr $ over2 * s * s) (advancePtr dTabPtr over2) ss (advancePtr wsumpc 1) (advancePtr npc 1)
      mkLayers _ _ _ [] _ _ = return ()
    --
    mkLayers nullPtr gcfDataPtr gcfTablePtr layerSizes wsump np
    -- FIXME!!!: Change all data types in such a way
    --   that they would bring their own finalizers with them !!!
    -- ATM we hack freeGCFSet slightly (see Data.hs)
    return $ GCFSet gcfp [] (CVector gcfTableSize $ castPtr gcfTablePtr)
  where
    uvwp = case visPositions vis of
             CVector _ p -> castPtr p
             _ -> error "Wrong uvw location for CPU GCF."
    f = fromIntegral
    t2 = gridTheta gp / 2
    wstep = gcfpStepW gcfp
    over = gcfpOver gcfp
    over2 = over * over
    gcfMaxSize = gcfpMaxSize gcfp
    size i = min gcfMaxSize
               (gcfpMinSize gcfp + gcfpGrowth gcfp * abs i)
    maxWPlane = max (round (visMaxW vis / wstep)) (round (-visMinW vis / wstep))
    numOfPlanes = 2*maxWPlane+1
    layerSizes = map size [-maxWPlane .. maxWPlane]
    gcfDataSize = over2 * sum layerSizes
    gcfTableSize = over2 * (1 + 2 * maxWPlane)
    l = over * gcfMaxSize
    pad = 4
    pitch = l + pad
    arenaSize = l * pitch<commit_msg>Fix 2 bugs in CPU GCF (NAN for unpopulated w-planes and wrong GCF layer size calculation).<commit_after>{-# LANGUAGE BangPatterns #-}

module Kernel.CPU.GCF (kernel) where

import Foreign.C
import Foreign.Ptr
import Foreign.Storable
import Foreign.Marshal.Array
import Data.Complex

import Kernel.CPU.FFT
import Data
import Vector

foreign import ccall mkGCFLayer ::
     FftPlan
  -> Ptr (Complex Double) -- destination buffer
  -> Ptr (Ptr (Complex Double)) -- table
  -> Ptr (Complex Double) -- work buffer
  -> CInt                 -- support size to extract
  -> CInt                 -- maximum support size
  -> CInt                 -- lead dim padding
  -> Double               -- Theta/2
  -> Double               -- w
  -> IO FftPlan

foreign import ccall calcAccums ::
     Ptr Double -- uvws
  -> Ptr Double -- sums (out)
  -> Ptr CInt   -- npts (out)
  -> Double     -- wstep
  -> CInt       -- # of points baselines * channels * timesteps
  -> CInt       -- # of planes
  -> IO ()

kernel ::
     GridPar
  -> GCFPar
  -> Vis
  -> IO GCFSet
kernel gp gcfp vis =
  allocaArray numOfPlanes $ \wsump ->
  allocaArray numOfPlanes $ \np ->
  allocaArray arenaSize $ \arenap -> do
    fftInitThreading
    gcfDataPtr <- mallocArray gcfDataSize
    gcfTablePtr <- mallocArray gcfTableSize

    calcAccums uvwp wsump np wstep (f $ visTimesteps vis * (length $ visBaselines vis)) (f numOfPlanes)
    -- I'm using an explicit recursion explicitly.
    -- Usually this generates much better code.
    -- OTOH, higher-order functions could potentially be fused.
    let
      mkLayers !p0 !destPtr !dTabPtr (s:ss) wsumpc npc wmid = do
        wsumc <- peek wsumpc
        nc <- peek npc
        let avg = if nc > 0 then wsumc/fromIntegral nc else wmid
        !p <- mkGCFLayer p0 destPtr dTabPtr arenap (f s) (f gcfMaxSize) (f pad) t2 avg
        mkLayers p (advancePtr destPtr $ over2 * s * s) (advancePtr dTabPtr over2) ss (advancePtr wsumpc 1) (advancePtr npc 1) (wmid + wstep)
      mkLayers _ _ _ [] _ _ _ = return ()
    --
    mkLayers nullPtr gcfDataPtr gcfTablePtr layerSupps wsump np (- wstep * fromIntegral maxWPlane)
    -- FIXME!!!: Change all data types in such a way
    --   that they would bring their own finalizers with them !!!
    -- ATM we hack freeGCFSet slightly (see Data.hs)
    return $ GCFSet gcfp [] (CVector gcfTableSize $ castPtr gcfTablePtr)
  where
    uvwp = case visPositions vis of
             CVector _ p -> castPtr p
             _ -> error "Wrong uvw location for CPU GCF."
    f = fromIntegral
    t2 = gridTheta gp / 2
    wstep = gcfpStepW gcfp
    over = gcfpOver gcfp
    over2 = over * over
    gcfMaxSize = gcfpMaxSize gcfp
    supp i = min gcfMaxSize
               (gcfpMinSize gcfp + gcfpGrowth gcfp * abs i)
    maxWPlane = max (round (visMaxW vis / wstep)) (round (-visMinW vis / wstep))
    numOfPlanes = 2*maxWPlane+1
    layerSupps = map supp [-maxWPlane .. maxWPlane]
    gcfDataSize = over2 * sum (map (\s -> s * s) layerSupps)
    gcfTableSize = over2 * (1 + 2 * maxWPlane)
    l = over * gcfMaxSize
    pad = 4
    pitch = l + pad
    arenaSize = l * pitch
{'fraction_non_alphanumeric': 0.09196354598177299, 'fraction_numerical': 0.006296603148301574, 'mean_word_length': 2.241675617615467, 'pattern_counts': {'":': 0, '<': 13, '<?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': '4103288', 'n_tokens_mistral': 2080, 'n_tokens_neox': 2009, 'n_words': 670}