futhark 0.25.26 → 0.25.27

raw patch · 33 files changed

+631/−365 lines, 33 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• CHANGELOG.md +28/−0
• futhark.cabal +1/−1
• rts/c/ispc_util.h +11/−11
• rts/c/scalar.h +25/−13
• rts/c/scalar_f16.h +6/−6
• rts/c/uniform.h +4/−5
• rts/python/scalar.py +12/−0
• src/Futhark/AD/Derivatives.hs +20/−2
• src/Futhark/AD/Rev/Scan.hs +225/−72
• src/Futhark/Analysis/PrimExp.hs +14/−1
• src/Futhark/Analysis/PrimExp/Convert.hs +2/−2
• src/Futhark/CLI/Eval.hs +2/−2
• src/Futhark/CLI/Fmt.hs +5/−5
• src/Futhark/CLI/Run.hs +5/−3
• src/Futhark/CodeGen/Backends/GenericC/Code.hs +1/−1
• src/Futhark/CodeGen/Backends/MulticoreC.hs +1/−1
• src/Futhark/CodeGen/Backends/MulticoreISPC.hs +2/−2
• src/Futhark/IR/Prop.hs +1/−1
• src/Futhark/IR/SOACS/Simplify.hs +19/−13
• src/Futhark/IR/TypeCheck.hs +1/−1
• src/Futhark/Optimise/Simplify/Rules.hs +1/−1
• src/Futhark/Optimise/Simplify/Rules/Simple.hs +0/−2
• src/Futhark/Optimise/Sink.hs +3/−4
• src/Futhark/Pass/ExtractKernels/DistributeNests.hs +1/−6
• src/Futhark/Pass/ExtractKernels/Distribution.hs +2/−2
• src/Language/Futhark/Core.hs +5/−0
• src/Language/Futhark/Interpreter.hs +151/−158
• src/Language/Futhark/Interpreter/AD.hs +37/−31
• src/Language/Futhark/Primitive.hs +27/−2
• src/Language/Futhark/Prop.hs +12/−11
• src/Language/Futhark/TypeChecker/Names.hs +1/−1
• src/Language/Futhark/TypeChecker/Unify.hs +2/−2
• unittests/Futhark/AD/DerivativesTests.hs +4/−3

@@ -5,6 +5,34 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html). +## [0.25.27]++### Added++* Improved reverse-mode AD of `scan` with complicated operators. Work+  by Peter Adema and Sophus Valentin Willumsgaard.++### Fixed++* `futhark eval`: any errors in the provided .fut file would cause a+  "file not found" error message.++* Handling of module-dependent size expressions in type abbreviations+  (#2209).++* A `let`-bound size would mistakenly be in scope of the bound+  expression (#2210).++* An overzealous floating-point simplification rule.++* Corrected AD of `x**y` where `x==0` (#2216).++* `futhark fmt`: correct file name in parse errors.++* A bug in the "sink" optimisation pass could cause compiler crashes.++* Compile errors with newer versions of `ispc`.+ ## [0.25.26]  ### Fixed

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.25.26+version:        0.25.27 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to

@@ -27,17 +27,17 @@ make_extract(float16) make_extract(float) make_extract(double)-make_extract(int8* uniform)-make_extract(int16* uniform)-make_extract(int32* uniform)-make_extract(int64* uniform)-make_extract(uint8* uniform)-make_extract(uint16* uniform)-make_extract(uint32* uniform)-make_extract(uint64* uniform)-make_extract(float16* uniform)-make_extract(float* uniform)-make_extract(double* uniform)+/* make_extract(int8* uniform) */+/* make_extract(int16* uniform) */+/* make_extract(int32* uniform) */+/* make_extract(int64* uniform) */+/* make_extract(uint8* uniform) */+/* make_extract(uint16* uniform) */+/* make_extract(uint32* uniform) */+/* make_extract(uint64* uniform) */+/* make_extract(float16* uniform) */+/* make_extract(float* uniform) */+/* make_extract(double* uniform) */ make_extract(struct futhark_context) make_extract(struct memblock) 

@@ -68,7 +68,7 @@   return x * y; } -#if ISPC+#if defined(ISPC)  SCALAR_FUN_ATTR uint8_t udiv8(uint8_t x, uint8_t y) {   // This strange pattern is used to prevent the ISPC compiler from@@ -1309,7 +1309,7 @@ SCALAR_FUN_ATTR uint16_t futrts_smul_hi16(int16_t a, int16_t b) { return ((int32_t)a) * ((int32_t)b) >> 16; } SCALAR_FUN_ATTR uint32_t futrts_smul_hi32(int32_t a, int32_t b) { return __mulhi(a, b); } SCALAR_FUN_ATTR uint64_t futrts_smul_hi64(int64_t a, int64_t b) { return __mul64hi(a, b); }-#elif ISPC+#elif defined(ISPC) SCALAR_FUN_ATTR uint8_t futrts_umul_hi8(uint8_t a, uint8_t b) { return ((uint16_t)a) * ((uint16_t)b) >> 8; } SCALAR_FUN_ATTR uint16_t futrts_umul_hi16(uint16_t a, uint16_t b) { return ((uint32_t)a) * ((uint32_t)b) >> 16; } SCALAR_FUN_ATTR uint32_t futrts_umul_hi32(uint32_t a, uint32_t b) { return ((uint64_t)a) * ((uint64_t)b) >> 32; }@@ -1430,7 +1430,7 @@   return __clzll(x); } -#elif ISPC+#elif defined(ISPC)  SCALAR_FUN_ATTR int32_t futrts_clzz8(int8_t x) {   return count_leading_zeros((int32_t)(uint8_t)x)-24;@@ -1518,7 +1518,7 @@   return y == 0 ? 64 : y - 1; } -#elif ISPC+#elif defined(ISPC)  SCALAR_FUN_ATTR int32_t futrts_ctzz8(int8_t x) {   return x == 0 ? 8 : count_trailing_zeros((int32_t)x);@@ -1628,7 +1628,7 @@   return pow(x, y); } -#elif ISPC+#elif defined(ISPC)  SCALAR_FUN_ATTR float fabs32(float x) {   return abs(x);@@ -1645,7 +1645,7 @@ SCALAR_FUN_ATTR float fpow32(float a, float b) {   float ret;   foreach_active (i) {-      uniform float r = __stdlib_powf(extract(a, i), extract(b, i));+      uniform float r = pow(extract(a, i), extract(b, i));       ret = insert(ret, i, r);   }   return ret;@@ -1674,7 +1674,7 @@   return isnan(x); } -#if ISPC+#if defined(ISPC)  SCALAR_FUN_ATTR bool futrts_isinf32(float x) {   return !isnan(x) && isnan(x - x);@@ -1905,7 +1905,7 @@   return fma(a, b, c); } -#elif ISPC+#elif defined(ISPC)  SCALAR_FUN_ATTR float futrts_log32(float x) {   return futrts_isfinite32(x) || (futrts_isinf32(x) && x < 0)? log(x) : x;@@ -2107,7 +2107,7 @@ }  SCALAR_FUN_ATTR float futrts_ldexp32(float x, int32_t y) {-  return x * pow((double)2.0, (double)y);+  return x * pow((uniform float)2.0, (float)y); }  SCALAR_FUN_ATTR float futrts_copysign32(float x, float y) {@@ -2267,7 +2267,7 @@ } #endif -#if ISPC+#if defined(ISPC) SCALAR_FUN_ATTR int32_t futrts_to_bits32(float x) {   return intbits(x); }@@ -2306,7 +2306,7 @@ SCALAR_FUN_ATTR double futrts_from_bits64(int64_t x); SCALAR_FUN_ATTR int64_t futrts_to_bits64(double x); -#if ISPC+#if defined(ISPC) SCALAR_FUN_ATTR bool futrts_isinf64(float x) {   return !isnan(x) && isnan(x - x); }@@ -2386,7 +2386,7 @@ SCALAR_FUN_ATTR double fpow64(double a, double b) {   float ret;   foreach_active (i) {-      uniform float r = __stdlib_powf(extract(a, i), extract(b, i));+      uniform float r = pow(extract(a, i), extract(b, i));       ret = insert(ret, i, r);   }   return ret;@@ -2673,7 +2673,7 @@ }  SCALAR_FUN_ATTR double futrts_ldexp64(double x, int32_t y) {-  return x * pow((double)2.0, (double)y);+  return x * pow((uniform double)2.0, (double)y); }  SCALAR_FUN_ATTR double futrts_copysign64(double x, double y) {@@ -3047,5 +3047,17 @@ #endif  #endif++#define futrts_cond_f16(x,y,z) ((x) ? (y) : (z))+#define futrts_cond_f32(x,y,z) ((x) ? (y) : (z))+#define futrts_cond_f64(x,y,z) ((x) ? (y) : (z))++#define futrts_cond_i8(x,y,z) ((x) ? (y) : (z))+#define futrts_cond_i16(x,y,z) ((x) ? (y) : (z))+#define futrts_cond_i32(x,y,z) ((x) ? (y) : (z))+#define futrts_cond_i64(x,y,z) ((x) ? (y) : (z))++#define futrts_cond_bool(x,y,z) ((x) ? (y) : (z))+#define futrts_cond_unit(x,y,z) ((x) ? (y) : (z))  // End of scalar.h.

@@ -23,7 +23,7 @@ // compiler will have to be real careful! typedef float f16; -#elif ISPC+#elif defined(ISPC) typedef float16 f16;  #else@@ -154,7 +154,7 @@   return pow(x, y); } -#elif ISPC+#elif defined(ISPC) SCALAR_FUN_ATTR f16 fabs16(f16 x) {   return abs(x); }@@ -190,7 +190,7 @@ } #endif -#if ISPC+#if defined(ISPC) SCALAR_FUN_ATTR bool futrts_isinf16(float x) {   return !futrts_isnan16(x) && futrts_isnan16(x - x); }@@ -345,7 +345,7 @@ SCALAR_FUN_ATTR f16 futrts_fma16(f16 a, f16 b, f16 c) {   return fma(a, b, c); }-#elif ISPC+#elif defined(ISPC)  SCALAR_FUN_ATTR f16 futrts_log16(f16 x) {   return futrts_isfinite16(x) || (futrts_isinf16(x) && x < 0) ? log(x) : x;@@ -664,7 +664,7 @@ SCALAR_FUN_ATTR f16 futrts_from_bits16(int16_t x) {   return __ushort_as_half(x); }-#elif ISPC+#elif defined(ISPC)  SCALAR_FUN_ATTR int16_t futrts_to_bits16(f16 x) {   varying int16_t y = *((varying int16_t * uniform)&x);@@ -916,7 +916,7 @@   return (double) x; } -#if ISPC+#if defined(ISPC) SCALAR_FUN_ATTR f16 fpconv_f64_f16(double x) {   return (f16) ((float)x); }

@@ -1,10 +1,9 @@- // Start of uniform.h  // Uniform versions of all library functions as to // improve performance in ISPC when in an uniform context. -#if ISPC+#if defined(ISPC)  static inline uniform uint8_t add8(uniform uint8_t x, uniform uint8_t y) {   return x + y;@@ -839,7 +838,7 @@ }  static inline uniform float fpow32(uniform float x, uniform float y) {-  return __stdlib_powf(x, y);+  return pow(x, y); }  static inline uniform bool futrts_isnan32(uniform float x) {@@ -1181,7 +1180,7 @@ }  static inline uniform double fpow64(uniform double x, uniform double y) {-  return __stdlib_powf(x, y);+  return pow(x, y); }  static inline uniform double futrts_log64(uniform double x) {@@ -1445,7 +1444,7 @@ }  static inline uniform f16 fpconv_f64_f16(uniform double x) {-  return (uniform f16) ((uniform float)x); +  return (uniform f16) ((uniform float)x); }  #endif

@@ -1028,4 +1028,16 @@  futhark_copysign16 = futhark_copysign32 = futhark_copysign64 = np.copysign ++def futhark_cond(x, y, z):+    return y if x else z+++futhark_cond_f16 = futhark_cond_f32 = futhark_cond_f64 = futhark_cond+futhark_cond_i18 = futhark_cond_i16 = futhark_cond_i32 = futhark_cond_i64 = (+    futhark_cond+)+futhark_cond_bool = futhark_cond_unit = futhark_cond++ # End of scalar.py.

@@ -8,7 +8,7 @@  import Data.Bifunctor (bimap) import Futhark.Analysis.PrimExp.Convert-import Futhark.IR.Syntax.Core (Name, VName)+import Futhark.IR.Syntax.Core (Name, VName, nameToText) import Futhark.Util.IntegralExp import Prelude hiding (quot) @@ -134,7 +134,10 @@ pdBinOp (FPow ft) a b =   floatBinOp derivs derivs derivs ft a b   where-    derivs x y = (y * (x ** (y - 1)), (x ** y) * log x)+    derivs x y =+      ( y * (x ** (y - 1)),+        condExp (x .<=. 0) 0 ((x ** y) * log x)+      ) pdBinOp (FMax ft) a b =   floatBinOp derivs derivs derivs ft a b   where@@ -375,6 +378,21 @@   Just [untyped $ 1 * isF32 (UnOpExp (FSignum Float32) y), fConst Float32 0] pdBuiltin "copysign64" [_x, y] =   Just [untyped $ 1 * isF64 (UnOpExp (FSignum Float64) y), fConst Float64 0]+pdBuiltin h [x, _y, _z]+  | Just t <- isCondFun $ nameToText h =+      Just+        [ boolToT t false,+          boolToT t $ isBool x,+          boolToT t $ bNot $ isBool x+        ]+  where+    boolToT t = case t of+      IntType it ->+        ConvOpExp (BToI it) . untyped+      FloatType ft ->+        ConvOpExp (SIToFP Int32 ft) . ConvOpExp (BToI Int32) . untyped+      Bool -> untyped+      Unit -> const $ ValueExp UnitValue -- More problematic derivatives follow below. pdBuiltin "umul_hi8" [x, y] = Just [y, x] pdBuiltin "umul_hi16" [x, y] = Just [y, x]

@@ -33,9 +33,7 @@ vectorAdd = zipWith (~+~)  orderArgs :: Special -> [a] -> [[a]]-orderArgs s lst =-  let d = div (length lst) $ specialScans s-   in chunk d lst+orderArgs s lst = chunk (div (length lst) $ specialScans s) lst  -- computes `d(x op y)/dx` or d(x op y)/dy mkScanAdjointLam :: VjpOps -> Lambda SOACS -> FirstOrSecond -> [SubExp] -> ADM (Lambda SOACS)@@ -55,19 +53,19 @@ --       `xs` is  the input  of scan --       `ys_adj` is the known adjoint of ys --       `j` draw values from `iota n`-mkScanFusedMapLam ::-  VjpOps ->-  SubExp ->-  Lambda SOACS ->-  [VName] ->-  [VName] ->-  [VName] ->-  Special ->-  Int ->-  ADM (Lambda SOACS)+mkScanFusedMapLam :: -- i and j above are probably swapped in the code below+  VjpOps -> -- (ops) helper functions+  SubExp -> -- (w) ~length of arrays e.g. xs+  Lambda SOACS -> -- (scn_lam) the scan to be differentiated ('scan' turned into a lambda)+  [VName] -> -- (xs) input of the scan (actually as)+  [VName] -> -- (ys) output of the scan+  [VName] -> -- (ys_adj) adjoint of ys+  Special -> -- (s) information about which special case we're working with for the scan derivative+  Int -> -- (d) dimension of the input (number of elements in the input tuple)+  ADM (Lambda SOACS) -- output: some kind of codegen for the lambda mkScanFusedMapLam ops w scn_lam xs ys ys_adj s d = do   let sc = specialCase s-  let k = specialSubSize s+      k = specialSubSize s   ys_ts <- traverse lookupType ys   idmat <- identityM (length ys) $ rowType $ head ys_ts   lams <- traverse (mkScanAdjointLam ops scn_lam WrtFirst) idmat@@ -82,7 +80,7 @@             y_s <- forM ys_adj $ \y_ ->               letSubExp (baseString y_ ++ "_j") =<< eIndex y_ [eSubExp j]             let zso = orderArgs s y_s-            let ido = orderArgs s $ case_jac k sc idmat+            let ido = orderArgs s $ caseJac k sc idmat             pure $ subExpsRes $ concat $ zipWith (++) zso $ fmap concat ido         )         ( buildBody_ $ do@@ -96,41 +94,38 @@             lam_rs <- traverse (`eLambda` args) lams              let yso = orderArgs s $ subExpsRes y_s-            let jaco = orderArgs s $ case_jac k sc $ transpose lam_rs+            let jaco = orderArgs s $ caseJac k sc $ transpose lam_rs              pure $ concat $ zipWith (++) yso $ fmap concat jaco         )   where-    case_jac :: Int -> SpecialCase -> [[a]] -> [[a]]-    case_jac _ Generic jac = jac-    case_jac k ZeroQuadrant jac =-      concat-        $ zipWith-          (\i -> map (take k . drop (i * k)))-          [0 .. d `div` k]-        $ chunk k jac-    case_jac k MatrixMul jac =+    caseJac :: Int -> Maybe SpecialCase -> [[a]] -> [[a]]+    caseJac _ Nothing jac = jac+    caseJac k (Just ZeroQuadrant) jac =+      concat $+        zipWith (\i -> map (take k . drop (i * k))) [0 .. d `div` k] $+          chunk k jac+    caseJac k (Just MatrixMul) jac =       take k <$> take k jac  -- a1 a2 b -> a2 + b * a1 linFunT0 :: [PrimExp VName] -> [PrimExp VName] -> [[PrimExp VName]] -> Special -> PrimType -> [PrimExp VName] linFunT0 a1 a2 b s pt =   let t = case specialCase s of-        MatrixMul ->+        Just MatrixMul ->           concatMap (\v -> matrixVecMul b v pt) $ chunk (specialSubSize s) a1         _ -> matrixVecMul b a1 pt    in a2 `vectorAdd` t  -- \(a1, b1) (a2, b2) -> (a2 + b2 * a1, b2 * b1)-mkScanLinFunO :: Type -> Special -> ADM (Scan SOACS)+mkScanLinFunO :: Type -> Special -> ADM (Scan SOACS) -- a is an instance of y_bar, b is a Jacobian (a 'c' in the 2023 paper) mkScanLinFunO t s = do   let pt = elemType t   neu_elm <- mkNeutral $ specialNeutral s-  let (as, bs) = specialParams s-  (a1s, b1s, a2s, b2s) <- mkParams (as, bs)-  let pet = primExpFromSubExp pt . Var-  let (_, n) = specialNeutral s-+  let (as, bs) = specialParams s -- input size, Jacobian element count+  (a1s, b1s, a2s, b2s) <- mkParams (as, bs) -- create sufficient free variables to bind every element of the vectors / matrices+  let pet = primExpFromSubExp pt . Var -- manifest variable names as expressions+  let (_, n) = specialNeutral s -- output size (one side of the Jacobian)   lam <- mkLambda (map (\v -> Param mempty v (rowType t)) (a1s ++ b1s ++ a2s ++ b2s)) . fmap subExpsRes $ do     let [a1s', b1s', a2s', b2s'] = (fmap . fmap) pet [a1s, b1s, a2s, b2s]     let (b1sm, b2sm) = (chunk n b1s', chunk n b2s')@@ -170,9 +165,8 @@       j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))        dj <--        traverse-          (\dd -> letExp (baseString dd ++ "_dj") =<< eIndex dd [eSubExp j])-          ds+        forM ds $ \dd ->+          letExp (baseString dd ++ "_dj") =<< eIndex dd [eSubExp j]        fmap varsRes . letTupExp "scan_contribs"         =<< eIf@@ -192,21 +186,40 @@   iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64   letTupExp "scan_contribs" $ Op $ Screma w (iota : xs) $ mapSOAC map_lam -data SpecialCase-  = Generic-  | ZeroQuadrant-  | MatrixMul-  deriving (Show)+-- | Scan special cases.+data SpecialCase = ZeroQuadrant | MatrixMul deriving (Show) +-- | Metadata for how to perform the scan for the return sweep. data Special = Special-  { specialNeutral :: (Int, Int),+  { -- | Size of one of the two dimensions of the Jacobian (e.g. 3 if+    --  it's 3x3, must be square because scan must be a->a->a). It's+    --  the size of the special neutral element, not the element itself+    specialNeutral :: (Int, Int),+    -- | Size of input (nr params); Flat size of Jacobian (dim1 *+    -- dim2)). Number of params for the special lambda.     specialParams :: (Int, Int),+    -- | The number of scans to do, 1 in most cases, k in the+    -- ZeroQuadrant (block diagonal?) case.     specialScans :: Int,+    -- | Probably: the size of submatrices for the ZeroQuadrant (block+    -- diagonal?) case, or 1 otherwise.     specialSubSize :: Int,-    specialCase :: SpecialCase+    -- | Which case.+    specialCase :: Maybe SpecialCase   }   deriving (Show) +-- | The different ways to handle scans. The best one is chosen+-- heuristically by looking at the operator.+data ScanAlgo+  = -- | Construct and compose the Jacobians; the approach presented+    -- in *Reverse-Mode AD of Multi-Reduce and Scan in Futhark*.+    GenericIFL23 Special+  | -- | The approach from *Parallelism-preserving automatic+    -- differentiation for second-order array languages*.+    GenericPPAD+  deriving (Show)+ subMats :: Int -> [[Exp SOACS]] -> Exp SOACS -> Maybe Int subMats d mat zero =   let sub_d = filter (\x -> d `mod` x == 0) [1 .. (d `div` 2)]@@ -215,25 +228,30 @@    in if null tmp then Nothing else Just $ snd $ head tmp   where     ok m (row, i) =-      all (\(v, j) -> v == zero || (i `div` m == j `div` m)) $+      all (\(v, j) -> v == zero || i `div` m == j `div` m) $         zip row [0 .. d - 1] -cases :: Int -> Type -> [[Exp SOACS]] -> Special-cases d t mat-  | Just k <- subMats d mat $ zeroExp t =-      let nonZeros = zipWith (\i -> map (take k . drop (i * k))) [0 .. d `div` k] $ chunk k mat-       in if all (== head nonZeros) $ tail nonZeros-            then Special (d, k) (d, k * k) 1 k MatrixMul-            else Special (k, k) (k, k * k) (d `div` k) k ZeroQuadrant-cases d _ _ = Special (d, d) (d, d * d) 1 d Generic+cases :: Int -> Type -> [[Exp SOACS]] -> ScanAlgo+cases d t mat = case subMats d mat $ zeroExp t of+  Just k ->+    let nonZeros = zipWith (\i -> map (take k . drop (i * k))) [0 .. d `div` k] $ chunk k mat+     in if all (== head nonZeros) $ tail nonZeros+          then GenericIFL23 $ Special (d, k) (d, k * k) 1 k $ Just MatrixMul+          else GenericIFL23 $ Special (k, k) (k, k * k) (d `div` k) k $ Just ZeroQuadrant+  Nothing ->+    case d of+      1 -> GenericIFL23 $ Special (d, d) (d, d * d) 1 d Nothing+      _ -> GenericPPAD -identifyCase :: VjpOps -> Lambda SOACS -> ADM Special+-- | construct and optimise a temporary lambda, that calculates the+-- Jacobian of the scan op. Figure out if the Jacobian has some+-- special shape, discarding the temporary lambda.+identifyCase :: VjpOps -> Lambda SOACS -> ADM ScanAlgo identifyCase ops lam = do   let t = lambdaReturnType lam   let d = length t   idmat <- identityM d $ head t   lams <- traverse (mkScanAdjointLam ops lam WrtFirst) idmat-   par1 <- traverse (newParam "tmp1") t   par2 <- traverse (newParam "tmp2") t   jac_lam <- mkLambda (par1 ++ par2) $ do@@ -246,22 +264,161 @@   let jac = chunk d $ fmap (BasicOp . SubExp . resSubExp) $ bodyResult $ lambdaBody simp   pure $ cases d (head t) jac +scanRight :: [VName] -> SubExp -> Scan SOACS -> ADM [VName]+scanRight as w scan = do+  as_types <- mapM lookupType as+  let arg_type_row = map rowType as_types++  par_a1 <- zipWithM (\x -> newParam (baseString x <> "_par_a1")) as arg_type_row+  par_a2 <- zipWithM (\x -> newParam (baseString x <> "_par_a2")) as arg_type_row+  -- Just the original operator but with par_a1 and par_a2 swapped.+  rev_op <- mkLambda (par_a1 <> par_a2) $ do+    op <- renameLambda $ scanLambda scan+    eLambda op (map (toExp . paramName) (par_a2 <> par_a1))+  -- same neutral element+  let e = scanNeutral scan+  let rev_scan = Scan rev_op e++  iota <-+    letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64+  -- flip the input array (this code is inspired from the code in+  -- diffScanAdd, but made to work with [VName] instead VName)+  map_scan <- revArrLam as+  -- perform the scan+  scan_res <-+    letTupExp "adj_ctrb_scan" . Op . Screma w [iota] $+      scanomapSOAC [rev_scan] map_scan+  -- flip the output array again+  rev_lam <- revArrLam scan_res+  letTupExp "reverse_scan_result" $ Op $ Screma w [iota] $ mapSOAC rev_lam+  where+    revArrLam :: [VName] -> ADM (Lambda SOACS)+    revArrLam arrs = do+      par_i <- newParam "i" $ Prim int64+      mkLambda [par_i] . forM arrs $ \arr ->+        fmap varRes . letExp "ys_bar_rev"+          =<< eIndex arr [toExp (pe64 w - le64 (paramName par_i) - 1)]++mkPPADOpLifted :: VjpOps -> [VName] -> Scan SOACS -> ADM (Lambda SOACS)+mkPPADOpLifted ops as scan = do+  as_types <- mapM lookupType as+  let arg_type_row = map rowType as_types+  par_x1 <- zipWithM (\x -> newParam (baseString x ++ "_par_x1")) as arg_type_row+  par_x2_unused <- zipWithM (\x -> newParam (baseString x ++ "_par_x2_unused")) as arg_type_row+  par_a1 <- zipWithM (\x -> newParam (baseString x ++ "_par_a1")) as arg_type_row+  par_a2 <- zipWithM (\x -> newParam (baseString x ++ "_par_a2")) as arg_type_row+  par_y1_h <- zipWithM (\x -> newParam (baseString x ++ "_par_y1_h")) as arg_type_row+  par_y2_h <- zipWithM (\x -> newParam (baseString x ++ "_par_y2_h")) as arg_type_row++  add_lams <- mapM addLambda arg_type_row++  mkLambda+    (par_x1 ++ par_a1 ++ par_y1_h ++ par_x2_unused ++ par_a2 ++ par_y2_h)+    (op_lift par_x1 par_a1 par_y1_h par_a2 par_y2_h add_lams)+  where+    op_lift px1 pa1 py1 pa2 py2 adds = do+      op_bar_1 <- mkScanAdjointLam ops (scanLambda scan) WrtFirst (Var . paramName <$> py2)+      let op_bar_args = toExp . Var . paramName <$> px1 ++ pa1+      z_term <- map resSubExp <$> eLambda op_bar_1 op_bar_args+      let z =+            mapM+              (\(z_t, y_1, add) -> head <$> eLambda add [toExp z_t, toExp y_1])+              (zip3 z_term (Var . paramName <$> py1) adds)++      let x1 = subExpsRes <$> mapM (toSubExp "x1" . Var . paramName) px1+      op <- renameLambda $ scanLambda scan+      let a3 = eLambda op (toExp . paramName <$> pa1 ++ pa2)++      concat <$> sequence [x1, a3, z]++asLiftPPAD :: [VName] -> SubExp -> [SubExp] -> ADM [VName]+asLiftPPAD as w e = do+  par_i <- newParam "i" $ Prim int64+  lmb <- mkLambda [par_i] $ do+    forM (zip as e) $ \(arr, arr_e) -> do+      a_lift <-+        letExp "a_lift"+          =<< eIf+            ( do+                nm1 <- toSubExp "n_minus_one" $ pe64 w - 1+                pure $ BasicOp $ CmpOp (CmpSlt Int64) (Var $ paramName par_i) nm1+            )+            ( buildBody_ $ (\x -> [subExpRes x]) <$> (letSubExp "val" =<< eIndex arr [toExp $ le64 (paramName par_i) + 1])+            )+            (buildBody_ $ pure [subExpRes arr_e])+      pure $ varRes a_lift++  iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64+  letTupExp "as_lift" $ Op $ Screma w [iota] $ mapSOAC lmb++ysRightPPAD :: [VName] -> SubExp -> [SubExp] -> ADM [VName]+ysRightPPAD ys w e = do+  par_i <- newParam "i" $ Prim int64+  lmb <- mkLambda [par_i] $ do+    forM (zip ys e) $ \(arr, arr_e) -> do+      a_lift <-+        letExp "y_right"+          =<< eIf+            ( pure $ BasicOp $ CmpOp (CmpEq int64) (Var $ paramName par_i) (constant (0 :: Int64))+            )+            (buildBody_ $ pure [subExpRes arr_e])+            ( buildBody_ $ (\x -> [subExpRes x]) <$> (letSubExp "val" =<< eIndex arr [toExp $ le64 (paramName par_i) - 1])+            )+      pure $ varRes a_lift++  iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64+  letTupExp "ys_right" $ Op $ Screma w [iota] $ mapSOAC lmb++finalMapPPAD :: VjpOps -> [VName] -> Scan SOACS -> ADM (Lambda SOACS)+finalMapPPAD ops as scan = do+  as_types <- mapM lookupType as+  let arg_type_row = map rowType as_types+  par_y_right <- zipWithM (\x -> newParam (baseString x ++ "_par_y_right")) as arg_type_row+  par_a <- zipWithM (\x -> newParam (baseString x ++ "_par_a")) as arg_type_row+  par_r_adj <- zipWithM (\x -> newParam (baseString x ++ "_par_r_adj")) as arg_type_row++  mkLambda (par_y_right ++ par_a ++ par_r_adj) $ do+    op_bar_2 <- mkScanAdjointLam ops (scanLambda scan) WrtSecond (Var . paramName <$> par_r_adj)+    eLambda op_bar_2 $ toExp . Var . paramName <$> par_y_right ++ par_a+ diffScan :: VjpOps -> [VName] -> SubExp -> [VName] -> Scan SOACS -> ADM () diffScan ops ys w as scan = do-  sc <- identifyCase ops (scanLambda scan)+  -- ys ~ results of scan, w ~ size of input array, as ~ (unzipped)+  -- arrays, scan ~ scan: operator with ne+  scan_case <- identifyCase ops $ scanLambda scan   let d = length as-  ys_adj <- mapM lookupAdjVal ys+  ys_adj <- mapM lookupAdjVal ys -- ys_bar   as_ts <- mapM lookupType as-  map1_lam <- mkScanFusedMapLam ops w (scanLambda scan) as ys ys_adj sc d-  scans_lin_fun_o <- mkScanLinFunO (head as_ts) sc-  scan_lams <- mkScans (specialScans sc) scans_lin_fun_o-  iota <--    letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64-  r_scan <--    letTupExp "adj_ctrb_scan" . Op . Screma w [iota] $-      scanomapSOAC scan_lams map1_lam -  as_contribs <- mkScanFinalMap ops w (scanLambda scan) as ys (splitScanRes sc r_scan d)+  as_contribs <- case scan_case of+    GenericPPAD -> do+      let e = scanNeutral scan+      as_lift <- asLiftPPAD as w e++      let m = ys ++ as_lift ++ ys_adj++      op_lft <- mkPPADOpLifted ops as scan+      a_zero <- mapM (fmap Var . letExp "rscan_zero" . zeroExp . rowType) as_ts+      let lft_scan = Scan op_lft $ e ++ e ++ a_zero+      rs_adj <- (!! 2) . chunk d <$> scanRight m w lft_scan++      ys_right <- ysRightPPAD ys w e++      final_lmb <- finalMapPPAD ops as scan+      letTupExp "as_bar" $ Op $ Screma w (ys_right ++ as ++ rs_adj) $ mapSOAC final_lmb+    GenericIFL23 sc -> do+      -- IFL23+      map1_lam <- mkScanFusedMapLam ops w (scanLambda scan) as ys ys_adj sc d+      scans_lin_fun_o <- mkScanLinFunO (head as_ts) sc+      scan_lams <- mkScans (specialScans sc) scans_lin_fun_o+      iota <-+        letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64+      r_scan <-+        letTupExp "adj_ctrb_scan" . Op . Screma w [iota] $+          scanomapSOAC scan_lams map1_lam+      mkScanFinalMap ops w (scanLambda scan) as ys (splitScanRes sc r_scan d)+  -- Goal: calculate as_contribs in new way+  -- zipWithM_ updateAdj as as_contribs -- as_bar += new adjoint   zipWithM_ updateAdj as as_contribs   where     mkScans :: Int -> Scan SOACS -> ADM [Scan SOACS]@@ -269,7 +426,6 @@       replicateM d $ do         lam' <- renameLambda $ scanLambda s         pure $ Scan lam' $ scanNeutral s-     splitScanRes sc res d =       concatMap (take (div d $ specialScans sc)) (orderArgs sc res) @@ -289,9 +445,8 @@     let rear = [1, 0] ++ drop 2 [0 .. rank - 1]      transp_as <--      traverse-        (\a -> letExp (baseString a ++ "_transp") $ BasicOp $ Rearrange rear a)-        as+      forM as $ \a ->+        letExp (baseString a ++ "_transp") $ BasicOp $ Rearrange rear a      ts <- traverse lookupType transp_as     let n = arraysSize 0 ts@@ -309,10 +464,8 @@       letTupExp "trans_ys" . Op $         Screma n (transp_as ++ subExpVars ne) (mapSOAC map_lam) -    zipWithM-      (\y x -> auxing aux $ letBindNames [y] $ BasicOp $ Rearrange rear x)-      ys-      transp_ys+    forM (zip ys transp_ys) $ \(y, x) ->+      auxing aux $ letBindNames [y] $ BasicOp $ Rearrange rear x    foldr (vjpStm ops) m stmts 

@@ -59,6 +59,7 @@     fMax16,     fMax32,     fMax64,+    condExp,      -- * Untyped construction     (~*~),@@ -73,6 +74,7 @@ import Control.Monad import Data.Map qualified as M import Data.Set qualified as S+import Data.Text qualified as T import Data.Traversable import Futhark.IR.Prop.Names import Futhark.Util.IntegralExp@@ -93,7 +95,7 @@   | CmpOpExp CmpOp (PrimExp v) (PrimExp v)   | UnOpExp UnOp (PrimExp v)   | ConvOpExp ConvOp (PrimExp v)-  | FunExp String [PrimExp v] PrimType+  | FunExp T.Text [PrimExp v] PrimType   deriving (Eq, Ord, Show)  instance Functor PrimExp where@@ -699,6 +701,17 @@ -- | 64-bit float maximum. fMax64 :: TPrimExp Double v -> TPrimExp Double v -> TPrimExp Double v fMax64 x y = isF64 $ BinOpExp (FMax Float64) (untyped x) (untyped y)++-- | Conditional expression.+condExp :: TPrimExp Bool v -> TPrimExp t v -> TPrimExp t v -> TPrimExp t v+condExp x y z =+  TPrimExp $+    FunExp+      (condFun t)+      [untyped x, untyped y, untyped z]+      t+  where+    t = primExpType $ untyped y  -- | Convert result of some integer expression to have the same type -- as another, using sign extension.

@@ -45,7 +45,7 @@   toExp (ValueExp v) =     pure $ BasicOp $ SubExp $ Constant v   toExp (FunExp h args t) =-    Apply (nameFromString h)+    Apply (nameFromText h)       <$> args'       <*> pure [(primRetType t, mempty)]       <*> pure (Safe, mempty, [])@@ -79,7 +79,7 @@ primExpFromExp f (Apply fname args ts _)   | isBuiltInFunction fname,     [Prim t] <- map (declExtTypeOf . fst) ts =-      FunExp (nameToString fname) <$> mapM (primExpFromSubExpM f . fst) args <*> pure t+      FunExp (nameToText fname) <$> mapM (primExpFromSubExpM f . fst) args <*> pure t primExpFromExp _ _ = fail "Not a PrimExp"  -- | Like 'primExpFromExp', but for a t'SubExp'.

@@ -38,8 +38,8 @@   let InterpreterConfig _ file = cfg   maybe_new_state <- newFutharkiState cfg file   (src, env, ctx) <- case maybe_new_state of-    Left _ -> do-      hPutStrLn stderr $ fromJust file <> ": file not found."+    Left reason -> do+      hPutDocLn stderr reason       exitWith $ ExitFailure 2     Right s -> pure s   mapM_ (runExpr src env ctx) exprs

@@ -32,11 +32,11 @@ main :: String -> [String] -> IO () main = mainWithOptions initialFmtCfg fmtOptions "[FILES]" $ \args cfg ->   case args of-    [] -> Just $ putDoc =<< onInput =<< T.getContents+    [] -> Just $ putDoc =<< onInput "<stdin>" =<< T.getContents     files ->       Just $ forM_ files $ \file -> do         file_s <- T.readFile file-        doc <- onInput file_s+        doc <- onInput file file_s         if cfgCheck cfg           then unless (docText doc == file_s) $ do             T.hPutStrLn stderr $ T.pack file <> ": not formatted correctly."@@ -44,9 +44,9 @@             exitFailure           else withFile file WriteMode $ \h -> hPutDoc h doc   where-    onInput s = do-      case fmtToDoc "<stdin>" s of+    onInput fname s = do+      case fmtToDoc fname s of         Left (SyntaxError loc err) -> do-          T.hPutStr stderr $ locText loc <> ":\n" <> prettyText err+          T.hPutStrLn stderr $ locText loc <> ":\n" <> prettyText err           exitFailure         Right fmt -> pure fmt

@@ -123,9 +123,11 @@       hPutDoc stderr $         prettyWarnings ws -  ictx <--    foldM (\ctx -> badOnLeft I.prettyInterpreterError <=< runInterpreter' . I.interpretImport ctx) I.initialCtx $-      map (fmap fileProg) imports+  let loadImport ctx =+        badOnLeft I.prettyInterpreterError+          <=< runInterpreter' . I.interpretImport ctx++  ictx <- foldM loadImport I.initialCtx $ map (fmap fileProg) imports   let (tenv, ienv) =         let (iname, fm) = last imports          in ( fileScope fm,

@@ -108,7 +108,7 @@     _ -> [C.cexp|$id:(prettyString bop)($exp:x', $exp:y')|] compilePrimExp f (FunExp h args _) = do   args' <- mapM (compilePrimExp f) args-  pure [C.cexp|$id:(funName (nameFromString h))($args:args')|]+  pure [C.cexp|$id:(funName (nameFromText h))($args:args')|]  -- | Compile prim expression to C expression. compileExp :: Exp -> CompilerM op s C.Exp

@@ -138,7 +138,7 @@   where     field name (ct, Prim _) =       [C.cstms|$id:struct.$id:(closureRetvalStructField name)=(($ty:ct*)&$id:name);-               $escstm:("#if ISPC")+               $escstm:("#if defined(ISPC)")                $id:struct.$id:(closureRetvalStructField name)+= programIndex;                $escstm:("#endif")|]     field name (_, MemBlock) =

@@ -467,7 +467,7 @@     _ -> [C.cexp|$id:(prettyString bop)($exp:x', $exp:y')|] compileExp (FunExp h args _) = do   args' <- mapM compileExp args-  pure [C.cexp|$id:(funName (nameFromString h))($args:args')|]+  pure [C.cexp|$id:(funName (nameFromText h))($args:args')|]  -- | Compile a block of code with ISPC specific semantics, falling back -- to generic C when this semantics is not needed.@@ -803,7 +803,7 @@   aos_name <- newVName "aos"   GC.items     [C.citems|-    $escstm:("#if ISPC")+    $escstm:("#if defined(ISPC)")     $tyqual:uniform struct $id:fstruct $id:aos_name[programCount];     $id:aos_name[programIndex] = $id:(fstruct <> "_");     $escstm:("foreach_active (i)")

@@ -65,7 +65,7 @@ builtInFunctions :: M.Map Name (PrimType, [PrimType]) builtInFunctions = M.fromList $ map namify $ M.toList primFuns   where-    namify (k, (paramts, ret, _)) = (nameFromString k, (ret, paramts))+    namify (k, (paramts, ret, _)) = (nameFromText k, (ret, paramts))  -- | If the expression is a t'BasicOp', return it, otherwise 'Nothing'. asBasicOp :: Exp rep -> Maybe BasicOp

@@ -537,30 +537,36 @@         let redlam_deps = dataDependencies $ lambdaBody redlam,         let redlam_res = bodyResult $ lambdaBody redlam,         let redlam_params = lambdaParams redlam,+        let (redlam_xparams, redlam_yparams) =+              splitAt (length nes) redlam_params,         let used_after =               map snd . filter ((`UT.used` used) . patElemName . fst) $-                zip red_pes redlam_params,+                zip (red_pes <> red_pes) redlam_params,         let necessary =               findNecessaryForReturned                 (`elem` used_after)                 (zip redlam_params $ map resSubExp $ redlam_res <> redlam_res)                 redlam_deps,-        let alive_mask = map ((`nameIn` necessary) . paramName) redlam_params,-        not $ and (take (length nes) alive_mask) = Simplify $ do+        let alive_mask =+              zipWith+                (||)+                (map ((`nameIn` necessary) . paramName) redlam_xparams)+                (map ((`nameIn` necessary) . paramName) redlam_yparams),+        not $ and alive_mask = Simplify $ do           let fixDeadToNeutral lives ne = if lives then Nothing else Just ne               dead_fix = zipWith fixDeadToNeutral alive_mask nes-              (used_red_pes, _, used_nes) =-                unzip3 . filter (\(_, x, _) -> paramName x `nameIn` necessary) $-                  zip3 red_pes redlam_params nes+              (used_red_pes, used_nes) =+                unzip . map snd . filter fst $ zip alive_mask $ zip red_pes nes -          let maplam' = removeLambdaResults (take (length nes) alive_mask) maplam-          redlam' <- removeLambdaResults (take (length nes) alive_mask) <$> fixLambdaParams redlam (dead_fix ++ dead_fix)+          when (used_nes == nes) cannotSimplify -          auxing aux $-            letBind (Pat $ used_red_pes ++ map_pes) $-              Op $-                Screma w arrs $-                  mkOp redlam' used_nes maplam'+          let maplam' = removeLambdaResults alive_mask maplam+          redlam' <-+            removeLambdaResults alive_mask+              <$> fixLambdaParams redlam (dead_fix ++ dead_fix)++          auxing aux . letBind (Pat $ used_red_pes ++ map_pes) . Op $+            Screma w arrs (mkOp redlam' used_nes maplam')     removeDeadReduction' _ _ _ _ = Skip removeDeadReduction _ _ _ _ = Skip 

@@ -1424,7 +1424,7 @@ checkPrimExp (FunExp h args t) = do   (h_ts, h_ret, _) <-     maybe-      (bad $ TypeError $ "Unknown function: " <> T.pack h)+      (bad $ TypeError $ "Unknown function: " <> h)       pure       $ M.lookup h primFuns   when (length h_ts /= length args) . bad . TypeError $

@@ -97,7 +97,7 @@ constantFoldPrimFun :: (BuilderOps rep) => TopDownRuleGeneric rep constantFoldPrimFun _ (Let pat (StmAux cs attrs _) (Apply fname args _ _))   | Just args' <- mapM (isConst . fst) args,-    Just (_, _, fun) <- M.lookup (nameToString fname) primFuns,+    Just (_, _, fun) <- M.lookup (nameToText fname) primFuns,     Just result <- fun args' =       Simplify $         certifying cs $

@@ -118,8 +118,6 @@   | isCt1 e1 = resIsSubExp e2   | isCt1 e2 = resIsSubExp e1 simplifyBinOp _ _ (BinOp FMul {} e1 e2)-  | isCt0 e1 = resIsSubExp e1-  | isCt0 e2 = resIsSubExp e2   | isCt1 e1 = resIsSubExp e2   | isCt1 e2 = resIsSubExp e1 simplifyBinOp look _ (BinOp (SMod t _) e1 e2)

@@ -77,12 +77,11 @@ multiplicity stm =   case stmExp stm of     Match cond cases defbody _ ->-      foldl comb mempty $-        free 1 cond-          : free 1 defbody-          : map (free 1 . caseBody) cases+      foldl' comb mempty $+        free 1 cond : free 1 defbody : map (free 1 . caseBody) cases     Op {} -> free 2 stm     Loop {} -> free 2 stm+    WithAcc {} -> free 2 stm     _ -> free 1 stm   where     free k x = M.fromList $ map (,k) $ namesToList $ freeIn x

@@ -806,12 +806,7 @@   traverse renameStm <=< runBuilder_ $ do     addStms k_stms -    let pat =-          Pat . rearrangeShape perm $-            patElems $-              loopNestingPat $-                fst nest-+    let pat = Pat . rearrangeShape perm $ patElems $ loopNestingPat $ fst nest     letBind pat $ Op $ segOp k   where     findInput kernel_inps a =

@@ -513,8 +513,8 @@         expandTarget       )   where-    isIdentity (patElem, SubExpRes cs (Var v))-      | v `notNameIn` bound = Left (patElem, (cs, v))+    isIdentity (patElem, SubExpRes _ (Var v))+      | v `notNameIn` bound = Left (patElem, (mempty, v))     isIdentity x = Right x  removeIdentityMappingFromNesting ::

@@ -29,6 +29,7 @@     baseTag,     baseName,     baseString,+    baseText,     quote,      -- * Number re-export@@ -195,6 +196,10 @@ -- | Return the base 'Name' converted to a string. baseString :: VName -> String baseString = nameToString . baseName++-- | Return the base 'Name' converted to a text.+baseText :: VName -> T.Text+baseText = nameToText . baseName  instance Eq VName where   VName _ x == VName _ y = x == y

@@ -34,12 +34,12 @@ import Control.Monad.Trans.Maybe import Data.Array import Data.Bifunctor+import Data.Bitraversable import Data.List   ( find,     foldl',     genericLength,     genericTake,-    isPrefixOf,     transpose,   ) import Data.List qualified as L@@ -157,11 +157,32 @@ valueStructType :: ValueType -> StructType valueStructType = first $ flip sizeFromInteger mempty . toInteger +-- | An expression along with an environment in which to evaluate that+-- expression. Used to represent non-interpreted size expressions,+-- which may still be in reference to some environment.+data SizeClosure = SizeClosure Env Size+  deriving (Show)++instance Pretty SizeClosure where+  pretty (SizeClosure _ e) = pretty e++instance Pretty (F.Shape SizeClosure) where+  pretty = mconcat . map (braces . pretty) . shapeDims++-- | A type where the sizes are unevaluated expressions.+type EvalType = TypeBase SizeClosure NoUniqueness++structToEval :: Env -> StructType -> EvalType+structToEval env = first (SizeClosure env)++evalToStruct :: EvalType -> StructType+evalToStruct = first (\(SizeClosure _ e) -> e)+ resolveTypeParams ::   [VName] ->   StructType ->-  StructType ->-  ([(VName, ([VName], StructType))], [(VName, Exp)])+  EvalType ->+  ([(VName, ([VName], EvalType))], [(VName, SizeClosure)]) resolveTypeParams names orig_t1 orig_t2 =   execState (match mempty orig_t1 orig_t2) mempty   where@@ -194,27 +215,36 @@           match bound t1' t2'     match _ _ _ = pure mempty -    matchDims bound e1 e2+    matchDims bound e1 (SizeClosure env e2)       | e1 == anySize || e2 == anySize = pure mempty-      | otherwise = matchExps bound e1 e2+      | otherwise = matchExps bound env e1 e2 -    matchExps bound (Var (QualName _ d1) _ _) e+    matchExps bound env (Var (QualName _ d1) _ _) e       | d1 `elem` names,-        not $ any (`elem` bound) $ fvVars $ freeInExp e =-          addDim d1 e-    matchExps bound e1 e2+        not $ any problematic $ fvVars $ freeInExp e =+          addDim d1 (SizeClosure env e)+      where+        problematic v = v `elem` bound || v `elem` names+    matchExps bound env e1 e2       | Just es <- similarExps e1 e2 =-          mapM_ (uncurry $ matchExps bound) es-    matchExps _ _ _ = pure mempty+          mapM_ (uncurry $ matchExps bound env) es+    matchExps _ _ _ _ = pure mempty +evalWithExts :: Env -> Exp -> EvalM Value+evalWithExts env e = do+  size_env <- extEnv+  eval (size_env <> env) e+ evalResolved ::-  Eval ->-  ([(VName, ([VName], StructType))], [(VName, Exp)]) ->+  ([(VName, ([VName], EvalType))], [(VName, SizeClosure)]) ->   EvalM Env-evalResolved eval' (ts, ds) = do-  ts' <- mapM (traverse $ \(bound, t) -> evalType eval' (S.fromList bound) t) ts-  ds' <- mapM (traverse $ fmap asInt64 . eval') ds+evalResolved (ts, ds) = do+  ts' <- mapM (traverse $ \(bound, t) -> first onDim <$> evalType (S.fromList bound) t) ts+  ds' <- mapM (traverse $ \(SizeClosure env e) -> asInt64 <$> evalWithExts env e) ds   pure $ typeEnv (M.fromList ts') <> i64Env (M.fromList ds')+  where+    onDim (Left x) = sizeFromInteger (toInteger x) mempty+    onDim (Right (SizeClosure _ e)) = e -- FIXME  resolveExistentials :: [VName] -> StructType -> ValueShape -> M.Map VName Int64 resolveExistentials names = match@@ -237,24 +267,22 @@       | d1 `elem` names = M.singleton d1 d2     matchDims _ _ = mempty -checkShape :: Shape (Maybe Int64) -> ValueShape -> Maybe ValueShape-checkShape (ShapeDim Nothing shape1) (ShapeDim d2 shape2) =-  ShapeDim d2 <$> checkShape shape1 shape2-checkShape (ShapeDim (Just d1) shape1) (ShapeDim d2 shape2) = do+checkShape :: Shape Int64 -> ValueShape -> Maybe ValueShape+checkShape (ShapeDim d1 shape1) (ShapeDim d2 shape2) = do   guard $ d1 == d2   ShapeDim d2 <$> checkShape shape1 shape2 checkShape (ShapeDim d1 shape1) ShapeLeaf =   -- This case is for handling polymorphism, when a function doesn't   -- know that the array it produced actually has more dimensions.-  ShapeDim (fromMaybe 0 d1) <$> checkShape shape1 ShapeLeaf+  ShapeDim d1 <$> checkShape shape1 ShapeLeaf checkShape (ShapeRecord shapes1) (ShapeRecord shapes2) =   ShapeRecord <$> sequence (M.intersectionWith checkShape shapes1 shapes2) checkShape (ShapeRecord shapes1) ShapeLeaf =-  Just $ fromMaybe 0 <$> ShapeRecord shapes1+  Just $ ShapeRecord shapes1 checkShape (ShapeSum shapes1) (ShapeSum shapes2) =   ShapeSum <$> sequence (M.intersectionWith (zipWithM checkShape) shapes1 shapes2) checkShape (ShapeSum shapes1) ShapeLeaf =-  Just $ fromMaybe 0 <$> ShapeSum shapes1+  Just $ ShapeSum shapes1 checkShape _ shape2 =   Just shape2 @@ -303,12 +331,9 @@ lookupVar :: QualName VName -> Env -> Maybe TermBinding lookupVar = lookupInEnv envTerm -lookupType :: QualName VName -> Env -> Maybe T.TypeBinding+lookupType :: QualName VName -> Env -> Maybe (Env, T.TypeBinding) lookupType = lookupInEnv envType --- | An expression evaluator that embeds an environment.-type Eval = Exp -> EvalM Value- -- | A TermValue with a 'Nothing' type annotation is an intrinsic or -- an existential. data TermBinding@@ -316,7 +341,7 @@   | -- | A polymorphic value that must be instantiated.  The     --  'StructType' provided is un-evaluated, but parts of it can be     --  evaluated using the provided 'Eval' function.-    TermPoly (Maybe T.BoundV) (StructType -> Eval -> EvalM Value)+    TermPoly (Maybe T.BoundV) (EvalType -> EvalM Value)   | TermModule Module  instance Show TermBinding where@@ -335,7 +360,7 @@ -- | The actual type- and value environment. data Env = Env   { envTerm :: M.Map VName TermBinding,-    envType :: M.Map VName T.TypeBinding+    envType :: M.Map VName (Env, T.TypeBinding)   }   deriving (Show) @@ -375,7 +400,7 @@       envType = M.map tbind m     }   where-    tbind = T.TypeAbbr Unlifted [] . RetType []+    tbind = (mempty,) . T.TypeAbbr Unlifted [] . RetType []  i64Env :: M.Map VName Int64 -> Env i64Env = valEnv . M.map f@@ -610,7 +635,7 @@  -- | Expand type based on information that was not available at -- type-checking time (the structure of abstract types).-expandType :: (Pretty u) => Env -> TypeBase Size u -> TypeBase Size u+expandType :: (Pretty u) => Env -> TypeBase Size u -> TypeBase SizeClosure u expandType _ (Scalar (Prim pt)) = Scalar $ Prim pt expandType env (Scalar (Record fs)) = Scalar $ Record $ fmap (expandType env) fs expandType env (Scalar (Arrow u p d t1 (RetType dims t2))) =@@ -618,23 +643,22 @@ expandType env t@(Array u shape _) =   let et = stripArray (shapeRank shape) t       et' = expandType env et-   in second (const u) (arrayOf shape $ toStruct et')+      shape' = fmap (SizeClosure env) shape+   in second (const u) (arrayOf shape' $ toStruct et') expandType env (Scalar (TypeVar u tn args)) =   case lookupType tn env of-    Just (T.TypeAbbr _ ps (RetType ext t')) ->+    Just (tn_env, T.TypeAbbr _ ps (RetType ext t')) ->       let (substs, types) = mconcat $ zipWith matchPtoA ps args-          onDim (Var v _ _)+          onDim (SizeClosure _ (Var v _ _))             | Just e <- M.lookup (qualLeaf v) substs =-                e+                SizeClosure env e           -- The next case can occur when a type with existential size           -- has been hidden by a module ascription,           -- e.g. tests/modules/sizeparams4.fut.-          onDim e-            | any (`elem` ext) $ fvVars $ freeInExp e = anySize+          onDim (SizeClosure _ e)+            | any (`elem` ext) $ fvVars $ freeInExp e = SizeClosure mempty anySize           onDim d = d-       in if null ps-            then bimap onDim (const u) t'-            else expandType (Env mempty types <> env) $ bimap onDim (const u) t'+       in bimap onDim (const u) $ expandType (Env mempty types <> tn_env) t'     Nothing ->       -- This case only happens for built-in abstract types,       -- e.g. accumulators.@@ -643,37 +667,38 @@     matchPtoA (TypeParamDim p _) (TypeArgDim e) =       (M.singleton p e, mempty)     matchPtoA (TypeParamType l p _) (TypeArgType t') =-      let t'' = expandType env t'-       in (mempty, M.singleton p $ T.TypeAbbr l [] $ RetType [] t'')+      let t'' = evalToStruct $ expandType env t' -- FIXME, we are throwing away the closure here.+       in (mempty, M.singleton p (mempty, T.TypeAbbr l [] $ RetType [] t''))     matchPtoA _ _ = mempty-    expandArg (TypeArgDim s) = TypeArgDim s+    expandArg (TypeArgDim s) = TypeArgDim $ SizeClosure env s     expandArg (TypeArgType t) = TypeArgType $ expandType env t expandType env (Scalar (Sum cs)) = Scalar $ Sum $ (fmap . fmap) (expandType env) cs -evalWithExts :: Env -> EvalM Eval-evalWithExts env = do-  size_env <- extEnv-  pure $ eval $ size_env <> env- -- | Evaluate all possible sizes, except those that contain free -- variables in the set of names.-evalType :: Eval -> S.Set VName -> StructType -> EvalM StructType-evalType eval' outer_bound t = do-  let evalDim bound _ e-        | canBeEvaluated bound e = do-            x <- asInteger <$> eval' e-            pure $ sizeFromInteger x mempty-      evalDim _ _ d = pure d+evalType :: S.Set VName -> EvalType -> EvalM (TypeBase (Either Int64 SizeClosure) NoUniqueness)+evalType outer_bound t = do+  let evalDim bound _ (SizeClosure env e)+        | canBeEvaluated bound e =+            Left . asInt64 <$> evalWithExts env e+      evalDim _ _ e = pure $ Right e   traverseDims evalDim t   where     canBeEvaluated bound e =       let free = fvVars $ freeInExp e        in not $ any (`S.member` bound) free || any (`S.member` outer_bound) free +-- | Evaluate all sizes, and it better work. This implies it must be a+-- size-dependent function type, or one that has existentials.+evalTypeFully :: EvalType -> EvalM ValueType+evalTypeFully t = do+  let evalDim (SizeClosure env e) = asInt64 <$> evalWithExts env e+  bitraverse evalDim pure t+ evalTermVar :: Env -> QualName VName -> StructType -> EvalM Value evalTermVar env qv t =   case lookupVar qv env of-    Just (TermPoly _ v) -> v (expandType env t) =<< evalWithExts env+    Just (TermPoly _ v) -> v $ expandType env t     Just (TermValue _ v) -> pure v     x -> do       ss <- map (locText . srclocOf) <$> stacktrace@@ -685,15 +710,7 @@           <> show x  typeValueShape :: Env -> StructType -> EvalM ValueShape-typeValueShape env t = do-  eval' <- evalWithExts env-  t' <- evalType eval' mempty $ expandType env t-  case traverse dim $ typeShape t' of-    Nothing -> error $ "typeValueShape: failed to fully evaluate type " <> prettyString t'-    Just shape -> pure shape-  where-    dim (IntLit x _ _) = Just $ fromIntegral x-    dim _ = Nothing+typeValueShape env t = typeShape <$> evalTypeFully (expandType env t)  -- Sometimes type instantiation is not quite enough - then we connect -- up the missing sizes here.  In particular used for eta-expanded@@ -703,7 +720,7 @@ linkMissingSizes missing_sizes p v env =   env <> i64Env (resolveExistentials missing_sizes p_t (valueShape v))   where-    p_t = expandType env $ patternStructType p+    p_t = evalToStruct $ expandType env $ patternStructType p  evalFunction :: Env -> [VName] -> [Pat ParamType] -> Exp -> ResType -> EvalM Value -- We treat zero-parameter lambdas as simply an expression to@@ -747,20 +764,19 @@       fmap (TermValue (Just $ T.BoundV [] ftype))         . returned env (retType ret) retext         =<< evalFunction env [] ps fbody (retType ret)-    else pure . TermPoly (Just $ T.BoundV [] ftype) $ \ftype' ->+    else pure . TermPoly (Just $ T.BoundV [] ftype) $ \ftype' -> do       let resolved = resolveTypeParams (map typeParamName tparams) ftype ftype'-       in \eval' -> do-            tparam_env <- evalResolved eval' resolved-            let env' = tparam_env <> env-                -- In some cases (abstract lifted types) there may be-                -- missing sizes that were not fixed by the type-                -- instantiation.  These will have to be set by looking-                -- at the actual function arguments.-                missing_sizes =-                  filter (`M.notMember` envTerm env') $-                    map typeParamName (filter isSizeParam tparams)-            returned env (retType ret) retext-              =<< evalFunction env' missing_sizes ps fbody (retType ret)+      tparam_env <- evalResolved resolved+      let env' = tparam_env <> env+          -- In some cases (abstract lifted types) there may be+          -- missing sizes that were not fixed by the type+          -- instantiation.  These will have to be set by looking+          -- at the actual function arguments.+          missing_sizes =+            filter (`M.notMember` envTerm env') $+              map typeParamName (filter isSizeParam tparams)+      returned env (retType ret) retext+        =<< evalFunction env' missing_sizes ps fbody (retType ret)  evalArg :: Env -> Exp -> Maybe VName -> EvalM Value evalArg env e ext = do@@ -770,12 +786,12 @@     _ -> pure ()   pure v -returned :: Env -> TypeBase Size als -> [VName] -> Value -> EvalM Value+returned :: Env -> TypeBase Size u -> [VName] -> Value -> EvalM Value returned _ _ [] v = pure v returned env ret retext v = do   mapM_ (uncurry putExtSize . second (ValuePrim . SignedValue . Int64Value))     . M.toList-    $ resolveExistentials retext (expandType env $ toStruct ret)+    $ resolveExistentials retext (evalToStruct $ expandType env $ toStruct ret)     $ valueShape v   pure v @@ -831,7 +847,7 @@ evalAppExp env (LetPat sizes p e body _) = do   v <- eval env e   env' <- matchPat env p v-  let p_t = expandType env $ patternStructType p+  let p_t = evalToStruct $ expandType env $ patternStructType p       v_s = valueShape v       env'' = env' <> i64Env (resolveExistentials (map sizeName sizes) p_t v_s)   eval env'' body@@ -896,10 +912,7 @@   where     withLoopParams v =       let sparams' =-            resolveExistentials-              sparams-              (patternStructType pat)-              (valueShape v)+            resolveExistentials sparams (patternStructType pat) (valueShape v)        in matchPat (i64Env sparams' <> env) pat v      inc = (`P.doAdd` Int64Value 1)@@ -980,16 +993,14 @@   -- Probably will not ever be used.   pure $ toArray' ShapeLeaf $ map ValuePrim vs eval env (AppExp e (Info (AppRes t retext))) = do-  let t' = expandType env $ toStruct t   v <- evalAppExp env e-  returned env t' retext v+  returned env (toStruct t) retext v eval env (Var qv (Info t) _) = evalTermVar env qv (toStruct t) eval env (Ascript e _ _) = eval env e eval env (Coerce e te (Info t) loc) = do   v <- eval env e-  eval' <- evalWithExts env-  t' <- evalType eval' mempty $ expandType env $ toStruct t-  case checkShape (structTypeShape t') (valueShape v) of+  t' <- evalTypeFully $ expandType env $ toStruct t+  case checkShape (typeShape t') (valueShape v) of     Just _ -> pure v     Nothing ->       bad loc env . docText $@@ -1115,8 +1126,10 @@       (k, v) <- M.toList m       k' <- replace k       pure (k', f v)-    onModule (Module (Env terms types)) =-      Module $ Env (replaceM onTerm terms) (replaceM onType types)+    onEnv (Env terms types) =+      Env (replaceM onTerm terms) (replaceM (bimap onEnv onType) types)+    onModule (Module env) =+      Module $ onEnv env     onModule (ModuleFun f) =       ModuleFun $ \m -> onModule <$> f (substituteInModule substs m)     onTerm (TermValue t v) = TermValue t v@@ -1154,15 +1167,11 @@ evalModExp env (ModDecs ds _) = do   Env terms types <- foldM evalDec env ds   -- Remove everything that was present in the original Env.-  pure-    ( Env-        (terms `M.difference` envTerm env)-        (types `M.difference` envType env),-      Module $+  let env' =         Env           (terms `M.difference` envTerm env)           (types `M.difference` envType env)-    )+  pure (env', Module env') evalModExp env (ModVar qv _) =   (mempty,) <$> evalModuleVar env qv evalModExp env (ModAscript me _ (Info substs) _) =@@ -1197,8 +1206,7 @@ evalDec env (ValDec (ValBind _ v _ (Info ret) tparams ps fbody _ _ _)) = localExts $ do   binding <- evalFunctionBinding env tparams ps ret fbody   sizes <- extEnv-  pure $-    env {envTerm = M.insert v binding $ envTerm env} <> sizes+  pure $ env {envTerm = M.insert v binding $ envTerm env} <> sizes evalDec env (OpenDec me _) = do   (me_env, me') <- evalModExp env me   case me' of@@ -1209,7 +1217,7 @@ evalDec env (LocalDec d _) = evalDec env d evalDec env ModTypeDec {} = pure env evalDec env (TypeDec (TypeBind v l ps _ (Info (RetType dims t)) _ _)) = do-  let abbr = T.TypeAbbr l ps . RetType dims $ expandType env t+  let abbr = (env, T.TypeAbbr l ps . RetType dims $ evalToStruct $ expandType env t)   pure env {envType = M.insert v abbr $ envType env} evalDec env (ModDec (ModBind v ps ret body _ loc)) = do   (mod_env, mod) <- evalModExp env $ wrapInLambda ps@@ -1255,7 +1263,7 @@   Ctx     ( Env         ( M.insert-            (VName (nameFromString "intrinsics") 0)+            (VName (nameFromText "intrinsics") 0)             (TermModule (Module $ Env terms types))             terms         )@@ -1263,8 +1271,8 @@     )     mempty   where-    terms = M.mapMaybeWithKey (const . def . baseString) intrinsics-    types = M.mapMaybeWithKey (const . tdef . baseString) intrinsics+    terms = M.mapMaybeWithKey (const . def . baseText) intrinsics+    types = M.mapMaybeWithKey (const . tdef . baseName) intrinsics      sintOp f =       [ (getS, putS, P.doBinOp (f Int8), adBinOp $ AD.OpBin (f Int8)),@@ -1472,6 +1480,7 @@               <+> dquotes (prettyValue v)               <> "." +    def :: T.Text -> Maybe TermBinding     def "!" =       Just $         unopDef@@ -1578,13 +1587,13 @@               ++ floatCmp P.FCmpLe               ++ boolCmp P.CmpLle     def s-      | Just bop <- find ((s ==) . prettyString) P.allBinOps =+      | Just bop <- find ((s ==) . prettyText) P.allBinOps =           Just $ tbopDef (AD.OpBin bop) $ P.doBinOp bop-      | Just unop <- find ((s ==) . prettyString) P.allCmpOps =+      | Just unop <- find ((s ==) . prettyText) P.allCmpOps =           Just $ tbopDef (AD.OpCmp unop) $ \x y -> P.BoolValue <$> P.doCmpOp unop x y-      | Just cop <- find ((s ==) . prettyString) P.allConvOps =+      | Just cop <- find ((s ==) . prettyText) P.allConvOps =           Just $ unopDef [(getV, Just . putV, P.doConvOp cop, adUnOp $ AD.OpConv cop)]-      | Just unop <- find ((s ==) . prettyString) P.allUnOps =+      | Just unop <- find ((s ==) . prettyText) P.allUnOps =           Just $ unopDef [(getV, Just . putV, P.doUnOp unop, adUnOp $ AD.OpUn unop)]       | Just (pts, _, f) <- M.lookup s P.primFuns =           case length pts of@@ -1600,42 +1609,33 @@                         pure $ ValuePrim res                   _ ->                     error $ "Cannot apply " <> prettyString s ++ " to " <> show x-      | "sign_" `isPrefixOf` s =+      | "sign_" `T.isPrefixOf` s =           Just $             fun1 $ \x ->               case x of                 (ValuePrim (UnsignedValue x')) ->                   pure $ ValuePrim $ SignedValue x'                 _ -> error $ "Cannot sign: " <> show x-      | "unsign_" `isPrefixOf` s =+      | "unsign_" `T.isPrefixOf` s =           Just $             fun1 $ \x ->               case x of                 (ValuePrim (SignedValue x')) ->                   pure $ ValuePrim $ UnsignedValue x'                 _ -> error $ "Cannot unsign: " <> show x-    def s-      | "map_stream" `isPrefixOf` s =-          Just $ fun2 stream-    def s | "reduce_stream" `isPrefixOf` s =-      Just $ fun3 $ \_ f arg -> stream f arg     def "map" = Just $-      TermPoly Nothing $ \t eval' -> do-        t' <- evalType eval' mempty t+      TermPoly Nothing $ \t -> do+        t' <- evalTypeFully t         pure $ ValueFun $ \f -> pure . ValueFun $ \xs ->           case unfoldFunType t' of             ([_, _], ret_t)-              | Just rowshape <- typeRowShape ret_t ->+              | rowshape <- typeShape $ stripArray 1 ret_t ->                   toArray' rowshape <$> mapM (apply noLoc mempty f) (snd $ fromArray xs)-              | otherwise ->-                  error $ "Bad return type: " <> prettyString ret_t             _ ->               error $                 "Invalid arguments to map intrinsic:\n"                   ++ unlines [prettyString t, show f, show xs]-      where-        typeRowShape = sequenceA . structTypeShape . stripArray 1-    def s | "reduce" `isPrefixOf` s = Just $+    def s | "reduce" `T.isPrefixOf` s = Just $       fun3 $ \f ne xs ->         foldM (apply2 noLoc mempty f) ne $ snd $ fromArray xs     def "scan" = Just $@@ -1994,17 +1994,14 @@         let o' = fst $ valueAccum (\a b -> (b : a, b)) [] o          -- For each output..-        let m =-              fromMaybe (error "vjp: differentiation failed") $-                zipWithM-                  ( \on sn -> case on of-                      -- If it is a VJP variable of the correct depth, run deriveTape on it- and its corresponding seed-                      (ValueAD d (AD.VJP (AD.VJPValue t))) | d == depth -> (putAD $ AD.tapePrimal t,) <$> AD.deriveTape t sn-                      -- Otherwise, its partial derivatives are all 0-                      _ -> Just (on, M.empty)-                  )-                  o'-                  s'+        let m = flip map (zip o' s') $ \(on, sn) -> case on of+              -- If it is a VJP variable of the correct depth, run+              -- deriveTapqe on it- and its corresponding seed+              (ValueAD d (AD.VJP (AD.VJPValue t)))+                | d == depth ->+                    (putAD $ AD.tapePrimal t, AD.deriveTape t sn)+              -- Otherwise, its partial derivatives are all 0+              _ -> (on, M.empty)          -- Add together every derivative         let drvs = M.map (Just . putAD) $ M.unionsWith add $ map snd m@@ -2095,17 +2092,13 @@         expectJust _ (Just v) = v         expectJust s Nothing = error s     def "acc" = Nothing-    def s | nameFromString s `M.member` namesToPrimTypes = Nothing-    def s = error $ "Missing intrinsic: " ++ s+    def s | nameFromText s `M.member` namesToPrimTypes = Nothing+    def s = error $ "Missing intrinsic: " ++ T.unpack s +    tdef :: Name -> Maybe (Env, T.TypeBinding)     tdef s = do-      t <- nameFromString s `M.lookup` namesToPrimTypes-      pure $ T.TypeAbbr Unlifted [] $ RetType [] $ Scalar $ Prim t--    stream f arg@(ValueArray _ xs) =-      let n = ValuePrim $ SignedValue $ Int64Value $ arrayLength xs-       in apply2 noLoc mempty f n arg-    stream _ arg = error $ "Cannot stream: " <> show arg+      t <- s `M.lookup` namesToPrimTypes+      pure (mempty, T.TypeAbbr Unlifted [] $ RetType [] $ Scalar $ Prim t)  intrinsicVal :: Name -> Value intrinsicVal name =@@ -2190,11 +2183,15 @@ -- horribly if these are ill-typed. interpretFunction :: Ctx -> VName -> [V.Value] -> Either T.Text (F ExtOp Value) interpretFunction ctx fname vs = do-  ft <- case lookupVar (qualName fname) $ ctxEnv ctx of-    Just (TermValue (Just (T.BoundV _ t)) _) ->-      updateType (map valueType vs) t-    Just (TermPoly (Just (T.BoundV _ t)) _) ->-      updateType (map valueType vs) t+  let env = ctxEnv ctx++  (ft, mkf) <- case lookupVar (qualName fname) env of+    Just (TermValue (Just (T.BoundV _ t)) v) -> do+      ft <- updateType (map valueType vs) t+      pure (ft, pure v)+    Just (TermPoly (Just (T.BoundV _ t)) v) -> do+      ft <- updateType (map valueType vs) t+      pure (ft, v (structToEval env ft))     _ ->       Left $ "Unknown function `" <> nameToText (toName fname) <> "`." @@ -2202,18 +2199,14 @@    checkEntryArgs fname vs ft -  Right $-    runEvalM (ctxImports ctx) $ do-      -- XXX: We are providing a dummy type here.  This is OK as long-      -- as the function we invoke is monomorphic, which is what we-      -- require of entry points.  This is to avoid reimplementing-      -- type inference machinery here.-      f <- evalTermVar (ctxEnv ctx) (qualName fname) (Scalar (Prim Bool))-      foldM (apply noLoc mempty) f vs'+  Right $ runEvalM (ctxImports ctx) $ do+    f <- mkf+    foldM (apply noLoc mempty) f vs'   where     updateType (vt : vts) (Scalar (Arrow als pn d pt (RetType dims rt))) = do       checkInput vt pt-      Scalar . Arrow als pn d (valueStructType vt) . RetType dims . toRes Nonunique <$> updateType vts (toStruct rt)+      Scalar . Arrow als pn d (valueStructType vt) . RetType dims . toRes Nonunique+        <$> updateType vts (toStruct rt)     updateType _ t =       Right t 

@@ -16,12 +16,13 @@  import Control.Monad (foldM, zipWithM) import Data.Either (isRight)-import Data.List (find)+import Data.List (find, foldl') import Data.Map qualified as M import Data.Maybe (fromMaybe)+import Data.Text qualified as T import Futhark.AD.Derivatives (pdBinOp, pdBuiltin, pdUnOp) import Futhark.Analysis.PrimExp (PrimExp (..))-import Language.Futhark.Core (VName (..), nameFromString)+import Language.Futhark.Core (VName (..), nameFromString, nameFromText) import Language.Futhark.Primitive  -- Mathematical operations subject to AD.@@ -29,7 +30,7 @@   = OpBin BinOp   | OpCmp CmpOp   | OpUn UnOp-  | OpFn String+  | OpFn T.Text   | OpConv ConvOp   deriving (Show) @@ -133,7 +134,7 @@   let (a, b) = pdBinOp op x y   [a, b] lookupPDs (OpUn op) [x] = Just [pdUnOp op x]-lookupPDs (OpFn fn) p = pdBuiltin (nameFromString fn) p+lookupPDs (OpFn fn) p = pdBuiltin (nameFromText fn) p lookupPDs _ _ = Nothing  -- Shared AD logic--@@ -216,17 +217,20 @@           -- least one variable of depth > 0           error "find isRight" -calculatePDs :: Op -> [ADValue] -> Maybe [ADValue]-calculatePDs op p = do+calculatePDs :: Op -> [ADValue] -> [ADValue]+calculatePDs op p =   -- Create a unique VName for each operand   let n = map (\i -> VName (nameFromString $ "x" ++ show i) i) [1 .. length p]-  -- Put the operands in the environment-  let m = M.fromList $ zip n p+      -- Put the operands in the environment+      m = M.fromList $ zip n p -  -- Look up, and calculate the partial derivative-  -- of the operation with respect to each operand-  pde <- lookupPDs op $ map (`LeafExp` opReturnType op) n-  mapM (evalPrimExp m) pde+      -- Look up, and calculate the partial derivative+      -- of the operation with respect to each operand+      pde =+        fromMaybe (error "lookupPDs failed") $+          lookupPDs op $+            map (`LeafExp` opReturnType op) n+   in map (fromMaybe (error "evalPrimExp failed") . evalPrimExp m) pde  -- VJP / Reverse mode automatic differentiation-- -- In reverse mode AD, the entire computation@@ -267,29 +271,31 @@ -- | This calculates every partial derivative of a 'Tape'. The result -- is a map of the partial derivatives, each key corresponding to the -- ID of a free variable (see TapeID).-deriveTape :: Tape -> ADValue -> Maybe (M.Map Int ADValue)-deriveTape (TapeID i _) s = Just $ M.fromList [(i, s)]-deriveTape (TapeConst _) _ = Just M.empty-deriveTape (TapeOp op p _) s = do+deriveTape :: Tape -> ADValue -> M.Map Int ADValue+deriveTape (TapeID i _) s = M.fromList [(i, s)]+deriveTape (TapeConst _) _ = M.empty+deriveTape (TapeOp op p _) s =   -- Calculate the new sensitivities-  s'' <- case op of-    OpConv op' -> do-      -- In case of type conversion, simply convert the sensitivity-      s' <- doOp (OpConv $ flipConvOp op') [s]-      Just [s']-    _ -> do-      pds <- calculatePDs op $ map tapePrimal p-      mapM (mul s) pds+  let s'' = case op of+        OpConv op' ->+          -- In case of type conversion, simply convert the sensitivity+          [ fromMaybe (error "deriveTape: doOp failed") $+              doOp (OpConv $ flipConvOp op') [s]+          ]+        _ ->+          map (mul s) $ calculatePDs op $ map tapePrimal p -  -- Propagate the new sensitivities-  pd <- zipWithM deriveTape p s''-  -- Add up the results-  Just $ foldl (M.unionWith add) M.empty pd+      -- Propagate the new sensitivities+      pd = zipWith deriveTape p s''+   in -- Add up the results+      foldl' (M.unionWith add) M.empty pd   where     add x y =-      fromMaybe (error "deriveTape: addition failed") $+      fromMaybe (error "deriveTape: add failed") $         doOp (OpBin $ addFor $ opReturnType op) [x, y]-    mul x y = doOp (OpBin $ mulFor $ opReturnType op) [x, y]+    mul x y =+      fromMaybe (error "deriveTape: mul failed") $+        doOp (OpBin $ mulFor $ opReturnType op) [x, y]  -- JVP / Forward mode automatic differentiation-- @@ -311,7 +317,7 @@     _ -> do       -- Calculate the new derivative using the chain       -- rule-      pds <- calculatePDs op $ map primal' p+      let pds = calculatePDs op $ map primal' p       vs <- zipWithM mul pds $ map derivative p       foldM add (Constant $ blankPrimValue $ opReturnType op) vs   where

@@ -92,6 +92,8 @@      -- * Primitive functions     primFuns,+    condFun,+    isCondFun,      -- * Utility     zeroIsh,@@ -128,6 +130,7 @@   ) import Data.Fixed (mod') -- Weird location. import Data.Int (Int16, Int32, Int64, Int8)+import Data.List qualified as L import Data.Map qualified as M import Data.Text qualified as T import Data.Word (Word16, Word32, Word64, Word8)@@ -1181,17 +1184,28 @@ wordToDouble :: Word64 -> Double wordToDouble = G.runGet G.getDoublele . P.runPut . P.putWord64le +-- | @condFun t@ is the name of the ternary conditional function that+-- accepts operands of type @[Bool, t, t]@, and returns either the+-- first or second @t@ based on the truth value of the @Bool@.+condFun :: PrimType -> T.Text+condFun t = "cond_" <> prettyText t++-- | Is this the name of a condition function as per 'condFun', and+-- for which type?+isCondFun :: T.Text -> Maybe PrimType+isCondFun v = L.find (\t -> condFun t == v) allPrimTypes+ -- | A mapping from names of primitive functions to their parameter -- types, their result type, and a function for evaluating them. primFuns ::   M.Map-    String+    T.Text     ( [PrimType],       PrimType,       [PrimValue] -> Maybe PrimValue     ) primFuns =-  M.fromList+  M.fromList $     [ f16 "sqrt16" sqrt,       f32 "sqrt32" sqrt,       f64 "sqrt64" sqrt,@@ -1529,6 +1543,17 @@       f32_3 "fma32" (\a b c -> a * b + c),       f64_3 "fma64" (\a b c -> a * b + c)     ]+      <> [ ( condFun t,+             ( [Bool, t, t],+               t,+               \case+                 [BoolValue b, tv, fv] ->+                   Just $ if b then tv else fv+                 _ -> Nothing+             )+           )+           | t <- allPrimTypes+         ]   where     i8 s f = (s, ([IntType Int8], IntType Int32, i8PrimFun f))     i16 s f = (s, ([IntType Int16], IntType Int32, i16PrimFun f))

@@ -140,6 +140,7 @@ import Data.Maybe import Data.Ord import Data.Set qualified as S+import Data.Text qualified as T import Futhark.Util (maxinum) import Futhark.Util.Pretty import Language.Futhark.Primitive qualified as Primitive@@ -1095,7 +1096,7 @@      intrinsicStart = 1 + baseTag (fst $ last primOp) -    [a, b, n, m, k, l, p, q] = zipWith VName (map nameFromString ["a", "b", "n", "m", "k", "l", "p", "q"]) [0 ..]+    [a, b, n, m, k, l, p, q] = zipWith VName (map nameFromText ["a", "b", "n", "m", "k", "l", "p", "q"]) [0 ..]      t_a u = TypeVar u (qualName a) []     array_a u s = Array u s $ t_a mempty@@ -1121,37 +1122,37 @@     accType u t =       TypeVar u (qualName (fst intrinsicAcc)) [TypeArgType t] -    namify i (x, y) = (VName (nameFromString x) i, y)+    namify i (x, y) = (VName (nameFromText x) i, y)      primFun (name, (ts, t, _)) =       (name, IntrinsicMonoFun (map unPrim ts) $ unPrim t) -    unOpFun bop = (prettyString bop, IntrinsicMonoFun [t] t)+    unOpFun bop = (prettyText bop, IntrinsicMonoFun [t] t)       where         t = unPrim $ Primitive.unOpType bop -    binOpFun bop = (prettyString bop, IntrinsicMonoFun [t, t] t)+    binOpFun bop = (prettyText bop, IntrinsicMonoFun [t, t] t)       where         t = unPrim $ Primitive.binOpType bop -    cmpOpFun bop = (prettyString bop, IntrinsicMonoFun [t, t] Bool)+    cmpOpFun bop = (prettyText bop, IntrinsicMonoFun [t, t] Bool)       where         t = unPrim $ Primitive.cmpOpType bop -    convOpFun cop = (prettyString cop, IntrinsicMonoFun [unPrim ft] $ unPrim tt)+    convOpFun cop = (prettyText cop, IntrinsicMonoFun [unPrim ft] $ unPrim tt)       where         (ft, tt) = Primitive.convOpType cop -    signFun t = ("sign_" ++ prettyString t, IntrinsicMonoFun [Unsigned t] $ Signed t)+    signFun t = ("sign_" <> prettyText t, IntrinsicMonoFun [Unsigned t] $ Signed t) -    unsignFun t = ("unsign_" ++ prettyString t, IntrinsicMonoFun [Signed t] $ Unsigned t)+    unsignFun t = ("unsign_" <> prettyText t, IntrinsicMonoFun [Signed t] $ Unsigned t)      unPrim (Primitive.IntType t) = Signed t     unPrim (Primitive.FloatType t) = FloatType t     unPrim Primitive.Bool = Bool     unPrim Primitive.Unit = Bool -    intrinsicPrim t = (prettyString t, IntrinsicType Unlifted [] $ Scalar $ Prim t)+    intrinsicPrim t = (prettyText t, IntrinsicType Unlifted [] $ Scalar $ Prim t)      anyIntType =       map Signed [minBound .. maxBound]@@ -1161,10 +1162,10 @@         ++ map FloatType [minBound .. maxBound]     anyPrimType = Bool : anyNumberType -    mkIntrinsicBinOp :: BinOp -> Maybe (String, Intrinsic)+    mkIntrinsicBinOp :: BinOp -> Maybe (T.Text, Intrinsic)     mkIntrinsicBinOp op = do       op' <- intrinsicBinOp op-      pure (prettyString op, op')+      pure (prettyText op, op')      binOp ts = Just $ IntrinsicOverloadedFun ts [Nothing, Nothing] Nothing     ordering = Just $ IntrinsicOverloadedFun anyPrimType [Nothing, Nothing] (Just Bool)

@@ -338,8 +338,8 @@       resolvePat p $ \p' -> CasePat p' <$> resolveExp body <*> pure cloc resolveAppExp (LetPat sizes p e1 e2 loc) = do   checkForDuplicateNames (map sizeBinderToParam sizes) [p]+  e1' <- resolveExp e1   resolveSizes sizes $ \sizes' -> do-    e1' <- resolveExp e1     resolvePat p $ \p' -> do       e2' <- resolveExp e2       pure $ LetPat sizes' p' e1' e2' loc

@@ -474,9 +474,9 @@             | uncurry (<) $ swap ord (uniqueness b2) (uniqueness b1) -> do                 unifyError usage mempty bcs . withIndexLink "unify-return-uniqueness" $                   "Return types"-                    </> indent 2 (pretty d1 <> pretty b1)+                    </> indent 2 (pretty b1)                     </> "and"-                    </> indent 2 (pretty d2 <> pretty b2)+                    </> indent 2 (pretty b2)                     </> "have incompatible uniqueness."             | otherwise -> do                 -- Introduce the existentials as size variables so they

@@ -1,9 +1,10 @@ module Futhark.AD.DerivativesTests (tests) where  import Data.Map qualified as M+import Data.Text qualified as T import Futhark.AD.Derivatives import Futhark.Analysis.PrimExp-import Futhark.IR.Syntax.Core (nameFromString)+import Futhark.IR.Syntax.Core (nameFromText) import Futhark.Util.Pretty (prettyString) import Test.Tasty import Test.Tasty.HUnit@@ -23,8 +24,8 @@     blank = ValueExp . blankPrimValue      primFunTest (f, (ts, ret, _)) =-      testCase f $-        case pdBuiltin (nameFromString f) (map blank ts) of+      testCase (T.unpack f) $+        case pdBuiltin (nameFromText f) (map blank ts) of           Nothing -> assertFailure "pdBuiltin gives Nothing"           Just v -> map primExpType v @?= replicate (length ts) ret