futhark 0.18.3 → 0.18.4

raw patch · 57 files changed

+2864/−2035 lines, 57 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• docs/language-reference.rst +7/−7
• docs/man/futhark-c.rst +17/−4
• docs/man/futhark-cuda.rst +24/−12
• docs/man/futhark-multicore.rst +18/−4
• docs/man/futhark-opencl.rst +17/−6
• docs/usage.rst +1/−2
• futhark.cabal +6/−1
• rts/c/cuda.h +6/−8
• rts/c/opencl.h +5/−4
• rts/c/timing.h +0/−18
• rts/c/util.h +0/−1
• rts/c/values.h +68/−64
• src/Futhark/Actions.hs +38/−69
• src/Futhark/Analysis/PrimExp.hs +12/−12
• src/Futhark/CLI/Dev.hs +18/−1
• src/Futhark/CLI/REPL.hs +12/−3
• src/Futhark/CLI/Test.hs +8/−6
• src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs +2/−2
• src/Futhark/CodeGen/Backends/COpenCL.hs +1/−1
• src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs +2/−2
• src/Futhark/CodeGen/Backends/GenericC.hs +74/−549
• src/Futhark/CodeGen/Backends/GenericC/CLI.hs +463/−0
• src/Futhark/CodeGen/Backends/GenericPython.hs +28/−12
• src/Futhark/CodeGen/Backends/MulticoreC.hs +0/−7
• src/Futhark/CodeGen/Backends/PyOpenCL.hs +2/−1
• src/Futhark/CodeGen/Backends/SequentialC.hs +1/−8
• src/Futhark/CodeGen/Backends/SimpleRep.hs +80/−1
• src/Futhark/CodeGen/ImpCode/Kernels.hs +1/−1
• src/Futhark/CodeGen/ImpGen.hs +2/−9
• src/Futhark/CodeGen/ImpGen/Kernels.hs +10/−5
• src/Futhark/CodeGen/ImpGen/Kernels/Base.hs +22/−7
• src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs +42/−487
• src/Futhark/CodeGen/ImpGen/Kernels/SegScan/SinglePass.hs +489/−0
• src/Futhark/CodeGen/ImpGen/Kernels/SegScan/TwoPass.hs +506/−0
• src/Futhark/CodeGen/ImpGen/Multicore/Base.hs +1/−12
• src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs +9/−0
• src/Futhark/IR/Mem/IxFun.hs +8/−3
• src/Futhark/Internalise.hs +24/−184
• src/Futhark/Internalise/Defunctionalise.hs +378/−392
• src/Futhark/Internalise/FreeVars.hs +134/−0
• src/Futhark/Internalise/LiftLambdas.hs +174/−0
• src/Futhark/Internalise/Monad.hs +3/−18
• src/Futhark/Internalise/Monomorphise.hs +55/−61
• src/Futhark/Optimise/Simplify.hs +2/−2
• src/Futhark/Optimise/Simplify/Engine.hs +3/−3
• src/Futhark/Optimise/Simplify/Rules.hs +9/−4
• src/Futhark/Pass/ExtractKernels/DistributeNests.hs +17/−16
• src/Futhark/Test.hs +2/−0
• src/Futhark/Test/Values.hs +6/−2
• src/Futhark/TypeCheck.hs +6/−1
• src/Futhark/Util.hs +12/−1
• src/Language/Futhark/Pretty.hs +3/−3
• src/Language/Futhark/Prop.hs +21/−5
• src/Language/Futhark/Traversals.hs +5/−5
• src/Language/Futhark/TypeChecker.hs +3/−4
• src/Language/Futhark/TypeChecker/Terms.hs +3/−2
• src/Language/Futhark/TypeChecker/Types.hs +4/−3

@@ -649,10 +649,10 @@ ............  Construct a signed integer array whose first element is ``x`` and-which proceeds stride of ``y-x`` until reaching ``z`` (inclusive).-The ``..y`` part can be elided in which case a stride of 1 is used.  A-run-time error occurs if ``z`` is lesser than ``x`` or ``y``, or if-``x`` and ``y`` are the same value.+which proceeds with a stride of ``y-x`` until reaching ``z``+(inclusive).  The ``..y`` part can be elided in which case a stride of+1 is used.  A run-time error occurs if ``z`` is less than ``x`` or+``y``, or if ``x`` and ``y`` are the same value.  In the general case, the size of the array produced by a range is unknown (see `Size types`_).  In a few cases, the size is known@@ -668,9 +668,9 @@ ............  Construct a signed integer array whose first elements is ``x``, and-which proceeds upwards with a stride of ``y`` until reaching ``z``+which proceeds upwards with a stride of ``y-x`` until reaching ``z`` (exclusive).  The ``..y`` part can be elided in which case a stride of-1 is used.  A run-time error occurs if ``z`` is lesser than ``x`` or+1 is used.  A run-time error occurs if ``z`` is less than ``x`` or ``y``, or if ``x`` and ``y`` are the same value.    * ``0..1..<n`` has size ``n``@@ -683,7 +683,7 @@ ...............  Construct a signed integer array whose first elements is ``x``, and-which proceeds downwards with a stride of ``y`` until reaching ``z``+which proceeds downwards with a stride of ``y-x`` until reaching ``z`` (exclusive).  The ``..y`` part can be elided in which case a stride of -1 is used.  A run-time error occurs if ``z`` is greater than ``x`` or ``y``, or if ``x`` and ``y`` are the same value.

@@ -15,10 +15,10 @@ ===========  ``futhark c`` translates a Futhark program to sequential C code, and-either compiles that C code with gcc(1) to an executable binary-program, or produces a ``.h`` and ``.c`` file that can be linked with-other code..  The standard Futhark optimisation pipeline is used, and-GCC is invoked with ``-O3``, ``-lm``, and ``-std=c99``.+either compiles that C code with a C compiler (see below) to an+executable binary program, or produces a ``.h`` and ``.c`` file that+can be linked with other code..  The standard Futhark optimisation+pipeline is used, and  The resulting program will read the arguments to the entry point (``main`` by default) from standard input and print its return value@@ -56,6 +56,19 @@  --Werror   Treat warnings as errors.++ENVIRONMENT VARIABLES+=====================++``CC``++  The C compiler used to compile the program.  Defaults to ``cc`` if+  unset.++``CFLAGS``++  Space-separated list of options passed to the C compiler.  Defaults+  to ``-O3 -std=c99`` if unset.  EXECUTABLE OPTIONS ==================

@@ -16,12 +16,10 @@   ``futhark cuda`` translates a Futhark program to C code invoking CUDA-kernels, and either compiles that C code with gcc(1) to an executable-binary program, or produces a ``.h`` and ``.c`` file that can be-linked with other code. The standard Futhark optimisation pipeline is-used, and GCC is invoked with ``-O``, ``-lm``, and ``-std=c99``. The-resulting program will otherwise behave exactly as one compiled with-``futhark c``.+kernels, and either compiles that C code with a C compiler to an+executable binary program, or produces a ``.h`` and ``.c`` file that+can be linked with other code. The standard Futhark optimisation+pipeline is used.  ``futhark cuda`` uses ``-lcuda -lcudart -lnvrtc`` to link.  If using ``--library``, you will need to do the same when linking the final@@ -63,6 +61,19 @@ --Werror   Treat warnings as errors. +ENVIRONMENT VARIABLES+=====================++``CC``++  The C compiler used to compile the program.  Defaults to ``cc`` if+  unset.++``CFLAGS``++  Space-separated list of options passed to the C compiler.  Defaults+  to ``-O -std=c99`` if unset.+ EXECUTABLE OPTIONS ================== @@ -140,12 +151,13 @@ ENVIRONMENT =========== -If run without ``--library``, ``futhark cuda`` will invoke ``gcc(1)``-to compile the generated C program into a binary.  This only works if-``gcc`` can find the necessary CUDA libraries.  On most systems, CUDA-is installed in ``/usr/local/cuda``, which is not part of the default-``gcc`` search path.  You may need to set the following environment-variables before running ``futhark cuda``::+If run without ``--library``, ``futhark cuda`` will invoke a C+compiler to compile the generated C program into a binary.  This only+works if the C compiler can find the necessary CUDA libraries.  On+most systems, CUDA is installed in ``/usr/local/cuda``, which is+usually not part of the default compiler search path.  You may need to+set the following environment variables before running ``futhark+cuda``::    LIBRARY_PATH=/usr/local/cuda/lib64   LD_LIBRARY_PATH=/usr/local/cuda/lib64/

@@ -15,10 +15,10 @@ ===========  ``futhark multicore`` translates a Futhark program to multithreaded C-code, and either compiles that C code with gcc(1) to an executable-binary program, or produces a ``.h`` and ``.c`` file that can be-linked with other code..  The standard Futhark optimisation pipeline-is used, and GCC is invoked with ``-O3 -lm -std=c11 -pthread``.+code, and either compiles that C code with a C compiler to an+executable binary program, or produces a ``.h`` and ``.c`` file that+can be linked with other code.  The standard Futhark optimisation+pipeline is used.  The resulting program will read the arguments to the entry point (``main`` by default) from standard input and print its return value@@ -56,6 +56,20 @@  --Werror   Treat warnings as errors.+++ENVIRONMENT VARIABLES+=====================++``CC``++  The C compiler used to compile the program.  Defaults to ``cc`` if+  unset.++``CFLAGS``++  Space-separated list of options passed to the C compiler.  Defaults+  to ``-O3 -std=c99 -pthread`` if unset.  EXECUTABLE OPTIONS ==================

@@ -16,12 +16,10 @@   ``futhark opencl`` translates a Futhark program to C code invoking-OpenCL kernels, and either compiles that C code with gcc(1) to an-executable binary program, or produces a ``.h`` and ``.c`` file that-can be linked with other code. The standard Futhark optimisation-pipeline is used, and GCC is invoked with ``-O``, ``-lm``, and-``-std=c99``. The resulting program will otherwise behave exactly as-one compiled with ``futhark c``.+OpenCL kernels, and either compiles that C code with a C compiler to+an executable binary program, or produces a ``.h`` and ``.c`` file+that can be linked with other code. The standard Futhark optimisation+pipeline is used.  ``futhark opencl`` uses ``-lOpenCL`` to link (``-framework OpenCL`` on macOS).  If using ``--library``, you will need to do the same when@@ -58,6 +56,19 @@  --Werror   Treat warnings as errors.++ENVIRONMENT VARIABLES+=====================++``CC``++  The C compiler used to compile the program.  Defaults to ``cc`` if+  unset.++``CFLAGS``++  Space-separated list of options passed to the C compiler.  Defaults+  to ``-O -std=c99`` if unset.  EXECUTABLE OPTIONS ==================

@@ -30,7 +30,7 @@   $ futhark c prog.fut  This makes use of the ``futhark c`` compiler, but any other will work-as well.  The compiler will automatically invoke ``gcc`` to produce an+as well.  The compiler will automatically invoke ``cc`` to produce an executable binary called ``prog``.  If we had used ``futhark py`` instead of ``futhark c``, the ``prog`` file would instead have contained Python code, along with a `shebang`_ for easy execution.  In@@ -66,7 +66,6 @@ ``futhark run`` and ``futhark repl``.  The latter is an interactive prompt, useful for experimenting with Futhark expressions.  Be aware that the interpreter runs code very slowly.-  .. _executable-options: 

@@ -1,7 +1,7 @@ cabal-version: 2.4  name:           futhark-version:        0.18.3+version:        0.18.4 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -93,6 +93,7 @@       Futhark.CodeGen.Backends.COpenCL       Futhark.CodeGen.Backends.COpenCL.Boilerplate       Futhark.CodeGen.Backends.GenericC+      Futhark.CodeGen.Backends.GenericC.CLI       Futhark.CodeGen.Backends.GenericC.Options       Futhark.CodeGen.Backends.GenericPython       Futhark.CodeGen.Backends.GenericPython.AST@@ -117,6 +118,8 @@       Futhark.CodeGen.ImpGen.Kernels.SegMap       Futhark.CodeGen.ImpGen.Kernels.SegRed       Futhark.CodeGen.ImpGen.Kernels.SegScan+      Futhark.CodeGen.ImpGen.Kernels.SegScan.SinglePass+      Futhark.CodeGen.ImpGen.Kernels.SegScan.TwoPass       Futhark.CodeGen.ImpGen.Kernels.ToOpenCL       Futhark.CodeGen.ImpGen.Kernels.Transpose       Futhark.CodeGen.ImpGen.Multicore@@ -178,7 +181,9 @@       Futhark.Internalise.Bindings       Futhark.Internalise.Defunctionalise       Futhark.Internalise.Defunctorise+      Futhark.Internalise.FreeVars       Futhark.Internalise.Lambdas+      Futhark.Internalise.LiftLambdas       Futhark.Internalise.Monad       Futhark.Internalise.Monomorphise       Futhark.Internalise.TypesValues

@@ -47,7 +47,7 @@   int num_sizes;   const char **size_names;   const char **size_vars;-  size_t *size_values;+  int64_t *size_values;   const char **size_classes; }; @@ -55,7 +55,7 @@                              int num_sizes,                              const char *size_names[],                              const char *size_vars[],-                             size_t *size_values,+                             int64_t *size_values,                              const char *size_classes[]) {   cfg->debugging = 0;   cfg->logging = 0;@@ -388,12 +388,10 @@   }    for (int i = 0; i < ctx->cfg.num_sizes; i++) {-    const char *size_class, *size_name;-    size_t *size_value, max_value = 0, default_value = 0;--    size_class = ctx->cfg.size_classes[i];-    size_value = &ctx->cfg.size_values[i];-    size_name = ctx->cfg.size_names[i];+    const char *size_class = ctx->cfg.size_classes[i];+    int64_t *size_value = &ctx->cfg.size_values[i];+    const char* size_name = ctx->cfg.size_names[i];+    int64_t max_value = 0, default_value = 0;      if (strstr(size_class, "group_size") == size_class) {       max_value = ctx->max_block_size;

@@ -46,7 +46,7 @@   int num_sizes;   const char **size_names;   const char **size_vars;-  size_t *size_values;+  int64_t *size_values;   const char **size_classes; }; @@ -54,7 +54,7 @@                                int num_sizes,                                const char *size_names[],                                const char *size_vars[],-                               size_t *size_values,+                               int64_t *size_values,                                const char *size_classes[]) {   cfg->debugging = 0;   cfg->logging = 0;@@ -598,9 +598,10 @@   // or set them to the default.   for (int i = 0; i < ctx->cfg.num_sizes; i++) {     const char *size_class = ctx->cfg.size_classes[i];-    size_t *size_value = &ctx->cfg.size_values[i];+    int64_t *size_value = &ctx->cfg.size_values[i];     const char* size_name = ctx->cfg.size_names[i];-    size_t max_value = 0, default_value = 0;+    int64_t max_value = 0, default_value = 0;+     if (strstr(size_class, "group_size") == size_class) {       max_value = max_group_size;       default_value = ctx->cfg.default_group_size;

@@ -32,24 +32,6 @@   return time.tv_sec * 1000000000 + time.tv_nsec; } --static inline uint64_t rdtsc() {-  unsigned int hi, lo;-  __asm__ __volatile__("rdtsc" : "=a"(lo), "=d"(hi));-  return  ((uint64_t) lo) | (((uint64_t) hi) << 32);-}--static inline void rdtsc_wait(uint64_t n) {-  const uint64_t start = rdtsc();-  while (rdtsc() < (start + n)) {-    __asm__("PAUSE");-  }-}-static inline void spin_for(uint64_t nb_cycles) {-  rdtsc_wait(nb_cycles);-}-- #endif  // End of timing.h.

@@ -32,7 +32,6 @@ static inline void check_err(int errval, int sets_errno, const char *fun, int line,                             const char *msg, ...) {   if (errval) {-    char str[256];     char errnum[10];      va_list vl;

@@ -15,14 +15,14 @@   str_reader elem_reader; }; -static void skipspaces() {+static void skipspaces(FILE *f) {   int c;   do {-    c = getchar();+    c = getc(f);   } while (isspace(c));    if (c != EOF) {-    ungetc(c, stdin);+    ungetc(c, f);   } } @@ -31,13 +31,13 @@ }  // Produces an empty token only on EOF.-static void next_token(char *buf, int bufsize) {+static void next_token(FILE *f, char *buf, int bufsize) {  start:-  skipspaces();+  skipspaces(f);    int i = 0;   while (i < bufsize) {-    int c = getchar();+    int c = getc(f);     buf[i] = (char)c;      if (c == EOF) {@@ -45,7 +45,7 @@       return;     } else if (c == '-' && i == 1 && buf[0] == '-') {       // Line comment, so skip to end of line and start over.-      for (; c != '\n' && c != EOF; c = getchar());+      for (; c != '\n' && c != EOF; c = getc(f));       goto start;     } else if (!constituent((char)c)) {       if (i == 0) {@@ -55,7 +55,7 @@         buf[i+1] = 0;         return;       } else {-        ungetc(c, stdin);+        ungetc(c, f);         buf[i] = 0;         return;       }@@ -67,8 +67,8 @@   buf[bufsize-1] = 0; } -static int next_token_is(char *buf, int bufsize, const char* expected) {-  next_token(buf, bufsize);+static int next_token_is(FILE *f, char *buf, int bufsize, const char* expected) {+  next_token(f, buf, bufsize);   return strcmp(buf, expected) == 0; } @@ -101,7 +101,8 @@   return ret; } -static int read_str_array_elems(char *buf, int bufsize,+static int read_str_array_elems(FILE *f,+                                char *buf, int bufsize,                                 struct array_reader *reader, int64_t dims) {   int ret;   int first = 1;@@ -110,7 +111,7 @@   int64_t *elems_read_in_dim = (int64_t*) calloc((size_t)dims, sizeof(int64_t));    while (1) {-    next_token(buf, bufsize);+    next_token(f, buf, bufsize);      if (strcmp(buf, "]") == 0) {       if (knows_dimsize[cur_dim]) {@@ -130,7 +131,7 @@         elems_read_in_dim[cur_dim]++;       }     } else if (strcmp(buf, ",") == 0) {-      next_token(buf, bufsize);+      next_token(f, buf, bufsize);       if (strcmp(buf, "[") == 0) {         if (cur_dim == dims - 1) {           ret = 1;@@ -180,7 +181,7 @@   return ret; } -static int read_str_empty_array(char *buf, int bufsize,+static int read_str_empty_array(FILE *f, char *buf, int bufsize,                                 const char *type_name, int64_t *shape, int64_t dims) {   if (strlen(buf) == 0) {     // EOF@@ -191,32 +192,32 @@     return 1;   } -  if (!next_token_is(buf, bufsize, "(")) {+  if (!next_token_is(f, buf, bufsize, "(")) {     return 1;   }    for (int i = 0; i < dims; i++) {-    if (!next_token_is(buf, bufsize, "[")) {+    if (!next_token_is(f, buf, bufsize, "[")) {       return 1;     } -    next_token(buf, bufsize);+    next_token(f, buf, bufsize);      if (sscanf(buf, "%"SCNu64, (uint64_t*)&shape[i]) != 1) {       return 1;     } -    if (!next_token_is(buf, bufsize, "]")) {+    if (!next_token_is(f, buf, bufsize, "]")) {       return 1;     }   } -  if (!next_token_is(buf, bufsize, type_name)) {+  if (!next_token_is(f, buf, bufsize, type_name)) {     return 1;   }  -  if (!next_token_is(buf, bufsize, ")")) {+  if (!next_token_is(f, buf, bufsize, ")")) {     return 1;   } @@ -231,7 +232,8 @@   return 1; } -static int read_str_array(int64_t elem_size, str_reader elem_reader,+static int read_str_array(FILE *f,+                          int64_t elem_size, str_reader elem_reader,                           const char *type_name,                           void **data, int64_t *shape, int64_t dims) {   int ret;@@ -240,13 +242,13 @@    int dims_seen;   for (dims_seen = 0; dims_seen < dims; dims_seen++) {-    if (!next_token_is(buf, sizeof(buf), "[")) {+    if (!next_token_is(f, buf, sizeof(buf), "[")) {       break;     }   }    if (dims_seen == 0) {-    return read_str_empty_array(buf, sizeof(buf), type_name, shape, dims);+    return read_str_empty_array(f, buf, sizeof(buf), type_name, shape, dims);   }    if (dims_seen != dims) {@@ -260,7 +262,7 @@   reader.elems = (char*) realloc(*data, (size_t)(elem_size*reader.n_elems_space));   reader.elem_reader = elem_reader; -  ret = read_str_array_elems(buf, sizeof(buf), &reader, dims);+  ret = read_str_array_elems(f, buf, sizeof(buf), &reader, dims);    *data = reader.elems; @@ -469,8 +471,8 @@ } #endif -static int read_byte(void* dest) {-  int num_elems_read = fread(dest, 1, 1, stdin);+static int read_byte(FILE *f, void* dest) {+  int num_elems_read = fread(dest, 1, 1, f);   return num_elems_read == 1 ? 0 : 1; } @@ -529,12 +531,12 @@ // General value interface.  All endian business taken care of at // lower layers. -static int read_is_binary() {-  skipspaces();-  int c = getchar();+static int read_is_binary(FILE *f) {+  skipspaces(f);+  int c = getc(f);   if (c == 'b') {     int8_t bin_version;-    int ret = read_byte(&bin_version);+    int ret = read_byte(f, &bin_version);      if (ret != 0) { futhark_panic(1, "binary-input: could not read version.\n"); } @@ -545,14 +547,14 @@      return 1;   }-  ungetc(c, stdin);+  ungetc(c, f);   return 0; } -static const struct primtype_info_t* read_bin_read_type_enum() {+static const struct primtype_info_t* read_bin_read_type_enum(FILE *f) {   char read_binname[4]; -  int num_matched = scanf("%4c", read_binname);+  int num_matched = fscanf(f, "%4c", read_binname);   if (num_matched != 1) { futhark_panic(1, "binary-input: Couldn't read element type.\n"); }    const struct primtype_info_t **type = primtypes;@@ -568,9 +570,9 @@   return NULL; } -static void read_bin_ensure_scalar(const struct primtype_info_t *expected_type) {+static void read_bin_ensure_scalar(FILE *f, const struct primtype_info_t *expected_type) {   int8_t bin_dims;-  int ret = read_byte(&bin_dims);+  int ret = read_byte(f, &bin_dims);   if (ret != 0) { futhark_panic(1, "binary-input: Couldn't get dims.\n"); }    if (bin_dims != 0) {@@ -578,7 +580,7 @@           bin_dims);   } -  const struct primtype_info_t *bin_type = read_bin_read_type_enum();+  const struct primtype_info_t *bin_type = read_bin_read_type_enum(f);   if (bin_type != expected_type) {     futhark_panic(1, "binary-input: Expected scalar of type %s but got scalar of type %s.\n",           expected_type->type_name,@@ -588,11 +590,12 @@  //// High-level interface -static int read_bin_array(const struct primtype_info_t *expected_type, void **data, int64_t *shape, int64_t dims) {+static int read_bin_array(FILE *f,+                          const struct primtype_info_t *expected_type, void **data, int64_t *shape, int64_t dims) {   int ret;    int8_t bin_dims;-  ret = read_byte(&bin_dims);+  ret = read_byte(f, &bin_dims);   if (ret != 0) { futhark_panic(1, "binary-input: Couldn't get dims.\n"); }    if (bin_dims != dims) {@@ -600,7 +603,7 @@           dims, bin_dims);   } -  const struct primtype_info_t *bin_primtype = read_bin_read_type_enum();+  const struct primtype_info_t *bin_primtype = read_bin_read_type_enum(f);   if (expected_type != bin_primtype) {     futhark_panic(1, "binary-input: Expected %iD-array with element type '%s' but got %iD-array with element type '%s'.\n",           dims, expected_type->type_name, dims, bin_primtype->type_name);@@ -609,7 +612,7 @@   int64_t elem_count = 1;   for (int i=0; i<dims; i++) {     int64_t bin_shape;-    ret = fread(&bin_shape, sizeof(bin_shape), 1, stdin);+    ret = fread(&bin_shape, sizeof(bin_shape), 1, f);     if (ret != 1) {       futhark_panic(1, "binary-input: Couldn't read size for dimension %i of array.\n", i);     }@@ -628,7 +631,7 @@   }   *data = tmp; -  int64_t num_elems_read = (int64_t)fread(*data, (size_t)elem_size, (size_t)elem_count, stdin);+  int64_t num_elems_read = (int64_t)fread(*data, (size_t)elem_size, (size_t)elem_count, f);   if (num_elems_read != elem_count) {     futhark_panic(1, "binary-input: tried to read %i elements of an array, but only got %i elements.\n",           elem_count, num_elems_read);@@ -643,18 +646,18 @@   return 0; } -static int read_array(const struct primtype_info_t *expected_type, void **data, int64_t *shape, int64_t dims) {-  if (!read_is_binary()) {-    return read_str_array(expected_type->size, (str_reader)expected_type->read_str, expected_type->type_name, data, shape, dims);+static int read_array(FILE *f, const struct primtype_info_t *expected_type, void **data, int64_t *shape, int64_t dims) {+  if (!read_is_binary(f)) {+    return read_str_array(f, expected_type->size, (str_reader)expected_type->read_str, expected_type->type_name, data, shape, dims);   } else {-    return read_bin_array(expected_type, data, shape, dims);+    return read_bin_array(f, expected_type, data, shape, dims);   } } -static int end_of_input() {-  skipspaces();+static int end_of_input(FILE *f) {+  skipspaces(f);   char token[2];-  next_token(token, sizeof(token));+  next_token(f, token, sizeof(token));   if (strcmp(token, "") == 0) {     return 0;   } else {@@ -679,30 +682,30 @@     }      if (len*slice_size == 0) {-      printf("empty(");+      fprintf(out, "empty(");       for (int64_t i = 0; i < rank; i++) {-        printf("[%"PRIi64"]", shape[i]);+        fprintf(out, "[%"PRIi64"]", shape[i]);       }-      printf("%s", elem_type->type_name);-      printf(")");+      fprintf(out, "%s", elem_type->type_name);+      fprintf(out, ")");     } else if (rank==1) {-      putchar('[');+      fputc('[', out);       for (int64_t i = 0; i < len; i++) {         elem_type->write_str(out, (void*) (data + i * elem_size));         if (i != len-1) {-          printf(", ");+          fprintf(out, ", ");         }       }-      putchar(']');+      fputc(']', out);     } else {-      putchar('[');+      fputc('[', out);       for (int64_t i = 0; i < len; i++) {         write_str_array(out, elem_type, data + i * slice_size * elem_size, shape+1, rank-1);         if (i != len-1) {-          printf(", ");+          fprintf(out, ", ");         }       }-      putchar(']');+      fputc(']', out);     }   }   return 0;@@ -721,7 +724,7 @@   fputc('b', out);   fputc((char)BINARY_FORMAT_VERSION, out);   fwrite(&rank, sizeof(int8_t), 1, out);-  fputs(elem_type->binname, out);+  fwrite(elem_type->binname, 4, 1, out);   if (shape != NULL) {     fwrite(shape, sizeof(int64_t), (size_t)rank, out);   }@@ -752,15 +755,16 @@   } } -static int read_scalar(const struct primtype_info_t *expected_type, void *dest) {-  if (!read_is_binary()) {+static int read_scalar(FILE *f,+                       const struct primtype_info_t *expected_type, void *dest) {+  if (!read_is_binary(f)) {     char buf[100];-    next_token(buf, sizeof(buf));+    next_token(f, buf, sizeof(buf));     return expected_type->read_str(buf, dest);   } else {-    read_bin_ensure_scalar(expected_type);+    read_bin_ensure_scalar(f, expected_type);     int64_t elem_size = expected_type->size;-    int num_elems_read = fread(dest, (size_t)elem_size, 1, stdin);+    int num_elems_read = fread(dest, (size_t)elem_size, 1, f);     if (IS_BIG_ENDIAN) {       flip_bytes(elem_size, (unsigned char*) dest);     }

@@ -19,6 +19,7 @@ import Control.Monad import Control.Monad.IO.Class import qualified Data.ByteString.Lazy.Char8 as ByteString+import Data.Maybe (fromMaybe) import Futhark.Analysis.Alias import Futhark.Analysis.Metrics import qualified Futhark.CodeGen.Backends.CCUDA as CCUDA@@ -34,7 +35,7 @@ import Futhark.IR.MCMem (MCMem) import Futhark.IR.Prop.Aliases import Futhark.IR.SeqMem (SeqMem)-import Futhark.Util (runProgramWithExitCode)+import Futhark.Util (runProgramWithExitCode, unixEnvironment) import Language.SexpGrammar as Sexp import System.Exit import System.FilePath@@ -106,11 +107,42 @@         Left s ->           error $ "Couldn't encode program: " ++ s +cmdCC :: String+cmdCC = fromMaybe "cc" $ lookup "CC" unixEnvironment++cmdCFLAGS :: [String] -> [String]+cmdCFLAGS def = maybe def words $ lookup "CFLAGS" unixEnvironment++runCC :: String -> String -> [String] -> [String] -> FutharkM ()+runCC cpath outpath cflags_def ldflags = do+  ret <-+    liftIO $+      runProgramWithExitCode+        cmdCC+        ( [cpath, "-o", outpath]+            ++ cmdCFLAGS cflags_def+            +++            -- The default LDFLAGS are always added.+            ldflags+        )+        mempty+  case ret of+    Left err ->+      externalErrorS $ "Failed to run " ++ cmdCC ++ ": " ++ show err+    Right (ExitFailure code, _, gccerr) ->+      externalErrorS $+        cmdCC ++ " failed with code "+          ++ show code+          ++ ":\n"+          ++ gccerr+    Right (ExitSuccess, _, _) ->+      return ()+ -- | The @futhark c@ action. compileCAction :: FutharkConfig -> CompilerMode -> FilePath -> Action SeqMem compileCAction fcfg mode outpath =   Action-    { actionName = "Compile to to sequential C",+    { actionName = "Compile to sequential C",       actionDescription = "Compile to sequential C",       actionProcedure = helper     }@@ -127,23 +159,7 @@           liftIO $ writeFile cpath impl         ToExecutable -> do           liftIO $ writeFile cpath $ SequentialC.asExecutable cprog-          ret <--            liftIO $-              runProgramWithExitCode-                "gcc"-                [cpath, "-O3", "-std=c99", "-lm", "-o", outpath]-                mempty-          case ret of-            Left err ->-              externalErrorS $ "Failed to run gcc: " ++ show err-            Right (ExitFailure code, _, gccerr) ->-              externalErrorS $-                "gcc failed with code "-                  ++ show code-                  ++ ":\n"-                  ++ gccerr-            Right (ExitSuccess, _, _) ->-              return ()+          runCC cpath outpath ["-O3", "-std=c99"] ["-lm"]  -- | The @futhark opencl@ action. compileOpenCLAction :: FutharkConfig -> CompilerMode -> FilePath -> Action KernelsMem@@ -173,23 +189,7 @@           liftIO $ writeFile cpath impl         ToExecutable -> do           liftIO $ writeFile cpath $ COpenCL.asExecutable cprog-          ret <--            liftIO $-              runProgramWithExitCode-                "gcc"-                ([cpath, "-O", "-std=c99", "-lm", "-o", outpath] ++ extra_options)-                mempty-          case ret of-            Left err ->-              externalErrorS $ "Failed to run gcc: " ++ show err-            Right (ExitFailure code, _, gccerr) ->-              externalErrorS $-                "gcc failed with code "-                  ++ show code-                  ++ ":\n"-                  ++ gccerr-            Right (ExitSuccess, _, _) ->-              return ()+          runCC cpath outpath ["-O", "-std=c99"] ("-lm" : extra_options)  -- | The @futhark cuda@ action. compileCUDAAction :: FutharkConfig -> CompilerMode -> FilePath -> Action KernelsMem@@ -216,21 +216,7 @@           liftIO $ writeFile cpath impl         ToExecutable -> do           liftIO $ writeFile cpath $ CCUDA.asExecutable cprog-          let args =-                [cpath, "-O", "-std=c99", "-lm", "-o", outpath]-                  ++ extra_options-          ret <- liftIO $ runProgramWithExitCode "gcc" args mempty-          case ret of-            Left err ->-              externalErrorS $ "Failed to run gcc: " ++ show err-            Right (ExitFailure code, _, gccerr) ->-              externalErrorS $-                "gcc failed with code "-                  ++ show code-                  ++ ":\n"-                  ++ gccerr-            Right (ExitSuccess, _, _) ->-              return ()+          runCC cpath outpath ["-O", "-std=c99"] ("-lm" : extra_options)  -- | The @futhark multicore@ action. compileMulticoreAction :: FutharkConfig -> CompilerMode -> FilePath -> Action MCMem@@ -253,21 +239,4 @@           liftIO $ writeFile cpath impl         ToExecutable -> do           liftIO $ writeFile cpath $ MulticoreC.asExecutable cprog-          -- let debug_flags = ["-g", "-fno-omit-frame-pointer", "-fsanitize=address", "-fsanitize=integer", "-fsanitize=undefined", "-fno-sanitize-recover=null"]-          ret <--            liftIO $-              runProgramWithExitCode-                "gcc"-                [cpath, "-O3", "-pthread", "-std=c11", "-lm", "-o", outpath]-                mempty-          case ret of-            Left err ->-              externalErrorS $ "Failed to run gcc: " ++ show err-            Right (ExitFailure code, _, gccerr) ->-              externalErrorS $-                "gcc failed with code "-                  ++ show code-                  ++ ":\n"-                  ++ gccerr-            Right (ExitSuccess, _, _) ->-              return ()+          runCC cpath outpath ["-O", "-std=c99"] ["-lm", "-pthread"]

@@ -41,6 +41,8 @@     (.&.),     (.|.),     (.^.),+    (.>>.),+    (.<<.),     bNot,     sMax32,     sMin32,@@ -396,20 +398,18 @@ x .>. y = y .<. x x .>=. y = y .<=. x --- | Lifted bitwise operators.-(.&.), (.|.), (.^.) :: TPrimExp t v -> TPrimExp t v -> TPrimExp t v-TPrimExp x .&. TPrimExp y =-  TPrimExp $-    constFoldPrimExp $-      BinOpExp (And $ primExpIntType x) x y-TPrimExp x .|. TPrimExp y =-  TPrimExp $-    constFoldPrimExp $-      BinOpExp (Or $ primExpIntType x) x y-TPrimExp x .^. TPrimExp y =+-- | Lifted bitwise operators.  The right-shift is logical, *not* arithmetic.+(.&.), (.|.), (.^.), (.>>.), (.<<.) :: TPrimExp t v -> TPrimExp t v -> TPrimExp t v+bitPrimExp :: (IntType -> BinOp) -> TPrimExp t v -> TPrimExp t v -> TPrimExp t v+bitPrimExp op (TPrimExp x) (TPrimExp y) =   TPrimExp $     constFoldPrimExp $-      BinOpExp (Xor $ primExpIntType x) x y+      BinOpExp (op $ primExpIntType x) x y+(.&.) = bitPrimExp And+(.|.) = bitPrimExp Or+(.^.) = bitPrimExp Xor+(.>>.) = bitPrimExp LShr+(.<<.) = bitPrimExp Shl  infix 4 .==., .<., .>., .<=., .>=. 

@@ -26,6 +26,7 @@ import qualified Futhark.IR.SeqMem as SeqMem import Futhark.Internalise.Defunctionalise as Defunctionalise import Futhark.Internalise.Defunctorise as Defunctorise+import Futhark.Internalise.LiftLambdas as LiftLambdas import Futhark.Internalise.Monomorphise as Monomorphise import Futhark.Optimise.CSE import Futhark.Optimise.DoubleBuffer@@ -69,7 +70,9 @@     Defunctorise   | -- | Defunctorise and monomorphise.     Monomorphise-  | -- | Defunctorise, monomorphise, and defunctionalise.+  | -- | Defunctorise, monomorphise, and lambda-lift.+    LiftLambdas+  | -- | Defunctorise, monomorphise, lambda-lift, and defunctionalise.     Defunctionalise  data Config = Config@@ -453,6 +456,11 @@       "Monomorphise the program.",     Option       []+      ["lift-lambdas"]+      (NoArg $ Right $ \opts -> opts {futharkPipeline = LiftLambdas})+      "Lambda-lift the program.",+    Option+      []       ["defunctionalise"]       (NoArg $ Right $ \opts -> opts {futharkPipeline = Defunctionalise})       "Defunctionalise the program.",@@ -579,6 +587,14 @@               flip evalState src $                 Defunctorise.transformProg imports                   >>= Monomorphise.transformProg+        LiftLambdas -> do+          (_, imports, src) <- readProgram file+          liftIO $+            p $+              flip evalState src $+                Defunctorise.transformProg imports+                  >>= Monomorphise.transformProg+                  >>= LiftLambdas.transformProg         Defunctionalise -> do           (_, imports, src) <- readProgram file           liftIO $@@ -586,6 +602,7 @@               flip evalState src $                 Defunctorise.transformProg imports                   >>= Monomorphise.transformProg+                  >>= LiftLambdas.transformProg                   >>= Defunctionalise.transformProg         Pipeline {} ->           case splitExtensions file of

@@ -87,9 +87,18 @@         if quit then return () else toploop s    maybe_init_state <- liftIO $ newFutharkiState 0 maybe_prog-  case maybe_init_state of-    Left err -> error $ "Failed to initialise interpreter state: " ++ err-    Right init_state -> Haskeline.runInputT Haskeline.defaultSettings $ toploop init_state+  s <- case maybe_init_state of+    Left prog_err -> do+      noprog_init_state <- liftIO $ newFutharkiState 0 Nothing+      case noprog_init_state of+        Left err ->+          error $ "Failed to initialise interpreter state: " ++ err+        Right s -> do+          liftIO $ putStrLn prog_err+          return s {futharkiLoaded = maybe_prog}+    Right s ->+      return s+  Haskeline.runInputT Haskeline.defaultSettings $ toploop s    putStrLn "Leaving 'futhark repl'." 

@@ -97,18 +97,20 @@ optimisedProgramMetrics programs pipeline program =   case pipeline of     SOACSPipeline ->-      check "-s"+      check ["-s"]     KernelsPipeline ->-      check "--kernels"+      check ["--kernels"]     SequentialCpuPipeline ->-      check "--cpu"+      check ["--cpu"]     GpuPipeline ->-      check "--gpu"+      check ["--gpu"]+    NoPipeline ->+      check []   where     check opt = do       futhark <- io $ maybe getExecutablePath return $ configFuthark programs-      (code, output, err) <--        io $ readProcessWithExitCode futhark ["dev", opt, "--metrics", program] ""+      let opts = ["dev"] ++ opt ++ ["--metrics", program]+      (code, output, err) <- io $ readProcessWithExitCode futhark opts ""       let output' = T.decodeUtf8 output       case code of         ExitSuccess

@@ -120,14 +120,14 @@  generateConfigFuns :: M.Map Name SizeClass -> GC.CompilerM OpenCL () String generateConfigFuns sizes = do-  let size_decls = map (\k -> [C.csdecl|size_t $id:k;|]) $ M.keys sizes+  let size_decls = map (\k -> [C.csdecl|typename int64_t $id:k;|]) $ M.keys sizes       num_sizes = M.size sizes   GC.earlyDecl [C.cedecl|struct sizes { $sdecls:size_decls };|]   cfg <- GC.publicDef "context_config" GC.InitDecl $ \s ->     ( [C.cedecl|struct $id:s;|],       [C.cedecl|struct $id:s { struct cuda_config cu_cfg;                               int profiling;-                              size_t sizes[$int:num_sizes];+                              typename int64_t sizes[$int:num_sizes];                               int num_nvrtc_opts;                               const char **nvrtc_opts;                             };|]

@@ -318,7 +318,7 @@   GC.stm [C.cstm|$id:v = ctx->sizes.$id:key <= $exp:x';|]   GC.stm     [C.cstm|if (ctx->logging) {-    fprintf(stderr, "Compared %s <= %d.\n", $string:(pretty key), $exp:x');+    fprintf(stderr, "Compared %s <= %ld: %s.\n", $string:(pretty key), (long)$exp:x', $id:v ? "true" : "false");     }|] callKernel (GetSizeMax v size_class) =   let field = "max_" ++ pretty size_class

@@ -114,12 +114,12 @@               }|]     ) -  let size_decls = map (\k -> [C.csdecl|size_t $id:k;|]) $ M.keys sizes+  let size_decls = map (\k -> [C.csdecl|typename int64_t $id:k;|]) $ M.keys sizes   GC.earlyDecl [C.cedecl|struct sizes { $sdecls:size_decls };|]   cfg <- GC.publicDef "context_config" GC.InitDecl $ \s ->     ( [C.cedecl|struct $id:s;|],       [C.cedecl|struct $id:s { struct opencl_config opencl;-                              size_t sizes[$int:num_sizes];+                              typename int64_t sizes[$int:num_sizes];                               int num_build_opts;                               const char **build_opts;                             };|]

@@ -5,7 +5,6 @@ {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE Trustworthy #-} {-# LANGUAGE TupleSections #-}-{-# OPTIONS_GHC -fno-warn-orphans #-}  -- | C code generator framework. module Futhark.CodeGen.Backends.GenericC@@ -82,23 +81,19 @@ import Control.Monad.Identity import Control.Monad.RWS import Data.Bifunctor (first)-import Data.Bits (shiftR, xor)-import Data.Char (isAlphaNum, isDigit, ord) import qualified Data.DList as DL import Data.FileEmbed-import Data.List (unzip4) import Data.Loc import qualified Data.Map.Strict as M import Data.Maybe+import Futhark.CodeGen.Backends.GenericC.CLI import Futhark.CodeGen.Backends.GenericC.Options import Futhark.CodeGen.Backends.SimpleRep import Futhark.CodeGen.ImpCode import Futhark.IR.Prop (isBuiltInFunction) import Futhark.MonadFreshNames-import Futhark.Util (zEncodeString) import qualified Language.C.Quote.OpenCL as C import qualified Language.C.Syntax as C-import Text.Printf  data CompilerState s = CompilerState   { compArrayStructs :: [((C.Type, Int), (C.Type, [C.Definition]))],@@ -441,36 +436,6 @@ cacheMem :: C.ToExp a => a -> CompilerM op s (Maybe VName) cacheMem a = asks $ M.lookup (C.toExp a noLoc) . envCachedMem -instance C.ToIdent Name where-  toIdent = C.toIdent . zEncodeString . nameToString--instance C.ToIdent VName where-  toIdent = C.toIdent . zEncodeString . pretty--instance C.ToExp VName where-  toExp v _ = [C.cexp|$id:v|]--instance C.ToExp IntValue where-  toExp (Int8Value v) = C.toExp v-  toExp (Int16Value v) = C.toExp v-  toExp (Int32Value v) = C.toExp v-  toExp (Int64Value v) = C.toExp v--instance C.ToExp FloatValue where-  toExp (Float32Value v) = C.toExp v-  toExp (Float64Value v) = C.toExp v--instance C.ToExp PrimValue where-  toExp (IntValue v) = C.toExp v-  toExp (FloatValue v) = C.toExp v-  toExp (BoolValue True) = C.toExp (1 :: Int8)-  toExp (BoolValue False) = C.toExp (0 :: Int8)-  toExp Checked = C.toExp (1 :: Int8)--instance C.ToExp SubExp where-  toExp (Var v) = C.toExp v-  toExp (Constant c) = C.toExp c- -- | Construct a publicly visible definition using the specified name -- as the template.  The first returned definition is put in the -- header file, and the second is the implementation.  Returns the public@@ -530,9 +495,6 @@ decl :: C.InitGroup -> CompilerM op s () decl x = item [C.citem|$decl:x;|] -addrOf :: C.Exp -> C.Exp-addrOf e = [C.cexp|&$exp:e|]- -- | Public names must have a consitent prefix. publicName :: String -> CompilerM op s String publicName s = return $ "futhark_" ++ s@@ -810,21 +772,6 @@       freeRawMem mem space mem_s       allocRawMem mem size space [C.cexp|desc|] -primTypeInfo :: PrimType -> Signedness -> C.Exp-primTypeInfo (IntType it) t = case (it, t) of-  (Int8, TypeUnsigned) -> [C.cexp|u8_info|]-  (Int16, TypeUnsigned) -> [C.cexp|u16_info|]-  (Int32, TypeUnsigned) -> [C.cexp|u32_info|]-  (Int64, TypeUnsigned) -> [C.cexp|u64_info|]-  (Int8, _) -> [C.cexp|i8_info|]-  (Int16, _) -> [C.cexp|i16_info|]-  (Int32, _) -> [C.cexp|i32_info|]-  (Int64, _) -> [C.cexp|i64_info|]-primTypeInfo (FloatType Float32) _ = [C.cexp|f32_info|]-primTypeInfo (FloatType Float64) _ = [C.cexp|f64_info|]-primTypeInfo Bool _ = [C.cexp|bool_info|]-primTypeInfo Cert _ = [C.cexp|bool_info|]- copyMemoryDefaultSpace ::   C.Exp ->   C.Exp ->@@ -840,38 +787,6 @@  --- Entry points. -arrayName :: PrimType -> Signedness -> Int -> String-arrayName pt signed rank =-  prettySigned (signed == TypeUnsigned) pt ++ "_" ++ show rank ++ "d"--opaqueName :: String -> [ValueDesc] -> String-opaqueName s _-  | valid = "opaque_" ++ s-  where-    valid =-      head s /= '_'-        && not (isDigit $ head s)-        && all ok s-    ok c = isAlphaNum c || c == '_'-opaqueName s vds = "opaque_" ++ hash (zipWith xor [0 ..] $ map ord (s ++ concatMap p vds))-  where-    p (ScalarValue pt signed _) =-      show (pt, signed)-    p (ArrayValue _ space pt signed dims) =-      show (space, pt, signed, length dims)--    -- FIXME: a stupid hash algorithm; may have collisions.-    hash =-      printf "%x" . foldl xor 0-        . map-          ( iter . (* 0x45d9f3b)-              . iter-              . (* 0x45d9f3b)-              . iter-              . fromIntegral-          )-    iter x = ((x :: Word32) `shiftR` 16) `xor` x- criticalSection :: Operations op s -> [C.BlockItem] -> [C.BlockItem] criticalSection ops x =   [C.citems|lock_lock(&ctx->lock);@@ -1115,45 +1030,68 @@       ct <- valueDescToCType vd       return [C.csdecl|$ty:ct *$id:(tupleField i);|] -externalValueToCType :: ExternalValue -> CompilerM op s C.Type-externalValueToCType (TransparentValue vd) = valueDescToCType vd-externalValueToCType (OpaqueValue desc vds) = opaqueToCType desc vds+allTrue :: [C.Exp] -> C.Exp+allTrue [] = [C.cexp|true|]+allTrue [x] = x+allTrue (x : xs) = [C.cexp|$exp:x && $exp:(allTrue xs)|] -prepareEntryInputs :: [ExternalValue] -> CompilerM op s [C.Param]-prepareEntryInputs = zipWithM prepare [(0 :: Int) ..]+prepareEntryInputs ::+  [ExternalValue] ->+  CompilerM op s ([(C.Param, C.Exp)], [C.BlockItem])+prepareEntryInputs args = collect' $ zipWithM prepare [(0 :: Int) ..] args   where+    arg_names = namesFromList $ concatMap evNames args+    evNames (OpaqueValue _ vds) = map vdName vds+    evNames (TransparentValue vd) = [vdName vd]+    vdName (ArrayValue v _ _ _ _) = v+    vdName (ScalarValue _ _ v) = v+     prepare pno (TransparentValue vd) = do       let pname = "in" ++ show pno-      ty <- prepareValue [C.cexp|$id:pname|] vd-      return [C.cparam|const $ty:ty $id:pname|]+      (ty, check) <- prepareValue [C.cexp|$id:pname|] vd+      return+        ( [C.cparam|const $ty:ty $id:pname|],+          allTrue check+        )     prepare pno (OpaqueValue desc vds) = do       ty <- opaqueToCType desc vds       let pname = "in" ++ show pno           field i ScalarValue {} = [C.cexp|$id:pname->$id:(tupleField i)|]           field i ArrayValue {} = [C.cexp|$id:pname->$id:(tupleField i)|]-      zipWithM_ prepareValue (zipWith field [0 ..] vds) vds-      return [C.cparam|const $ty:ty *$id:pname|]+      checks <- map snd <$> zipWithM prepareValue (zipWith field [0 ..] vds) vds+      return+        ( [C.cparam|const $ty:ty *$id:pname|],+          allTrue $ concat checks+        )      prepareValue src (ScalarValue pt signed name) = do       let pt' = signedPrimTypeToCType signed pt       stm [C.cstm|$id:name = $exp:src;|]-      return pt'+      return (pt', [])     prepareValue src vd@(ArrayValue mem _ _ _ shape) = do       ty <- valueDescToCType vd        stm [C.cstm|$exp:mem = $exp:src->mem;|]        let rank = length shape-          maybeCopyDim (Var d) i =-            Just [C.cstm|$id:d = $exp:src->shape[$int:i];|]-          maybeCopyDim _ _ = Nothing+          maybeCopyDim (Var d) i+            | not $ d `nameIn` arg_names =+              ( Just [C.cstm|$id:d = $exp:src->shape[$int:i];|],+                [C.cexp|$id:d == $exp:src->shape[$int:i]|]+              )+          maybeCopyDim x i =+            ( Nothing,+              [C.cexp|$exp:x == $exp:src->shape[$int:i]|]+            ) -      stms $ catMaybes $ zipWith maybeCopyDim shape [0 .. rank -1]+      let (sets, checks) =+            unzip $ zipWith maybeCopyDim shape [0 .. rank -1]+      stms $ catMaybes sets -      return [C.cty|$ty:ty*|]+      return ([C.cty|$ty:ty*|], checks) -prepareEntryOutputs :: [ExternalValue] -> CompilerM op s [C.Param]-prepareEntryOutputs = zipWithM prepare [(0 :: Int) ..]+prepareEntryOutputs :: [ExternalValue] -> CompilerM op s ([C.Param], [C.BlockItem])+prepareEntryOutputs = collect' . zipWithM prepare [(0 :: Int) ..]   where     prepare pno (TransparentValue vd) = do       let pname = "out" ++ show pno@@ -1196,10 +1134,11 @@       stms $ zipWith maybeCopyDim shape [0 .. rank -1]  onEntryPoint ::+  [C.BlockItem] ->   Name ->   Function op ->-  CompilerM op s (C.Definition, C.Definition, C.Initializer)-onEntryPoint fname function@(Function _ outputs inputs _ results args) = do+  CompilerM op s C.Definition+onEntryPoint get_consts fname (Function _ outputs inputs _ results args) = do   let out_args = map (\p -> [C.cexp|&$id:(paramName p)|]) outputs       in_args = map (\p -> [C.cexp|$id:(paramName p)|]) inputs @@ -1209,12 +1148,11 @@   let entry_point_name = nameToString fname   entry_point_function_name <- publicName $ "entry_" ++ entry_point_name -  (entry_point_input_params, unpack_entry_inputs) <--    collect' $ prepareEntryInputs args-  (entry_point_output_params, pack_entry_outputs) <--    collect' $ prepareEntryOutputs results+  (inputs', unpack_entry_inputs) <- prepareEntryInputs args+  let (entry_point_input_params, entry_point_input_checks) = unzip inputs' -  (cli_entry_point, cli_init) <- cliEntryPoint fname function+  (entry_point_output_params, pack_entry_outputs) <-+    prepareEntryOutputs results    ctx_ty <- contextType @@ -1229,17 +1167,26 @@         [C.citems|          $items:unpack_entry_inputs -         int ret = $id:(funName fname)(ctx, $args:out_args, $args:in_args);+         if (!($exp:(allTrue entry_point_input_checks))) {+           ret = 1;+           if (!ctx->error) {+             ctx->error = msgprintf("Error: entry point arguments have invalid sizes.");+           }+         } else {+           ret = $id:(funName fname)(ctx, $args:out_args, $args:in_args); -         if (ret == 0) {-           $items:pack_entry_outputs+           if (ret == 0) {+             $items:get_consts++             $items:pack_entry_outputs+           }          }         |]    ops <- asks envOperations    return-    ( [C.cedecl|+    [C.cedecl|        int $id:entry_point_function_name            ($ty:ctx_ty *ctx,             $params:entry_point_output_params,@@ -1247,13 +1194,12 @@          $items:inputdecls          $items:outputdecls +         int ret = 0;+          $items:(criticalSection ops critical)           return ret;-       }|],-      cli_entry_point,-      cli_init-    )+       }|]   where     stubParam (MemParam name space) =       declMem name space@@ -1261,321 +1207,6 @@       let ty' = primTypeToCType ty       decl [C.cdecl|$ty:ty' $id:name;|] ---- CLI interface------ Our strategy for CLI entry points is to parse everything into--- host memory ('DefaultSpace') and copy the result into host memory--- after the entry point has returned.  We have some ad-hoc frobbery--- to copy the host-level memory blocks to another memory space if--- necessary.  This will break if the Futhark entry point uses--- non-trivial index functions for its input or output.------ The idea here is to keep the nastyness in the wrapper, whilst not--- messing up anything else.--printPrimStm :: (C.ToExp a, C.ToExp b) => a -> b -> PrimType -> Signedness -> C.Stm-printPrimStm dest val bt ept =-  [C.cstm|write_scalar($exp:dest, binary_output, &$exp:(primTypeInfo bt ept), &$exp:val);|]---- | Return a statement printing the given external value.-printStm :: ExternalValue -> C.Exp -> CompilerM op s C.Stm-printStm (OpaqueValue desc _) _ =-  return [C.cstm|printf("#<opaque %s>", $string:desc);|]-printStm (TransparentValue (ScalarValue bt ept _)) e =-  return $ printPrimStm [C.cexp|stdout|] e bt ept-printStm (TransparentValue (ArrayValue _ _ bt ept shape)) e = do-  values_array <- publicName $ "values_" ++ name-  shape_array <- publicName $ "shape_" ++ name-  let num_elems = cproduct [[C.cexp|$id:shape_array(ctx, $exp:e)[$int:i]|] | i <- [0 .. rank -1]]-  return-    [C.cstm|{-      $ty:bt' *arr = calloc(sizeof($ty:bt'), $exp:num_elems);-      assert(arr != NULL);-      assert($id:values_array(ctx, $exp:e, arr) == 0);-      write_array(stdout, binary_output, &$exp:(primTypeInfo bt ept), arr,-                  $id:shape_array(ctx, $exp:e), $int:rank);-      free(arr);-    }|]-  where-    rank = length shape-    bt' = primTypeToCType bt-    name = arrayName bt ept rank--readPrimStm :: C.ToExp a => a -> Int -> PrimType -> Signedness -> C.Stm-readPrimStm place i t ept =-  [C.cstm|if (read_scalar(&$exp:(primTypeInfo t ept),&$exp:place) != 0) {-        futhark_panic(1, "Error when reading input #%d of type %s (errno: %s).\n",-              $int:i,-              $exp:(primTypeInfo t ept).type_name,-              strerror(errno));-      }|]--readInputs :: [ExternalValue] -> CompilerM op s [(C.Stm, C.Stm, C.Stm, C.Exp)]-readInputs = zipWithM readInput [0 ..]--readInput :: Int -> ExternalValue -> CompilerM op s (C.Stm, C.Stm, C.Stm, C.Exp)-readInput i (OpaqueValue desc _) = do-  stm [C.cstm|futhark_panic(1, "Cannot read input #%d of type %s\n", $int:i, $string:desc);|]-  return ([C.cstm|;|], [C.cstm|;|], [C.cstm|;|], [C.cexp|NULL|])-readInput i (TransparentValue (ScalarValue t ept _)) = do-  dest <- newVName "read_value"-  item [C.citem|$ty:(primTypeToCType t) $id:dest;|]-  stm $ readPrimStm dest i t ept-  return ([C.cstm|;|], [C.cstm|;|], [C.cstm|;|], [C.cexp|$id:dest|])-readInput i (TransparentValue vd@(ArrayValue _ _ t ept dims)) = do-  dest <- newVName "read_value"-  shape <- newVName "read_shape"-  arr <- newVName "read_arr"-  ty <- valueDescToCType vd-  item [C.citem|$ty:ty *$id:dest;|]--  let t' = signedPrimTypeToCType ept t-      rank = length dims-      name = arrayName t ept rank-      dims_exps = [[C.cexp|$id:shape[$int:j]|] | j <- [0 .. rank -1]]-      dims_s = concat $ replicate rank "[]"--  new_array <- publicName $ "new_" ++ name-  free_array <- publicName $ "free_" ++ name--  items-    [C.citems|-     typename int64_t $id:shape[$int:rank];-     $ty:t' *$id:arr = NULL;-     errno = 0;-     if (read_array(&$exp:(primTypeInfo t ept),-                    (void**) &$id:arr,-                    $id:shape,-                    $int:(length dims))-         != 0) {-       futhark_panic(1, "Cannot read input #%d of type %s%s (errno: %s).\n",-                 $int:i,-                 $string:dims_s,-                 $exp:(primTypeInfo t ept).type_name,-                 strerror(errno));-     }|]--  return-    ( [C.cstm|assert(($exp:dest = $id:new_array(ctx, $id:arr, $args:dims_exps)) != 0);|],-      [C.cstm|assert($id:free_array(ctx, $exp:dest) == 0);|],-      [C.cstm|free($id:arr);|],-      [C.cexp|$id:dest|]-    )--prepareOutputs :: [ExternalValue] -> CompilerM op s [(C.Exp, C.Stm)]-prepareOutputs = mapM prepareResult-  where-    prepareResult ev = do-      ty <- externalValueToCType ev-      result <- newVName "result"--      case ev of-        TransparentValue ScalarValue {} -> do-          item [C.citem|$ty:ty $id:result;|]-          return ([C.cexp|$id:result|], [C.cstm|;|])-        TransparentValue (ArrayValue _ _ t ept dims) -> do-          let name = arrayName t ept $ length dims-          free_array <- publicName $ "free_" ++ name-          item [C.citem|$ty:ty *$id:result;|]-          return-            ( [C.cexp|$id:result|],-              [C.cstm|assert($id:free_array(ctx, $exp:result) == 0);|]-            )-        OpaqueValue desc vds -> do-          free_opaque <- publicName $ "free_" ++ opaqueName desc vds-          item [C.citem|$ty:ty *$id:result;|]-          return-            ( [C.cexp|$id:result|],-              [C.cstm|assert($id:free_opaque(ctx, $exp:result) == 0);|]-            )--printResult :: [(ExternalValue, C.Exp)] -> CompilerM op s [C.Stm]-printResult vs = fmap concat $-  forM vs $ \(v, e) -> do-    p <- printStm v e-    return [p, [C.cstm|printf("\n");|]]--cliEntryPoint ::-  Name ->-  FunctionT a ->-  CompilerM op s (C.Definition, C.Initializer)-cliEntryPoint fname (Function _ _ _ _ results args) = do-  ((pack_input, free_input, free_parsed, input_args), input_items) <--    collect' $ unzip4 <$> readInputs args--  ((output_vals, free_outputs), output_decls) <--    collect' $ unzip <$> prepareOutputs results-  printstms <- printResult $ zip results output_vals--  ctx_ty <- contextType-  sync_ctx <- publicName "context_sync"-  error_ctx <- publicName "context_get_error"--  let entry_point_name = nameToString fname-      cli_entry_point_function_name = "futrts_cli_entry_" ++ entry_point_name-  entry_point_function_name <- publicName $ "entry_" ++ entry_point_name--  pause_profiling <- publicName "context_pause_profiling"-  unpause_profiling <- publicName "context_unpause_profiling"--  let run_it =-        [C.citems|-                  int r;-                  // Run the program once.-                  $stms:pack_input-                  if ($id:sync_ctx(ctx) != 0) {-                    futhark_panic(1, "%s", $id:error_ctx(ctx));-                  };-                  // Only profile last run.-                  if (profile_run) {-                    $id:unpause_profiling(ctx);-                  }-                  t_start = get_wall_time();-                  r = $id:entry_point_function_name(ctx,-                                                    $args:(map addrOf output_vals),-                                                    $args:input_args);-                  if (r != 0) {-                    futhark_panic(1, "%s", $id:error_ctx(ctx));-                  }-                  if ($id:sync_ctx(ctx) != 0) {-                    futhark_panic(1, "%s", $id:error_ctx(ctx));-                  };-                  if (profile_run) {-                    $id:pause_profiling(ctx);-                  }-                  t_end = get_wall_time();-                  long int elapsed_usec = t_end - t_start;-                  if (time_runs && runtime_file != NULL) {-                    fprintf(runtime_file, "%lld\n", (long long) elapsed_usec);-                    fflush(runtime_file);-                  }-                  $stms:free_input-                |]--  return-    ( [C.cedecl|static void $id:cli_entry_point_function_name($ty:ctx_ty *ctx) {-    typename int64_t t_start, t_end;-    int time_runs = 0, profile_run = 0;--    // We do not want to profile all the initialisation.-    $id:pause_profiling(ctx);--    // Declare and read input.-    set_binary_mode(stdin);-    $items:input_items--    if (end_of_input() != 0) {-      futhark_panic(1, "Expected EOF on stdin after reading input for %s.\n", $string:(quote (pretty fname)));-    }--    $items:output_decls--    // Warmup run-    if (perform_warmup) {-      $items:run_it-      $stms:free_outputs-    }-    time_runs = 1;-    // Proper run.-    for (int run = 0; run < num_runs; run++) {-      // Only profile last run.-      profile_run = run == num_runs -1;-      $items:run_it-      if (run < num_runs-1) {-        $stms:free_outputs-      }-    }--    // Free the parsed input.-    $stms:free_parsed--    // Print the final result.-    if (binary_output) {-      set_binary_mode(stdout);-    }-    $stms:printstms--    $stms:free_outputs-  }-                |],-      [C.cinit|{ .name = $string:entry_point_name,-                      .fun = $id:cli_entry_point_function_name }|]-    )--genericOptions :: [Option]-genericOptions =-  [ Option-      { optionLongName = "write-runtime-to",-        optionShortName = Just 't',-        optionArgument = RequiredArgument "FILE",-        optionDescription = "Print the time taken to execute the program to the indicated file, an integral number of microseconds.",-        optionAction = set_runtime_file-      },-    Option-      { optionLongName = "runs",-        optionShortName = Just 'r',-        optionArgument = RequiredArgument "INT",-        optionDescription = "Perform NUM runs of the program.",-        optionAction = set_num_runs-      },-    Option-      { optionLongName = "debugging",-        optionShortName = Just 'D',-        optionArgument = NoArgument,-        optionDescription = "Perform possibly expensive internal correctness checks and verbose logging.",-        optionAction = [C.cstm|futhark_context_config_set_debugging(cfg, 1);|]-      },-    Option-      { optionLongName = "log",-        optionShortName = Just 'L',-        optionArgument = NoArgument,-        optionDescription = "Print various low-overhead logging information to stderr while running.",-        optionAction = [C.cstm|futhark_context_config_set_logging(cfg, 1);|]-      },-    Option-      { optionLongName = "entry-point",-        optionShortName = Just 'e',-        optionArgument = RequiredArgument "NAME",-        optionDescription = "The entry point to run. Defaults to main.",-        optionAction = [C.cstm|if (entry_point != NULL) entry_point = optarg;|]-      },-    Option-      { optionLongName = "binary-output",-        optionShortName = Just 'b',-        optionArgument = NoArgument,-        optionDescription = "Print the program result in the binary output format.",-        optionAction = [C.cstm|binary_output = 1;|]-      },-    Option-      { optionLongName = "help",-        optionShortName = Just 'h',-        optionArgument = NoArgument,-        optionDescription = "Print help information and exit.",-        optionAction =-          [C.cstm|{-                   printf("Usage: %s [OPTION]...\nOptions:\n\n%s\nFor more information, consult the Futhark User's Guide or the man pages.\n",-                          fut_progname, option_descriptions);-                   exit(0);-                  }|]-      }-  ]-  where-    set_runtime_file =-      [C.cstm|{-          runtime_file = fopen(optarg, "w");-          if (runtime_file == NULL) {-            futhark_panic(1, "Cannot open %s: %s\n", optarg, strerror(errno));-          }-        }|]-    set_num_runs =-      [C.cstm|{-          num_runs = atoi(optarg);-          perform_warmup = 1;-          if (num_runs <= 0) {-            futhark_panic(1, "Need a positive number of runs, not %s\n", optarg);-          }-        }|]- -- | The result of compilation to C is four parts, which can be put -- together in various ways.  The obvious way is to concatenate all of -- them, which yields a CLI program.  Another is to compile the@@ -1638,11 +1269,8 @@   m CParts compileProg backend ops extra header_extra spaces options prog = do   src <- getNameSource-  let ((prototypes, definitions, entry_points), endstate) =+  let ((prototypes, definitions, entry_point_decls), endstate) =         runCompilerM ops src () compileProg'-      (entry_point_decls, cli_entry_point_decls, entry_point_inits) =-        unzip3 entry_points-      option_parser = generateOptionParser "parse_options" $ genericOptions ++ options    let headerdefs =         [C.cunit|@@ -1694,103 +1322,9 @@ $esc:timing_h |] -  let clidefs =-        [C.cunit|-$esc:("#include <string.h>")-$esc:("#include <inttypes.h>")-$esc:("#include <errno.h>")-$esc:("#include <ctype.h>")-$esc:("#include <errno.h>")-$esc:("#include <getopt.h>")--$esc:values_h--$esc:("#define __private")--static int binary_output = 0;-static typename FILE *runtime_file;-static int perform_warmup = 0;-static int num_runs = 1;-// If the entry point is NULL, the program will terminate after doing initialisation and such.-static const char *entry_point = "main";--$esc:tuning_h--$func:option_parser--$edecls:cli_entry_point_decls--typedef void entry_point_fun(struct futhark_context*);--struct entry_point_entry {-  const char *name;-  entry_point_fun *fun;-};--int main(int argc, char** argv) {-  fut_progname = argv[0];--  struct entry_point_entry entry_points[] = {-    $inits:entry_point_inits-  };--  struct futhark_context_config *cfg = futhark_context_config_new();-  assert(cfg != NULL);--  int parsed_options = parse_options(cfg, argc, argv);-  argc -= parsed_options;-  argv += parsed_options;--  if (argc != 0) {-    futhark_panic(1, "Excess non-option: %s\n", argv[0]);-  }--  struct futhark_context *ctx = futhark_context_new(cfg);-  assert (ctx != NULL);--  char* error = futhark_context_get_error(ctx);-  if (error != NULL) {-    futhark_panic(1, "%s", error);-  }--  if (entry_point != NULL) {-    int num_entry_points = sizeof(entry_points) / sizeof(entry_points[0]);-    entry_point_fun *entry_point_fun = NULL;-    for (int i = 0; i < num_entry_points; i++) {-      if (strcmp(entry_points[i].name, entry_point) == 0) {-        entry_point_fun = entry_points[i].fun;-        break;-      }-    }--    if (entry_point_fun == NULL) {-      fprintf(stderr, "No entry point '%s'.  Select another with --entry-point.  Options are:\n",-                      entry_point);-      for (int i = 0; i < num_entry_points; i++) {-        fprintf(stderr, "%s\n", entry_points[i].name);-      }-      return 1;-    }--    entry_point_fun(ctx);--    if (runtime_file != NULL) {-      fclose(runtime_file);-    }--    char *report = futhark_context_report(ctx);-    fputs(report, stderr);-    free(report);-  }--  futhark_context_free(ctx);-  futhark_context_config_free(cfg);-  return 0;-}-                        |]-   let early_decls = DL.toList $ compEarlyDecls endstate   let lib_decls = DL.toList $ compLibDecls endstate+  let clidefs = cliDefs options $ Functions entry_funs   let libdefs =         [C.cunit| $esc:("#ifdef _MSC_VER\n#define inline __inline\n#endif")@@ -1800,7 +1334,7 @@ $esc:("#include <errno.h>") $esc:("#include <assert.h>") -$esc:(header_extra)+$esc:header_extra  $esc:lock_h @@ -1823,9 +1357,10 @@    return $ CParts (pretty headerdefs) (pretty utildefs) (pretty clidefs) (pretty libdefs)   where-    compileProg' = do-      let Definitions consts (Functions funs) = prog+    Definitions consts (Functions funs) = prog+    entry_funs = filter (functionEntry . snd) funs +    compileProg' = do       (memstructs, memfuns, memreport) <- unzip3 <$> mapM defineMemorySpace spaces        get_consts <- compileConstants consts@@ -1837,7 +1372,7 @@        mapM_ earlyDecl memstructs       entry_points <--        mapM (uncurry onEntryPoint) $ filter (functionEntry . snd) funs+        mapM (uncurry (onEntryPoint get_consts)) $ filter (functionEntry . snd) funs        extra @@ -1859,10 +1394,8 @@         ++ cFloat64Funs      util_h = $(embedStringFile "rts/c/util.h")-    values_h = $(embedStringFile "rts/c/values.h")     timing_h = $(embedStringFile "rts/c/timing.h")     lock_h = $(embedStringFile "rts/c/lock.h")-    tuning_h = $(embedStringFile "rts/c/tuning.h")  commonLibFuns :: [C.BlockItem] -> CompilerM op s () commonLibFuns memreport = do@@ -2051,10 +1584,10 @@       return ([C.cparam|$ty:ty *$id:p_name|], [C.cexp|$id:p_name|])  compilePrimValue :: PrimValue -> C.Exp-compilePrimValue (IntValue (Int8Value k)) = [C.cexp|$int:k|]-compilePrimValue (IntValue (Int16Value k)) = [C.cexp|$int:k|]+compilePrimValue (IntValue (Int8Value k)) = [C.cexp|(typename int8_t)$int:k|]+compilePrimValue (IntValue (Int16Value k)) = [C.cexp|(typename int16_t)$int:k|] compilePrimValue (IntValue (Int32Value k)) = [C.cexp|$int:k|]-compilePrimValue (IntValue (Int64Value k)) = [C.cexp|$int:k|]+compilePrimValue (IntValue (Int64Value k)) = [C.cexp|(typename int64_t)$int:k|] compilePrimValue (FloatValue (Float64Value x))   | isInfinite x =     if x > 0 then [C.cexp|INFINITY|] else [C.cexp|-INFINITY|]@@ -2439,11 +1972,3 @@ assignmentOperator Sub {} = Just $ \d e -> [C.cexp|$id:d -= $exp:e|] assignmentOperator Mul {} = Just $ \d e -> [C.cexp|$id:d *= $exp:e|] assignmentOperator _ = Nothing---- | Return an expression multiplying together the given expressions.--- If an empty list is given, the expression @1@ is returned.-cproduct :: [C.Exp] -> C.Exp-cproduct [] = [C.cexp|1|]-cproduct (e : es) = foldl mult e es-  where-    mult x y = [C.cexp|$exp:x * $exp:y|]

@@ -0,0 +1,463 @@+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE Trustworthy #-}+{-# LANGUAGE TupleSections #-}++module Futhark.CodeGen.Backends.GenericC.CLI+  ( cliDefs,+  )+where++import Data.FileEmbed+import Data.List (unzip5)+import Futhark.CodeGen.Backends.GenericC.Options+import Futhark.CodeGen.Backends.SimpleRep+import Futhark.CodeGen.ImpCode+import qualified Language.C.Quote.OpenCL as C+import qualified Language.C.Syntax as C++genericOptions :: [Option]+genericOptions =+  [ Option+      { optionLongName = "write-runtime-to",+        optionShortName = Just 't',+        optionArgument = RequiredArgument "FILE",+        optionDescription = "Print the time taken to execute the program to the indicated file, an integral number of microseconds.",+        optionAction = set_runtime_file+      },+    Option+      { optionLongName = "runs",+        optionShortName = Just 'r',+        optionArgument = RequiredArgument "INT",+        optionDescription = "Perform NUM runs of the program.",+        optionAction = set_num_runs+      },+    Option+      { optionLongName = "debugging",+        optionShortName = Just 'D',+        optionArgument = NoArgument,+        optionDescription = "Perform possibly expensive internal correctness checks and verbose logging.",+        optionAction = [C.cstm|futhark_context_config_set_debugging(cfg, 1);|]+      },+    Option+      { optionLongName = "log",+        optionShortName = Just 'L',+        optionArgument = NoArgument,+        optionDescription = "Print various low-overhead logging information to stderr while running.",+        optionAction = [C.cstm|futhark_context_config_set_logging(cfg, 1);|]+      },+    Option+      { optionLongName = "entry-point",+        optionShortName = Just 'e',+        optionArgument = RequiredArgument "NAME",+        optionDescription = "The entry point to run. Defaults to main.",+        optionAction = [C.cstm|if (entry_point != NULL) entry_point = optarg;|]+      },+    Option+      { optionLongName = "binary-output",+        optionShortName = Just 'b',+        optionArgument = NoArgument,+        optionDescription = "Print the program result in the binary output format.",+        optionAction = [C.cstm|binary_output = 1;|]+      },+    Option+      { optionLongName = "help",+        optionShortName = Just 'h',+        optionArgument = NoArgument,+        optionDescription = "Print help information and exit.",+        optionAction =+          [C.cstm|{+                   printf("Usage: %s [OPTION]...\nOptions:\n\n%s\nFor more information, consult the Futhark User's Guide or the man pages.\n",+                          fut_progname, option_descriptions);+                   exit(0);+                  }|]+      }+  ]+  where+    set_runtime_file =+      [C.cstm|{+          runtime_file = fopen(optarg, "w");+          if (runtime_file == NULL) {+            futhark_panic(1, "Cannot open %s: %s\n", optarg, strerror(errno));+          }+        }|]+    set_num_runs =+      [C.cstm|{+          num_runs = atoi(optarg);+          perform_warmup = 1;+          if (num_runs <= 0) {+            futhark_panic(1, "Need a positive number of runs, not %s\n", optarg);+          }+        }|]++valueDescToCType :: ValueDesc -> C.Type+valueDescToCType (ScalarValue pt signed _) =+  signedPrimTypeToCType signed pt+valueDescToCType (ArrayValue _ _ pt signed shape) =+  let name = "futhark_" ++ arrayName pt signed (length shape)+   in [C.cty|struct $id:name|]++opaqueToCType :: String -> [ValueDesc] -> C.Type+opaqueToCType desc vds =+  let name = "futhark_" ++ opaqueName desc vds+   in [C.cty|struct $id:name|]++externalValueToCType :: ExternalValue -> C.Type+externalValueToCType (TransparentValue vd) = valueDescToCType vd+externalValueToCType (OpaqueValue desc vds) = opaqueToCType desc vds++primTypeInfo :: PrimType -> Signedness -> C.Exp+primTypeInfo (IntType it) t = case (it, t) of+  (Int8, TypeUnsigned) -> [C.cexp|u8_info|]+  (Int16, TypeUnsigned) -> [C.cexp|u16_info|]+  (Int32, TypeUnsigned) -> [C.cexp|u32_info|]+  (Int64, TypeUnsigned) -> [C.cexp|u64_info|]+  (Int8, _) -> [C.cexp|i8_info|]+  (Int16, _) -> [C.cexp|i16_info|]+  (Int32, _) -> [C.cexp|i32_info|]+  (Int64, _) -> [C.cexp|i64_info|]+primTypeInfo (FloatType Float32) _ = [C.cexp|f32_info|]+primTypeInfo (FloatType Float64) _ = [C.cexp|f64_info|]+primTypeInfo Bool _ = [C.cexp|bool_info|]+primTypeInfo Cert _ = [C.cexp|bool_info|]++readPrimStm :: C.ToIdent a => a -> Int -> PrimType -> Signedness -> C.Stm+readPrimStm place i t ept =+  [C.cstm|if (read_scalar(stdin, &$exp:(primTypeInfo t ept), &$id:place) != 0) {+            futhark_panic(1, "Error when reading input #%d of type %s (errno: %s).\n",+                          $int:i,+                          $exp:(primTypeInfo t ept).type_name,+                          strerror(errno));+          }|]++readInput :: Int -> ExternalValue -> ([C.BlockItem], C.Stm, C.Stm, C.Stm, C.Exp)+readInput i (OpaqueValue desc _) =+  ( [C.citems|futhark_panic(1, "Cannot read input #%d of type %s\n", $int:i, $string:desc);|],+    [C.cstm|;|],+    [C.cstm|;|],+    [C.cstm|;|],+    [C.cexp|NULL|]+  )+readInput i (TransparentValue (ScalarValue t ept _)) =+  let dest = "read_value_" ++ show i+   in ( [C.citems|$ty:(primTypeToCType t) $id:dest;+                  $stm:(readPrimStm dest i t ept);|],+        [C.cstm|;|],+        [C.cstm|;|],+        [C.cstm|;|],+        [C.cexp|$id:dest|]+      )+readInput i (TransparentValue (ArrayValue _ _ t ept dims)) =+  let dest = "read_value_" ++ show i+      shape = "read_shape_" ++ show i+      arr = "read_arr_" ++ show i++      name = arrayName t ept rank+      arr_ty_name = "futhark_" ++ name+      ty = [C.cty|struct $id:arr_ty_name|]+      rank = length dims+      dims_exps = [[C.cexp|$id:shape[$int:j]|] | j <- [0 .. rank -1]]+      dims_s = concat $ replicate rank "[]"+      t' = signedPrimTypeToCType ept t++      new_array = "futhark_new_" ++ name+      free_array = "futhark_free_" ++ name++      items =+        [C.citems|+           $ty:ty *$id:dest;+           typename int64_t $id:shape[$int:rank];+           $ty:t' *$id:arr = NULL;+           errno = 0;+           if (read_array(stdin,+                          &$exp:(primTypeInfo t ept),+                          (void**) &$id:arr,+                          $id:shape,+                          $int:(length dims))+               != 0) {+             futhark_panic(1, "Cannot read input #%d of type %s%s (errno: %s).\n",+                           $int:i,+                           $string:dims_s,+                           $exp:(primTypeInfo t ept).type_name,+                           strerror(errno));+           }|]+   in ( items,+        [C.cstm|assert(($id:dest = $id:new_array(ctx, $id:arr, $args:dims_exps)) != NULL);|],+        [C.cstm|assert($id:free_array(ctx, $id:dest) == 0);|],+        [C.cstm|free($id:arr);|],+        [C.cexp|$id:dest|]+      )++readInputs :: [ExternalValue] -> [([C.BlockItem], C.Stm, C.Stm, C.Stm, C.Exp)]+readInputs = zipWith readInput [0 ..]++prepareOutputs :: [ExternalValue] -> [(C.BlockItem, C.Exp, C.Stm)]+prepareOutputs = zipWith prepareResult [(0 :: Int) ..]+  where+    prepareResult i ev = do+      let ty = externalValueToCType ev+          result = "result_" ++ show i++      case ev of+        TransparentValue ScalarValue {} ->+          ( [C.citem|$ty:ty $id:result;|],+            [C.cexp|$id:result|],+            [C.cstm|;|]+          )+        TransparentValue (ArrayValue _ _ t ept dims) ->+          let name = arrayName t ept $ length dims+              free_array = "futhark_free_" ++ name+           in ( [C.citem|$ty:ty *$id:result;|],+                [C.cexp|$id:result|],+                [C.cstm|assert($id:free_array(ctx, $id:result) == 0);|]+              )+        OpaqueValue desc vds ->+          let free_opaque = "futhark_free_" ++ opaqueName desc vds+           in ( [C.citem|$ty:ty *$id:result;|],+                [C.cexp|$id:result|],+                [C.cstm|assert($id:free_opaque(ctx, $id:result) == 0);|]+              )++printPrimStm :: (C.ToExp a, C.ToExp b) => a -> b -> PrimType -> Signedness -> C.Stm+printPrimStm dest val bt ept =+  [C.cstm|write_scalar($exp:dest, binary_output, &$exp:(primTypeInfo bt ept), &$exp:val);|]++-- | Return a statement printing the given external value.+printStm :: ExternalValue -> C.Exp -> C.Stm+printStm (OpaqueValue desc _) _ =+  [C.cstm|printf("#<opaque %s>", $string:desc);|]+printStm (TransparentValue (ScalarValue bt ept _)) e =+  printPrimStm [C.cexp|stdout|] e bt ept+printStm (TransparentValue (ArrayValue _ _ bt ept shape)) e =+  let values_array = "futhark_values_" ++ name+      shape_array = "futhark_shape_" ++ name+      num_elems = cproduct [[C.cexp|$id:shape_array(ctx, $exp:e)[$int:i]|] | i <- [0 .. rank -1]]+   in [C.cstm|{+        $ty:bt' *arr = calloc(sizeof($ty:bt'), $exp:num_elems);+        assert(arr != NULL);+        assert($id:values_array(ctx, $exp:e, arr) == 0);+        write_array(stdout, binary_output, &$exp:(primTypeInfo bt ept), arr,+                    $id:shape_array(ctx, $exp:e), $int:rank);+        free(arr);+      }|]+  where+    rank = length shape+    bt' = primTypeToCType bt+    name = arrayName bt ept rank++printResult :: [(ExternalValue, C.Exp)] -> [C.Stm]+printResult = concatMap f+  where+    f (v, e) = [printStm v e, [C.cstm|printf("\n");|]]++cliEntryPoint ::+  Name ->+  FunctionT a ->+  (C.Definition, C.Initializer)+cliEntryPoint fname (Function _ _ _ _ results args) =+  let (input_items, pack_input, free_input, free_parsed, input_args) =+        unzip5 $ readInputs args++      (output_decls, output_vals, free_outputs) =+        unzip3 $ prepareOutputs results++      printstms = printResult $ zip results output_vals++      ctx_ty = [C.cty|struct futhark_context|]+      sync_ctx = "futhark_context_sync" :: Name+      error_ctx = "futhark_context_get_error" :: Name++      entry_point_name = nameToString fname+      cli_entry_point_function_name = "futrts_cli_entry_" ++ entry_point_name+      entry_point_function_name = "futhark_entry_" ++ entry_point_name++      pause_profiling = "futhark_context_pause_profiling" :: Name+      unpause_profiling = "futhark_context_unpause_profiling" :: Name++      addrOf e = [C.cexp|&$exp:e|]++      run_it =+        [C.citems|+                  int r;+                  // Run the program once.+                  $stms:pack_input+                  if ($id:sync_ctx(ctx) != 0) {+                    futhark_panic(1, "%s", $id:error_ctx(ctx));+                  };+                  // Only profile last run.+                  if (profile_run) {+                    $id:unpause_profiling(ctx);+                  }+                  t_start = get_wall_time();+                  r = $id:entry_point_function_name(ctx,+                                                    $args:(map addrOf output_vals),+                                                    $args:input_args);+                  if (r != 0) {+                    futhark_panic(1, "%s", $id:error_ctx(ctx));+                  }+                  if ($id:sync_ctx(ctx) != 0) {+                    futhark_panic(1, "%s", $id:error_ctx(ctx));+                  };+                  if (profile_run) {+                    $id:pause_profiling(ctx);+                  }+                  t_end = get_wall_time();+                  long int elapsed_usec = t_end - t_start;+                  if (time_runs && runtime_file != NULL) {+                    fprintf(runtime_file, "%lld\n", (long long) elapsed_usec);+                    fflush(runtime_file);+                  }+                  $stms:free_input+                |]+   in ( [C.cedecl|+  static void $id:cli_entry_point_function_name($ty:ctx_ty *ctx) {+    typename int64_t t_start, t_end;+    int time_runs = 0, profile_run = 0;++    // We do not want to profile all the initialisation.+    $id:pause_profiling(ctx);++    // Declare and read input.+    set_binary_mode(stdin);+    $items:(mconcat input_items)++    if (end_of_input(stdin) != 0) {+      futhark_panic(1, "Expected EOF on stdin after reading input for %s.\n", $string:(quote (pretty fname)));+    }++    $items:output_decls++    // Warmup run+    if (perform_warmup) {+      $items:run_it+      $stms:free_outputs+    }+    time_runs = 1;+    // Proper run.+    for (int run = 0; run < num_runs; run++) {+      // Only profile last run.+      profile_run = run == num_runs -1;+      $items:run_it+      if (run < num_runs-1) {+        $stms:free_outputs+      }+    }++    // Free the parsed input.+    $stms:free_parsed++    // Print the final result.+    if (binary_output) {+      set_binary_mode(stdout);+    }+    $stms:printstms++    $stms:free_outputs+  }|],+        [C.cinit|{ .name = $string:entry_point_name,+                      .fun = $id:cli_entry_point_function_name }|]+      )++{-# NOINLINE cliDefs #-}+cliDefs :: [Option] -> Functions a -> [C.Definition]+cliDefs options (Functions funs) =+  let values_h = $(embedStringFile "rts/c/values.h")+      tuning_h = $(embedStringFile "rts/c/tuning.h")++      option_parser =+        generateOptionParser "parse_options" $ genericOptions ++ options+      (cli_entry_point_decls, entry_point_inits) =+        unzip $ map (uncurry cliEntryPoint) funs+   in [C.cunit|+$esc:("#include <string.h>")+$esc:("#include <inttypes.h>")+$esc:("#include <errno.h>")+$esc:("#include <ctype.h>")+$esc:("#include <errno.h>")+$esc:("#include <getopt.h>")++$esc:values_h++static int binary_output = 0;+static typename FILE *runtime_file;+static int perform_warmup = 0;+static int num_runs = 1;+// If the entry point is NULL, the program will terminate after doing initialisation and such.+static const char *entry_point = "main";++$esc:tuning_h++$func:option_parser++$edecls:cli_entry_point_decls++typedef void entry_point_fun(struct futhark_context*);++struct entry_point_entry {+  const char *name;+  entry_point_fun *fun;+};++int main(int argc, char** argv) {+  fut_progname = argv[0];++  struct entry_point_entry entry_points[] = {+    $inits:entry_point_inits+  };++  struct futhark_context_config *cfg = futhark_context_config_new();+  assert(cfg != NULL);++  int parsed_options = parse_options(cfg, argc, argv);+  argc -= parsed_options;+  argv += parsed_options;++  if (argc != 0) {+    futhark_panic(1, "Excess non-option: %s\n", argv[0]);+  }++  struct futhark_context *ctx = futhark_context_new(cfg);+  assert (ctx != NULL);++  char* error = futhark_context_get_error(ctx);+  if (error != NULL) {+    futhark_panic(1, "%s", error);+  }++  if (entry_point != NULL) {+    int num_entry_points = sizeof(entry_points) / sizeof(entry_points[0]);+    entry_point_fun *entry_point_fun = NULL;+    for (int i = 0; i < num_entry_points; i++) {+      if (strcmp(entry_points[i].name, entry_point) == 0) {+        entry_point_fun = entry_points[i].fun;+        break;+      }+    }++    if (entry_point_fun == NULL) {+      fprintf(stderr, "No entry point '%s'.  Select another with --entry-point.  Options are:\n",+                      entry_point);+      for (int i = 0; i < num_entry_points; i++) {+        fprintf(stderr, "%s\n", entry_points[i].name);+      }+      return 1;+    }++    entry_point_fun(ctx);++    if (runtime_file != NULL) {+      fclose(runtime_file);+    }++    char *report = futhark_context_report(ctx);+    fputs(report, stderr);+    free(report);+  }++  futhark_context_free(ctx);+  futhark_context_config_free(cfg);+  return 0;+}|]

@@ -53,7 +53,7 @@ import Futhark.CodeGen.Backends.GenericPython.Options import qualified Futhark.CodeGen.ImpCode as Imp import Futhark.IR.Primitive hiding (Bool)-import Futhark.IR.Prop (isBuiltInFunction)+import Futhark.IR.Prop (isBuiltInFunction, subExpVars) import Futhark.IR.Syntax (Space (..)) import Futhark.MonadFreshNames import Futhark.Util (zEncodeString)@@ -459,9 +459,9 @@ compileName :: VName -> String compileName = zEncodeString . pretty -compileDim :: Imp.DimSize -> PyExp-compileDim (Imp.Constant v) = compilePrimValue v-compileDim (Imp.Var v) = Var $ compileName v+compileDim :: Imp.DimSize -> CompilerM op s PyExp+compileDim (Imp.Constant v) = pure $ compilePrimValue v+compileDim (Imp.Var v) = compileVar v  unpackDim :: PyExp -> Imp.DimSize -> Int32 -> CompilerM op s () unpackDim arr_name (Imp.Constant c) i = do@@ -470,12 +470,18 @@   let constant_i = Integer $ toInteger i   stm $     Assert (BinOp "==" constant_c (Index shape_name $ IdxExp constant_i)) $-      String "constant dimension wrong"+      String "Entry point arguments have invalid sizes." unpackDim arr_name (Imp.Var var) i = do   let shape_name = Field arr_name "shape"       src = Index shape_name $ IdxExp $ Integer $ toInteger i   var' <- compileVar var-  stm $ Assign var' $ simpleCall "np.int64" [src]+  stm $+    If+      (BinOp "==" var' None)+      [Assign var' $ simpleCall "np.int64" [src]]+      [ Assert (BinOp "==" var' src) $+          String "Error: entry point arguments have invalid sizes."+      ]  entryPointOutput :: Imp.ExternalValue -> CompilerM op s PyExp entryPointOutput (Imp.OpaqueValue desc vs) =@@ -492,7 +498,8 @@ entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ bt ept dims)) = do   mem' <- compileVar mem   let cast = Cast mem' (compilePrimTypeExt bt ept)-  return $ simpleCall "createArray" [cast, Tuple $ map compileDim dims]+  dims' <- mapM compileDim dims+  return $ simpleCall "createArray" [cast, Tuple dims']  badInput :: Int -> PyExp -> String -> PyStmt badInput i e t =@@ -551,6 +558,15 @@           "Actual Futhark type: {}"         ] +declEntryPointInputSizes :: [Imp.ExternalValue] -> CompilerM op s ()+declEntryPointInputSizes = mapM_ onSize . concatMap sizes+  where+    sizes (Imp.TransparentValue v) = valueSizes v+    sizes (Imp.OpaqueValue _ vs) = concatMap valueSizes vs+    valueSizes (Imp.ArrayValue _ _ _ _ dims) = subExpVars dims+    valueSizes Imp.ScalarValue {} = []+    onSize v = stm $ Assign (Var (compileName v)) None+ entryPointInput :: (Int, Imp.ExternalValue, PyExp) -> CompilerM op s () entryPointInput (i, Imp.OpaqueValue desc vs, e) = do   let type_is_ok =@@ -750,11 +766,11 @@         return $ Just $ compileName name'       _ -> return Nothing -  prepareIn <--    collect $-      mapM_ entryPointInput $-        zip3 [0 ..] args $-          map (Var . extValueDescName) args+  prepareIn <- collect $ do+    declEntryPointInputSizes args+    mapM_ entryPointInput $+      zip3 [0 ..] args $+        map (Var . extValueDescName) args   (res, prepareOut) <- collect' $ mapM entryPointOutput results    let argexps_lib = map (compileName . Imp.paramName) inputs

@@ -207,7 +207,6 @@ operations =   GC.defaultOperations     { GC.opsCompiler = compileOp,-      GC.opsCopy = copyMulticoreMemory,       GC.opsCritical =         -- The thread entering an API function is always considered         -- the "first worker" - note that this might differ from the@@ -218,12 +217,6 @@           []         )     }--copyMulticoreMemory :: GC.Copy Multicore ()-copyMulticoreMemory destmem destidx DefaultSpace srcmem srcidx DefaultSpace nbytes =-  GC.copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes-copyMulticoreMemory _ _ destspace _ _ srcspace _ =-  error $ "Cannot copy to " ++ show destspace ++ " from " ++ show srcspace  closureFreeStructField :: VName -> Name closureFreeStructField v =

@@ -410,11 +410,12 @@ packArrayOutput :: Py.EntryOutput Imp.OpenCL () packArrayOutput mem "device" bt ept dims = do   mem' <- Py.compileVar mem+  dims' <- mapM Py.compileDim dims   return $     Call       (Var "cl.array.Array")       [ Arg $ Var "self.queue",-        Arg $ Tuple $ map Py.compileDim dims,+        Arg $ Tuple dims',         Arg $ Var $ Py.compilePrimTypeExt bt ept,         ArgKeyword "data" mem'       ]

@@ -29,8 +29,7 @@     operations :: GC.Operations Imp.Sequential ()     operations =       GC.defaultOperations-        { GC.opsCompiler = const $ return (),-          GC.opsCopy = copySequentialMemory+        { GC.opsCompiler = const $ return ()         }      generateContext = do@@ -130,9 +129,3 @@                                  return 0;                                }|]         )--copySequentialMemory :: GC.Copy Imp.Sequential ()-copySequentialMemory destmem destidx DefaultSpace srcmem srcidx DefaultSpace nbytes =-  GC.copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes-copySequentialMemory _ _ destspace _ _ srcspace _ =-  error $ "Cannot copy to " ++ show destspace ++ " from " ++ show srcspace

@@ -1,5 +1,6 @@ {-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE Trustworthy #-}+{-# OPTIONS_GHC -fno-warn-orphans #-}  -- | Simple C runtime representation. module Futhark.CodeGen.Backends.SimpleRep@@ -9,6 +10,9 @@     intTypeToCType,     primTypeToCType,     signedPrimTypeToCType,+    arrayName,+    opaqueName,+    cproduct,      -- * Primitive value operations     cIntOps,@@ -20,11 +24,14 @@   ) where +import Data.Bits (shiftR, xor)+import Data.Char (isAlphaNum, isDigit, ord) import Futhark.CodeGen.ImpCode import Futhark.Util (zEncodeString) import Futhark.Util.Pretty (prettyOneLine) import qualified Language.C.Quote.C as C import qualified Language.C.Syntax as C+import Text.Printf  -- | The C type corresponding to a signed integer type. intTypeToCType :: IntType -> C.Type@@ -76,6 +83,78 @@ defaultMemBlockType :: C.Type defaultMemBlockType = [C.cty|char*|] +-- | The name of exposed array type structs.+arrayName :: PrimType -> Signedness -> Int -> String+arrayName pt signed rank =+  prettySigned (signed == TypeUnsigned) pt ++ "_" ++ show rank ++ "d"++-- | The name of exposed opaque types.+opaqueName :: String -> [ValueDesc] -> String+opaqueName s _+  | valid = "opaque_" ++ s+  where+    valid =+      head s /= '_'+        && not (isDigit $ head s)+        && all ok s+    ok c = isAlphaNum c || c == '_'+opaqueName s vds = "opaque_" ++ hash (zipWith xor [0 ..] $ map ord (s ++ concatMap p vds))+  where+    p (ScalarValue pt signed _) =+      show (pt, signed)+    p (ArrayValue _ space pt signed dims) =+      show (space, pt, signed, length dims)++    -- FIXME: a stupid hash algorithm; may have collisions.+    hash =+      printf "%x" . foldl xor 0+        . map+          ( iter . (* 0x45d9f3b)+              . iter+              . (* 0x45d9f3b)+              . iter+              . fromIntegral+          )+    iter x = ((x :: Word32) `shiftR` 16) `xor` x++-- | Return an expression multiplying together the given expressions.+-- If an empty list is given, the expression @1@ is returned.+cproduct :: [C.Exp] -> C.Exp+cproduct [] = [C.cexp|1|]+cproduct (e : es) = foldl mult e es+  where+    mult x y = [C.cexp|$exp:x * $exp:y|]++instance C.ToIdent Name where+  toIdent = C.toIdent . zEncodeString . nameToString++instance C.ToIdent VName where+  toIdent = C.toIdent . zEncodeString . pretty++instance C.ToExp VName where+  toExp v _ = [C.cexp|$id:v|]++instance C.ToExp IntValue where+  toExp (Int8Value v) = C.toExp v+  toExp (Int16Value v) = C.toExp v+  toExp (Int32Value v) = C.toExp v+  toExp (Int64Value v) = C.toExp v++instance C.ToExp FloatValue where+  toExp (Float32Value v) = C.toExp v+  toExp (Float64Value v) = C.toExp v++instance C.ToExp PrimValue where+  toExp (IntValue v) = C.toExp v+  toExp (FloatValue v) = C.toExp v+  toExp (BoolValue True) = C.toExp (1 :: Int8)+  toExp (BoolValue False) = C.toExp (0 :: Int8)+  toExp Checked = C.toExp (1 :: Int8)++instance C.ToExp SubExp where+  toExp (Var v) = C.toExp v+  toExp (Constant c) = C.toExp c+ cIntOps :: [C.Definition] cIntOps =   concatMap (`map` [minBound .. maxBound]) ops@@ -509,7 +588,7 @@      return x == 0 ? 32 :  __builtin_ctz(x);    }    static typename int32_t $id:(funName' "ctz64") (typename int64_t x) {-     return x == 0 ? 64 : __builtin_ctzl(x);+     return x == 0 ? 64 : __builtin_ctzll(x);    } $esc:("#endif")                 |]

@@ -82,7 +82,7 @@   = ScalarUse VName PrimType   | MemoryUse VName   | ConstUse VName KernelConstExp-  deriving (Eq, Show)+  deriving (Eq, Ord, Show)  instance Pretty KernelConst where   ppr (SizeConst key) = text "get_size" <> parens (ppr key)

@@ -498,9 +498,6 @@  data ArrayDecl = ArrayDecl VName PrimType MemLocation -fparamSizes :: Typed dec => Param dec -> S.Set VName-fparamSizes = S.fromList . subExpVars . arrayDims . paramType- compileInParams ::   Mem lore =>   [FParam lore] ->@@ -511,7 +508,6 @@         splitAt (length params - sum (map entryPointSize orig_epts)) params   (inparams, arrayds) <- partitionEithers <$> mapM compileInParam (ctx_params ++ val_params)   let findArray x = find (isArrayDecl x) arrayds-      sizes = mconcat $ map fparamSizes $ ctx_params ++ val_params        summaries = M.fromList $ mapMaybe memSummary params         where@@ -529,11 +525,8 @@           (Just (ArrayDecl _ bt (MemLocation mem shape _)), _) -> do             memspace <- findMemInfo mem             Just $ Imp.ArrayValue mem memspace bt signedness shape-          (_, Prim bt)-            | paramName fparam `S.member` sizes ->-              Nothing-            | otherwise ->-              Just $ Imp.ScalarValue bt signedness $ paramName fparam+          (_, Prim bt) ->+            Just $ Imp.ScalarValue bt signedness $ paramName fparam           _ ->             Nothing 

@@ -34,8 +34,8 @@ import Futhark.IR.KernelsMem import qualified Futhark.IR.Mem.IxFun as IxFun import Futhark.MonadFreshNames-import Futhark.Util.IntegralExp (IntegralExp, divUp, quot)-import Prelude hiding (quot)+import Futhark.Util.IntegralExp (IntegralExp, divUp, quot, rem)+import Prelude hiding (quot, rem)  callKernelOperations :: Operations KernelsMem HostEnv Imp.HostOp callKernelOperations =@@ -76,8 +76,8 @@ compileProgOpenCL,   compileProgCUDA ::     MonadFreshNames m => Prog KernelsMem -> m (Warnings, Imp.Program)-compileProgOpenCL = compileProg $ HostEnv openclAtomics-compileProgCUDA = compileProg $ HostEnv cudaAtomics+compileProgOpenCL = compileProg $ HostEnv openclAtomics OpenCL+compileProgCUDA = compileProg $ HostEnv cudaAtomics CUDA  opCompiler ::   Pattern KernelsMem ->@@ -155,7 +155,7 @@ checkLocalMemoryReqs :: Imp.Code -> CallKernelGen (Maybe (Imp.TExp Bool)) checkLocalMemoryReqs code = do   scope <- askScope-  let alloc_sizes = map (sum . localAllocSizes . Imp.kernelBody) $ getKernels code+  let alloc_sizes = map (sum . map alignedSize . localAllocSizes . Imp.kernelBody) $ getKernels code    -- If any of the sizes involve a variable that is not known at this   -- point, then we cannot check the requirements.@@ -178,6 +178,11 @@     localAllocSizes = foldMap localAllocSize     localAllocSize (Imp.LocalAlloc _ size) = [size]     localAllocSize _ = []++    -- These allocations will actually be padded to an 8-byte aligned+    -- size, so we should take that into account when checking whether+    -- they fit.+    alignedSize x = x + ((8 - (x `rem` 8)) `rem` 8)  expCompiler :: ExpCompiler KernelsMem HostEnv Imp.HostOp -- We generate a simple kernel for itoa and replicate.

@@ -8,6 +8,7 @@     CallKernelGen,     InKernelGen,     HostEnv (..),+    Target (..),     KernelEnv (..),     computeThreadChunkSize,     groupReduce,@@ -34,7 +35,7 @@ where  import Control.Monad.Except-import Data.List (elemIndex, find, nub, zip4)+import Data.List (elemIndex, find, zip4) import qualified Data.Map.Strict as M import Data.Maybe import qualified Data.Set as S@@ -45,13 +46,21 @@ import qualified Futhark.IR.Mem.IxFun as IxFun import Futhark.MonadFreshNames import Futhark.Transform.Rename-import Futhark.Util (chunks, dropLast, mapAccumLM, maybeNth, takeLast)+import Futhark.Util (chunks, dropLast, mapAccumLM, maybeNth, nubOrd, takeLast) import Futhark.Util.IntegralExp (divUp, quot, rem) import Prelude hiding (quot, rem) -newtype HostEnv = HostEnv-  {hostAtomics :: AtomicBinOp}+-- | Which target are we ultimately generating code for?  While most+-- of the kernels code is the same, there are some cases where we+-- generate special code based on the ultimate low-level API we are+-- targeting.+data Target = CUDA | OpenCL +data HostEnv = HostEnv+  { hostAtomics :: AtomicBinOp,+    hostTarget :: Target+  }+ data KernelEnv = KernelEnv   { kernelAtomics :: AtomicBinOp,     kernelConstants :: KernelConstants@@ -301,7 +310,13 @@ whenActive :: SegLevel -> SegSpace -> InKernelGen () -> InKernelGen () whenActive lvl space m   | SegNoVirtFull <- segVirt lvl = m-  | otherwise = sWhen (isActive $ unSegSpace space) m+  | otherwise = do+    group_size <- kernelGroupSize . kernelConstants <$> askEnv+    -- XXX: the following check is too naive - we should also handle+    -- the multi-dimensional case.+    if [group_size] == map (toInt64Exp . snd) (unSegSpace space)+      then m+      else sWhen (isActive $ unSegSpace space) m  compileGroupOp :: OpCompiler KernelsMem KernelEnv Imp.KernelOp compileGroupOp pat (Alloc size space) =@@ -734,7 +749,7 @@ computeKernelUses kernel_body bound_in_kernel = do   let actually_free = freeIn kernel_body `namesSubtract` namesFromList bound_in_kernel   -- Compute the variables that we need to pass to the kernel.-  nub <$> readsFromSet actually_free+  nubOrd <$> readsFromSet actually_free  readsFromSet :: Names -> CallKernelGen [Imp.KernelUse] readsFromSet free =@@ -1316,7 +1331,7 @@   InKernelGen () ->   CallKernelGen () sKernelFailureTolerant tol ops constants name m = do-  HostEnv atomics <- askEnv+  HostEnv atomics _ <- askEnv   body <- makeAllMemoryGlobal $ subImpM_ (KernelEnv atomics constants) ops m   uses <- computeKernelUses body mempty   emit $

@@ -1,479 +1,48 @@-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE TypeFamilies #-}---- | Code generation for segmented and non-segmented scans.  Uses a--- fairly inefficient two-pass algorithm. module Futhark.CodeGen.ImpGen.Kernels.SegScan (compileSegScan) where -import Control.Monad.Except-import Control.Monad.State-import Data.List (delete, find, foldl', zip4)-import Data.Maybe import qualified Futhark.CodeGen.ImpCode.Kernels as Imp-import Futhark.CodeGen.ImpGen+import Futhark.CodeGen.ImpGen hiding (compileProg) import Futhark.CodeGen.ImpGen.Kernels.Base+import qualified Futhark.CodeGen.ImpGen.Kernels.SegScan.SinglePass as SinglePass+import qualified Futhark.CodeGen.ImpGen.Kernels.SegScan.TwoPass as TwoPass import Futhark.IR.KernelsMem-import qualified Futhark.IR.Mem.IxFun as IxFun-import Futhark.Transform.Rename-import Futhark.Util (takeLast)-import Futhark.Util.IntegralExp (divUp, quot, rem)-import Prelude hiding (quot, rem) --- Aggressively try to reuse memory for different SegBinOps, because--- we will run them sequentially after another.-makeLocalArrays ::-  Count GroupSize SubExp ->-  SubExp ->-  [SegBinOp KernelsMem] ->-  InKernelGen [[VName]]-makeLocalArrays (Count group_size) num_threads scans = do-  (arrs, mems_and_sizes) <- runStateT (mapM onScan scans) mempty-  let maxSize sizes = Imp.bytes $ foldl' sMax64 1 $ map Imp.unCount sizes-  forM_ mems_and_sizes $ \(sizes, mem) ->-    sAlloc_ mem (maxSize sizes) (Space "local")-  return arrs-  where-    onScan (SegBinOp _ scan_op nes _) = do-      let (scan_x_params, _scan_y_params) =-            splitAt (length nes) $ lambdaParams scan_op-      (arrs, used_mems) <- fmap unzip $-        forM scan_x_params $ \p ->-          case paramDec p of-            MemArray pt shape _ (ArrayIn mem _) -> do-              let shape' = Shape [num_threads] <> shape-              arr <--                lift $-                  sArray "scan_arr" pt shape' $-                    ArrayIn mem $ IxFun.iota $ map pe64 $ shapeDims shape'-              return (arr, [])-            _ -> do-              let pt = elemType $ paramType p-                  shape = Shape [group_size]-              (sizes, mem') <- getMem pt shape-              arr <- lift $ sArrayInMem "scan_arr" pt shape mem'-              return (arr, [(sizes, mem')])-      modify (<> concat used_mems)-      return arrs--    getMem pt shape = do-      let size = typeSize $ Array pt shape NoUniqueness-      mems <- get-      case (find ((size `elem`) . fst) mems, mems) of-        (Just mem, _) -> do-          modify $ delete mem-          return mem-        (Nothing, (size', mem) : mems') -> do-          put mems'-          return (size : size', mem)-        (Nothing, []) -> do-          mem <- lift $ sDeclareMem "scan_arr_mem" $ Space "local"-          return ([size], mem)--type CrossesSegment = Maybe (Imp.TExp Int64 -> Imp.TExp Int64 -> Imp.TExp Bool)--localArrayIndex :: KernelConstants -> Type -> Imp.TExp Int64-localArrayIndex constants t =-  if primType t-    then sExt64 (kernelLocalThreadId constants)-    else sExt64 (kernelGlobalThreadId constants)--barrierFor :: Lambda KernelsMem -> (Bool, Imp.Fence, InKernelGen ())-barrierFor scan_op = (array_scan, fence, sOp $ Imp.Barrier fence)+-- The single-pass scan does not support multiple operators, so jam+-- them together here.+combineScans :: [SegBinOp KernelsMem] -> SegBinOp KernelsMem+combineScans ops =+  SegBinOp+    { segBinOpComm = mconcat (map segBinOpComm ops),+      segBinOpLambda = lam',+      segBinOpNeutral = concatMap segBinOpNeutral ops,+      segBinOpShape = mempty -- Assumed+    }   where-    array_scan = not $ all primType $ lambdaReturnType scan_op-    fence-      | array_scan = Imp.FenceGlobal-      | otherwise = Imp.FenceLocal--xParams, yParams :: SegBinOp KernelsMem -> [LParam KernelsMem]-xParams scan =-  take (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan))-yParams scan =-  drop (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan))--writeToScanValues ::-  [VName] ->-  ([PatElem KernelsMem], SegBinOp KernelsMem, [KernelResult]) ->-  InKernelGen ()-writeToScanValues gtids (pes, scan, scan_res)-  | shapeRank (segBinOpShape scan) > 0 =-    forM_ (zip pes scan_res) $ \(pe, res) ->-      copyDWIMFix-        (patElemName pe)-        (map Imp.vi64 gtids)-        (kernelResultSubExp res)-        []-  | otherwise =-    forM_ (zip (yParams scan) scan_res) $ \(p, res) ->-      copyDWIMFix (paramName p) [] (kernelResultSubExp res) []--readToScanValues ::-  [Imp.TExp Int64] ->-  [PatElem KernelsMem] ->-  SegBinOp KernelsMem ->-  InKernelGen ()-readToScanValues is pes scan-  | shapeRank (segBinOpShape scan) > 0 =-    forM_ (zip (yParams scan) pes) $ \(p, pe) ->-      copyDWIMFix (paramName p) [] (Var (patElemName pe)) is-  | otherwise =-    return ()+    lams = map segBinOpLambda ops+    xParams lam = take (length (lambdaReturnType lam)) (lambdaParams lam)+    yParams lam = drop (length (lambdaReturnType lam)) (lambdaParams lam)+    lam' =+      Lambda+        { lambdaParams = concatMap xParams lams ++ concatMap yParams lams,+          lambdaReturnType = concatMap lambdaReturnType lams,+          lambdaBody =+            Body+              ()+              (mconcat (map (bodyStms . lambdaBody) lams))+              (concatMap (bodyResult . lambdaBody) lams)+        } -readCarries ::-  Imp.TExp Int64 ->-  [Imp.TExp Int64] ->-  [Imp.TExp Int64] ->-  [PatElem KernelsMem] ->-  SegBinOp KernelsMem ->-  InKernelGen ()-readCarries chunk_offset dims' vec_is pes scan-  | shapeRank (segBinOpShape scan) > 0 = do-    ltid <- kernelLocalThreadId . kernelConstants <$> askEnv-    -- We may have to reload the carries from the output of the-    -- previous chunk.-    sIf-      (chunk_offset .>. 0 .&&. ltid .==. 0)-      ( do-          let is = unflattenIndex dims' $ chunk_offset - 1-          forM_ (zip (xParams scan) pes) $ \(p, pe) ->-            copyDWIMFix (paramName p) [] (Var (patElemName pe)) (is ++ vec_is)-      )-      ( forM_ (zip (xParams scan) (segBinOpNeutral scan)) $ \(p, ne) ->-          copyDWIMFix (paramName p) [] ne []-      )+canBeSinglePass :: SegSpace -> [SegBinOp KernelsMem] -> Maybe (SegBinOp KernelsMem)+canBeSinglePass space ops+  | [_] <- unSegSpace space,+    all ok ops =+    Just $ combineScans ops   | otherwise =-    return ()---- | Produce partially scanned intervals; one per workgroup.-scanStage1 ::-  Pattern KernelsMem ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->-  SegSpace ->-  [SegBinOp KernelsMem] ->-  KernelBody KernelsMem ->-  CallKernelGen (TV Int32, Imp.TExp Int64, CrossesSegment)-scanStage1 (Pattern _ all_pes) num_groups group_size space scans kbody = do-  let num_groups' = fmap toInt64Exp num_groups-      group_size' = fmap toInt64Exp group_size-  num_threads <- dPrimV "num_threads" $ sExt32 $ unCount num_groups' * unCount group_size'--  let (gtids, dims) = unzip $ unSegSpace space-      dims' = map toInt64Exp dims-  let num_elements = product dims'-      elems_per_thread = num_elements `divUp` sExt64 (tvExp num_threads)-      elems_per_group = unCount group_size' * elems_per_thread--  let crossesSegment =-        case reverse dims' of-          segment_size : _ : _ -> Just $ \from to ->-            (to - from) .>. (to `rem` segment_size)-          _ -> Nothing--  sKernelThread "scan_stage1" num_groups' group_size' (segFlat space) $ do-    constants <- kernelConstants <$> askEnv-    all_local_arrs <- makeLocalArrays group_size (tvSize num_threads) scans--    -- The variables from scan_op will be used for the carry and such-    -- in the big chunking loop.-    forM_ scans $ \scan -> do-      dScope Nothing $ scopeOfLParams $ lambdaParams $ segBinOpLambda scan-      forM_ (zip (xParams scan) (segBinOpNeutral scan)) $ \(p, ne) ->-        copyDWIMFix (paramName p) [] ne []--    sFor "j" elems_per_thread $ \j -> do-      chunk_offset <--        dPrimV "chunk_offset" $-          sExt64 (kernelGroupSize constants) * j-            + sExt64 (kernelGroupId constants) * elems_per_group-      flat_idx <--        dPrimV "flat_idx" $-          tvExp chunk_offset + sExt64 (kernelLocalThreadId constants)-      -- Construct segment indices.-      zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx--      let per_scan_pes = segBinOpChunks scans all_pes--          in_bounds =-            foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims'--          when_in_bounds = compileStms mempty (kernelBodyStms kbody) $ do-            let (all_scan_res, map_res) =-                  splitAt (segBinOpResults scans) $ kernelBodyResult kbody-                per_scan_res =-                  segBinOpChunks scans all_scan_res--            sComment "write to-scan values to parameters" $-              mapM_ (writeToScanValues gtids) $-                zip3 per_scan_pes scans per_scan_res--            sComment "write mapped values results to global memory" $-              forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(pe, se) ->-                copyDWIMFix-                  (patElemName pe)-                  (map Imp.vi64 gtids)-                  (kernelResultSubExp se)-                  []--      sComment "threads in bounds read input" $-        sWhen in_bounds when_in_bounds--      forM_ (zip3 per_scan_pes scans all_local_arrs) $-        \(pes, scan@(SegBinOp _ scan_op nes vec_shape), local_arrs) ->-          sComment "do one intra-group scan operation" $ do-            let rets = lambdaReturnType scan_op-                scan_x_params = xParams scan-                (array_scan, fence, barrier) = barrierFor scan_op--            when array_scan barrier--            sLoopNest vec_shape $ \vec_is -> do-              sComment "maybe restore some to-scan values to parameters, or read neutral" $-                sIf-                  in_bounds-                  ( do-                      readToScanValues (map Imp.vi64 gtids ++ vec_is) pes scan-                      readCarries (tvExp chunk_offset) dims' vec_is pes scan-                  )-                  ( forM_ (zip (yParams scan) (segBinOpNeutral scan)) $ \(p, ne) ->-                      copyDWIMFix (paramName p) [] ne []-                  )--              sComment "combine with carry and write to local memory" $-                compileStms mempty (bodyStms $ lambdaBody scan_op) $-                  forM_ (zip3 rets local_arrs (bodyResult $ lambdaBody scan_op)) $-                    \(t, arr, se) ->-                      copyDWIMFix arr [localArrayIndex constants t] se []--              let crossesSegment' = do-                    f <- crossesSegment-                    Just $ \from to ->-                      let from' = sExt64 from + tvExp chunk_offset-                          to' = sExt64 to + tvExp chunk_offset-                       in f from' to'--              sOp $ Imp.ErrorSync fence--              -- We need to avoid parameter name clashes.-              scan_op_renamed <- renameLambda scan_op-              groupScan-                crossesSegment'-                (sExt64 $ tvExp num_threads)-                (sExt64 $ kernelGroupSize constants)-                scan_op_renamed-                local_arrs--              sComment "threads in bounds write partial scan result" $-                sWhen in_bounds $-                  forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) ->-                    copyDWIMFix-                      (patElemName pe)-                      (map Imp.vi64 gtids ++ vec_is)-                      (Var arr)-                      [localArrayIndex constants t]--              barrier--              let load_carry =-                    forM_ (zip local_arrs scan_x_params) $ \(arr, p) ->-                      copyDWIMFix-                        (paramName p)-                        []-                        (Var arr)-                        [ if primType $ paramType p-                            then sExt64 (kernelGroupSize constants) - 1-                            else-                              (sExt64 (kernelGroupId constants) + 1)-                                * sExt64 (kernelGroupSize constants) - 1-                        ]-                  load_neutral =-                    forM_ (zip nes scan_x_params) $ \(ne, p) ->-                      copyDWIMFix (paramName p) [] ne []--              sComment "first thread reads last element as carry-in for next iteration" $ do-                crosses_segment <- dPrimVE "crosses_segment" $-                  case crossesSegment of-                    Nothing -> false-                    Just f ->-                      f-                        ( tvExp chunk_offset-                            + sExt64 (kernelGroupSize constants) -1-                        )-                        ( tvExp chunk_offset-                            + sExt64 (kernelGroupSize constants)-                        )-                should_load_carry <--                  dPrimVE "should_load_carry" $-                    kernelLocalThreadId constants .==. 0 .&&. bNot crosses_segment-                sWhen should_load_carry load_carry-                when array_scan barrier-                sUnless should_load_carry load_neutral--              barrier--  return (num_threads, elems_per_group, crossesSegment)--scanStage2 ::-  Pattern KernelsMem ->-  TV Int32 ->-  Imp.TExp Int64 ->-  Count NumGroups SubExp ->-  CrossesSegment ->-  SegSpace ->-  [SegBinOp KernelsMem] ->-  CallKernelGen ()-scanStage2 (Pattern _ all_pes) stage1_num_threads elems_per_group num_groups crossesSegment space scans = do-  let (gtids, dims) = unzip $ unSegSpace space-      dims' = map toInt64Exp dims--  -- Our group size is the number of groups for the stage 1 kernel.-  let group_size = Count $ unCount num_groups-      group_size' = fmap toInt64Exp group_size--  let crossesSegment' = do-        f <- crossesSegment-        Just $ \from to ->-          f-            ((sExt64 from + 1) * elems_per_group - 1)-            ((sExt64 to + 1) * elems_per_group - 1)--  sKernelThread "scan_stage2" 1 group_size' (segFlat space) $ do-    constants <- kernelConstants <$> askEnv-    per_scan_local_arrs <- makeLocalArrays group_size (tvSize stage1_num_threads) scans-    let per_scan_rets = map (lambdaReturnType . segBinOpLambda) scans-        per_scan_pes = segBinOpChunks scans all_pes--    flat_idx <--      dPrimV "flat_idx" $-        (sExt64 (kernelLocalThreadId constants) + 1) * elems_per_group - 1-    -- Construct segment indices.-    zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx--    forM_ (zip4 scans per_scan_local_arrs per_scan_rets per_scan_pes) $-      \(SegBinOp _ scan_op nes vec_shape, local_arrs, rets, pes) ->-        sLoopNest vec_shape $ \vec_is -> do-          let glob_is = map Imp.vi64 gtids ++ vec_is--              in_bounds =-                foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims'--              when_in_bounds = forM_ (zip3 rets local_arrs pes) $ \(t, arr, pe) ->-                copyDWIMFix-                  arr-                  [localArrayIndex constants t]-                  (Var $ patElemName pe)-                  glob_is--              when_out_of_bounds = forM_ (zip3 rets local_arrs nes) $ \(t, arr, ne) ->-                copyDWIMFix arr [localArrayIndex constants t] ne []-              (_, _, barrier) =-                barrierFor scan_op--          sComment "threads in bound read carries; others get neutral element" $-            sIf in_bounds when_in_bounds when_out_of_bounds--          barrier--          groupScan-            crossesSegment'-            (sExt64 $ tvExp stage1_num_threads)-            (sExt64 $ kernelGroupSize constants)-            scan_op-            local_arrs--          sComment "threads in bounds write scanned carries" $-            sWhen in_bounds $-              forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) ->-                copyDWIMFix-                  (patElemName pe)-                  glob_is-                  (Var arr)-                  [localArrayIndex constants t]--scanStage3 ::-  Pattern KernelsMem ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->-  Imp.TExp Int64 ->-  CrossesSegment ->-  SegSpace ->-  [SegBinOp KernelsMem] ->-  CallKernelGen ()-scanStage3 (Pattern _ all_pes) num_groups group_size elems_per_group crossesSegment space scans = do-  let num_groups' = fmap toInt64Exp num_groups-      group_size' = fmap toInt64Exp group_size-      (gtids, dims) = unzip $ unSegSpace space-      dims' = map toInt64Exp dims-  required_groups <--    dPrimVE "required_groups" $-      sExt32 $ product dims' `divUp` sExt64 (unCount group_size')--  sKernelThread "scan_stage3" num_groups' group_size' (segFlat space) $-    virtualiseGroups SegVirt required_groups $ \virt_group_id -> do-      constants <- kernelConstants <$> askEnv--      -- Compute our logical index.-      flat_idx <--        dPrimVE "flat_idx" $-          sExt64 virt_group_id * sExt64 (unCount group_size')-            + sExt64 (kernelLocalThreadId constants)-      zipWithM_ dPrimV_ gtids $ unflattenIndex dims' flat_idx--      -- Figure out which group this element was originally in.-      orig_group <- dPrimV "orig_group" $ flat_idx `quot` elems_per_group-      -- Then the index of the carry-in of the preceding group.-      carry_in_flat_idx <--        dPrimV "carry_in_flat_idx" $-          tvExp orig_group * elems_per_group - 1-      -- Figure out the logical index of the carry-in.-      let carry_in_idx = unflattenIndex dims' $ tvExp carry_in_flat_idx--      -- Apply the carry if we are not in the scan results for the first-      -- group, and are not the last element in such a group (because-      -- then the carry was updated in stage 2), and we are not crossing-      -- a segment boundary.-      let in_bounds =-            foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims'-          crosses_segment =-            fromMaybe false $-              crossesSegment-                <*> pure (tvExp carry_in_flat_idx)-                <*> pure flat_idx-          is_a_carry = flat_idx .==. (tvExp orig_group + 1) * elems_per_group - 1-          no_carry_in = tvExp orig_group .==. 0 .||. is_a_carry .||. crosses_segment--      let per_scan_pes = segBinOpChunks scans all_pes-      sWhen in_bounds $-        sUnless no_carry_in $-          forM_ (zip per_scan_pes scans) $-            \(pes, SegBinOp _ scan_op nes vec_shape) -> do-              dScope Nothing $ scopeOfLParams $ lambdaParams scan_op-              let (scan_x_params, scan_y_params) =-                    splitAt (length nes) $ lambdaParams scan_op--              sLoopNest vec_shape $ \vec_is -> do-                forM_ (zip scan_x_params pes) $ \(p, pe) ->-                  copyDWIMFix-                    (paramName p)-                    []-                    (Var $ patElemName pe)-                    (carry_in_idx ++ vec_is)--                forM_ (zip scan_y_params pes) $ \(p, pe) ->-                  copyDWIMFix-                    (paramName p)-                    []-                    (Var $ patElemName pe)-                    (map Imp.vi64 gtids ++ vec_is)--                compileBody' scan_x_params $ lambdaBody scan_op--                forM_ (zip scan_x_params pes) $ \(p, pe) ->-                  copyDWIMFix-                    (patElemName pe)-                    (map Imp.vi64 gtids ++ vec_is)-                    (Var $ paramName p)-                    []+    Nothing+  where+    ok op =+      segBinOpShape op == mempty+        && all primType (lambdaReturnType (segBinOpLambda op))  -- | Compile 'SegScan' instance to host-level code with calls to -- various kernels.@@ -486,25 +55,11 @@   CallKernelGen () compileSegScan pat lvl space scans kbody = sWhen (0 .<. n) $ do   emit $ Imp.DebugPrint "\n# SegScan" Nothing--  -- Since stage 2 involves a group size equal to the number of groups-  -- used for stage 1, we have to cap this number to the maximum group-  -- size.-  stage1_max_num_groups <- dPrim "stage1_max_num_groups" int64-  sOp $ Imp.GetSizeMax (tvVar stage1_max_num_groups) SizeGroup--  stage1_num_groups <--    fmap (Imp.Count . tvSize) $-      dPrimV "stage1_num_groups" $-        sMin64 (tvExp stage1_max_num_groups) $-          toInt64Exp $ Imp.unCount $ segNumGroups lvl--  (stage1_num_threads, elems_per_group, crossesSegment) <--    scanStage1 pat stage1_num_groups (segGroupSize lvl) space scans kbody--  emit $ Imp.DebugPrint "elems_per_group" $ Just $ untyped elems_per_group--  scanStage2 pat stage1_num_threads elems_per_group stage1_num_groups crossesSegment space scans-  scanStage3 pat (segNumGroups lvl) (segGroupSize lvl) elems_per_group crossesSegment space scans+  target <- hostTarget <$> askEnv+  case target of+    CUDA+      | Just scan' <- canBeSinglePass space scans ->+        SinglePass.compileSegScan pat lvl space scan' kbody+    _ -> TwoPass.compileSegScan pat lvl space scans kbody   where     n = product $ map toInt64Exp $ segSpaceDims space

@@ -0,0 +1,489 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE TypeFamilies #-}++-- | Code generation for segmented and non-segmented scans.  Uses a+-- fast single-pass algorithm, but which only works on NVIDIA GPUs and+-- with some constraints on the operator.  We use this when we can.+module Futhark.CodeGen.ImpGen.Kernels.SegScan.SinglePass (compileSegScan) where++import Control.Monad.Except+import Data.List (zip4)+import Data.Maybe+import qualified Futhark.CodeGen.ImpCode.Kernels as Imp+import Futhark.CodeGen.ImpGen+import Futhark.CodeGen.ImpGen.Kernels.Base+import Futhark.IR.KernelsMem+import qualified Futhark.IR.Mem.IxFun as IxFun+import Futhark.Transform.Rename+import Futhark.Util (takeLast)+import Futhark.Util.IntegralExp (IntegralExp, divUp, quot)+import Prelude hiding (quot)++xParams, yParams :: SegBinOp KernelsMem -> [LParam KernelsMem]+xParams scan =+  take (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan))+yParams scan =+  drop (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan))++alignTo :: IntegralExp a => a -> a -> a+alignTo x a = (x `divUp` a) * a++createLocalArrays ::+  Count GroupSize SubExp ->+  SubExp ->+  [PrimType] ->+  InKernelGen (VName, [VName], [VName], VName, VName, [VName])+createLocalArrays (Count groupSize) m types = do+  let groupSizeE = toInt64Exp groupSize+      workSize = toInt64Exp m * groupSizeE+      prefixArraysSize =+        foldl (\acc tySize -> alignTo acc tySize + tySize * groupSizeE) 0 $+          map primByteSize types+      maxTransposedArraySize =+        foldl1 sMax64 $ map (\ty -> workSize * primByteSize ty) types++      warpSize :: Num a => a+      warpSize = 32+      maxWarpExchangeSize =+        foldl (\acc tySize -> alignTo acc tySize + tySize * fromInteger warpSize) 0 $+          map primByteSize types+      maxLookbackSize = maxWarpExchangeSize + warpSize+      size = Imp.bytes $ maxLookbackSize `sMax64` prefixArraysSize `sMax64` maxTransposedArraySize++      varTE :: TV Int64 -> TPrimExp Int64 VName+      varTE = le64 . tvVar++  byteOffsets <-+    mapM (fmap varTE . dPrimV "byte_offsets") $+      scanl (\off tySize -> alignTo off tySize + toInt64Exp groupSize * tySize) 0 $+        map primByteSize types++  warpByteOffsets <-+    mapM (fmap varTE . dPrimV "warp_byte_offset") $+      scanl (\off tySize -> alignTo off tySize + warpSize * tySize) warpSize $+        map primByteSize types++  sComment "Allocate reused shared memeory" $ return ()++  localMem <- sAlloc "local_mem" size (Space "local")+  transposeArrayLength <- dPrimV "trans_arr_len" workSize++  sharedId <- sArrayInMem "shared_id" int32 (Shape [constant (1 :: Int32)]) localMem+  sharedReadOffset <- sArrayInMem "shared_read_offset" int32 (Shape [constant (1 :: Int32)]) localMem++  transposedArrays <-+    forM types $ \ty ->+      sArrayInMem+        "local_transpose_arr"+        ty+        (Shape [tvSize transposeArrayLength])+        localMem++  prefixArrays <-+    forM (zip byteOffsets types) $ \(off, ty) -> do+      let off' = off `quot` primByteSize ty+      sArray+        "local_prefix_arr"+        ty+        (Shape [groupSize])+        $ ArrayIn localMem $ IxFun.iotaOffset off' [pe64 groupSize]++  warpscan <- sArrayInMem "warpscan" int8 (Shape [constant (warpSize :: Int64)]) localMem+  warpExchanges <-+    forM (zip warpByteOffsets types) $ \(off, ty) -> do+      let off' = off `quot` primByteSize ty+      sArray+        "warp_exchange"+        ty+        (Shape [constant (warpSize :: Int64)])+        $ ArrayIn localMem $ IxFun.iotaOffset off' [warpSize]++  return (sharedId, transposedArrays, prefixArrays, sharedReadOffset, warpscan, warpExchanges)++-- | Compile 'SegScan' instance to host-level code with calls to a+-- single-pass kernel.+compileSegScan ::+  Pattern KernelsMem ->+  SegLevel ->+  SegSpace ->+  SegBinOp KernelsMem ->+  KernelBody KernelsMem ->+  CallKernelGen ()+compileSegScan pat lvl space scanOp kbody = do+  let Pattern _ all_pes = pat+      group_size = toInt64Exp <$> segGroupSize lvl+      n = product $ map toInt64Exp $ segSpaceDims space+      m :: Num a => a+      m = 9+      num_groups = Count (n `divUp` (unCount group_size * m))+      num_threads = unCount num_groups * unCount group_size+      (mapIdx, _) = head $ unSegSpace space+      scanOpNe = segBinOpNeutral scanOp+      tys = map (\(Prim pt) -> pt) $ lambdaReturnType $ segBinOpLambda scanOp+      statusX, statusA, statusP :: Num a => a+      statusX = 0+      statusA = 1+      statusP = 2+      makeStatusUsed flag used = tvExp flag .|. (tvExp used .<<. 2)+      unmakeStatusUsed :: TV Int8 -> TV Int8 -> TV Int8 -> InKernelGen ()+      unmakeStatusUsed flagUsed flag used = do+        used <-- tvExp flagUsed .>>. 2+        flag <-- tvExp flagUsed .&. 3++  -- Allocate the shared memory for output component+  numThreads <- dPrimV "numThreads" num_threads+  numGroups <- dPrimV "numGroups" $ unCount num_groups++  globalId <- sStaticArray "id_counter" (Space "device") int32 $ Imp.ArrayZeros 1+  statusFlags <- sAllocArray "status_flags" int8 (Shape [tvSize numGroups]) (Space "device")+  (aggregateArrays, incprefixArrays) <-+    fmap unzip $+      forM tys $ \ty ->+        (,) <$> sAllocArray "aggregates" ty (Shape [tvSize numGroups]) (Space "device")+          <*> sAllocArray "incprefixes" ty (Shape [tvSize numGroups]) (Space "device")++  sReplicate statusFlags $ intConst Int8 statusX++  sKernelThread "segscan" num_groups group_size (segFlat space) $ do+    constants <- kernelConstants <$> askEnv++    (sharedId, transposedArrays, prefixArrays, sharedReadOffset, warpscan, exchanges) <-+      createLocalArrays (segGroupSize lvl) (intConst Int64 m) tys++    dynamicId <- dPrim "dynamic_id" int32+    sWhen (kernelLocalThreadId constants .==. 0) $ do+      (globalIdMem, _, globalIdOff) <- fullyIndexArray globalId [0]+      sOp $+        Imp.Atomic DefaultSpace $+          Imp.AtomicAdd+            Int32+            (tvVar dynamicId)+            globalIdMem+            (Count $ unCount globalIdOff)+            (untyped (1 :: Imp.TExp Int32))+      copyDWIMFix sharedId [0] (tvSize dynamicId) []++    let localBarrier = Imp.Barrier Imp.FenceLocal+        localFence = Imp.MemFence Imp.FenceLocal+        globalFence = Imp.MemFence Imp.FenceGlobal++    sOp localBarrier+    copyDWIMFix (tvVar dynamicId) [] (Var sharedId) [0]+    sOp localBarrier++    blockOff <-+      dPrimV "blockOff" $+        sExt64 (tvExp dynamicId) * m * kernelGroupSize constants++    privateArrays <-+      forM tys $ \ty ->+        sAllocArray+          "private"+          ty+          (Shape [intConst Int64 m])+          (ScalarSpace [intConst Int64 m] ty)++    sComment "Load and map" $+      sFor "i" m $ \i -> do+        -- The map's input index+        dPrimV_ mapIdx $+          tvExp blockOff + sExt64 (kernelLocalThreadId constants)+            + i * kernelGroupSize constants+        -- Perform the map+        let in_bounds =+              compileStms mempty (kernelBodyStms kbody) $ do+                let (all_scan_res, map_res) = splitAt (segBinOpResults [scanOp]) $ kernelBodyResult kbody++                -- Write map results to their global memory destinations+                forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(dest, src) ->+                  copyDWIMFix (patElemName dest) [Imp.vi64 mapIdx] (kernelResultSubExp src) []++                -- Write to-scan results to private memory.+                forM_ (zip privateArrays $ map kernelResultSubExp all_scan_res) $ \(dest, src) ->+                  copyDWIMFix dest [i] src []++            out_of_bounds =+              forM_ (zip privateArrays scanOpNe) $ \(dest, ne) ->+                copyDWIMFix dest [i] ne []++        sIf (Imp.vi64 mapIdx .<. n) in_bounds out_of_bounds++    sComment "Transpose scan inputs" $ do+      forM_ (zip transposedArrays privateArrays) $ \(trans, priv) -> do+        sOp localBarrier+        sFor "i" m $ \i -> do+          sharedIdx <-+            dPrimVE "sharedIdx" $+              sExt64 (kernelLocalThreadId constants)+                + i * kernelGroupSize constants+          copyDWIMFix trans [sharedIdx] (Var priv) [i]+        sOp localBarrier+        sFor "i" m $ \i -> do+          sharedIdx <- dPrimV "sharedIdx" $ kernelLocalThreadId constants * m + i+          copyDWIMFix priv [sExt64 i] (Var trans) [sExt64 $ tvExp sharedIdx]+      sOp localBarrier++    sComment "Per thread scan" $+      -- We don't need to touch the first element, so only m-1+      -- iterations here.+      sFor "i" (m -1) $ \i -> do+        let xs = map paramName $ xParams scanOp+            ys = map paramName $ yParams scanOp++        forM_ (zip privateArrays $ zip3 xs ys tys) $ \(src, (x, y, ty)) -> do+          dPrim_ x ty+          dPrim_ y ty+          copyDWIMFix x [] (Var src) [i]+          copyDWIMFix y [] (Var src) [i + 1]++        compileStms mempty (bodyStms $ lambdaBody $ segBinOpLambda scanOp) $+          forM_ (zip privateArrays $ bodyResult $ lambdaBody $ segBinOpLambda scanOp) $ \(dest, res) ->+            copyDWIMFix dest [i + 1] res []++    sComment "Publish results in shared memory" $ do+      forM_ (zip prefixArrays privateArrays) $ \(dest, src) ->+        copyDWIMFix dest [sExt64 $ kernelLocalThreadId constants] (Var src) [m - 1]+      sOp localBarrier++    scanOp' <- renameLambda $ segBinOpLambda scanOp++    accs <- mapM (dPrim "acc") tys+    sComment "Scan results (with warp scan)" $ do+      groupScan+        Nothing -- TODO+        (tvExp numThreads)+        (kernelGroupSize constants)+        scanOp'+        prefixArrays++      sOp localBarrier++      let firstThread acc prefixes =+            copyDWIMFix (tvVar acc) [] (Var prefixes) [sExt64 (kernelGroupSize constants) - 1]+          notFirstThread acc prefixes =+            copyDWIMFix (tvVar acc) [] (Var prefixes) [sExt64 (kernelLocalThreadId constants) - 1]+      sIf+        (kernelLocalThreadId constants .==. 0)+        (zipWithM_ firstThread accs prefixArrays)+        (zipWithM_ notFirstThread accs prefixArrays)++      sOp localBarrier++    prefixes <- forM (zip scanOpNe tys) $ \(ne, ty) ->+      dPrimV "prefix" $ TPrimExp $ toExp' ty ne+    sComment "Perform lookback" $ do+      sWhen (tvExp dynamicId .==. 0 .&&. kernelLocalThreadId constants .==. 0) $ do+        everythingVolatile $+          forM_ (zip incprefixArrays accs) $ \(incprefixArray, acc) ->+            copyDWIMFix incprefixArray [tvExp dynamicId] (tvSize acc) []+        sOp globalFence+        everythingVolatile $+          copyDWIMFix statusFlags [tvExp dynamicId] (intConst Int8 statusP) []+        forM_ (zip scanOpNe accs) $ \(ne, acc) ->+          copyDWIMFix (tvVar acc) [] ne []+      -- end sWhen++      let warpSize = kernelWaveSize constants+      sWhen (bNot (tvExp dynamicId .==. 0) .&&. kernelLocalThreadId constants .<. warpSize) $ do+        sWhen (kernelLocalThreadId constants .==. 0) $ do+          everythingVolatile $+            forM_ (zip aggregateArrays accs) $ \(aggregateArray, acc) ->+              copyDWIMFix aggregateArray [tvExp dynamicId] (tvSize acc) []+          sOp globalFence+          everythingVolatile $+            copyDWIMFix statusFlags [tvExp dynamicId] (intConst Int8 statusA) []+          copyDWIMFix warpscan [0] (Var statusFlags) [tvExp dynamicId - 1]+        -- sWhen+        sOp localFence++        status <- dPrim "status" int8 :: InKernelGen (TV Int8)+        copyDWIMFix (tvVar status) [] (Var warpscan) [0]++        sIf+          (tvExp status .==. statusP)+          ( sWhen (kernelLocalThreadId constants .==. 0) $+              everythingVolatile $+                forM_ (zip prefixes incprefixArrays) $ \(prefix, incprefixArray) ->+                  copyDWIMFix (tvVar prefix) [] (Var incprefixArray) [tvExp dynamicId - 1]+          )+          ( do+              readOffset <-+                dPrimV "readOffset" $+                  sExt32 $ tvExp dynamicId - sExt64 (kernelWaveSize constants)+              let loopStop = warpSize * (-1)+              sWhile (tvExp readOffset .>. loopStop) $ do+                readI <- dPrimV "read_i" $ tvExp readOffset + kernelLocalThreadId constants+                aggrs <- forM (zip scanOpNe tys) $ \(ne, ty) ->+                  dPrimV "aggr" $ TPrimExp $ toExp' ty ne+                flag <- dPrimV "flag" statusX+                used <- dPrimV "used" (0 :: Imp.TExp Int8)+                everythingVolatile $+                  sWhen (tvExp readI .>=. 0) $ do+                    copyDWIMFix (tvVar flag) [] (Var statusFlags) [sExt64 $ tvExp readI]+                    sIf+                      (tvExp flag .==. statusP)+                      ( forM_ (zip incprefixArrays aggrs) $ \(incprefix, aggr) ->+                          copyDWIMFix (tvVar aggr) [] (Var incprefix) [sExt64 $ tvExp readI]+                      )+                      ( sWhen (tvExp flag .==. statusA) $ do+                          forM_ (zip aggrs aggregateArrays) $ \(aggr, aggregate) ->+                            copyDWIMFix (tvVar aggr) [] (Var aggregate) [sExt64 $ tvExp readI]+                          used <-- (1 :: Imp.TExp Int8)+                      )+                -- end sIf+                -- end sWhen+                forM_ (zip exchanges aggrs) $ \(exchange, aggr) ->+                  copyDWIMFix exchange [sExt64 $ kernelLocalThreadId constants] (tvSize aggr) []+                tmp <- dPrimV "tmp" $ makeStatusUsed flag used+                copyDWIMFix warpscan [sExt64 $ kernelLocalThreadId constants] (tvSize tmp) []+                sOp localFence++                (warpscanMem, warpscanSpace, warpscanOff) <-+                  fullyIndexArray warpscan [sExt64 warpSize - 1]+                flag <-- TPrimExp (Imp.index warpscanMem warpscanOff int8 warpscanSpace Imp.Volatile)+                sWhen (kernelLocalThreadId constants .==. 0) $ do+                  -- TODO: This is a single-threaded reduce+                  sIf+                    (bNot $ tvExp flag .==. statusP)+                    ( do+                        scanOp'' <- renameLambda scanOp'+                        let (agg1s, agg2s) = splitAt (length tys) $ map paramName $ lambdaParams scanOp''++                        forM_ (zip3 agg1s scanOpNe tys) $ \(agg1, ne, ty) ->+                          dPrimV_ agg1 $ TPrimExp $ toExp' ty ne+                        zipWithM_ dPrim_ agg2s tys++                        flag1 <- dPrimV "flag1" statusX+                        flag2 <- dPrim "flag2" int8+                        used1 <- dPrimV "used1" (0 :: Imp.TExp Int8)+                        used2 <- dPrim "used2" int8+                        sFor "i" warpSize $ \i -> do+                          copyDWIMFix (tvVar flag2) [] (Var warpscan) [sExt64 i]+                          unmakeStatusUsed flag2 flag2 used2+                          forM_ (zip agg2s exchanges) $ \(agg2, exchange) ->+                            copyDWIMFix agg2 [] (Var exchange) [sExt64 i]+                          sIf+                            (bNot $ tvExp flag2 .==. statusA)+                            ( do+                                flag1 <-- tvExp flag2+                                used1 <-- tvExp used2+                                forM_ (zip3 agg1s tys agg2s) $ \(agg1, ty, agg2) ->+                                  agg1 <~~ toExp' ty (Var agg2)+                            )+                            ( do+                                used1 <-- tvExp used1 + tvExp used2+                                compileStms mempty (bodyStms $ lambdaBody scanOp'') $+                                  forM_ (zip3 agg1s tys $ bodyResult $ lambdaBody scanOp'') $+                                    \(agg1, ty, res) -> agg1 <~~ toExp' ty res+                            )+                        flag <-- tvExp flag1+                        used <-- tvExp used1+                        forM_ (zip3 aggrs tys agg1s) $ \(aggr, ty, agg1) ->+                          tvVar aggr <~~ toExp' ty (Var agg1)+                    )+                    -- else+                    ( forM_ (zip aggrs exchanges) $ \(aggr, exchange) ->+                        copyDWIMFix (tvVar aggr) [] (Var exchange) [sExt64 warpSize - 1]+                    )+                  -- end sIf+                  sIf+                    (tvExp flag .==. statusP)+                    (readOffset <-- loopStop)+                    (readOffset <-- tvExp readOffset - zExt32 (tvExp used))+                  copyDWIMFix sharedReadOffset [0] (tvSize readOffset) []+                  scanOp''' <- renameLambda scanOp'+                  let (xs, ys) = splitAt (length tys) $ map paramName $ lambdaParams scanOp'''+                  forM_ (zip xs aggrs) $ \(x, aggr) -> dPrimV_ x (tvExp aggr)+                  forM_ (zip ys prefixes) $ \(y, prefix) -> dPrimV_ y (tvExp prefix)+                  compileStms mempty (bodyStms $ lambdaBody scanOp''') $+                    forM_ (zip3 prefixes tys $ bodyResult $ lambdaBody scanOp''') $+                      \(prefix, ty, res) -> prefix <-- TPrimExp (toExp' ty res)+                -- end sWhen+                sOp localFence+                copyDWIMFix (tvVar readOffset) [] (Var sharedReadOffset) [0]+          )+        -- end sWhile+        -- end sIf+        sWhen (kernelLocalThreadId constants .==. 0) $ do+          scanOp'''' <- renameLambda scanOp'+          let xs = map paramName $ take (length tys) $ lambdaParams scanOp''''+              ys = map paramName $ drop (length tys) $ lambdaParams scanOp''''+          forM_ (zip xs prefixes) $ \(x, prefix) -> dPrimV_ x $ tvExp prefix+          forM_ (zip ys accs) $ \(y, acc) -> dPrimV_ y $ tvExp acc+          compileStms mempty (bodyStms $ lambdaBody scanOp'''') $+            everythingVolatile $+              forM_ (zip incprefixArrays $ bodyResult $ lambdaBody scanOp'''') $+                \(incprefixArray, res) -> copyDWIMFix incprefixArray [tvExp dynamicId] res []+          sOp globalFence+          everythingVolatile $ copyDWIMFix statusFlags [tvExp dynamicId] (intConst Int8 statusP) []+          forM_ (zip exchanges prefixes) $ \(exchange, prefix) ->+            copyDWIMFix exchange [0] (tvSize prefix) []+          forM_ (zip3 accs tys scanOpNe) $ \(acc, ty, ne) ->+            tvVar acc <~~ toExp' ty ne+      -- end sWhen+      -- end sWhen++      sWhen (bNot $ tvExp dynamicId .==. 0) $ do+        sOp localBarrier+        forM_ (zip exchanges prefixes) $ \(exchange, prefix) ->+          copyDWIMFix (tvVar prefix) [] (Var exchange) [0]+        sOp localBarrier+    -- end sWhen+    -- end sComment++    scanOp''''' <- renameLambda scanOp'+    scanOp'''''' <- renameLambda scanOp'++    sComment "Distribute results" $ do+      let (xs, ys) = splitAt (length tys) $ map paramName $ lambdaParams scanOp'''''+          (xs', ys') = splitAt (length tys) $ map paramName $ lambdaParams scanOp''''''++      forM_ (zip4 (zip prefixes accs) (zip xs xs') (zip ys ys') tys) $+        \((prefix, acc), (x, x'), (y, y'), ty) -> do+          dPrim_ x ty+          dPrim_ y ty+          dPrimV_ x' $ tvExp prefix+          dPrimV_ y' $ tvExp acc++      compileStms mempty (bodyStms $ lambdaBody scanOp'''''') $+        forM_ (zip3 xs tys $ bodyResult $ lambdaBody scanOp'''''') $+          \(x, ty, res) -> x <~~ toExp' ty res++      sFor "i" m $ \i -> do+        forM_ (zip privateArrays ys) $ \(src, y) ->+          copyDWIMFix y [] (Var src) [i]++        compileStms mempty (bodyStms $ lambdaBody scanOp''''') $+          forM_ (zip privateArrays $ bodyResult $ lambdaBody scanOp''''') $+            \(dest, res) ->+              copyDWIMFix dest [i] res []++    sComment "Transpose scan output" $ do+      forM_ (zip transposedArrays privateArrays) $ \(trans, priv) -> do+        sOp localBarrier+        sFor "i" m $ \i -> do+          sharedIdx <-+            dPrimV "sharedIdx" $+              sExt64 (kernelLocalThreadId constants * m) + i+          copyDWIMFix trans [tvExp sharedIdx] (Var priv) [i]+        sOp localBarrier+        sFor "i" m $ \i -> do+          sharedIdx <-+            dPrimV "sharedIdx" $+              kernelLocalThreadId constants+                + sExt32 (kernelGroupSize constants * i)+          copyDWIMFix priv [i] (Var trans) [sExt64 $ tvExp sharedIdx]+      sOp localBarrier++    sComment "Write block scan results to global memory" $+      forM_ (zip (map patElemName all_pes) privateArrays) $ \(dest, src) ->+        sFor "i" m $ \i -> do+          dPrimV_ mapIdx $+            tvExp blockOff + kernelGroupSize constants * i+              + sExt64 (kernelLocalThreadId constants)+          sWhen (Imp.vi64 mapIdx .<. n) $+            copyDWIMFix dest [Imp.vi64 mapIdx] (Var src) [i]++    sComment "If this is the last block, reset the dynamicId" $+      sWhen (tvExp dynamicId .==. unCount num_groups - 1) $+        copyDWIMFix globalId [0] (constant (0 :: Int32)) []

@@ -0,0 +1,506 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE TypeFamilies #-}++-- | Code generation for segmented and non-segmented scans.  Uses a+-- fairly inefficient two-pass algorithm, but can handle anything.+module Futhark.CodeGen.ImpGen.Kernels.SegScan.TwoPass (compileSegScan) where++import Control.Monad.Except+import Control.Monad.State+import Data.List (delete, find, foldl', zip4)+import Data.Maybe+import qualified Futhark.CodeGen.ImpCode.Kernels as Imp+import Futhark.CodeGen.ImpGen+import Futhark.CodeGen.ImpGen.Kernels.Base+import Futhark.IR.KernelsMem+import qualified Futhark.IR.Mem.IxFun as IxFun+import Futhark.Transform.Rename+import Futhark.Util (takeLast)+import Futhark.Util.IntegralExp (divUp, quot, rem)+import Prelude hiding (quot, rem)++-- Aggressively try to reuse memory for different SegBinOps, because+-- we will run them sequentially after another.+makeLocalArrays ::+  Count GroupSize SubExp ->+  SubExp ->+  [SegBinOp KernelsMem] ->+  InKernelGen [[VName]]+makeLocalArrays (Count group_size) num_threads scans = do+  (arrs, mems_and_sizes) <- runStateT (mapM onScan scans) mempty+  let maxSize sizes = Imp.bytes $ foldl' sMax64 1 $ map Imp.unCount sizes+  forM_ mems_and_sizes $ \(sizes, mem) ->+    sAlloc_ mem (maxSize sizes) (Space "local")+  return arrs+  where+    onScan (SegBinOp _ scan_op nes _) = do+      let (scan_x_params, _scan_y_params) =+            splitAt (length nes) $ lambdaParams scan_op+      (arrs, used_mems) <- fmap unzip $+        forM scan_x_params $ \p ->+          case paramDec p of+            MemArray pt shape _ (ArrayIn mem _) -> do+              let shape' = Shape [num_threads] <> shape+              arr <-+                lift $+                  sArray "scan_arr" pt shape' $+                    ArrayIn mem $ IxFun.iota $ map pe64 $ shapeDims shape'+              return (arr, [])+            _ -> do+              let pt = elemType $ paramType p+                  shape = Shape [group_size]+              (sizes, mem') <- getMem pt shape+              arr <- lift $ sArrayInMem "scan_arr" pt shape mem'+              return (arr, [(sizes, mem')])+      modify (<> concat used_mems)+      return arrs++    getMem pt shape = do+      let size = typeSize $ Array pt shape NoUniqueness+      mems <- get+      case (find ((size `elem`) . fst) mems, mems) of+        (Just mem, _) -> do+          modify $ delete mem+          return mem+        (Nothing, (size', mem) : mems') -> do+          put mems'+          return (size : size', mem)+        (Nothing, []) -> do+          mem <- lift $ sDeclareMem "scan_arr_mem" $ Space "local"+          return ([size], mem)++type CrossesSegment = Maybe (Imp.TExp Int64 -> Imp.TExp Int64 -> Imp.TExp Bool)++localArrayIndex :: KernelConstants -> Type -> Imp.TExp Int64+localArrayIndex constants t =+  if primType t+    then sExt64 (kernelLocalThreadId constants)+    else sExt64 (kernelGlobalThreadId constants)++barrierFor :: Lambda KernelsMem -> (Bool, Imp.Fence, InKernelGen ())+barrierFor scan_op = (array_scan, fence, sOp $ Imp.Barrier fence)+  where+    array_scan = not $ all primType $ lambdaReturnType scan_op+    fence+      | array_scan = Imp.FenceGlobal+      | otherwise = Imp.FenceLocal++xParams, yParams :: SegBinOp KernelsMem -> [LParam KernelsMem]+xParams scan =+  take (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan))+yParams scan =+  drop (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan))++writeToScanValues ::+  [VName] ->+  ([PatElem KernelsMem], SegBinOp KernelsMem, [KernelResult]) ->+  InKernelGen ()+writeToScanValues gtids (pes, scan, scan_res)+  | shapeRank (segBinOpShape scan) > 0 =+    forM_ (zip pes scan_res) $ \(pe, res) ->+      copyDWIMFix+        (patElemName pe)+        (map Imp.vi64 gtids)+        (kernelResultSubExp res)+        []+  | otherwise =+    forM_ (zip (yParams scan) scan_res) $ \(p, res) ->+      copyDWIMFix (paramName p) [] (kernelResultSubExp res) []++readToScanValues ::+  [Imp.TExp Int64] ->+  [PatElem KernelsMem] ->+  SegBinOp KernelsMem ->+  InKernelGen ()+readToScanValues is pes scan+  | shapeRank (segBinOpShape scan) > 0 =+    forM_ (zip (yParams scan) pes) $ \(p, pe) ->+      copyDWIMFix (paramName p) [] (Var (patElemName pe)) is+  | otherwise =+    return ()++readCarries ::+  Imp.TExp Int64 ->+  [Imp.TExp Int64] ->+  [Imp.TExp Int64] ->+  [PatElem KernelsMem] ->+  SegBinOp KernelsMem ->+  InKernelGen ()+readCarries chunk_offset dims' vec_is pes scan+  | shapeRank (segBinOpShape scan) > 0 = do+    ltid <- kernelLocalThreadId . kernelConstants <$> askEnv+    -- We may have to reload the carries from the output of the+    -- previous chunk.+    sIf+      (chunk_offset .>. 0 .&&. ltid .==. 0)+      ( do+          let is = unflattenIndex dims' $ chunk_offset - 1+          forM_ (zip (xParams scan) pes) $ \(p, pe) ->+            copyDWIMFix (paramName p) [] (Var (patElemName pe)) (is ++ vec_is)+      )+      ( forM_ (zip (xParams scan) (segBinOpNeutral scan)) $ \(p, ne) ->+          copyDWIMFix (paramName p) [] ne []+      )+  | otherwise =+    return ()++-- | Produce partially scanned intervals; one per workgroup.+scanStage1 ::+  Pattern KernelsMem ->+  Count NumGroups SubExp ->+  Count GroupSize SubExp ->+  SegSpace ->+  [SegBinOp KernelsMem] ->+  KernelBody KernelsMem ->+  CallKernelGen (TV Int32, Imp.TExp Int64, CrossesSegment)+scanStage1 (Pattern _ all_pes) num_groups group_size space scans kbody = do+  let num_groups' = fmap toInt64Exp num_groups+      group_size' = fmap toInt64Exp group_size+  num_threads <- dPrimV "num_threads" $ sExt32 $ unCount num_groups' * unCount group_size'++  let (gtids, dims) = unzip $ unSegSpace space+      dims' = map toInt64Exp dims+  let num_elements = product dims'+      elems_per_thread = num_elements `divUp` sExt64 (tvExp num_threads)+      elems_per_group = unCount group_size' * elems_per_thread++  let crossesSegment =+        case reverse dims' of+          segment_size : _ : _ -> Just $ \from to ->+            (to - from) .>. (to `rem` segment_size)+          _ -> Nothing++  sKernelThread "scan_stage1" num_groups' group_size' (segFlat space) $ do+    constants <- kernelConstants <$> askEnv+    all_local_arrs <- makeLocalArrays group_size (tvSize num_threads) scans++    -- The variables from scan_op will be used for the carry and such+    -- in the big chunking loop.+    forM_ scans $ \scan -> do+      dScope Nothing $ scopeOfLParams $ lambdaParams $ segBinOpLambda scan+      forM_ (zip (xParams scan) (segBinOpNeutral scan)) $ \(p, ne) ->+        copyDWIMFix (paramName p) [] ne []++    sFor "j" elems_per_thread $ \j -> do+      chunk_offset <-+        dPrimV "chunk_offset" $+          sExt64 (kernelGroupSize constants) * j+            + sExt64 (kernelGroupId constants) * elems_per_group+      flat_idx <-+        dPrimV "flat_idx" $+          tvExp chunk_offset + sExt64 (kernelLocalThreadId constants)+      -- Construct segment indices.+      zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx++      let per_scan_pes = segBinOpChunks scans all_pes++          in_bounds =+            foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims'++          when_in_bounds = compileStms mempty (kernelBodyStms kbody) $ do+            let (all_scan_res, map_res) =+                  splitAt (segBinOpResults scans) $ kernelBodyResult kbody+                per_scan_res =+                  segBinOpChunks scans all_scan_res++            sComment "write to-scan values to parameters" $+              mapM_ (writeToScanValues gtids) $+                zip3 per_scan_pes scans per_scan_res++            sComment "write mapped values results to global memory" $+              forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(pe, se) ->+                copyDWIMFix+                  (patElemName pe)+                  (map Imp.vi64 gtids)+                  (kernelResultSubExp se)+                  []++      sComment "threads in bounds read input" $+        sWhen in_bounds when_in_bounds++      forM_ (zip3 per_scan_pes scans all_local_arrs) $+        \(pes, scan@(SegBinOp _ scan_op nes vec_shape), local_arrs) ->+          sComment "do one intra-group scan operation" $ do+            let rets = lambdaReturnType scan_op+                scan_x_params = xParams scan+                (array_scan, fence, barrier) = barrierFor scan_op++            when array_scan barrier++            sLoopNest vec_shape $ \vec_is -> do+              sComment "maybe restore some to-scan values to parameters, or read neutral" $+                sIf+                  in_bounds+                  ( do+                      readToScanValues (map Imp.vi64 gtids ++ vec_is) pes scan+                      readCarries (tvExp chunk_offset) dims' vec_is pes scan+                  )+                  ( forM_ (zip (yParams scan) (segBinOpNeutral scan)) $ \(p, ne) ->+                      copyDWIMFix (paramName p) [] ne []+                  )++              sComment "combine with carry and write to local memory" $+                compileStms mempty (bodyStms $ lambdaBody scan_op) $+                  forM_ (zip3 rets local_arrs (bodyResult $ lambdaBody scan_op)) $+                    \(t, arr, se) ->+                      copyDWIMFix arr [localArrayIndex constants t] se []++              let crossesSegment' = do+                    f <- crossesSegment+                    Just $ \from to ->+                      let from' = sExt64 from + tvExp chunk_offset+                          to' = sExt64 to + tvExp chunk_offset+                       in f from' to'++              sOp $ Imp.ErrorSync fence++              -- We need to avoid parameter name clashes.+              scan_op_renamed <- renameLambda scan_op+              groupScan+                crossesSegment'+                (sExt64 $ tvExp num_threads)+                (sExt64 $ kernelGroupSize constants)+                scan_op_renamed+                local_arrs++              sComment "threads in bounds write partial scan result" $+                sWhen in_bounds $+                  forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) ->+                    copyDWIMFix+                      (patElemName pe)+                      (map Imp.vi64 gtids ++ vec_is)+                      (Var arr)+                      [localArrayIndex constants t]++              barrier++              let load_carry =+                    forM_ (zip local_arrs scan_x_params) $ \(arr, p) ->+                      copyDWIMFix+                        (paramName p)+                        []+                        (Var arr)+                        [ if primType $ paramType p+                            then sExt64 (kernelGroupSize constants) - 1+                            else+                              (sExt64 (kernelGroupId constants) + 1)+                                * sExt64 (kernelGroupSize constants) - 1+                        ]+                  load_neutral =+                    forM_ (zip nes scan_x_params) $ \(ne, p) ->+                      copyDWIMFix (paramName p) [] ne []++              sComment "first thread reads last element as carry-in for next iteration" $ do+                crosses_segment <- dPrimVE "crosses_segment" $+                  case crossesSegment of+                    Nothing -> false+                    Just f ->+                      f+                        ( tvExp chunk_offset+                            + sExt64 (kernelGroupSize constants) -1+                        )+                        ( tvExp chunk_offset+                            + sExt64 (kernelGroupSize constants)+                        )+                should_load_carry <-+                  dPrimVE "should_load_carry" $+                    kernelLocalThreadId constants .==. 0 .&&. bNot crosses_segment+                sWhen should_load_carry load_carry+                when array_scan barrier+                sUnless should_load_carry load_neutral++              barrier++  return (num_threads, elems_per_group, crossesSegment)++scanStage2 ::+  Pattern KernelsMem ->+  TV Int32 ->+  Imp.TExp Int64 ->+  Count NumGroups SubExp ->+  CrossesSegment ->+  SegSpace ->+  [SegBinOp KernelsMem] ->+  CallKernelGen ()+scanStage2 (Pattern _ all_pes) stage1_num_threads elems_per_group num_groups crossesSegment space scans = do+  let (gtids, dims) = unzip $ unSegSpace space+      dims' = map toInt64Exp dims++  -- Our group size is the number of groups for the stage 1 kernel.+  let group_size = Count $ unCount num_groups+      group_size' = fmap toInt64Exp group_size++  let crossesSegment' = do+        f <- crossesSegment+        Just $ \from to ->+          f+            ((sExt64 from + 1) * elems_per_group - 1)+            ((sExt64 to + 1) * elems_per_group - 1)++  sKernelThread "scan_stage2" 1 group_size' (segFlat space) $ do+    constants <- kernelConstants <$> askEnv+    per_scan_local_arrs <- makeLocalArrays group_size (tvSize stage1_num_threads) scans+    let per_scan_rets = map (lambdaReturnType . segBinOpLambda) scans+        per_scan_pes = segBinOpChunks scans all_pes++    flat_idx <-+      dPrimV "flat_idx" $+        (sExt64 (kernelLocalThreadId constants) + 1) * elems_per_group - 1+    -- Construct segment indices.+    zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx++    forM_ (zip4 scans per_scan_local_arrs per_scan_rets per_scan_pes) $+      \(SegBinOp _ scan_op nes vec_shape, local_arrs, rets, pes) ->+        sLoopNest vec_shape $ \vec_is -> do+          let glob_is = map Imp.vi64 gtids ++ vec_is++              in_bounds =+                foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims'++              when_in_bounds = forM_ (zip3 rets local_arrs pes) $ \(t, arr, pe) ->+                copyDWIMFix+                  arr+                  [localArrayIndex constants t]+                  (Var $ patElemName pe)+                  glob_is++              when_out_of_bounds = forM_ (zip3 rets local_arrs nes) $ \(t, arr, ne) ->+                copyDWIMFix arr [localArrayIndex constants t] ne []+              (_, _, barrier) =+                barrierFor scan_op++          sComment "threads in bound read carries; others get neutral element" $+            sIf in_bounds when_in_bounds when_out_of_bounds++          barrier++          groupScan+            crossesSegment'+            (sExt64 $ tvExp stage1_num_threads)+            (sExt64 $ kernelGroupSize constants)+            scan_op+            local_arrs++          sComment "threads in bounds write scanned carries" $+            sWhen in_bounds $+              forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) ->+                copyDWIMFix+                  (patElemName pe)+                  glob_is+                  (Var arr)+                  [localArrayIndex constants t]++scanStage3 ::+  Pattern KernelsMem ->+  Count NumGroups SubExp ->+  Count GroupSize SubExp ->+  Imp.TExp Int64 ->+  CrossesSegment ->+  SegSpace ->+  [SegBinOp KernelsMem] ->+  CallKernelGen ()+scanStage3 (Pattern _ all_pes) num_groups group_size elems_per_group crossesSegment space scans = do+  let num_groups' = fmap toInt64Exp num_groups+      group_size' = fmap toInt64Exp group_size+      (gtids, dims) = unzip $ unSegSpace space+      dims' = map toInt64Exp dims+  required_groups <-+    dPrimVE "required_groups" $+      sExt32 $ product dims' `divUp` sExt64 (unCount group_size')++  sKernelThread "scan_stage3" num_groups' group_size' (segFlat space) $+    virtualiseGroups SegVirt required_groups $ \virt_group_id -> do+      constants <- kernelConstants <$> askEnv++      -- Compute our logical index.+      flat_idx <-+        dPrimVE "flat_idx" $+          sExt64 virt_group_id * sExt64 (unCount group_size')+            + sExt64 (kernelLocalThreadId constants)+      zipWithM_ dPrimV_ gtids $ unflattenIndex dims' flat_idx++      -- Figure out which group this element was originally in.+      orig_group <- dPrimV "orig_group" $ flat_idx `quot` elems_per_group+      -- Then the index of the carry-in of the preceding group.+      carry_in_flat_idx <-+        dPrimV "carry_in_flat_idx" $+          tvExp orig_group * elems_per_group - 1+      -- Figure out the logical index of the carry-in.+      let carry_in_idx = unflattenIndex dims' $ tvExp carry_in_flat_idx++      -- Apply the carry if we are not in the scan results for the first+      -- group, and are not the last element in such a group (because+      -- then the carry was updated in stage 2), and we are not crossing+      -- a segment boundary.+      let in_bounds =+            foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims'+          crosses_segment =+            fromMaybe false $+              crossesSegment+                <*> pure (tvExp carry_in_flat_idx)+                <*> pure flat_idx+          is_a_carry = flat_idx .==. (tvExp orig_group + 1) * elems_per_group - 1+          no_carry_in = tvExp orig_group .==. 0 .||. is_a_carry .||. crosses_segment++      let per_scan_pes = segBinOpChunks scans all_pes+      sWhen in_bounds $+        sUnless no_carry_in $+          forM_ (zip per_scan_pes scans) $+            \(pes, SegBinOp _ scan_op nes vec_shape) -> do+              dScope Nothing $ scopeOfLParams $ lambdaParams scan_op+              let (scan_x_params, scan_y_params) =+                    splitAt (length nes) $ lambdaParams scan_op++              sLoopNest vec_shape $ \vec_is -> do+                forM_ (zip scan_x_params pes) $ \(p, pe) ->+                  copyDWIMFix+                    (paramName p)+                    []+                    (Var $ patElemName pe)+                    (carry_in_idx ++ vec_is)++                forM_ (zip scan_y_params pes) $ \(p, pe) ->+                  copyDWIMFix+                    (paramName p)+                    []+                    (Var $ patElemName pe)+                    (map Imp.vi64 gtids ++ vec_is)++                compileBody' scan_x_params $ lambdaBody scan_op++                forM_ (zip scan_x_params pes) $ \(p, pe) ->+                  copyDWIMFix+                    (patElemName pe)+                    (map Imp.vi64 gtids ++ vec_is)+                    (Var $ paramName p)+                    []++-- | Compile 'SegScan' instance to host-level code with calls to+-- various kernels.+compileSegScan ::+  Pattern KernelsMem ->+  SegLevel ->+  SegSpace ->+  [SegBinOp KernelsMem] ->+  KernelBody KernelsMem ->+  CallKernelGen ()+compileSegScan pat lvl space scans kbody = do+  -- Since stage 2 involves a group size equal to the number of groups+  -- used for stage 1, we have to cap this number to the maximum group+  -- size.+  stage1_max_num_groups <- dPrim "stage1_max_num_groups" int64+  sOp $ Imp.GetSizeMax (tvVar stage1_max_num_groups) SizeGroup++  stage1_num_groups <-+    fmap (Imp.Count . tvSize) $+      dPrimV "stage1_num_groups" $+        sMin64 (tvExp stage1_max_num_groups) $+          toInt64Exp $ Imp.unCount $ segNumGroups lvl++  (stage1_num_threads, elems_per_group, crossesSegment) <-+    scanStage1 pat stage1_num_groups (segGroupSize lvl) space scans kbody++  emit $ Imp.DebugPrint "elems_per_group" $ Just $ untyped elems_per_group++  scanStage2 pat stage1_num_threads elems_per_group stage1_num_groups crossesSegment space scans+  scanStage3 pat (segNumGroups lvl) (segGroupSize lvl) elems_per_group crossesSegment space scans

@@ -1,6 +1,5 @@ module Futhark.CodeGen.ImpGen.Multicore.Base-  ( toParam,-    compileKBody,+  ( compileKBody,     extractAllocations,     compileThreadResult,     HostEnv (..),@@ -10,7 +9,6 @@     decideScheduling',     groupResultArrays,     renameSegBinOp,-    resultArrays,     freeParams,     renameHistOpLambda,     atomicUpdateLocking,@@ -130,15 +128,6 @@   let freeVars = freeVariables code names   ts <- mapM lookupType freeVars   zipWithM toParam freeVars ts---- | Arrays for storing group results.-resultArrays :: String -> [SegBinOp MCMem] -> MulticoreGen [[VName]]-resultArrays s segops =-  forM segops $ \(SegBinOp _ lam _ shape) ->-    forM (lambdaReturnType lam) $ \t -> do-      let pt = elemType t-          full_shape = shape <> arrayShape t-      sAllocArray s pt full_shape DefaultSpace  -- | Arrays for storing group results shared between threads groupResultArrays ::

@@ -35,6 +35,15 @@ lamBody :: SegBinOp MCMem -> Body MCMem lamBody = lambdaBody . segBinOpLambda +-- Arrays for storing worker results.+resultArrays :: String -> [SegBinOp MCMem] -> MulticoreGen [[VName]]+resultArrays s segops =+  forM segops $ \(SegBinOp _ lam _ shape) ->+    forM (lambdaReturnType lam) $ \t -> do+      let pt = elemType t+          full_shape = shape <> arrayShape t+      sAllocArray s pt full_shape DefaultSpace+ nonsegmentedScan ::   Pattern MCMem ->   SegSpace ->

@@ -8,6 +8,7 @@   ( IxFun (..),     index,     iota,+    iotaOffset,     permute,     rotate,     reshape,@@ -364,11 +365,15 @@                 (permuteInv (lmadPermutation lmad) inds)        in off + prod +-- | iota with offset.+iotaOffset :: IntegralExp num => num -> Shape num -> IxFun num+iotaOffset o ns =+  let rs = replicate (length ns) 0+   in IxFun (makeRotIota Inc o (zip rs ns) :| []) ns True+ -- | iota. iota :: IntegralExp num => Shape num -> IxFun num-iota ns =-  let rs = replicate (length ns) 0-   in IxFun (makeRotIota Inc 0 (zip rs ns) :| []) ns True+iota = iotaOffset 0  -- | Permute dimensions. permute ::

@@ -12,9 +12,7 @@ module Futhark.Internalise (internaliseProg) where  import Control.Monad.Reader-import Control.Monad.State-import Data.Bitraversable-import Data.List (find, intercalate, intersperse, nub, transpose)+import Data.List (find, intercalate, intersperse, transpose) import qualified Data.List.NonEmpty as NE import qualified Data.Map.Strict as M import qualified Data.Set as S@@ -24,6 +22,7 @@ import Futhark.Internalise.Defunctionalise as Defunctionalise import Futhark.Internalise.Defunctorise as Defunctorise import Futhark.Internalise.Lambdas+import Futhark.Internalise.LiftLambdas as LiftLambdas import Futhark.Internalise.Monad as I import Futhark.Internalise.Monomorphise as Monomorphise import Futhark.Internalise.TypesValues@@ -42,9 +41,10 @@ internaliseProg always_safe prog = do   prog_decs <- Defunctorise.transformProg prog   prog_decs' <- Monomorphise.transformProg prog_decs-  prog_decs'' <- Defunctionalise.transformProg prog_decs'+  prog_decs'' <- LiftLambdas.transformProg prog_decs'+  prog_decs''' <- Defunctionalise.transformProg prog_decs''   (consts, funs) <--    runInternaliseM always_safe (internaliseValBinds prog_decs'')+    runInternaliseM always_safe (internaliseValBinds prog_decs''')   I.renameProg $ I.Prog consts funs  internaliseAttr :: E.AttrInfo -> Attr@@ -57,28 +57,15 @@ internaliseValBinds :: [E.ValBind] -> InternaliseM () internaliseValBinds = mapM_ internaliseValBind -internaliseFunName :: VName -> [E.Pattern] -> InternaliseM Name-internaliseFunName ofname [] = return $ nameFromString $ pretty ofname ++ "f"-internaliseFunName ofname _ = do-  info <- lookupFunction' ofname-  -- In some rare cases involving local functions, the same function-  -- name may be re-used in multiple places.  We check whether the-  -- function name has already been used, and generate a new one if-  -- so.-  case info of-    Just _ -> nameFromString . pretty <$> newNameFromString (baseString ofname)-    Nothing -> return $ nameFromString $ pretty ofname+internaliseFunName :: VName -> Name+internaliseFunName = nameFromString . pretty  internaliseValBind :: E.ValBind -> InternaliseM () internaliseValBind fb@(E.ValBind entry fname retdecl (Info (rettype, _)) tparams params body _ attrs loc) = do   localConstsScope $     bindingParams tparams params $ \shapeparams params' -> do       let shapenames = map I.paramName shapeparams-          normal_params = shapenames ++ map I.paramName (concat params')-          normal_param_names = namesFromList normal_params -      fname' <- internaliseFunName fname params-       msg <- case retdecl of         Just dt ->           errorMsg@@ -95,27 +82,13 @@         ensureResultExtShape msg loc (map I.fromDecl rettype') $           mkBody body_stms body_res -      constants <- allConsts-      let free_in_fun =-            freeIn body'-              `namesSubtract` normal_param_names-              `namesSubtract` constants--      used_free_params <- forM (namesToList free_in_fun) $ \v -> do-        v_t <- lookupType v-        return $ Param v $ toDecl v_t Nonunique--      let free_shape_params =-            map (`Param` I.Prim int64) $-              concatMap (I.shapeVars . I.arrayShape . I.paramType) used_free_params-          free_params = nub $ free_shape_params ++ used_free_params-          all_params = free_params ++ shapeparams ++ concat params'+      let all_params = shapeparams ++ concat params'        let fd =             I.FunDef               Nothing               (internaliseAttrs attrs)-              fname'+              (internaliseFunName fname)               rettype'               all_params               body'@@ -126,9 +99,7 @@           bindFunction             fname             fd-            ( fname',-              map I.paramName free_params,-              shapenames,+            ( shapenames,               map declTypeOf $ concat params',               all_params,               applyRetType rettype' all_params@@ -140,93 +111,11 @@   where     zeroExts ts = generaliseExtTypes ts ts -allDimsFreshInType :: MonadFreshNames m => E.PatternType -> m E.PatternType-allDimsFreshInType = bitraverse onDim pure-  where-    onDim (E.NamedDim v) =-      E.NamedDim . E.qualName <$> newVName (baseString $ E.qualLeaf v)-    onDim _ =-      E.NamedDim . E.qualName <$> newVName "size"---- | Replace all named dimensions with a fresh name, and remove all--- constant dimensions.  The point is to remove the constraints, but--- keep the names around.  We use this for constructing the entry--- point parameters.-allDimsFreshInPat :: MonadFreshNames m => E.Pattern -> m E.Pattern-allDimsFreshInPat (PatternAscription p _ _) =-  allDimsFreshInPat p-allDimsFreshInPat (PatternParens p _) =-  allDimsFreshInPat p-allDimsFreshInPat (Id v (Info t) loc) =-  Id v <$> (Info <$> allDimsFreshInType t) <*> pure loc-allDimsFreshInPat (TuplePattern ps loc) =-  TuplePattern <$> mapM allDimsFreshInPat ps <*> pure loc-allDimsFreshInPat (RecordPattern ps loc) =-  RecordPattern <$> mapM (traverse allDimsFreshInPat) ps <*> pure loc-allDimsFreshInPat (Wildcard (Info t) loc) =-  Wildcard <$> (Info <$> allDimsFreshInType t) <*> pure loc-allDimsFreshInPat (PatternLit e (Info t) loc) =-  PatternLit e <$> (Info <$> allDimsFreshInType t) <*> pure loc-allDimsFreshInPat (PatternConstr c (Info t) pats loc) =-  PatternConstr c <$> (Info <$> allDimsFreshInType t)-    <*> mapM allDimsFreshInPat pats-    <*> pure loc--data EntryTrust-  = -- | This parameter or return value is an opaque type.  When a-    -- parameter, this implies that it must have been returned by a-    -- previous call to Futhark, and hence we can preserve (constant)-    -- size constraints.-    EntryTrusted-  | -- | The type is directly exposed.  Any size constraint cannot be-    -- trusted.-    EntryUntrusted--entryTrust :: EntryType -> EntryTrust-entryTrust t-  | E.Scalar (E.Prim E.Unsigned {}) <- E.entryType t =-    EntryUntrusted-  | E.Array _ _ (E.Prim E.Unsigned {}) _ <- E.entryType t =-    EntryUntrusted-  | E.Scalar E.Prim {} <- E.entryType t =-    EntryUntrusted-  | E.Array _ _ E.Prim {} _ <- E.entryType t =-    EntryUntrusted-  | otherwise =-    EntryTrusted--fixEntryParamSizes :: MonadFreshNames m => E.Pattern -> EntryTrust -> m E.Pattern-fixEntryParamSizes p EntryTrusted = pure p-fixEntryParamSizes p EntryUntrusted = allDimsFreshInPat p---- When we are returning a value from the entry point, we fully--- existentialise the return type.  This is because it might otherwise--- refer to sizes that are not in scope, because the generated entry--- point function does not keep the size parameters of the original--- entry point.-fullyExistential ::-  [[I.TypeBase ExtShape u]] ->-  [[I.TypeBase ExtShape u]]-fullyExistential tss =-  evalState (mapM (mapM (bitraverse (traverse onDim) pure)) tss) 0-  where-    onDim _ = do-      i <- get-      modify (+ 1)-      pure $ Ext i- generateEntryPoint :: E.EntryPoint -> E.ValBind -> InternaliseM () generateEntryPoint (E.EntryPoint e_paramts e_rettype) vb = localConstsScope $ do-  let (E.ValBind _ ofname _ (Info (rettype, _)) _ params _ _ attrs loc) = vb-  -- We replace all shape annotations, so there should be no constant-  -- parameters here.-  params_fresh <- zipWithM fixEntryParamSizes params $ map entryTrust e_paramts-  let tparams =-        map (`E.TypeParamDim` mempty) $-          S.toList $-            mconcat $ map E.patternDimNames params_fresh-  bindingParams tparams params_fresh $ \shapeparams params' -> do-    entry_rettype <- fullyExistential <$> internaliseEntryReturnType rettype+  let (E.ValBind _ ofname _ (Info (rettype, _)) tparams params _ _ attrs loc) = vb+  bindingParams tparams params $ \shapeparams params' -> do+    entry_rettype <- internaliseEntryReturnType rettype     let entry' = entryPoint (zip e_paramts params') (e_rettype, entry_rettype)         args = map (I.Var . I.paramName) $ concat params' @@ -305,20 +194,6 @@        in (d + 1, te')     withoutDims te = (0 :: Int, te) -internaliseIdent :: E.Ident -> InternaliseM I.VName-internaliseIdent (E.Ident name (Info tp) loc) =-  case tp of-    E.Scalar E.Prim {} -> return name-    _ ->-      error $-        "Futhark.Internalise.internaliseIdent: asked to internalise non-prim-typed ident '"-          ++ pretty name-          ++ " of type "-          ++ pretty tp-          ++ " at "-          ++ locStr loc-          ++ "."- internaliseBody :: String -> E.Exp -> InternaliseM Body internaliseBody desc e =   insertStmsM $ resultBody <$> internaliseExp (desc <> "_res") e@@ -339,17 +214,11 @@   fmap pure $     letSubExp desc $       I.BasicOp $ I.ArrayLit (map constant vs) $ I.Prim int8-internaliseExp _ (E.Var (E.QualName _ name) (Info t) loc) = do+internaliseExp _ (E.Var (E.QualName _ name) _ _) = do   subst <- lookupSubst name   case subst of     Just substs -> return substs-    Nothing -> do-      -- If this identifier is the name of a constant, we have to turn it-      -- into a call to the corresponding function.-      is_const <- lookupConst name-      case is_const of-        Just ses -> return ses-        Nothing -> (: []) . I.Var <$> internaliseIdent (E.Ident name (Info t) loc)+    Nothing -> pure [I.Var name] internaliseExp desc (E.Index e idxs (Info ret, Info retext) loc) = do   vs <- internaliseExpToVars "indexed" e   dims <- case vs of@@ -644,10 +513,8 @@   ses <- internalisePat desc pat e body (internaliseExp desc)   bindExtSizes (E.toStruct ret) retext ses   return ses-internaliseExp desc (E.LetFun ofname (tparams, params, retdecl, Info rettype, body) letbody _ loc) = do-  internaliseValBind $-    E.ValBind Nothing ofname retdecl (Info (rettype, [])) tparams params body Nothing mempty loc-  internaliseExp desc letbody+internaliseExp _ (E.LetFun ofname _ _ _ _) =+  error $ "Unexpected LetFun " ++ pretty ofname internaliseExp desc (E.DoLoop sparams mergepat mergeexp form loopbody (Info (ret, retext)) loc) = do   ses <- internaliseExp "loop_init" mergeexp   ((loopbody', (form', shapepat, mergepat', mergeinit')), initstms) <-@@ -725,7 +592,6 @@               I.ForLoop i Int64 w loopvars     handleForm mergeinit (E.For i num_iterations) = do       num_iterations' <- internaliseExp1 "upper_bound" num_iterations-      i' <- internaliseIdent i       num_iterations_t <- I.subExpType num_iterations'       it <- case num_iterations_t of         I.Prim (IntType it) -> return it@@ -734,7 +600,7 @@       bindingLoopParams sparams' mergepat $         \shapepat mergepat' ->           forLoop mergepat' shapepat mergeinit $-            I.ForLoop i' it num_iterations' []+            I.ForLoop (E.identName i) it num_iterations' []     handleForm mergeinit (E.While cond) =       bindingLoopParams sparams' mergepat $ \shapepat mergepat' -> do         mergeinit_ts <- mapM subExpType mergeinit@@ -926,34 +792,6 @@  -- Builtin operators are handled specially because they are -- overloaded.-internaliseExp desc (E.BinOp (op, _) _ (xe, _) (ye, _) _ _ loc)-  | Just internalise <- isOverloadedFunction op [xe, ye] loc =-    internalise desc--- User-defined operators are just the same as a function call.-internaliseExp-  desc-  ( E.BinOp-      (op, oploc)-      (Info t)-      (xarg, Info (xt, xext))-      (yarg, Info (yt, yext))-      _-      (Info retext)-      loc-    ) =-    internaliseExp desc $-      E.Apply-        ( E.Apply-            (E.Var op (Info t) oploc)-            xarg-            (Info (E.diet xt, xext))-            (Info $ foldFunType [E.fromStruct yt] t, Info [])-            loc-        )-        yarg-        (Info (E.diet yt, yext))-        (Info t, Info retext)-        loc internaliseExp desc (E.Project k e (Info rt) _) = do   n <- internalisedTypeSize $ rt `setAliases` ()   i' <- fmap sum $@@ -963,6 +801,8 @@           map snd $ takeWhile ((/= k) . fst) $ sortFields fs         t -> [t]   take n . drop i' <$> internaliseExp desc e+internaliseExp _ e@E.BinOp {} =+  error $ "internaliseExp: Unexpected BinOp " ++ pretty e internaliseExp _ e@E.Lambda {} =   error $ "internaliseExp: Unexpected lambda at " ++ locStr (srclocOf e) internaliseExp _ e@E.OpSection {} =@@ -1988,13 +1828,13 @@   SrcLoc ->   InternaliseM ([SubExp], [I.ExtType]) funcall desc (QualName _ fname) args loc = do-  (fname', closure, shapes, value_paramts, fun_params, rettype_fun) <-+  (shapes, value_paramts, fun_params, rettype_fun) <-     lookupFunction fname   argts <- mapM subExpType args    shapeargs <- argShapes shapes fun_params argts   let diets =-        replicate (length closure + length shapeargs) I.ObservePrim+        replicate (length shapeargs) I.ObservePrim           ++ map I.diet value_paramts   args' <-     ensureArgShapes@@ -2002,7 +1842,7 @@       loc       (map I.paramName fun_params)       (map I.paramType fun_params)-      (map I.Var closure ++ shapeargs ++ args)+      (shapeargs ++ args)   argts' <- mapM subExpType args'   case rettype_fun $ zip args' argts' of     Nothing ->@@ -2027,7 +1867,7 @@       ses <-         attributing attrs $           letTupExp' desc $-            I.Apply fname' (zip args' diets) ts (safety, loc, mempty)+            I.Apply (internaliseFunName fname) (zip args' diets) ts (safety, loc, mempty)       return (ses, map I.fromDecl ts)  -- Bind existential names defined by an expression, based on the

@@ -12,13 +12,14 @@ import Data.Bifunctor import Data.Bitraversable import Data.Foldable-import Data.List (nub, partition, sortOn, tails)+import Data.List (partition, sortOn, tails) import qualified Data.List.NonEmpty as NE import qualified Data.Map.Strict as M import Data.Maybe import qualified Data.Sequence as Seq import qualified Data.Set as S import Futhark.IR.Pretty ()+import qualified Futhark.Internalise.FreeVars as FV import Futhark.MonadFreshNames import Language.Futhark import Language.Futhark.Traversals@@ -26,7 +27,7 @@ -- | An expression or an extended 'Lambda' (with size parameters, -- which AST lambdas do not support). data ExtExp-  = ExtLambda [TypeParam] [Pattern] Exp (Aliasing, StructType) SrcLoc+  = ExtLambda [Pattern] Exp StructType SrcLoc   | ExtExp Exp   deriving (Show) @@ -34,20 +35,31 @@ -- defunctionalization of an expression, aside from the residual expression. data StaticVal   = Dynamic PatternType-  | -- | The 'VName's are shape parameters that are bound-    -- by the 'Pattern'.-    LambdaSV [VName] Pattern StructType ExtExp Env+  | LambdaSV Pattern StructType ExtExp Env   | RecordSV [(Name, StaticVal)]   | -- | The constructor that is actually present, plus     -- the others that are not.     SumSV Name [StaticVal] [(Name, [PatternType])]-  | DynamicFun (Exp, StaticVal) StaticVal+  | -- | The pair is the StaticVal and residual expression of this+    -- function as a whole, while the second StaticVal is its+    -- body. (Don't trust this too much, my understanding may have+    -- holes.)+    DynamicFun (Exp, StaticVal) StaticVal   | IntrinsicSV   deriving (Show) --- | Environment mapping variable names to their associated static value.-type Env = M.Map VName StaticVal+-- | The type is Just if this is a polymorphic binding that must be+-- instantiated.+data Binding = Binding (Maybe ([VName], StructType)) StaticVal+  deriving (Show) +bindingSV :: Binding -> StaticVal+bindingSV (Binding _ sv) = sv++-- | Environment mapping variable names to their associated static+-- value.+type Env = M.Map VName Binding+ localEnv :: Env -> DefM a -> DefM a localEnv env = local $ Arrow.second (env <>) @@ -57,96 +69,153 @@ localNewEnv env = local $ \(globals, old_env) ->   (globals, M.filterWithKey (\k _ -> k `S.member` globals) old_env <> env) -extendEnv :: VName -> StaticVal -> DefM a -> DefM a-extendEnv vn sv = localEnv (M.singleton vn sv)- askEnv :: DefM Env askEnv = asks snd  isGlobal :: VName -> DefM a -> DefM a isGlobal v = local $ Arrow.first (S.insert v) -replaceStaticValSizes :: M.Map VName VName -> StaticVal -> StaticVal-replaceStaticValSizes substs sv =+replaceTypeSizes ::+  M.Map VName SizeSubst ->+  TypeBase (DimDecl VName) als ->+  TypeBase (DimDecl VName) als+replaceTypeSizes substs = first onDim+  where+    onDim (NamedDim v) =+      case M.lookup (qualLeaf v) substs of+        Just (SubstNamed v') -> NamedDim v'+        Just (SubstConst d) -> ConstDim d+        Nothing -> NamedDim v+    onDim d = d++replaceStaticValSizes ::+  S.Set VName ->+  M.Map VName SizeSubst ->+  StaticVal ->+  StaticVal+replaceStaticValSizes globals orig_substs sv =   case sv of-    LambdaSV sizes param t e closure_env ->-      LambdaSV-        sizes-        (onAST param)-        (onType t)-        (onExtExp e)-        (onEnv closure_env)+    _ | M.null orig_substs -> sv+    LambdaSV param t e closure_env ->+      let substs =+            foldl' (flip M.delete) orig_substs $+              S.fromList (M.keys closure_env)+       in LambdaSV+            (onAST substs param)+            (replaceTypeSizes substs t)+            (onExtExp substs e)+            (onEnv orig_substs closure_env) --intentional     Dynamic t ->-      Dynamic $ onType t+      Dynamic $ replaceTypeSizes orig_substs t     RecordSV fs ->-      RecordSV $ map (fmap (replaceStaticValSizes substs)) fs+      RecordSV $ map (fmap (replaceStaticValSizes globals orig_substs)) fs     SumSV c svs ts ->-      SumSV c (map (replaceStaticValSizes substs) svs) $-        map (fmap (map onType)) ts+      SumSV c (map (replaceStaticValSizes globals orig_substs) svs) $+        map (fmap $ map $ replaceTypeSizes orig_substs) ts     DynamicFun (e, sv1) sv2 ->-      DynamicFun (onAST e, replaceStaticValSizes substs sv1) $-        replaceStaticValSizes substs sv2+      DynamicFun (onExp orig_substs e, replaceStaticValSizes globals orig_substs sv1) $+        replaceStaticValSizes globals orig_substs sv2     IntrinsicSV ->       IntrinsicSV   where-    onName v = fromMaybe v $ M.lookup v substs-    onQualName v = maybe v qualName $ M.lookup (qualLeaf v) substs--    tv =+    tv substs =       identityMapper-        { mapOnPatternType = pure . onType,-          mapOnStructType = pure . onType,-          mapOnQualName = pure . onQualName,-          mapOnExp = pure . onAST+        { mapOnPatternType = pure . replaceTypeSizes substs,+          mapOnStructType = pure . replaceTypeSizes substs,+          mapOnExp = pure . onExp substs,+          mapOnName = pure . onName substs         } -    onExtExp (ExtExp e) =-      ExtExp $ onAST e-    onExtExp (ExtLambda dims params e (als, t) loc) =-      ExtLambda dims (map onAST params) (onAST e) (als, onType t) loc+    onName substs v =+      case M.lookup v substs of+        Just (SubstNamed v') -> qualLeaf v'+        _ -> v -    onEnv =+    onExp substs (Var v t loc) =+      case M.lookup (qualLeaf v) substs of+        Just (SubstNamed v') ->+          Var v' t loc+        Just (SubstConst d) ->+          Literal (SignedValue (Int64Value (fromIntegral d))) loc+        Nothing ->+          Var v (replaceTypeSizes substs <$> t) loc+    onExp substs (Coerce e tdecl t loc) =+      Coerce (onExp substs e) tdecl' (first (fmap (replaceTypeSizes substs)) t) loc+      where+        tdecl' =+          TypeDecl+            { declaredType = onTypeExp substs $ declaredType tdecl,+              expandedType = replaceTypeSizes substs <$> expandedType tdecl+            }+    onExp substs e = onAST substs e++    onTypeExpDim substs d@(DimExpNamed v loc) =+      case M.lookup (qualLeaf v) substs of+        Just (SubstNamed v') ->+          DimExpNamed v' loc+        Just (SubstConst x) ->+          DimExpConst x loc+        Nothing ->+          d+    onTypeExpDim _ d = d++    onTypeArgExp substs (TypeArgExpDim d loc) =+      TypeArgExpDim (onTypeExpDim substs d) loc+    onTypeArgExp substs (TypeArgExpType te) =+      TypeArgExpType (onTypeExp substs te)++    onTypeExp substs (TEArray te d loc) =+      TEArray (onTypeExp substs te) (onTypeExpDim substs d) loc+    onTypeExp substs (TEUnique t loc) =+      TEUnique (onTypeExp substs t) loc+    onTypeExp substs (TEApply t1 t2 loc) =+      TEApply (onTypeExp substs t1) (onTypeArgExp substs t2) loc+    onTypeExp substs (TEArrow p t1 t2 loc) =+      TEArrow p (onTypeExp substs t1) (onTypeExp substs t2) loc+    onTypeExp substs (TETuple ts loc) =+      TETuple (map (onTypeExp substs) ts) loc+    onTypeExp substs (TERecord ts loc) =+      TERecord (map (fmap $ onTypeExp substs) ts) loc+    onTypeExp substs (TESum ts loc) =+      TESum (map (fmap $ map $ onTypeExp substs) ts) loc+    onTypeExp _ (TEVar v loc) =+      TEVar v loc++    onExtExp substs (ExtExp e) =+      ExtExp $ onExp substs e+    onExtExp substs (ExtLambda params e t loc) =+      ExtLambda (map (onAST substs) params) (onExp substs e) (replaceTypeSizes substs t) loc++    onEnv substs =       M.fromList-        . map (bimap onName $ replaceStaticValSizes substs)+        . map (second (onBinding substs))         . M.toList -    onAST :: ASTMappable x => x -> x-    onAST = runIdentity . astMap tv--    onType = first onDim-      where-        onDim (NamedDim v) =-          NamedDim $ maybe v qualName $ M.lookup (qualLeaf v) substs-        onDim d = d+    onBinding substs (Binding t bsv) =+      Binding+        (second (replaceTypeSizes substs) <$> t)+        (replaceStaticValSizes globals substs bsv) --- | Construct new sizes for a LambdaSV (if that is what it is).  This--- is needed because sizes must be unique when we substitute the--- closure for the LambdaSV into another function, because sizes float--- to the top (see issue #1147).-newSizesForLambda :: StaticVal -> DefM StaticVal-newSizesForLambda (LambdaSV sizes param t e closure_env) = do-  sizes' <- mapM newName sizes-  let substs = M.fromList $ zip sizes sizes'-  pure $ replaceStaticValSizes substs $ LambdaSV sizes' param t e closure_env-newSizesForLambda sv = pure sv+    onAST :: ASTMappable x => M.Map VName SizeSubst -> x -> x+    onAST substs = runIdentity . astMap (tv substs)  -- | Returns the defunctionalization environment restricted -- to the given set of variable names and types.-restrictEnvTo :: NameSet -> DefM Env-restrictEnvTo (NameSet m) = restrict <$> ask+restrictEnvTo :: FV.NameSet -> DefM Env+restrictEnvTo (FV.NameSet m) = restrict <$> ask   where     restrict (globals, env) = M.mapMaybeWithKey keep env       where-        keep k sv = do+        keep k (Binding t sv) = do           guard $ not $ k `S.member` globals-          u <- M.lookup k m-          Just $ restrict' u sv+          u <- uniqueness <$> M.lookup k m+          Just $ Binding t $ restrict' u sv     restrict' Nonunique (Dynamic t) =       Dynamic $ t `setUniqueness` Nonunique     restrict' _ (Dynamic t) =       Dynamic t-    restrict' u (LambdaSV dims pat t e env) =-      LambdaSV dims pat t e $ M.map (restrict' u) env+    restrict' u (LambdaSV pat t e env) =+      LambdaSV pat t e $ M.map (restrict'' u) env     restrict' u (RecordSV fields) =       RecordSV $ map (fmap $ restrict' u) fields     restrict' u (SumSV c svs fields) =@@ -154,6 +223,7 @@     restrict' u (DynamicFun (e, sv1) sv2) =       DynamicFun (e, restrict' u sv1) $ restrict' u sv2     restrict' _ IntrinsicSV = IntrinsicSV+    restrict'' u (Binding t sv) = Binding t $ restrict' u sv  -- | Defunctionalization monad.  The Reader environment tracks both -- the current Env as well as the set of globally defined dynamic@@ -180,11 +250,15 @@   return ((x, decs), const mempty)  -- | Looks up the associated static value for a given name in the environment.-lookupVar :: SrcLoc -> VName -> DefM StaticVal-lookupVar loc x = do+lookupVar :: StructType -> SrcLoc -> VName -> DefM StaticVal+lookupVar t loc x = do   env <- askEnv   case M.lookup x env of-    Just sv -> return sv+    Just (Binding (Just (dims, sv_t)) sv) -> do+      globals <- asks fst+      instStaticVal globals dims t sv_t sv+    Just (Binding Nothing sv) ->+      pure sv     Nothing -- If the variable is unknown, it may refer to the 'intrinsics'     -- module, which we will have to treat specially.       | baseTag x <= maxIntrinsicTag -> return IntrinsicSV@@ -220,94 +294,143 @@ patternArraySizes :: Pattern -> S.Set VName patternArraySizes = arraySizes . patternStructType +data SizeSubst+  = SubstNamed (QualName VName)+  | SubstConst Int+  deriving (Eq, Ord, Show)+ dimMapping ::   Monoid a =>   TypeBase (DimDecl VName) a ->   TypeBase (DimDecl VName) a ->-  M.Map VName VName+  M.Map VName SizeSubst dimMapping t1 t2 = execState (matchDims f t1 t2) mempty   where     f (NamedDim d1) (NamedDim d2) = do-      modify $ M.insert (qualLeaf d1) (qualLeaf d2)+      modify $ M.insert (qualLeaf d1) $ SubstNamed d2       return $ NamedDim d1+    f (NamedDim d1) (ConstDim d2) = do+      modify $ M.insert (qualLeaf d1) $ SubstConst d2+      return $ NamedDim d1     f d _ = return d +dimMapping' ::+  Monoid a =>+  TypeBase (DimDecl VName) a ->+  TypeBase (DimDecl VName) a ->+  M.Map VName VName+dimMapping' t1 t2 = M.mapMaybe f $ dimMapping t1 t2+  where+    f (SubstNamed d) = Just $ qualLeaf d+    f _ = Nothing++sizesToRename :: StaticVal -> S.Set VName+sizesToRename (DynamicFun (_, sv1) sv2) =+  sizesToRename sv1 <> sizesToRename sv2+sizesToRename IntrinsicSV =+  mempty+sizesToRename Dynamic {} =+  mempty+sizesToRename (RecordSV fs) =+  foldMap (sizesToRename . snd) fs+sizesToRename (SumSV _ svs _) =+  foldMap sizesToRename svs+sizesToRename (LambdaSV param _ _ _) =+  patternDimNames param+    <> S.map identName (S.filter couldBeSize $ patternIdents param)+  where+    couldBeSize ident =+      unInfo (identType ident) == Scalar (Prim (Signed Int64))++-- When we instantiate a polymorphic StaticVal, we rename all the+-- sizes to avoid name conflicts later on.  This is a bit of a hack...+instStaticVal ::+  MonadFreshNames m =>+  S.Set VName ->+  [VName] ->+  StructType ->+  StructType ->+  StaticVal ->+  m StaticVal+instStaticVal globals dims t sv_t sv = do+  fresh_substs <- mkSubsts $ S.toList $ S.fromList dims <> sizesToRename sv++  let dims' = map (onName fresh_substs) dims+      isDim k _ = k `elem` dims'+      dim_substs =+        M.filterWithKey isDim $ dimMapping (replaceTypeSizes fresh_substs sv_t) t+      replace (SubstNamed k) = fromMaybe (SubstNamed k) $ M.lookup (qualLeaf k) dim_substs+      replace k = k+      substs = M.map replace fresh_substs <> dim_substs++  pure $ replaceStaticValSizes globals substs sv+  where+    mkSubsts names =+      M.fromList . zip names . map (SubstNamed . qualName)+        <$> mapM newName names++    onName substs v =+      case M.lookup v substs of+        Just (SubstNamed v') -> qualLeaf v'+        _ -> v+ defuncFun ::-  [TypeParam] ->+  [VName] ->   [Pattern] ->   Exp ->-  (Aliasing, StructType) ->+  StructType ->   SrcLoc ->   DefM (Exp, StaticVal)-defuncFun tparams pats e0 (closure, ret) loc = do-  when (any isTypeParam tparams) $-    error $-      "Received a lambda with type parameters at " ++ locStr loc-        ++ ", but the defunctionalizer expects a monomorphic input program."+defuncFun tparams pats e0 ret loc = do   -- Extract the first parameter of the lambda and "push" the   -- remaining ones (if there are any) into the body of the lambda.-  let (dims, pat, ret', e0') = case pats of+  let (pat, ret', e0') = case pats of         [] -> error "Received a lambda with no parameters."-        [pat'] -> (map typeParamName tparams, pat', ret, ExtExp e0)+        [pat'] -> (pat', ret, ExtExp e0)         (pat' : pats') ->-          -- Split shape parameters into those that are determined by-          -- the first pattern, and those that are determined by later-          -- patterns.-          let bound_by_pat = (`S.member` patternArraySizes pat') . typeParamName-              (pat_dims, rest_dims) = partition bound_by_pat tparams-           in ( map typeParamName pat_dims,-                pat',-                foldFunType (map (toStruct . patternType) pats') ret,-                ExtLambda rest_dims pats' e0 (closure, ret) loc-              )+          ( pat',+            foldFunType (map (toStruct . patternType) pats') ret,+            ExtLambda pats' e0 ret loc+          )    -- Construct a record literal that closes over the environment of   -- the lambda.  Closed-over 'DynamicFun's are converted to their   -- closure representation.   let used =-        freeVars (Lambda pats e0 Nothing (Info (closure, ret)) loc)-          `without` mconcat (map oneName dims)+        FV.freeVars (Lambda pats e0 Nothing (Info (mempty, ret)) loc)+          `FV.without` S.fromList tparams   used_env <- restrictEnvTo used    -- The closure parts that are sizes are proactively turned into size   -- parameters.   let sizes_of_arrays =-        foldMap (arraySizes . toStruct . typeFromSV') used_env+        foldMap (arraySizes . toStruct . typeFromSV . bindingSV) used_env           <> patternArraySizes pat       notSize = not . (`S.member` sizes_of_arrays)       (fields, env) =-        unzip $-          map closureFromDynamicFun $-            filter (notSize . fst) $ M.toList used_env-      env' = M.fromList env-      closure_dims = S.toList sizes_of_arrays--  global <- asks fst+        second M.fromList $+          unzip $+            map closureFromDynamicFun $+              filter (notSize . fst) $ M.toList used_env    return     ( RecordLit fields loc,-      LambdaSV-        ( nub $-            filter (`S.notMember` global) $-              dims <> closure_dims-        )-        pat-        ret'-        e0'-        env'+      LambdaSV pat ret' e0' env     )   where-    closureFromDynamicFun (vn, DynamicFun (clsr_env, sv) _) =+    closureFromDynamicFun (vn, Binding _ (DynamicFun (clsr_env, sv) _)) =       let name = nameFromString $ pretty vn-       in (RecordFieldExplicit name clsr_env mempty, (vn, sv))-    closureFromDynamicFun (vn, sv) =+       in ( RecordFieldExplicit name clsr_env mempty,+            (vn, Binding Nothing sv)+          )+    closureFromDynamicFun (vn, Binding _ sv) =       let name = nameFromString $ pretty vn-          tp' = typeFromSV' sv+          tp' = typeFromSV sv        in ( RecordFieldExplicit               name               (Var (qualName vn) (Info tp') mempty)               mempty,-            (vn, sv)+            (vn, Binding Nothing sv)           )  -- | Defunctionalization of an expression. Returns the residual expression and@@ -337,8 +460,8 @@     defuncField (RecordFieldExplicit vn e loc') = do       (e', sv) <- defuncExp e       return (RecordFieldExplicit vn e' loc', (vn, sv))-    defuncField (RecordFieldImplicit vn _ loc') = do-      sv <- lookupVar loc' vn+    defuncField (RecordFieldImplicit vn (Info t) loc') = do+      sv <- lookupVar (toStruct t) loc' vn       case sv of         -- If the implicit field refers to a dynamic function, we         -- convert it to an explicit field with a record closing over@@ -352,7 +475,7 @@         -- The field may refer to a functional expression, so we get the         -- type from the static value and not the one from the AST.         _ ->-          let tp = Info $ typeFromSV' sv+          let tp = Info $ typeFromSV sv            in return (RecordFieldImplicit vn tp loc', (baseName vn, sv)) defuncExp (ArrayLit es t@(Info t') loc) = do   es' <- mapM defuncExp' es@@ -362,8 +485,8 @@   me' <- mapM defuncExp' me   incl' <- mapM defuncExp' incl   return (Range e1' me' incl' t loc, Dynamic t')-defuncExp e@(Var qn _ loc) = do-  sv <- lookupVar loc (qualLeaf qn)+defuncExp e@(Var qn (Info t) loc) = do+  sv <- lookupVar (toStruct t) loc (qualLeaf qn)   case sv of     -- If the variable refers to a dynamic function, we return its closure     -- representation (i.e., a record expression capturing the free variables@@ -375,7 +498,7 @@       (pats, body, tp) <- etaExpand (typeOf e) e       defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) mempty     _ ->-      let tp = typeFromSV' sv+      let tp = typeFromSV sv        in return (Var qn (Info tp) loc, sv) defuncExp (Ascript e0 tydecl loc)   | orderZero (typeOf e0) = do@@ -390,31 +513,17 @@ defuncExp (LetPat pat e1 e2 (Info t, retext) loc) = do   (e1', sv1) <- defuncExp e1   let env = matchPatternSV pat sv1-      pat' = updatePattern' pat sv1+      pat' = updatePattern pat sv1   (e2', sv2) <- localEnv env $ defuncExp e2   -- To maintain any sizes going out of scope, we need to compute the   -- old size substitution induced by retext and also apply it to the   -- newly computed body type.-  let mapping = dimMapping (typeOf e2) t+  let mapping = dimMapping' (typeOf e2) t       subst v = fromMaybe v $ M.lookup v mapping       t' = first (fmap subst) $ typeOf e2'   return (LetPat pat' e1' e2' (Info t', retext) loc, sv2)---- Local functions are handled by rewriting them to lambdas, so that--- the same machinery can be re-used.  But we may have to eta-expand--- first.-defuncExp (LetFun vn (dims, pats, _, Info ret, e1) e2 let_t loc)-  | Scalar Arrow {} <- ret = do-    (body_pats, e1', ret') <- etaExpand (fromStruct ret) e1-    let f = (dims, pats <> body_pats, Nothing, Info ret', e1')-    defuncExp $ LetFun vn f e2 let_t loc-  | otherwise = do-    (e1', sv1) <- defuncFun dims pats e1 (mempty, ret) loc-    (e2', sv2) <- localEnv (M.singleton vn sv1) $ defuncExp e2-    return-      ( LetPat (Id vn (Info (typeOf e1')) loc) e1' e2' (Info $ typeOf e2', Info []) loc,-        sv2-      )+defuncExp (LetFun vn _ _ _ _) =+  error $ "defuncExp: Unexpected LetFun: " ++ prettyName vn defuncExp (If e1 e2 e3 tp loc) = do   (e1', _) <- defuncExp e1   (e2', sv) <- defuncExp e2@@ -433,8 +542,8 @@ defuncExp (Negate e0 loc) = do   (e0', sv) <- defuncExp e0   return (Negate e0' loc, sv)-defuncExp (Lambda pats e0 _ (Info (closure, ret)) loc) =-  defuncFun [] pats e0 (closure, ret) loc+defuncExp (Lambda pats e0 _ (Info (_, ret)) loc) =+  defuncFun [] pats e0 ret loc -- Operator sections are expected to be converted to lambda-expressions -- by the monomorphizer, so they should no longer occur at this point. defuncExp OpSection {} = error "defuncExp: unexpected operator section."@@ -459,45 +568,24 @@   return (DoLoop sparams pat e1' form' e3' ret loc, sv)   where     envFromIdent (Ident vn (Info tp) _) =-      M.singleton vn $ Dynamic tp---- We handle BinOps by turning them into ordinary function applications.-defuncExp-  ( BinOp-      (qn, qnloc)-      (Info t)-      (e1, Info (pt1, ext1))-      (e2, Info (pt2, ext2))-      (Info ret)-      (Info retext)-      loc-    ) =-    defuncExp $-      Apply-        ( Apply-            (Var qn (Info t) qnloc)-            e1-            (Info (diet pt1, ext1))-            (Info (Scalar $ Arrow mempty Unnamed (fromStruct pt2) ret), Info [])-            loc-        )-        e2-        (Info (diet pt2, ext2))-        (Info ret, Info retext)-        loc+      M.singleton vn $ Binding Nothing $ Dynamic tp+defuncExp e@BinOp {} =+  error $ "defuncExp: unexpected binary operator: " ++ pretty e defuncExp (Project vn e0 tp@(Info tp') loc) = do   (e0', sv0) <- defuncExp e0   case sv0 of     RecordSV svs -> case lookup vn svs of-      Just sv -> return (Project vn e0' (Info $ typeFromSV' sv) loc, sv)+      Just sv -> return (Project vn e0' (Info $ typeFromSV sv) loc, sv)       Nothing -> error "Invalid record projection."     Dynamic _ -> return (Project vn e0' tp loc, Dynamic tp')     _ -> error $ "Projection of an expression with static value " ++ show sv0 defuncExp (LetWith id1 id2 idxs e1 body t loc) = do   e1' <- defuncExp' e1-  sv1 <- lookupVar (identSrcLoc id2) $ identName id2   idxs' <- mapM defuncDimIndex idxs-  (body', sv) <- extendEnv (identName id1) sv1 $ defuncExp body+  let id1_binding = Binding Nothing $ Dynamic $ unInfo $ identType id1+  (body', sv) <-+    localEnv (M.singleton (identName id1) id1_binding) $+      defuncExp body   return (LetWith id1 id2 idxs' e1' body' t loc, sv) defuncExp expr@(Index e0 idxs info loc) = do   e0' <- defuncExp' e0@@ -517,7 +605,7 @@   (e2', sv2) <- defuncExp e2   let sv = staticField sv1 sv2 fs   return-    ( RecordUpdate e1' fs e2' (Info $ typeFromSV' sv1) loc,+    ( RecordUpdate e1' fs e2' (Info $ typeFromSV sv1) loc,       sv     )   where@@ -540,7 +628,7 @@         SumSV name svs $           M.toList $             name `M.delete` M.map (map defuncType) all_fs-  return (Constr name es' (Info (typeFromSV' sv)) loc, sv)+  return (Constr name es' (Info (typeFromSV sv)) loc, sv)   where     defuncType ::       Monoid als =>@@ -580,13 +668,12 @@  defuncExtExp :: ExtExp -> DefM (Exp, StaticVal) defuncExtExp (ExtExp e) = defuncExp e-defuncExtExp (ExtLambda tparams pats e0 (closure, ret) loc) =-  traverse newSizesForLambda-    =<< defuncFun tparams pats e0 (closure, ret) loc+defuncExtExp (ExtLambda pats e0 ret loc) =+  defuncFun [] pats e0 ret loc  defuncCase :: StaticVal -> Case -> DefM (Case, StaticVal) defuncCase sv (CasePat p e loc) = do-  let p' = updatePattern' p sv+  let p' = updatePattern p sv       env = matchPatternSV p sv   (e', sv') <- localEnv env $ defuncExp e   return (CasePat p' e' loc, sv')@@ -652,22 +739,27 @@ -- | Defunctionalize a let-bound function, while preserving parameters -- that have order 0 types (i.e., non-functional). defuncLet ::-  [TypeParam] ->+  [VName] ->   [Pattern] ->   Exp ->   StructType ->-  DefM ([TypeParam], [Pattern], Exp, StaticVal)+  DefM ([VName], [Pattern], Exp, StaticVal) defuncLet dims ps@(pat : pats) body rettype   | patternOrderZero pat = do-    let bound_by_pat = (`S.member` patternDimNames pat) . typeParamName+    let bound_by_pat = (`S.member` patternDimNames pat)         -- Take care to not include more size parameters than necessary.         (pat_dims, rest_dims) = partition bound_by_pat dims-        env = envFromPattern pat <> envFromShapeParams pat_dims+        env = envFromPattern pat <> envFromDimNames pat_dims     (rest_dims', pats', body', sv) <- localEnv env $ defuncLet rest_dims pats body rettype-    closure <- defuncFun dims ps body (mempty, rettype) mempty-    return (pat_dims ++ rest_dims', pat : pats', body', DynamicFun closure sv)+    closure <- defuncFun dims ps body rettype mempty+    return+      ( pat_dims ++ rest_dims',+        pat : pats',+        body',+        DynamicFun closure sv+      )   | otherwise = do-    (e, sv) <- defuncFun dims ps body (mempty, rettype) mempty+    (e, sv) <- defuncFun dims ps body rettype mempty     return ([], [], e, sv) defuncLet _ [] body rettype = do   (body', sv) <- defuncExp body@@ -679,16 +771,21 @@       RecordSV $ M.toList $ M.intersectionWith imposeType (M.fromList fs1) fs2     imposeType sv _ = sv -sizesForAll :: MonadFreshNames m => [Pattern] -> m ([VName], [Pattern])-sizesForAll params = do-  (params', sizes) <- runStateT (mapM (astMap tv) params) []-  return (sizes, params')+sizesForAll :: MonadFreshNames m => S.Set VName -> [Pattern] -> m ([VName], [Pattern])+sizesForAll bound_sizes params = do+  (params', sizes) <- runStateT (mapM (astMap tv) params) mempty+  return (S.toList sizes, params')   where+    bound = bound_sizes <> foldMap patternNames params     tv = identityMapper {mapOnPatternType = bitraverse onDim pure}     onDim AnyDim = do       v <- lift $ newVName "size"-      modify (v :)+      modify $ S.insert v       pure $ NamedDim $ qualName v+    onDim (NamedDim d) = do+      unless (qualLeaf d `S.member` bound) $+        modify $ S.insert $ qualLeaf d+      pure $ NamedDim d     onDim d = pure d  -- | Defunctionalize an application expression at a given depth of application.@@ -702,12 +799,14 @@   (e2', sv2) <- defuncExp e2   let e' = Apply e1' e2' d t loc   case sv1 of-    LambdaSV dims pat e0_t e0 closure_env -> do+    LambdaSV pat e0_t e0 closure_env -> do       let env' = matchPatternSV pat sv2-          env_dim = envFromDimNames dims-      (e0', sv) <- localNewEnv (env' <> closure_env <> env_dim) $ defuncExtExp e0+          dims = mempty+      (e0', sv) <-+        localNewEnv (env' <> closure_env) $+          defuncExtExp e0 -      let closure_pat = buildEnvPattern closure_env+      let closure_pat = buildEnvPattern dims closure_env           pat' = updatePattern pat sv2        globals <- asks fst@@ -718,13 +817,14 @@       -- and a hack.  There is some piece we're missing.       let params = [closure_pat, pat']           params_for_rettype = params ++ svParams sv1 ++ svParams sv2-          svParams (LambdaSV _ sv_pat _ _ _) = [sv_pat]+          svParams (LambdaSV sv_pat _ _ _) = [sv_pat]           svParams _ = []           rettype = buildRetType closure_env params_for_rettype e0_t $ typeOf e0'            already_bound =             globals <> S.fromList dims               <> S.map identName (foldMap patternIdents params)+           more_dims =             S.toList $               S.filter (`S.notMember` already_bound) $@@ -734,15 +834,16 @@           -- into the name of the lifted function, to make the           -- result slightly more human-readable.           liftedName i (Var f _ _) =-            "lifted_" ++ show i ++ "_" ++ baseString (qualLeaf f)+            "defunc_" ++ show i ++ "_" ++ baseString (qualLeaf f)           liftedName i (Apply f _ _ _ _) =             liftedName (i + 1) f-          liftedName _ _ = "lifted"+          liftedName _ _ = "defunc"        -- Ensure that no parameter sizes are AnyDim.  The internaliser       -- expects this.  This is easy, because they are all       -- first-order.-      (missing_dims, params') <- sizesForAll params+      let bound_sizes = S.fromList (dims <> more_dims) <> globals+      (missing_dims, params') <- sizesForAll bound_sizes params        fname <- newNameFromString $ liftedName (0 :: Int) e1       liftValDec@@ -805,8 +906,7 @@             | orderZero ret = (Info ret, Info ext)             | otherwise = (Info restype, Info ext)           apply_e = Apply e1' e2' d callret loc-      sv' <- newSizesForLambda sv-      return (apply_e, sv')+      return (apply_e, sv)     -- Propagate the 'IntrinsicsSV' until we reach the outermost application,     -- where we construct a dynamic static value with the appropriate type.     IntrinsicSV@@ -829,22 +929,29 @@           ++ show sv1 defuncApply depth e@(Var qn (Info t) loc) = do   let (argtypes, _) = unfoldFunType t-  sv <- lookupVar loc (qualLeaf qn)+  sv <- lookupVar (toStruct t) loc (qualLeaf qn)+   case sv of     DynamicFun _ _-      | fullyApplied sv depth ->+      | fullyApplied sv depth -> do         -- We still need to update the types in case the dynamic         -- function returns a higher-order term.         let (argtypes', rettype) = dynamicFunType sv argtypes-         in return (Var qn (Info (foldFunType argtypes' rettype)) loc, sv)+        return (Var qn (Info (foldFunType argtypes' rettype)) loc, sv)       | otherwise -> do         fname <- newName $ qualLeaf qn-        let (dims, pats, e0, sv') = liftDynFun sv depth-            pats_names = S.map identName $ mconcat $ map patternIdents pats-            notInPats = (`S.notMember` pats_names)-            dims' = filter notInPats dims+        let (pats, e0, sv') = liftDynFun (pretty qn) sv depth             (argtypes', rettype) = dynamicFunType sv' argtypes-        liftValDec fname (fromStruct rettype) dims' pats e0+            dims' = mempty++        -- Ensure that no parameter sizes are AnyDim.  The internaliser+        -- expects this.  This is easy, because they are all+        -- first-order.+        globals <- asks fst+        let bound_sizes = S.fromList dims' <> globals+        (missing_dims, pats') <- sizesForAll bound_sizes pats++        liftValDec fname (fromStruct rettype) (dims' ++ missing_dims) pats' e0         return           ( Var               (qualName fname)@@ -853,7 +960,7 @@             sv'           )     IntrinsicSV -> return (e, IntrinsicSV)-    _ -> return (Var qn (Info (typeFromSV' sv)) loc, sv)+    _ -> return (Var qn (Info (typeFromSV sv)) loc, sv) defuncApply depth (Parens e _) = defuncApply depth e defuncApply _ expr = defuncExp expr @@ -869,16 +976,19 @@ -- dimensions, a list of parameters, a function body, and the -- appropriate static value for applying the function at the given -- depth of partial application.-liftDynFun :: StaticVal -> Int -> ([VName], [Pattern], Exp, StaticVal)-liftDynFun (DynamicFun (e, sv) _) 0 = ([], [], e, sv)-liftDynFun (DynamicFun clsr@(_, LambdaSV dims pat _ _ _) sv) d+liftDynFun :: String -> StaticVal -> Int -> ([Pattern], Exp, StaticVal)+liftDynFun _ (DynamicFun (e, sv) _) 0 = ([], e, sv)+liftDynFun s (DynamicFun clsr@(_, LambdaSV pat _ _ _) sv) d   | d > 0 =-    let (dims', pats, e', sv') = liftDynFun sv (d -1)-     in (nub $ dims ++ dims', pat : pats, e', DynamicFun clsr sv')-liftDynFun sv _ =+    let (pats, e', sv') = liftDynFun s sv (d -1)+     in (pat : pats, e', DynamicFun clsr sv')+liftDynFun s sv d =   error $-    "Tried to lift a StaticVal " ++ show sv-      ++ ", but expected a dynamic function."+    s+      ++ " Tried to lift a StaticVal "+      ++ take 100 (show sv)+      ++ ", but expected a dynamic function.\n"+      ++ pretty d  -- | Converts a pattern to an environment that binds the individual names of the -- pattern to their corresponding types wrapped in a 'Dynamic' static value.@@ -887,28 +997,16 @@   TuplePattern ps _ -> foldMap envFromPattern ps   RecordPattern fs _ -> foldMap (envFromPattern . snd) fs   PatternParens p _ -> envFromPattern p-  Id vn (Info t) _ -> M.singleton vn $ Dynamic t+  Id vn (Info t) _ -> M.singleton vn $ Binding Nothing $ Dynamic t   Wildcard _ _ -> mempty   PatternAscription p _ _ -> envFromPattern p   PatternLit {} -> mempty   PatternConstr _ _ ps _ -> foldMap envFromPattern ps --- | Create an environment that binds the shape parameters.-envFromShapeParams :: [TypeParamBase VName] -> Env-envFromShapeParams = envFromDimNames . map dim-  where-    dim (TypeParamDim vn _) = vn-    dim tparam =-      error $-        "The defunctionalizer expects a monomorphic input program,\n"-          ++ "but it received a type parameter "-          ++ pretty tparam-          ++ " at "-          ++ locStr (srclocOf tparam)-          ++ "."- envFromDimNames :: [VName] -> Env-envFromDimNames = M.fromList . flip zip (repeat $ Dynamic $ Scalar $ Prim $ Signed Int64)+envFromDimNames = M.fromList . flip zip (repeat d)+  where+    d = Binding Nothing $ Dynamic $ Scalar $ Prim $ Signed Int64  -- | Create a new top-level value declaration with the given function name, -- return type, list of parameters, and body expression.@@ -942,13 +1040,16 @@         }  -- | Given a closure environment, construct a record pattern that--- binds the closed over variables.-buildEnvPattern :: Env -> Pattern-buildEnvPattern env = RecordPattern (map buildField $ M.toList env) mempty+-- binds the closed over variables.  Insert wildcard for any patterns+-- that would otherwise clash with size parameters.+buildEnvPattern :: [VName] -> Env -> Pattern+buildEnvPattern sizes env = RecordPattern (map buildField $ M.toList env) mempty   where-    buildField (vn, sv) =+    buildField (vn, Binding _ sv) =       ( nameFromString (pretty vn),-        Id vn (Info $ snd $ typeFromSV sv) mempty+        if vn `elem` sizes+          then Wildcard (Info $ typeFromSV sv) mempty+          else Id vn (Info $ typeFromSV sv) mempty       )  -- | Given a closure environment pattern and the type of a term,@@ -961,11 +1062,12 @@ buildRetType :: Env -> [Pattern] -> StructType -> PatternType -> PatternType buildRetType env pats = comb   where-    bound = foldMap oneName (M.keys env) <> foldMap patternVars pats+    bound =+      S.fromList (M.keys env) <> S.map identName (foldMap patternIdents pats)     boundAsUnique v =       maybe False (unique . unInfo . identType) $         find ((== v) . identName) $ S.toList $ foldMap patternIdents pats-    problematic v = (v `member` bound) && not (boundAsUnique v)+    problematic v = (v `S.member` bound) && not (boundAsUnique v)     comb (Scalar (Record fs_annot)) (Scalar (Record fs_got)) =       Scalar $ Record $ M.intersectionWith comb fs_annot fs_got     comb (Scalar (Sum cs_annot)) (Scalar (Sum cs_got)) =@@ -980,43 +1082,34 @@     descend (Scalar (Record t)) = Scalar $ Record $ fmap descend t     descend t = t --- | Compute the corresponding type for a given static value.-typeFromSV :: StaticVal -> ([VName], PatternType)+-- | Compute the corresponding type for the *representation* of a+-- given static value (not the original possibly higher-order value).+typeFromSV :: StaticVal -> PatternType typeFromSV (Dynamic tp) =-  (mempty, tp)-typeFromSV (LambdaSV sizes _ _ _ env) =-  ( sizes <> env_sizes,-    Scalar $ Record $ M.fromList $ map (fmap snd) env'-  )-  where-    env' = map (bimap (nameFromString . pretty) typeFromSV) $ M.toList env-    env_sizes = concatMap (fst . snd) env'+  tp+typeFromSV (LambdaSV _ _ _ env) =+  Scalar $+    Record $+      M.fromList $+        map (bimap (nameFromString . pretty) (typeFromSV . bindingSV)) $+          M.toList env typeFromSV (RecordSV ls) =   let ts = map (fmap typeFromSV) ls-   in ( concatMap (fst . snd) ts,-        Scalar $ Record $ M.fromList $ map (fmap snd) ts-      )+   in Scalar $ Record $ M.fromList ts typeFromSV (DynamicFun (_, sv) _) =   typeFromSV sv typeFromSV (SumSV name svs fields) =-  let (sizes, svs') = unzip $ map typeFromSV svs-   in ( concat sizes,-        Scalar $ Sum $ M.insert name svs' $ M.fromList fields-      )+  let svs' = map typeFromSV svs+   in Scalar $ Sum $ M.insert name svs' $ M.fromList fields typeFromSV IntrinsicSV =   error "Tried to get the type from the static value of an intrinsic." -typeFromSV' :: StaticVal -> PatternType-typeFromSV' sv =-  let (sizes, t) = typeFromSV sv-   in unscopeType (S.fromList sizes) t- -- | Construct the type for a fully-applied dynamic function from its -- static value and the original types of its arguments. dynamicFunType :: StaticVal -> [PatternType] -> ([PatternType], PatternType) dynamicFunType (DynamicFun _ sv) (p : ps) =   let (ps', ret) = dynamicFunType sv ps in (p : ps', ret)-dynamicFunType sv _ = ([], typeFromSV' sv)+dynamicFunType sv _ = ([], typeFromSV sv)  -- | Match a pattern with its static value. Returns an environment with -- the identifier components of the pattern mapped to the corresponding@@ -1035,8 +1128,8 @@   -- the pattern wins out.  This is important when matching a   -- nonunique pattern with a unique value.   if orderZeroSV sv-    then M.singleton vn $ Dynamic t-    else M.singleton vn sv+    then M.singleton vn $ Binding Nothing $ Dynamic t+    else M.singleton vn $ Binding Nothing sv matchPatternSV (Wildcard _ _) _ = mempty matchPatternSV (PatternAscription pat _ _) sv = matchPatternSV pat sv matchPatternSV PatternLit {} _ = mempty@@ -1085,7 +1178,7 @@ updatePattern (PatternParens pat loc) sv =   PatternParens (updatePattern pat sv) loc updatePattern (Id vn (Info tp) loc) sv =-  Id vn (Info $ comb tp (snd (typeFromSV sv) `setUniqueness` Nonunique)) loc+  Id vn (Info $ comb tp (typeFromSV sv `setUniqueness` Nonunique)) loc   where     -- Preserve any original zeroth-order types.     comb (Scalar Arrow {}) t2 = t2@@ -1096,7 +1189,7 @@     comb t1 _ = t1 -- t1 must be array or prim. updatePattern pat@(Wildcard (Info tp) loc) sv   | orderZero tp = pat-  | otherwise = Wildcard (Info $ snd $ typeFromSV sv) loc+  | otherwise = Wildcard (Info $ typeFromSV sv) loc updatePattern (PatternAscription pat tydecl loc) sv   | orderZero . unInfo $ expandedType tydecl =     PatternAscription (updatePattern pat sv) tydecl loc@@ -1106,7 +1199,7 @@   | orderZero t = pat   | otherwise = PatternConstr c1 (Info t') ps' loc   where-    t' = snd (typeFromSV sv) `setUniqueness` Nonunique+    t' = typeFromSV sv `setUniqueness` Nonunique     ps' = zipWith updatePattern ps svs updatePattern (PatternConstr c1 _ ps loc) (Dynamic t) =   PatternConstr c1 (Info t) ps loc@@ -1117,135 +1210,12 @@       ++ "to reflect the static value "       ++ show sv --- Like updatePattern, but discard sizes.  This is used for--- let-bindings, where we might otherwise introduce sizes that are--- free.-updatePattern' :: Pattern -> StaticVal -> Pattern-updatePattern' pat sv =-  let pat' = updatePattern pat sv-      (sizes, _) = typeFromSV sv-      tr =-        identityMapper-          { mapOnPatternType =-              pure . unscopeType (S.fromList sizes)-          }-   in runIdentity $ astMap tr pat'- -- | Convert a record (or tuple) type to a record static value. This is used for -- "unwrapping" tuples and records that are nested in 'Dynamic' static values. svFromType :: PatternType -> StaticVal svFromType (Scalar (Record fs)) = RecordSV . M.toList $ M.map svFromType fs svFromType t = Dynamic t --- A set of names where we also track uniqueness.-newtype NameSet = NameSet (M.Map VName Uniqueness) deriving (Show)--instance Semigroup NameSet where-  NameSet x <> NameSet y = NameSet $ M.unionWith max x y--instance Monoid NameSet where-  mempty = NameSet mempty--without :: NameSet -> NameSet -> NameSet-without (NameSet x) (NameSet y) = NameSet $ x `M.difference` y--member :: VName -> NameSet -> Bool-member v (NameSet m) = v `M.member` m--ident :: Ident -> NameSet-ident v = NameSet $ M.singleton (identName v) (uniqueness $ unInfo $ identType v)--oneName :: VName -> NameSet-oneName v = NameSet $ M.singleton v Nonunique--names :: S.Set VName -> NameSet-names = foldMap oneName---- | Compute the set of free variables of an expression.-freeVars :: Exp -> NameSet-freeVars expr = case expr of-  Literal {} -> mempty-  IntLit {} -> mempty-  FloatLit {} -> mempty-  StringLit {} -> mempty-  Parens e _ -> freeVars e-  QualParens _ e _ -> freeVars e-  TupLit es _ -> foldMap freeVars es-  RecordLit fs _ -> foldMap freeVarsField fs-    where-      freeVarsField (RecordFieldExplicit _ e _) = freeVars e-      freeVarsField (RecordFieldImplicit vn t _) = ident $ Ident vn t mempty-  ArrayLit es t _ ->-    foldMap freeVars es-      <> names (typeDimNames $ unInfo t)-  Range e me incl _ _ ->-    freeVars e <> foldMap freeVars me-      <> foldMap freeVars incl-  Var qn (Info t) _ -> NameSet $ M.singleton (qualLeaf qn) $ uniqueness t-  Ascript e t _ -> freeVars e <> names (typeDimNames $ unInfo $ expandedType t)-  Coerce e t _ _ -> freeVars e <> names (typeDimNames $ unInfo $ expandedType t)-  LetPat pat e1 e2 _ _ ->-    freeVars e1-      <> ( (names (patternDimNames pat) <> freeVars e2)-             `without` patternVars pat-         )-  LetFun vn (tparams, pats, _, _, e1) e2 _ _ ->-    ( (freeVars e1 <> names (foldMap patternDimNames pats))-        `without` ( foldMap patternVars pats-                      <> foldMap (oneName . typeParamName) tparams-                  )-    )-      <> (freeVars e2 `without` oneName vn)-  If e1 e2 e3 _ _ -> freeVars e1 <> freeVars e2 <> freeVars e3-  Apply e1 e2 _ _ _ -> freeVars e1 <> freeVars e2-  Negate e _ -> freeVars e-  Lambda pats e0 _ _ _ ->-    (names (foldMap patternDimNames pats) <> freeVars e0)-      `without` foldMap patternVars pats-  OpSection {} -> mempty-  OpSectionLeft _ _ e _ _ _ -> freeVars e-  OpSectionRight _ _ e _ _ _ -> freeVars e-  ProjectSection {} -> mempty-  IndexSection idxs _ _ -> foldMap freeDimIndex idxs-  DoLoop sparams pat e1 form e3 _ _ ->-    let (e2fv, e2ident) = formVars form-     in freeVars e1 <> e2fv-          <> ( freeVars e3-                 `without` (names (S.fromList sparams) <> patternVars pat <> e2ident)-             )-    where-      formVars (For v e2) = (freeVars e2, ident v)-      formVars (ForIn p e2) = (freeVars e2, patternVars p)-      formVars (While e2) = (freeVars e2, mempty)-  BinOp (qn, _) _ (e1, _) (e2, _) _ _ _ ->-    oneName (qualLeaf qn)-      <> freeVars e1-      <> freeVars e2-  Project _ e _ _ -> freeVars e-  LetWith id1 id2 idxs e1 e2 _ _ ->-    ident id2 <> foldMap freeDimIndex idxs <> freeVars e1-      <> (freeVars e2 `without` ident id1)-  Index e idxs _ _ -> freeVars e <> foldMap freeDimIndex idxs-  Update e1 idxs e2 _ -> freeVars e1 <> foldMap freeDimIndex idxs <> freeVars e2-  RecordUpdate e1 _ e2 _ _ -> freeVars e1 <> freeVars e2-  Assert e1 e2 _ _ -> freeVars e1 <> freeVars e2-  Constr _ es _ _ -> foldMap freeVars es-  Attr _ e _ -> freeVars e-  Match e cs _ _ -> freeVars e <> foldMap caseFV cs-    where-      caseFV (CasePat p eCase _) =-        (names (patternDimNames p) <> freeVars eCase)-          `without` patternVars p--freeDimIndex :: DimIndexBase Info VName -> NameSet-freeDimIndex (DimFix e) = freeVars e-freeDimIndex (DimSlice me1 me2 me3) =-  foldMap (foldMap freeVars) [me1, me2, me3]---- | Extract all the variable names bound in a pattern.-patternVars :: Pattern -> NameSet-patternVars = mconcat . map ident . S.toList . patternIdents- -- | Defunctionalize a top-level value binding. Returns the -- transformed result as well as an environment that binds the name of -- the value binding to the static value of the transformed body.  The@@ -1268,9 +1238,16 @@         attrs         loc defuncValBind valbind@(ValBind _ name retdecl (Info (rettype, retext)) tparams params body _ _ _) = do-  (tparams', params', body', sv) <- defuncLet tparams params body rettype+  when (any isTypeParam tparams) $+    error $+      prettyName name ++ " has type parameters, "+        ++ "but the defunctionaliser expects a monomorphic input program."+  (tparams', params', body', sv) <-+    defuncLet (map typeParamName tparams) params body rettype   let rettype' = combineTypeShapes rettype $ anySizes $ toStruct $ typeOf body'-  (missing_dims, params'') <- sizesForAll params'+  globals <- asks fst+  let bound_sizes = S.fromList tparams' <> globals+  (missing_dims, params'') <- sizesForAll bound_sizes params'   return     ( valbind         { valBindRetDecl = retdecl,@@ -1282,12 +1259,19 @@                 retext               ),           valBindTypeParams =-            tparams'-              ++ map (`TypeParamDim` mempty) missing_dims,+            map (`TypeParamDim` mempty) $ tparams' ++ missing_dims,           valBindParams = params'',           valBindBody = body'         },-      M.singleton name sv,+      M.singleton name $+        Binding+          ( Just+              ( first+                  (map typeParamName)+                  (valBindTypeScheme valbind)+              )+          )+          sv,       case sv of         DynamicFun {} -> True         Dynamic {} -> True@@ -1305,6 +1289,8 @@         then isGlobal (valBindName valbind') $ defuncVals ds         else defuncVals ds   return $ defs <> Seq.singleton valbind' <> ds'++{-# NOINLINE transformProg #-}  -- | Transform a list of top-level value bindings. May produce new -- lifted function definitions, which are placed in front of the

@@ -0,0 +1,134 @@+-- | Facilities for computing free variables in an expression, which+-- we need for both lambda-lifting and defunctionalisation.+module Futhark.Internalise.FreeVars+  ( freeVars,+    without,+    member,+    ident,+    size,+    sizes,+    NameSet (..),+    patternVars,+  )+where++import qualified Data.Map.Strict as M+import qualified Data.Set as S+import Futhark.IR.Pretty ()+import Language.Futhark++-- | A set of names where we also track uniqueness.+newtype NameSet = NameSet {unNameSet :: M.Map VName StructType}+  deriving (Show)++instance Semigroup NameSet where+  NameSet x <> NameSet y = NameSet $ M.unionWith max x y++instance Monoid NameSet where+  mempty = NameSet mempty++-- | Set subtraction.+without :: NameSet -> S.Set VName -> NameSet+without (NameSet x) y = NameSet $ M.filterWithKey keep x+  where+    keep k _ = k `S.notMember` y++withoutM :: NameSet -> NameSet -> NameSet+withoutM (NameSet x) (NameSet y) = NameSet $ x `M.difference` y++-- | Is this name in the 'NameSet'?+member :: VName -> NameSet -> Bool+member v (NameSet m) = v `M.member` m++-- | A 'NameSet' with a single 'Nonunique' name.+ident :: Ident -> NameSet+ident v = NameSet $ M.singleton (identName v) (toStruct $ unInfo $ identType v)++size :: VName -> NameSet+size v = NameSet $ M.singleton v $ Scalar $ Prim $ Signed Int64++sizes :: S.Set VName -> NameSet+sizes = foldMap size++-- | Compute the set of free variables of an expression.+freeVars :: Exp -> NameSet+freeVars expr = case expr of+  Literal {} -> mempty+  IntLit {} -> mempty+  FloatLit {} -> mempty+  StringLit {} -> mempty+  Parens e _ -> freeVars e+  QualParens _ e _ -> freeVars e+  TupLit es _ -> foldMap freeVars es+  RecordLit fs _ -> foldMap freeVarsField fs+    where+      freeVarsField (RecordFieldExplicit _ e _) = freeVars e+      freeVarsField (RecordFieldImplicit vn t _) = ident $ Ident vn t mempty+  ArrayLit es t _ ->+    foldMap freeVars es <> sizes (typeDimNames $ unInfo t)+  Range e me incl _ _ ->+    freeVars e <> foldMap freeVars me <> foldMap freeVars incl+  Var qn (Info t) _ -> NameSet $ M.singleton (qualLeaf qn) $ toStruct t+  Ascript e t _ -> freeVars e <> sizes (typeDimNames $ unInfo $ expandedType t)+  Coerce e t _ _ -> freeVars e <> sizes (typeDimNames $ unInfo $ expandedType t)+  LetPat pat e1 e2 _ _ ->+    freeVars e1+      <> ( (sizes (patternDimNames pat) <> freeVars e2)+             `withoutM` patternVars pat+         )+  LetFun vn (tparams, pats, _, _, e1) e2 _ _ ->+    ( (freeVars e1 <> sizes (foldMap patternDimNames pats))+        `without` ( S.map identName (foldMap patternIdents pats)+                      <> S.fromList (map typeParamName tparams)+                  )+    )+      <> (freeVars e2 `without` S.singleton vn)+  If e1 e2 e3 _ _ -> freeVars e1 <> freeVars e2 <> freeVars e3+  Apply e1 e2 _ _ _ -> freeVars e1 <> freeVars e2+  Negate e _ -> freeVars e+  Lambda pats e0 _ _ _ ->+    (sizes (foldMap patternDimNames pats) <> freeVars e0)+      `withoutM` foldMap patternVars pats+  OpSection {} -> mempty+  OpSectionLeft _ _ e _ _ _ -> freeVars e+  OpSectionRight _ _ e _ _ _ -> freeVars e+  ProjectSection {} -> mempty+  IndexSection idxs _ _ -> foldMap freeDimIndex idxs+  DoLoop sparams pat e1 form e3 _ _ ->+    let (e2fv, e2ident) = formVars form+     in freeVars e1+          <> ( (e2fv <> freeVars e3)+                 `withoutM` (sizes (S.fromList sparams) <> patternVars pat <> e2ident)+             )+    where+      formVars (For v e2) = (freeVars e2, ident v)+      formVars (ForIn p e2) = (freeVars e2, patternVars p)+      formVars (While e2) = (freeVars e2, mempty)+  BinOp (qn, _) (Info qn_t) (e1, _) (e2, _) _ _ _ ->+    NameSet (M.singleton (qualLeaf qn) $ toStruct qn_t)+      <> freeVars e1+      <> freeVars e2+  Project _ e _ _ -> freeVars e+  LetWith id1 id2 idxs e1 e2 _ _ ->+    ident id2 <> foldMap freeDimIndex idxs <> freeVars e1+      <> (freeVars e2 `without` S.singleton (identName id1))+  Index e idxs _ _ -> freeVars e <> foldMap freeDimIndex idxs+  Update e1 idxs e2 _ -> freeVars e1 <> foldMap freeDimIndex idxs <> freeVars e2+  RecordUpdate e1 _ e2 _ _ -> freeVars e1 <> freeVars e2+  Assert e1 e2 _ _ -> freeVars e1 <> freeVars e2+  Constr _ es _ _ -> foldMap freeVars es+  Attr _ e _ -> freeVars e+  Match e cs _ _ -> freeVars e <> foldMap caseFV cs+    where+      caseFV (CasePat p eCase _) =+        (sizes (patternDimNames p) <> freeVars eCase)+          `withoutM` patternVars p++freeDimIndex :: DimIndexBase Info VName -> NameSet+freeDimIndex (DimFix e) = freeVars e+freeDimIndex (DimSlice me1 me2 me3) =+  foldMap (foldMap freeVars) [me1, me2, me3]++-- | Extract all the variable names bound in a pattern.+patternVars :: Pattern -> NameSet+patternVars = mconcat . map ident . S.toList . patternIdents

@@ -0,0 +1,174 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE Trustworthy #-}++-- | Lambda-lifting of typed, monomorphic Futhark programs without+-- modules.  After this pass, the program will no longer contain any+-- 'LetFun's or 'Lambda's.+module Futhark.Internalise.LiftLambdas (transformProg) where++import Control.Monad.Reader+import Control.Monad.State+import Data.Foldable+import Data.List (partition)+import qualified Data.Map.Strict as M+import Data.Maybe+import qualified Data.Set as S+import Futhark.IR.Pretty ()+import qualified Futhark.Internalise.FreeVars as FV+import Futhark.MonadFreshNames+import Language.Futhark+import Language.Futhark.Traversals++newtype Env = Env {envReplace :: M.Map VName Exp}++initialEnv :: Env+initialEnv = Env mempty++data LiftState = State+  { stateNameSource :: VNameSource,+    stateValBinds :: [ValBind],+    stateGlobal :: S.Set VName+  }++initialState :: VNameSource -> LiftState+initialState src = State src mempty $ S.fromList $ M.keys intrinsics++newtype LiftM a = LiftM (ReaderT Env (State LiftState) a)+  deriving (Functor, Applicative, Monad, MonadReader Env, MonadState LiftState)++instance MonadFreshNames LiftM where+  putNameSource src = modify $ \s -> s {stateNameSource = src}+  getNameSource = gets stateNameSource++runLiftM :: VNameSource -> LiftM () -> ([ValBind], VNameSource)+runLiftM src (LiftM m) =+  let s = execState (runReaderT m initialEnv) (initialState src)+   in (reverse (stateValBinds s), stateNameSource s)++addValBind :: ValBind -> LiftM ()+addValBind vb = modify $ \s ->+  s+    { stateValBinds = vb : stateValBinds s,+      stateGlobal = foldl' (flip S.insert) (stateGlobal s) names+    }+  where+    names = valBindName vb : snd (unInfo (valBindRetType vb))++replacing :: VName -> Exp -> LiftM a -> LiftM a+replacing v e = local $ \env ->+  env {envReplace = M.insert v e $ envReplace env}++existentials :: Exp -> S.Set VName+existentials e =+  let here = case e of+        Apply _ _ (Info (_, pdim)) (_, Info ext) _ ->+          S.fromList (maybeToList pdim ++ ext)+        If _ _ _ (_, Info ext) _ ->+          S.fromList ext+        LetPat _ _ _ (_, Info ext) _ ->+          S.fromList ext+        Coerce _ _ (_, Info ext) _ ->+          S.fromList ext+        Range _ _ _ (_, Info ext) _ ->+          S.fromList ext+        Index _ _ (_, Info ext) _ ->+          S.fromList ext+        Match _ _ (_, Info ext) _ ->+          S.fromList ext+        _ ->+          mempty++      m = identityMapper {mapOnExp = \e' -> modify (<> existentials e') >> pure e'}+   in execState (astMap m e) here++liftFunction :: VName -> [TypeParam] -> [Pattern] -> StructType -> Exp -> LiftM Exp+liftFunction fname tparams params ret funbody = do+  -- Find free variables+  global <- gets stateGlobal+  let bound =+        global+          <> foldMap patternNames params+          <> S.fromList (map typeParamName tparams)+          <> existentials funbody++      free =+        let immediate_free = FV.freeVars funbody `FV.without` bound+            sizes_in_free =+              foldMap typeDimNames $+                M.elems $ FV.unNameSet immediate_free+            sizes =+              FV.sizes $+                sizes_in_free+                  <> foldMap patternDimNames params+                  <> typeDimNames ret+         in M.toList $ FV.unNameSet $ immediate_free <> (sizes `FV.without` bound)++      -- Those parameters that correspond to sizes must come first.+      sizes_in_types =+        foldMap typeDimNames (ret : map snd free ++ map patternStructType params)+      isSize (v, _) = v `S.member` sizes_in_types+      (free_dims, free_nondims) = partition isSize free++      free_params = map mkParam $ free_dims ++ free_nondims++  addValBind $+    ValBind+      { valBindName = fname,+        valBindTypeParams = tparams,+        valBindParams = free_params ++ params,+        valBindRetDecl = Nothing,+        valBindRetType = Info (ret, mempty),+        valBindBody = funbody,+        valBindDoc = Nothing,+        valBindAttrs = mempty,+        valBindLocation = mempty,+        valBindEntryPoint = Nothing+      }++  return $+    apply+      (Var (qualName fname) (Info (augType $ free_dims ++ free_nondims)) mempty)+      $ free_dims ++ free_nondims+  where+    orig_type = funType params ret+    mkParam (v, t) = Id v (Info (fromStruct t)) mempty+    freeVar (v, t) = Var (qualName v) (Info (fromStruct t)) mempty+    augType rem_free = fromStruct $ funType (map mkParam rem_free) orig_type++    apply :: Exp -> [(VName, StructType)] -> Exp+    apply f [] = f+    apply f (p : rem_ps) =+      let inner_ret = (Info (fromStruct (augType rem_ps)), Info mempty)+          inner = Apply f (freeVar p) (Info (Observe, Nothing)) inner_ret mempty+       in apply inner rem_ps++transformExp :: Exp -> LiftM Exp+transformExp (LetFun fname (tparams, params, _, Info ret, funbody) body _ _) = do+  funbody' <- transformExp funbody+  fname' <- newVName $ "lifted_" ++ baseString fname+  lifted_call <- liftFunction fname' tparams params ret funbody'+  replacing fname lifted_call $ transformExp body+transformExp (Lambda params body _ (Info (_, ret)) _) = do+  body' <- transformExp body+  fname <- newVName "lifted_lambda"+  liftFunction fname [] params ret body'+transformExp e@(Var v _ _) =+  -- Note that function-typed variables can only occur in expressions,+  -- not in other places where VNames/QualNames can occur.+  asks (fromMaybe e . M.lookup (qualLeaf v) . envReplace)+transformExp e =+  astMap m e+  where+    m = identityMapper {mapOnExp = transformExp}++transformValBind :: ValBind -> LiftM ()+transformValBind vb = do+  e <- transformExp $ valBindBody vb+  addValBind $ vb {valBindBody = e}++{-# NOINLINE transformProg #-}+transformProg :: MonadFreshNames m => [ValBind] -> m [ValBind]+transformProg vbinds =+  modifyNameSource $ \namesrc ->+    runLiftM namesrc $ mapM_ transformValBind vbinds

@@ -11,7 +11,6 @@     throwError,     VarSubstitutions,     InternaliseEnv (..),-    Closure,     FunInfo,     substitutingVars,     lookupSubst,@@ -19,7 +18,6 @@     lookupFunction,     lookupFunction',     lookupConst,-    allConsts,     bindFunction,     bindConstant,     localConstsScope,@@ -38,14 +36,8 @@ import Futhark.Tools import Futhark.Util (takeLast) --- | Extra parameters to pass when calling this function.  This--- corresponds to the closure of a locally defined function.-type Closure = [VName]- type FunInfo =-  ( Name,-    Closure,-    [VName],+  ( [VName],     [DeclType],     [FParam],     [(SubExp, Type)] -> Maybe [DeclExtType]@@ -68,8 +60,7 @@   { stateNameSource :: VNameSource,     stateFunTable :: FunTable,     stateConstSubsts :: VarSubstitutions,-    stateConstScope :: Scope SOACS,-    stateConsts :: Names+    stateConstScope :: Scope SOACS   }  data InternaliseResult = InternaliseResult (Stms SOACS) [FunDef SOACS]@@ -139,7 +130,6 @@         { stateNameSource = src,           stateFunTable = mempty,           stateConstSubsts = mempty,-          stateConsts = mempty,           stateConstScope = mempty         } @@ -168,9 +158,6 @@ lookupConst :: VName -> InternaliseM (Maybe [SubExp]) lookupConst fname = gets $ M.lookup fname . stateConstSubsts -allConsts :: InternaliseM Names-allConsts = gets stateConsts- bindFunction :: VName -> FunDef SOACS -> FunInfo -> InternaliseM () bindFunction fname fd info = do   addFunDef fd@@ -182,13 +169,11 @@       substs =         takeLast (length (funDefRetType fd)) $           bodyResult $ funDefBody fd-      const_names = namesFromList $ M.keys $ scopeOf stms   addStms stms   modify $ \s ->     s       { stateConstSubsts = M.insert cname substs $ stateConstSubsts s,-        stateConstScope = scopeOf stms <> stateConstScope s,-        stateConsts = const_names <> stateConsts s+        stateConstScope = scopeOf stms <> stateConstScope s       }  localConstsScope :: InternaliseM a -> InternaliseM a

@@ -23,6 +23,8 @@ -- -- * Turns implicit record fields into explicit record fields. --+-- * Rewrite BinOp nodes to Apply nodes.+-- -- Note that these changes are unfortunately not visible in the AST -- representation. module Futhark.Internalise.Monomorphise (transformProg) where@@ -62,7 +64,6 @@       ( VName,         [TypeParam],         [Pattern],-        Maybe (TypeExp VName),         StructType,         [VName],         Exp,@@ -334,7 +335,7 @@     -- filter those that are monomorphic versions of the current let-bound     -- function and insert them at this point, and propagate the rest.     rr <- asks envRecordReplacements-    let funbind = PolyBinding rr (fname, tparams, params, retdecl, ret, [], body, mempty, loc)+    let funbind = PolyBinding rr (fname, tparams, params, ret, [], body, mempty, loc)     pass $ do       (e', bs) <- listen $ extendEnv fname funbind $ transformExp e       -- Do not remember this one for next time we monomorphise this@@ -364,46 +365,33 @@   return $ Lambda params e0' decl tp loc transformExp (OpSection qn t loc) =   transformExp $ Var qn t loc-transformExp-  ( OpSectionLeft-      fname-      (Info t)-      e-      (Info (xtype, xargext), Info ytype)-      (Info rettype, Info retext)-      loc-    ) = do-    fname' <- transformFName loc fname $ toStruct t-    e' <- transformExp e-    desugarBinOpSection-      fname'-      (Just e')-      Nothing-      t-      (xtype, xargext)-      (ytype, Nothing)-      (rettype, retext)-      loc-transformExp-  ( OpSectionRight-      fname-      (Info t)-      e-      (Info xtype, Info (ytype, yargext))-      (Info rettype)-      loc-    ) = do-    fname' <- transformFName loc fname $ toStruct t-    e' <- transformExp e-    desugarBinOpSection-      fname'-      Nothing-      (Just e')-      t-      (xtype, Nothing)-      (ytype, yargext)-      (rettype, [])-      loc+transformExp (OpSectionLeft fname (Info t) e arg ret loc) = do+  let (Info (xtype, xargext), Info ytype) = arg+      (Info rettype, Info retext) = ret+  fname' <- transformFName loc fname $ toStruct t+  e' <- transformExp e+  desugarBinOpSection+    fname'+    (Just e')+    Nothing+    t+    (xtype, xargext)+    (ytype, Nothing)+    (rettype, retext)+    loc+transformExp (OpSectionRight fname (Info t) e arg (Info rettype) loc) = do+  let (Info xtype, Info (ytype, yargext)) = arg+  fname' <- transformFName loc fname $ toStruct t+  e' <- transformExp e+  desugarBinOpSection+    fname'+    Nothing+    (Just e')+    t+    (xtype, Nothing)+    (ytype, yargext)+    (rettype, [])+    loc transformExp (ProjectSection fields (Info t) loc) =   desugarProjectSection fields t loc transformExp (IndexSection idxs (Info t) loc) =@@ -420,16 +408,14 @@   -- sizes for them.   (pat_sizes, pat') <- sizesForPat pat   return $ DoLoop (sparams ++ pat_sizes) pat' e1' form' e3' ret loc-transformExp (BinOp (fname, oploc) (Info t) (e1, d1) (e2, d2) tp ext loc) = do+transformExp (BinOp (fname, _) (Info t) (e1, d1) (e2, d2) tp ext loc) = do   fname' <- transformFName loc fname $ toStruct t   e1' <- transformExp e1   e2' <- transformExp e2-  case fname' of-    Var fname'' _ _-      | orderZero (typeOf e1'),-        orderZero (typeOf e2') ->-        return $ BinOp (fname'', oploc) (Info t) (e1', d1) (e2', d2) tp ext loc-    _ -> do+  if orderZero (typeOf e1')+    && orderZero (typeOf e2')+    then return $ applyOp fname' e1' e2'+    else do       -- We have to flip the arguments to the function, because       -- operator application is left-to-right, while function       -- application is outside-in.  This matters when the arguments@@ -538,8 +524,8 @@   SrcLoc ->   MonoM Exp desugarBinOpSection op e_left e_right t (xtype, xext) (ytype, yext) (rettype, retext) loc = do-  (e1, p1) <- makeVarParam e_left $ fromStruct xtype-  (e2, p2) <- makeVarParam e_right $ fromStruct ytype+  (wrap_left, e1, p1) <- makeVarParam e_left $ fromStruct xtype+  (wrap_right, e2, p2) <- makeVarParam e_right $ fromStruct ytype   let apply_left =         Apply           op@@ -555,16 +541,24 @@           (Info rettype, Info retext)           loc       rettype' = toStruct rettype-  return $ Lambda (p1 ++ p2) body Nothing (Info (mempty, rettype')) loc+  return $ wrap_left $ wrap_right $ Lambda (p1 ++ p2) body Nothing (Info (mempty, rettype')) loc   where-    makeVarParam (Just e) _ = return (e, [])-    makeVarParam Nothing argtype = do+    patAndVar argtype = do       x <- newNameFromString "x"-      return-        ( Var (qualName x) (Info argtype) mempty,-          [Id x (Info $ fromStruct argtype) mempty]+      pure+        ( Id x (Info argtype) mempty,+          Var (qualName x) (Info argtype) mempty         ) +    makeVarParam (Just e) argtype = do+      (pat, var_e) <- patAndVar argtype+      let wrap body =+            LetPat pat e body (Info (typeOf body), Info mempty) mempty+      return (wrap, var_e, [])+    makeVarParam Nothing argtype = do+      (pat, var_e) <- patAndVar argtype+      return (id, var_e, [pat])+ desugarProjectSection :: [Name] -> PatternType -> SrcLoc -> MonoM Exp desugarProjectSection fields (Scalar (Arrow _ _ t1 t2)) loc = do   p <- newVName "project_p"@@ -714,7 +708,7 @@   PolyBinding ->   MonoType ->   MonoM (VName, InferSizeArgs, ValBind)-monomorphiseBinding entry (PolyBinding rr (name, tparams, params, retdecl, rettype, retext, body, attrs, loc)) t =+monomorphiseBinding entry (PolyBinding rr (name, tparams, params, rettype, retext, body, attrs, loc)) t =   replaceRecordReplacements rr $ do     let bind_t = foldFunType (map patternStructType params) rettype     (substs, t_shape_params) <- typeSubstsM loc (noSizes bind_t) $ noNamedParams t@@ -774,8 +768,8 @@       ValBind         { valBindEntryPoint = Nothing,           valBindName = name',-          valBindRetDecl = retdecl,           valBindRetType = Info rettype',+          valBindRetDecl = Nothing,           valBindTypeParams = tparams',           valBindParams = params'',           valBindBody = body'',@@ -851,8 +845,8 @@   PatternConstr n (Info tp) ps loc -> PatternConstr n (Info $ f tp) ps loc  toPolyBinding :: ValBind -> PolyBinding-toPolyBinding (ValBind _ name retdecl (Info (rettype, retext)) tparams params body _ attrs loc) =-  PolyBinding mempty (name, tparams, params, retdecl, rettype, retext, body, attrs, loc)+toPolyBinding (ValBind _ name _ (Info (rettype, retext)) tparams params body _ attrs loc) =+  PolyBinding mempty (name, tparams, params, rettype, retext, body, attrs, loc)  -- Remove all type variables and type abbreviations from a value binding. removeTypeVariables :: Bool -> ValBind -> MonoM ValBind

@@ -44,11 +44,11 @@ simplifyProg simpl rules blockers (Prog consts funs) = do   (consts_vtable, consts') <-     simplifyConsts-      (UT.usages $ foldMap (freeIn . funDefBody) funs)+      (UT.usages $ foldMap freeIn funs)       (mempty, consts)    funs' <- parPass (simplifyFun' consts_vtable) funs-  let funs_uses = UT.usages $ foldMap (freeIn . funDefBody) funs'+  let funs_uses = UT.usages $ foldMap freeIn funs'    (_, consts'') <- simplifyConsts funs_uses (mempty, consts') 

@@ -66,7 +66,7 @@ import Control.Monad.Reader import Control.Monad.State.Strict import Data.Either-import Data.List (find, foldl', mapAccumL, nub)+import Data.List (find, foldl', mapAccumL) import Data.Maybe import qualified Futhark.Analysis.SymbolTable as ST import qualified Futhark.Analysis.UsageTable as UT@@ -75,7 +75,7 @@ import Futhark.IR.Prop.Aliases import Futhark.Optimise.Simplify.Lore import Futhark.Optimise.Simplify.Rule-import Futhark.Util (splitFromEnd)+import Futhark.Util (nubOrd, splitFromEnd)  data HoistBlockers lore = HoistBlockers   { -- | Blocker for hoisting out of parallel loops.@@ -1011,7 +1011,7 @@     inspect _ = mempty  instance Simplifiable Certificates where-  simplify (Certificates ocs) = Certificates . nub . concat <$> mapM check ocs+  simplify (Certificates ocs) = Certificates . nubOrd . concat <$> mapM check ocs     where       check idd = do         vv <- ST.lookupSubExp idd <$> askVtable

@@ -488,25 +488,30 @@ simplifyBinOp _ _ (BinOp FAdd {} e1 e2)   | isCt0 e1 = subExpRes e2   | isCt0 e2 = subExpRes e1-simplifyBinOp look _ (BinOp Sub {} e1 e2)+simplifyBinOp look _ (BinOp sub@(Sub t _) e1 e2)   | isCt0 e2 = subExpRes e1   -- Cases for simplifying (a+b)-b and permutations.++  -- (e1_a+e1_b)-e1_a == e1_b   | Var v1 <- e1,     Just (BasicOp (BinOp Add {} e1_a e1_b), cs) <- look v1,     e1_a == e2 =     Just (SubExp e1_b, cs)+  -- (e1_a+e1_b)-e1_b == e1_a   | Var v1 <- e1,     Just (BasicOp (BinOp Add {} e1_a e1_b), cs) <- look v1,     e1_b == e2 =     Just (SubExp e1_a, cs)+  -- e2_a-(e2_a+e2_b) == 0-e2_b   | Var v2 <- e2,     Just (BasicOp (BinOp Add {} e2_a e2_b), cs) <- look v2,     e2_a == e1 =-    Just (SubExp e2_b, cs)-  | Var v2 <- e1,+    Just (BinOp sub (intConst t 0) e2_b, cs)+  -- e2_b-(e2_a+e2_b) == 0-e2_a+  | Var v2 <- e2,     Just (BasicOp (BinOp Add {} e2_a e2_b), cs) <- look v2,     e2_b == e1 =-    Just (SubExp e2_a, cs)+    Just (BinOp sub (intConst t 0) e2_a, cs) simplifyBinOp _ _ (BinOp FSub {} e1 e2)   | isCt0 e2 = subExpRes e1 simplifyBinOp _ _ (BinOp Mul {} e1 e2)

@@ -551,17 +551,17 @@               lambdaBody = mkBody (oneStm tmpbnd) [Var tmp]             }     maybeDistributeStm newbnd acc-maybeDistributeStm bnd@(Let _ aux (BasicOp Copy {})) acc =-  distributeSingleUnaryStm acc bnd $ \_ outerpat arr ->+maybeDistributeStm stm@(Let _ aux (BasicOp (Copy stm_arr))) acc =+  distributeSingleUnaryStm acc stm stm_arr $ \_ outerpat arr ->     return $ oneStm $ Let outerpat aux $ BasicOp $ Copy arr -- Opaques are applied to the full array, because otherwise they can -- drastically inhibit parallelisation in some cases.-maybeDistributeStm bnd@(Let (Pattern [] [pe]) aux (BasicOp Opaque {})) acc+maybeDistributeStm stm@(Let (Pattern [] [pe]) aux (BasicOp (Opaque (Var stm_arr)))) acc   | not $ primType $ typeOf pe =-    distributeSingleUnaryStm acc bnd $ \_ outerpat arr ->+    distributeSingleUnaryStm acc stm stm_arr $ \_ outerpat arr ->       return $ oneStm $ Let outerpat aux $ BasicOp $ Copy arr-maybeDistributeStm bnd@(Let _ aux (BasicOp (Rearrange perm _))) acc =-  distributeSingleUnaryStm acc bnd $ \nest outerpat arr -> do+maybeDistributeStm stm@(Let _ aux (BasicOp (Rearrange perm stm_arr))) acc =+  distributeSingleUnaryStm acc stm stm_arr $ \nest outerpat arr -> do     let r = length (snd nest) + 1         perm' = [0 .. r -1] ++ map (+ r) perm     -- We need to add a copy, because the original map nest@@ -573,14 +573,14 @@         [ Let (Pattern [] [PatElem arr' arr_t]) aux $ BasicOp $ Copy arr,           Let outerpat aux $ BasicOp $ Rearrange perm' arr'         ]-maybeDistributeStm bnd@(Let _ aux (BasicOp (Reshape reshape _))) acc =-  distributeSingleUnaryStm acc bnd $ \nest outerpat arr -> do+maybeDistributeStm stm@(Let _ aux (BasicOp (Reshape reshape stm_arr))) acc =+  distributeSingleUnaryStm acc stm stm_arr $ \nest outerpat arr -> do     let reshape' =           map DimNew (kernelNestWidths nest)             ++ map DimNew (newDims reshape)     return $ oneStm $ Let outerpat aux $ BasicOp $ Reshape reshape' arr-maybeDistributeStm stm@(Let _ aux (BasicOp (Rotate rots _))) acc =-  distributeSingleUnaryStm acc stm $ \nest outerpat arr -> do+maybeDistributeStm stm@(Let _ aux (BasicOp (Rotate rots stm_arr))) acc =+  distributeSingleUnaryStm acc stm stm_arr $ \nest outerpat arr -> do     let rots' = map (const $ intConst Int64 0) (kernelNestWidths nest) ++ rots     return $ oneStm $ Let outerpat aux $ BasicOp $ Rotate rots' arr maybeDistributeStm stm@(Let pat aux (BasicOp (Update arr slice (Var v)))) acc@@ -644,15 +644,16 @@   (MonadFreshNames m, LocalScope lore m, DistLore lore) =>   DistAcc lore ->   Stm SOACS ->+  VName ->   (KernelNest -> PatternT Type -> VName -> DistNestT lore m (Stms lore)) ->   DistNestT lore m (DistAcc lore)-distributeSingleUnaryStm acc bnd f =-  distributeSingleStm acc bnd >>= \case+distributeSingleUnaryStm acc stm stm_arr f =+  distributeSingleStm acc stm >>= \case     Just (kernels, res, nest, acc')-      | res == map Var (patternNames $ stmPattern bnd),+      | res == map Var (patternNames $ stmPattern stm),         (outer, _) <- nest,         [(arr_p, arr)] <- loopNestingParamsAndArrs outer,-        boundInKernelNest nest `namesIntersection` freeIn bnd+        boundInKernelNest nest `namesIntersection` freeIn stm           == oneName (paramName arr_p),         perfectlyMapped arr nest -> do         addPostStms kernels@@ -660,13 +661,13 @@         localScope (typeEnvFromDistAcc acc') $ do           postStm =<< f nest outerpat arr           return acc'-    _ -> addStmToAcc bnd acc+    _ -> addStmToAcc stm acc   where     perfectlyMapped arr (outer, nest)       | [(p, arr')] <- loopNestingParamsAndArrs outer,         arr == arr' =         case nest of-          [] -> True+          [] -> paramName p == stm_arr           x : xs -> perfectlyMapped (paramName p) (x, xs)       | otherwise =         False

@@ -123,6 +123,7 @@   | SOACSPipeline   | SequentialCpuPipeline   | GpuPipeline+  | NoPipeline   deriving (Show)  -- | A structure test specifies a compilation pipeline, as well as@@ -441,6 +442,7 @@   lexstr "distributed" $> KernelsPipeline     <|> lexstr "gpu" $> GpuPipeline     <|> lexstr "cpu" $> SequentialCpuPipeline+    <|> lexstr "internalised" $> NoPipeline     <|> pure SOACSPipeline  parseMetrics :: Parser AstMetrics

@@ -248,8 +248,12 @@   ms <- readRankedArrayOfST r rv s   case ms of     Just (i, shape, arr, t)-      | Just t' <- symbol ',' t ->-        readArrayElemsST (j + 1) r rv (i, shape, arr, t')+      | Just t' <- symbol ',' t -> do+        next <- readArrayElemsST (j + 1) r rv (i, shape, arr, t')+        -- Not OK to have zero values after a comma.+        case next of+          Just (0, _) -> return Nothing+          _ -> return next       | otherwise -> return $ Just (j, (i, shape, arr, t))     _ ->       return $ Just (0, s)

@@ -876,8 +876,13 @@   void $ checkArrIdent e checkBasicOp (Manifest perm arr) =   checkBasicOp $ Rearrange perm arr -- Basically same thing!-checkBasicOp (Assert e _ _) =+checkBasicOp (Assert e (ErrorMsg parts) _) = do   require [Prim Bool] e+  mapM_ checkPart parts+  where+    checkPart ErrorString {} = return ()+    checkPart (ErrorInt32 x) = require [Prim int32] x+    checkPart (ErrorInt64 x) = require [Prim int64] x  matchLoopResultExt ::   Checkable lore =>

@@ -25,6 +25,7 @@     focusNth,     unixEnvironment,     isEnvVarSet,+    isEnvVarAtLeast,     fancyTerminal,     runProgramWithExitCode,     directoryContents,@@ -69,6 +70,7 @@ import System.IO (hIsTerminalDevice, stdout) import System.IO.Unsafe import System.Process.ByteString+import Text.Read (readMaybe)  -- | Like 'nub', but without the quadratic runtime. nubOrd :: Ord a => [a] -> [a]@@ -161,7 +163,7 @@ unixEnvironment = unsafePerformIO getEnvironment  -- | Is an environment variable set to 0 or 1?  If 0, return False; if--- 1, True; otherwise the default value.+-- 1, True; otherwise default. isEnvVarSet :: String -> Bool -> Bool isEnvVarSet name default_val = fromMaybe default_val $ do   val <- lookup name unixEnvironment@@ -169,6 +171,15 @@     "0" -> return False     "1" -> return True     _ -> Nothing++-- | True if the environment variable, viewed as an integer, has at+-- least this numeric value.  Returns False if variable is unset or+-- not numeric.+isEnvVarAtLeast :: String -> Int -> Bool+isEnvVarAtLeast s x =+  case readMaybe =<< lookup s unixEnvironment of+    Just y -> y >= x+    _ -> False  {-# NOINLINE fancyTerminal #-} 

@@ -51,7 +51,7 @@ -- the base name. instance IsName VName where   pprName-    | isEnvVarSet "FUTHARK_COMPILER_DEBUGGING" False =+    | isEnvVarAtLeast "FUTHARK_COMPILER_DEBUGGING" 1 =       \(VName vn i) -> ppr vn <> text "_" <> text (show i)     | otherwise = ppr . baseName @@ -222,7 +222,7 @@     where       inst = case unAnnot t of         Just t'-          | isEnvVarSet "FUTHARK_COMPILER_DEBUGGING" False ->+          | isEnvVarAtLeast "FUTHARK_COMPILER_DEBUGGING" 2 ->             text "@" <> parens (align $ ppr t')         _ -> mempty   pprPrec _ (Parens e _) = align $ parens $ ppr e@@ -248,7 +248,7 @@     where       info' = case unAnnot info of         Just t-          | isEnvVarSet "FUTHARK_COMPILER_DEBUGGING" False ->+          | isEnvVarAtLeast "FUTHARK_COMPILER_DEBUGGING" 2 ->             text "@" <> parens (align $ ppr t)         _ -> mempty   pprPrec _ (StringLit s _) =

@@ -26,6 +26,8 @@      -- * Queries on expressions     typeOf,+    valBindTypeScheme,+    funType,      -- * Queries on patterns and params     patternIdents,@@ -111,13 +113,13 @@ import Data.Bitraversable (bitraverse) import Data.Char import Data.Foldable-import Data.List (genericLength, isPrefixOf, nub, sortOn)+import Data.List (genericLength, isPrefixOf, sortOn) import qualified Data.Map.Strict as M import Data.Maybe import Data.Ord import qualified Data.Set as S import qualified Futhark.IR.Primitive as Primitive-import Futhark.Util (maxinum)+import Futhark.Util (maxinum, nubOrd) import Futhark.Util.Pretty import Language.Futhark.Syntax @@ -138,13 +140,13 @@ nestedDims t =   case t of     Array _ _ a ds ->-      nub $ nestedDims (Scalar a) <> shapeDims ds+      nubOrd $ nestedDims (Scalar a) <> shapeDims ds     Scalar (Record fs) ->-      nub $ foldMap nestedDims fs+      nubOrd $ foldMap nestedDims fs     Scalar Prim {} ->       mempty     Scalar (Sum cs) ->-      nub $ foldMap (foldMap nestedDims) cs+      nubOrd $ foldMap (foldMap nestedDims) cs     Scalar (Arrow _ v t1 t2) ->       filter (notV v) $ nestedDims t1 <> nestedDims t2     Scalar (TypeVar _ _ _ targs) ->@@ -646,6 +648,20 @@   let (ps, r) = unfoldFunType t2    in (t1 : ps, r) unfoldFunType t = ([], t)++-- | The type scheme of a value binding, comprising the type+-- parameters and the actual type.+valBindTypeScheme :: ValBindBase Info VName -> ([TypeParamBase VName], StructType)+valBindTypeScheme vb =+  ( valBindTypeParams vb,+    funType (valBindParams vb) (fst (unInfo (valBindRetType vb)))+  )++-- | The type of a function with the given parameters and return type.+funType :: [PatternBase Info VName] -> StructType -> StructType+funType params ret = foldr (arrow . patternParam) ret params+  where+    arrow (xp, xt) yt = Scalar $ Arrow () xp xt yt  -- | The type names mentioned in a type. typeVars :: Monoid as => TypeBase dim as -> S.Set VName

@@ -212,7 +212,7 @@     Attr attr <$> mapOnExp tv e <*> pure loc  instance ASTMappable (LoopFormBase Info VName) where-  astMap tv (For i bound) = For <$> astMap tv i <*> astMap tv bound+  astMap tv (For i bound) = For <$> astMap tv i <*> mapOnExp tv bound   astMap tv (ForIn pat e) = ForIn <$> astMap tv pat <*> mapOnExp tv e   astMap tv (While e) = While <$> mapOnExp tv e @@ -252,12 +252,12 @@   astMap = traverse . mapOnName  instance ASTMappable (DimIndexBase Info VName) where-  astMap tv (DimFix j) = DimFix <$> astMap tv j+  astMap tv (DimFix j) = DimFix <$> mapOnExp tv j   astMap tv (DimSlice i j stride) =     DimSlice-      <$> maybe (return Nothing) (fmap Just . astMap tv) i-      <*> maybe (return Nothing) (fmap Just . astMap tv) j-      <*> maybe (return Nothing) (fmap Just . astMap tv) stride+      <$> maybe (return Nothing) (fmap Just . mapOnExp tv) i+      <*> maybe (return Nothing) (fmap Just . mapOnExp tv) j+      <*> maybe (return Nothing) (fmap Just . mapOnExp tv) stride  instance ASTMappable Alias where   astMap tv (AliasBound v) = AliasBound <$> mapOnName tv v

@@ -626,16 +626,15 @@             </> "\nwill have an opaque type, so the result will likely not be usable."     _ -> return () -  let arrow (xp, xt) yt = Scalar $ Arrow () xp xt yt+  let vb = ValBind entry' fname' maybe_tdecl' (Info (rettype, retext)) tparams' params' body' doc attrs loc   return     ( mempty         { envVtable =-            M.singleton fname' $-              BoundV tparams' $ foldr (arrow . patternParam) rettype params',+            M.singleton fname' $ uncurry BoundV $ valBindTypeScheme vb,           envNameMap =             M.singleton (Term, fname) $ qualName fname'         },-      ValBind entry' fname' maybe_tdecl' (Info (rettype, retext)) tparams' params' body' doc attrs loc+      vb     )  nastyType :: Monoid als => TypeBase dim als -> Bool

@@ -27,12 +27,13 @@ import Data.Bifunctor import Data.Char (isAscii) import Data.Either-import Data.List (find, foldl', isPrefixOf, nub, sort)+import Data.List (find, foldl', isPrefixOf, sort) import qualified Data.List.NonEmpty as NE import qualified Data.Map.Strict as M import Data.Maybe import qualified Data.Set as S import Futhark.IR.Primitive (intByteSize)+import Futhark.Util (nubOrd) import Futhark.Util.Pretty hiding (bool, group, space) import Language.Futhark hiding (unscopeType) import Language.Futhark.Semantic (includeToString)@@ -1806,7 +1807,7 @@           mapM_ dimToInit $ M.toList init_substs'            mergepat'' <- applySubst (`M.lookup` init_substs') <$> updateTypes mergepat'-          return (nub sparams, mergepat'')+          return (nubOrd sparams, mergepat'')      -- First we do a basic check of the loop body to figure out which of     -- the merge parameters are being consumed.  For this, we first need

@@ -26,9 +26,10 @@ import Control.Monad.Reader import Control.Monad.State import Data.Bifunctor-import Data.List (foldl', nub, sort)+import Data.List (foldl', sort) import qualified Data.Map.Strict as M import Data.Maybe+import Futhark.Util (nubOrd) import Futhark.Util.Pretty import Language.Futhark import Language.Futhark.Traversals@@ -144,7 +145,7 @@ checkTypeExp t@(TERecord fs loc) = do   -- Check for duplicate field names.   let field_names = map fst fs-  unless (sort field_names == sort (nub field_names)) $+  unless (sort field_names == sort (nubOrd field_names)) $     typeError loc mempty $ "Duplicate record fields in" <+> ppr t <> "."    fs_ts_ls <- traverse checkTypeExp $ M.fromList fs@@ -265,7 +266,7 @@           <+> ppr p <> "." checkTypeExp t@(TESum cs loc) = do   let constructors = map fst cs-  unless (sort constructors == sort (nub constructors)) $+  unless (sort constructors == sort (nubOrd constructors)) $     typeError loc mempty $ "Duplicate constructors in" <+> ppr t    unless (length constructors < 256) $