From bacd70ca6ba028e935bb512aeee713943901acdd Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Mon, 3 Mar 2025 23:47:42 +0100
Subject: Compile: replicate1i
---
src/Compile.hs | 28 +++++++++++++++++++++++++---
1 file changed, 25 insertions(+), 3 deletions(-)
(limited to 'src')
diff --git a/src/Compile.hs b/src/Compile.hs
index 3b02c60..7a846d4 100644
--- a/src/Compile.hs
+++ b/src/Compile.hs
@@ -20,7 +20,7 @@ import Data.Set (Set)
import Data.Some
import qualified Data.Vector as V
import Foreign
-import System.IO (hPutStrLn, stderr)
+-- import System.IO (hPutStrLn, stderr)
import Prelude hiding ((^))
import qualified Prelude
@@ -664,12 +664,28 @@ compile' env = \case
return (CELit name)
EReplicate1Inner _ elen earg -> do
+ let STArr n t = typeOf earg
lenname <- genName' "replen"
emit . SVarDecl True (repSTy tIx) lenname =<< compile' env elen
argname <- genName' "reparg"
emit . SVarDecl True (repSTy (typeOf earg)) argname =<< compile' env earg
- _
+ shszname <- genName' "shsz"
+ emit $ SVarDecl True (repSTy tIx) shszname (compileArrShapeSize n argname)
+
+ resname <- allocArray "rep" (SS n) t
+ (CEBinop (CELit shszname) "*" (CELit lenname))
+ ([CELit (argname ++ ".buf->sh[" ++ show i ++ "]") | i <- [0 .. fromSNat n - 1]]
+ ++ [CELit lenname])
+
+ ivar <- genName' "i"
+ jvar <- genName' "j"
+ emit $ SLoop (repSTy tIx) ivar (CELit "0") (CELit shszname) $
+ pure $ SLoop (repSTy tIx) jvar (CELit "0") (CELit lenname) $
+ pure $ SAsg (resname ++ ".buf->a[" ++ ivar ++ " * " ++ lenname ++ " + " ++ jvar ++ "]")
+ (CELit (argname ++ ".buf->a[" ++ ivar ++ "]"))
+
+ return (CELit resname)
-- EMaximum1Inner _ e -> error "TODO" -- EMaximum1Inner ext (compile' e)
@@ -737,7 +753,13 @@ compile' env = \case
EPlus{} -> error "Compile: monoid operations should have been eliminated"
EOneHot{} -> error "Compile: monoid operations should have been eliminated"
- _ -> error "Compile: not implemented"
+ EFold1Inner{} -> error "Compile: not implemented: EFold1Inner"
+ EMaximum1Inner{} -> error "Compile: not implemented: EMaximum1Inner"
+ EMinimum1Inner{} -> error "Compile: not implemented: EMinimum1Inner"
+ EIdx1{} -> error "Compile: not implemented: EIdx1"
+ ECustom{} -> error "Compile: not implemented: ECustom"
+ EWith{} -> error "Compile: not implemented: EWith"
+ EAccum{} -> error "Compile: not implemented: EAccum"
data Increment = Increment | Decrement
deriving (Show)
--
cgit v1.2.3-70-g09d2