diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,177 @@ +# Makefile for Sphinx documentation +# + +# You can set these variables from the command line. +SPHINXOPTS = +SPHINXBUILD = sphinx-build +PAPER = +BUILDDIR = _build + +# User-friendly check for sphinx-build +ifeq ($(shell which $(SPHINXBUILD) >/dev/null 2>&1; echo $$?), 1) +$(error The '$(SPHINXBUILD)' command was not found. Make sure you have Sphinx installed, then set the SPHINXBUILD environment variable to point to the full path of the '$(SPHINXBUILD)' executable. Alternatively you can add the directory with the executable to your PATH. If you don't have Sphinx installed, grab it from http://sphinx-doc.org/) +endif + +# Internal variables. +PAPEROPT_a4 = -D latex_paper_size=a4 +PAPEROPT_letter = -D latex_paper_size=letter +ALLSPHINXOPTS = -d $(BUILDDIR)/doctrees $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) . +# the i18n builder cannot share the environment and doctrees with the others +I18NSPHINXOPTS = $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) . + +.PHONY: help clean html dirhtml singlehtml pickle json htmlhelp qthelp devhelp epub latex latexpdf text man changes linkcheck doctest gettext + +help: + @echo "Please use \`make ' where is one of" + @echo " html to make standalone HTML files" + @echo " dirhtml to make HTML files named index.html in directories" + @echo " singlehtml to make a single large HTML file" + @echo " pickle to make pickle files" + @echo " json to make JSON files" + @echo " htmlhelp to make HTML files and a HTML help project" + @echo " qthelp to make HTML files and a qthelp project" + @echo " devhelp to make HTML files and a Devhelp project" + @echo " epub to make an epub" + @echo " latex to make LaTeX files, you can set PAPER=a4 or PAPER=letter" + @echo " latexpdf to make LaTeX files and run them through pdflatex" + @echo " latexpdfja to make LaTeX files and run them through platex/dvipdfmx" + @echo " text to make text files" + @echo " man to make manual pages" + @echo " texinfo to make Texinfo files" + @echo " info to make Texinfo files and run them through makeinfo" + @echo " gettext to make PO message catalogs" + @echo " changes to make an overview of all changed/added/deprecated items" + @echo " xml to make Docutils-native XML files" + @echo " pseudoxml to make pseudoxml-XML files for display purposes" + @echo " linkcheck to check all external links for integrity" + @echo " doctest to run all doctests embedded in the documentation (if enabled)" + +clean: + rm -rf $(BUILDDIR)/* + +html: + $(SPHINXBUILD) -b html $(ALLSPHINXOPTS) $(BUILDDIR)/html + @echo + @echo "Build finished. The HTML pages are in $(BUILDDIR)/html." + +dirhtml: + $(SPHINXBUILD) -b dirhtml $(ALLSPHINXOPTS) $(BUILDDIR)/dirhtml + @echo + @echo "Build finished. The HTML pages are in $(BUILDDIR)/dirhtml." + +singlehtml: + $(SPHINXBUILD) -b singlehtml $(ALLSPHINXOPTS) $(BUILDDIR)/singlehtml + @echo + @echo "Build finished. The HTML page is in $(BUILDDIR)/singlehtml." + +pickle: + $(SPHINXBUILD) -b pickle $(ALLSPHINXOPTS) $(BUILDDIR)/pickle + @echo + @echo "Build finished; now you can process the pickle files." + +json: + $(SPHINXBUILD) -b json $(ALLSPHINXOPTS) $(BUILDDIR)/json + @echo + @echo "Build finished; now you can process the JSON files." + +htmlhelp: + $(SPHINXBUILD) -b htmlhelp $(ALLSPHINXOPTS) $(BUILDDIR)/htmlhelp + @echo + @echo "Build finished; now you can run HTML Help Workshop with the" \ + ".hhp project file in $(BUILDDIR)/htmlhelp." + +qthelp: + $(SPHINXBUILD) -b qthelp $(ALLSPHINXOPTS) $(BUILDDIR)/qthelp + @echo + @echo "Build finished; now you can run "qcollectiongenerator" with the" \ + ".qhcp project file in $(BUILDDIR)/qthelp, like this:" + @echo "# qcollectiongenerator $(BUILDDIR)/qthelp/Futhark.qhcp" + @echo "To view the help file:" + @echo "# assistant -collectionFile $(BUILDDIR)/qthelp/Futhark.qhc" + +devhelp: + $(SPHINXBUILD) -b devhelp $(ALLSPHINXOPTS) $(BUILDDIR)/devhelp + @echo + @echo "Build finished." + @echo "To view the help file:" + @echo "# mkdir -p $$HOME/.local/share/devhelp/Futhark" + @echo "# ln -s $(BUILDDIR)/devhelp $$HOME/.local/share/devhelp/Futhark" + @echo "# devhelp" + +epub: + $(SPHINXBUILD) -b epub $(ALLSPHINXOPTS) $(BUILDDIR)/epub + @echo + @echo "Build finished. The epub file is in $(BUILDDIR)/epub." + +latex: + $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex + @echo + @echo "Build finished; the LaTeX files are in $(BUILDDIR)/latex." + @echo "Run \`make' in that directory to run these through (pdf)latex" \ + "(use \`make latexpdf' here to do that automatically)." + +latexpdf: + $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex + @echo "Running LaTeX files through pdflatex..." + $(MAKE) -C $(BUILDDIR)/latex all-pdf + @echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex." + +latexpdfja: + $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex + @echo "Running LaTeX files through platex and dvipdfmx..." + $(MAKE) -C $(BUILDDIR)/latex all-pdf-ja + @echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex." + +text: + $(SPHINXBUILD) -b text $(ALLSPHINXOPTS) $(BUILDDIR)/text + @echo + @echo "Build finished. The text files are in $(BUILDDIR)/text." + +man: + $(SPHINXBUILD) -b man $(ALLSPHINXOPTS) $(BUILDDIR)/man + @echo + @echo "Build finished. The manual pages are in $(BUILDDIR)/man." + +texinfo: + $(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(BUILDDIR)/texinfo + @echo + @echo "Build finished. The Texinfo files are in $(BUILDDIR)/texinfo." + @echo "Run \`make' in that directory to run these through makeinfo" \ + "(use \`make info' here to do that automatically)." + +info: + $(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(BUILDDIR)/texinfo + @echo "Running Texinfo files through makeinfo..." + make -C $(BUILDDIR)/texinfo info + @echo "makeinfo finished; the Info files are in $(BUILDDIR)/texinfo." + +gettext: + $(SPHINXBUILD) -b gettext $(I18NSPHINXOPTS) $(BUILDDIR)/locale + @echo + @echo "Build finished. The message catalogs are in $(BUILDDIR)/locale." + +changes: + $(SPHINXBUILD) -b changes $(ALLSPHINXOPTS) $(BUILDDIR)/changes + @echo + @echo "The overview file is in $(BUILDDIR)/changes." + +linkcheck: + $(SPHINXBUILD) -b linkcheck $(ALLSPHINXOPTS) $(BUILDDIR)/linkcheck + @echo + @echo "Link check complete; look for any errors in the above output " \ + "or in $(BUILDDIR)/linkcheck/output.txt." + +doctest: + $(SPHINXBUILD) -b doctest $(ALLSPHINXOPTS) $(BUILDDIR)/doctest + @echo "Testing of doctests in the sources finished, look at the " \ + "results in $(BUILDDIR)/doctest/output.txt." + +xml: + $(SPHINXBUILD) -b xml $(ALLSPHINXOPTS) $(BUILDDIR)/xml + @echo + @echo "Build finished. The XML files are in $(BUILDDIR)/xml." + +pseudoxml: + $(SPHINXBUILD) -b pseudoxml $(ALLSPHINXOPTS) $(BUILDDIR)/pseudoxml + @echo + @echo "Build finished. The pseudo-XML files are in $(BUILDDIR)/pseudoxml." diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,321 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +# +# Futhark documentation build configuration file, created by +# sphinx-quickstart on Tue Mar 24 14:21:12 2015. +# +# This file is execfile()d with the current directory set to its +# containing dir. +# +# Note that not all possible configuration values are present in this +# autogenerated file. +# +# All configuration values have a default; values that are commented out +# serve to show the default. + +import sys +import os +import re +from pygments.lexer import RegexLexer +from pygments import token +from sphinx.highlighting import lexers + +# If extensions (or modules to document with autodoc) are in another directory, +# add these directories to sys.path here. If the directory is relative to the +# documentation root, use os.path.abspath to make it absolute, like shown here. +#sys.path.insert(0, os.path.abspath('.')) + +# -- General configuration ------------------------------------------------ + +# If your documentation needs a minimal Sphinx version, state it here. +#needs_sphinx = '1.0' + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom +# ones. +extensions = [ + 'sphinx.ext.todo', + 'sphinx.ext.mathjax', +] + +# Add any paths that contain templates here, relative to this directory. +templates_path = ['_templates'] + +# The suffix of source filenames. +source_suffix = '.rst' + +# The encoding of source files. +#source_encoding = 'utf-8-sig' + +# The master toctree document. +master_doc = 'index' + +# General information about the project. +project = 'Futhark' +copyright = '2013-2019, DIKU, University of Copenhagen' + +# The version info for the project you're documenting, acts as replacement for +# |version| and |release|, also used in various other places throughout the +# built documents. +# +# The short X.Y version. + +# No reason for a full YAML parser; let's just hack it. +def get_version(): + # Get lines + lines = open('../package.yaml', 'r').read().split('\n') + # Find version line. + version_line = lines[1] + return re.search('version: "(.*)"', version_line).group(1) + +version = get_version() +# The full version, including alpha/beta/rc tags. +release = version + +# The language for content autogenerated by Sphinx. Refer to documentation +# for a list of supported languages. +#language = None + +# There are two options for replacing |today|: either, you set today to some +# non-false value, then it is used: +#today = '' +# Else, today_fmt is used as the format for a strftime call. +#today_fmt = '%B %d, %Y' + +# List of patterns, relative to source directory, that match files and +# directories to ignore when looking for source files. +exclude_patterns = ['_build'] + +# The reST default role (used for this markup: `text`) to use for all +# documents. +#default_role = None + +# If true, '()' will be appended to :func: etc. cross-reference text. +#add_function_parentheses = True + +# If true, the current module name will be prepended to all description +# unit titles (such as .. function::). +#add_module_names = True + +# If true, sectionauthor and moduleauthor directives will be shown in the +# output. They are ignored by default. +#show_authors = false + +# The name of the Pygments (syntax highlighting) style to use. +pygments_style = 'sphinx' + +class FutharkLexer(RegexLexer): + name = 'Futhark' + + tokens = { + 'root': [ + (r'if|then|else|let|loop|in|val|for|do|with|local|open|include|import|type|entry|module|while|unsafe|module', token.Keyword), + (r"[a-zA-Z_][a-zA-Z0-9_']*", token.Name), + (r"-- .*", token.Comment), + (r'.', token.Text) + ] + } + + +lexers['futhark'] = FutharkLexer() + +highlight_language = 'futhark' + +# A list of ignored prefixes for module index sorting. +#modindex_common_prefix = [] + +# If true, keep warnings as "system message" paragraphs in the built documents. +#keep_warnings = false + + +# -- Options for HTML output ---------------------------------------------- + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +html_theme = 'futhark' +html_theme_path = ['_theme'] + +# Theme options are theme-specific and customize the look and feel of a theme +# further. For a list of options available for each theme, see the +# documentation. +html_theme_options = {} + +# Add any paths that contain custom themes here, relative to this directory. +#html_theme_path = [] + +# The name for this set of Sphinx documents. If None, it defaults to +# " v documentation". +#html_title = None + +# A shorter title for the navigation bar. Default is the same as html_title. +#html_short_title = None + +# The name of an image file (relative to this directory) to place at the top +# of the sidebar. +#html_logo = None + +# The name of an image file (within the static path) to use as favicon of the +# docs. This file should be a Windows icon file (.ico) being 16x16 or 32x32 +# pixels large. +#html_favicon = None + +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +html_static_path = ['_static'] + +# Add any extra paths that contain custom files (such as robots.txt or +# .htaccess) here, relative to this directory. These files are copied +# directly to the root of the documentation. +#html_extra_path = [] + +# If not '', a 'Last updated on:' timestamp is inserted at every page bottom, +# using the given strftime format. +#html_last_updated_fmt = '%b %d, %Y' + +# If true, SmartyPants will be used to convert quotes and dashes to +# typographically correct entities. +#html_use_smartypants = True + +# Custom sidebar templates, maps document names to template names. +html_sidebars = { '**': ['globaltoc.html', 'relations.html', 'sourcelink.html', 'searchbox.html'] } + +# Additional templates that should be rendered to pages, maps page names to +# template names. +#html_additional_pages = {} + +# If false, no module index is generated. +#html_domain_indices = True + +# If false, no index is generated. +#html_use_index = True + +# If true, the index is split into individual pages for each letter. +#html_split_index = false + +# If true, links to the reST sources are added to the pages. +#html_show_sourcelink = True + +# If true, "Created using Sphinx" is shown in the HTML footer. Default is True. +#html_show_sphinx = True + +# If true, "(C) Copyright ..." is shown in the HTML footer. Default is True. +#html_show_copyright = True + +# If true, an OpenSearch description file will be output, and all pages will +# contain a tag referring to it. The value of this option must be the +# base URL from which the finished HTML is served. +#html_use_opensearch = '' + +# This is the file name suffix for HTML files (e.g. ".xhtml"). +#html_file_suffix = None + +# Output file base name for HTML help builder. +htmlhelp_basename = 'Futharkdoc' + + +# -- Options for LaTeX output --------------------------------------------- + +latex_elements = { +# The paper size ('letterpaper' or 'a4paper'). +#'papersize': 'letterpaper', + +# The font size ('10pt', '11pt' or '12pt'). +#'pointsize': '10pt', + +# Additional stuff for the LaTeX preamble. +#'preamble': '', +} + +# Grouping the document tree into LaTeX files. List of tuples +# (source start file, target name, title, +# author, documentclass [howto, manual, or own class]). +latex_documents = [ + ('index', 'Futhark.tex', 'Futhark Documentation', + 'DIKU', 'manual'), +] + +# The name of an image file (relative to this directory) to place at the top of +# the title page. +#latex_logo = None + +# For "manual" documents, if this is true, then toplevel headings are parts, +# not chapters. +#latex_use_parts = false + +# If true, show page references after internal links. +#latex_show_pagerefs = false + +# If true, show URL addresses after external links. +#latex_show_urls = false + +# Documents to append as an appendix to all manuals. +#latex_appendices = [] + +# If false, no module index is generated. +#latex_domain_indices = True + + +# -- Options for manual page output --------------------------------------- + +# One entry per manual page. List of tuples +# (source start file, name, description, authors, manual section). +man_pages = [ + ('man/futhark', 'futhark', 'a parallel functional array language', + ['Troels Henriksen (athas@sigkill.dk)'], 1), + ('man/futhark-c', 'futhark-c', 'compile Futhark to sequential C', + ['Troels Henriksen (athas@sigkill.dk)'], 1), + ('man/futhark-opencl', 'futhark-opencl', 'compile Futhark to OpenCL', + ['Troels Henriksen (athas@sigkill.dk)'], 1), + ('man/futhark-cuda', 'futhark-cuda', 'compile Futhark to CUDA', + ['Troels Henriksen (athas@sigkill.dk)'], 1), + ('man/futhark-python', 'futhark-python', 'compile Futhark to sequential Python', + ['Troels Henriksen (athas@sigkill.dk)'], 1), + ('man/futhark-pyopencl', 'futhark-pyopencl', 'compile Futhark to Python and OpenCL', + ['Troels Henriksen (athas@sigkill.dk)'], 1), + ('man/futhark-csharp', 'futhark-csharp', 'compile Futhark to sequential C#', + ['Troels Henriksen (athas@sigkill.dk)'], 1), + ('man/futhark-csopencl', 'futhark-csopencl', 'compile Futhark to C# and OpenCL', + ['Troels Henriksen (athas@sigkill.dk)'], 1), + ('man/futhark-run', 'futhark-run', 'interpret Futhark program', + ['Troels Henriksen (athas@sigkill.dk)'], 1), + ('man/futhark-repl', 'futhark-repl', 'interactive Futhark read-eval-print-loop', + ['Troels Henriksen (athas@sigkill.dk)'], 1), + ('man/futhark-test', 'futhark-test', 'test Futhark programs', + ['Troels Henriksen (athas@sigkill.dk)'], 1), + ('man/futhark-bench', 'futhark-bench', 'benchmark Futhark programs', + ['Troels Henriksen (athas@sigkill.dk)'], 1), + ('man/futhark-doc', 'futhark-doc', 'generate documentation for Futhark code', + ['Troels Henriksen (athas@sigkill.dk)'], 1), + ('man/futhark-dataset', 'futhark-dataset', 'generate random data sets', + ['Troels Henriksen (athas@sigkill.dk)'], 1), + ('man/futhark-pkg', 'futhark-pkg', 'manage Futhark packages', + ['Troels Henriksen (athas@sigkill.dk)'], 1) +] + +# If true, show URL addresses after external links. +#man_show_urls = false + + +# -- Options for Texinfo output ------------------------------------------- + +# Grouping the document tree into Texinfo files. List of tuples +# (source start file, target name, title, author, +# dir menu entry, description, category) +texinfo_documents = [ + ('index', 'Futhark', 'Futhark Documentation', + 'DIKU', 'Futhark', 'One line description of project.', + 'Miscellaneous'), +] + +# Documents to append as an appendix to all manuals. +#texinfo_appendices = [] + +# If false, no module index is generated. +#texinfo_domain_indices = True + +# How to display URL addresses: 'footnote', 'no', or 'inline'. +#texinfo_show_urls = 'footnote' + +# If true, do not generate a @detailmenu in the "Top" node's menu. +#texinfo_no_detailmenu = false diff --git a/docs/index.rst b/docs/index.rst new file mode 100644 --- /dev/null +++ b/docs/index.rst @@ -0,0 +1,57 @@ +.. Futhark documentation master file, created by + sphinx-quickstart on Tue Mar 24 14:21:12 2015. + You can adapt this file completely to your liking, but it should at least + contain the root `toctree` directive. + +Futhark User's Guide +==================== + +Welcome to the documentation for the Futhark compiler and language. +For a basic introduction, please see `the Futhark website +`_. To get started, read the page on +:ref:`installation`. Once the compiler has been installed, you might +want to take a look at :ref:`usage`. This User's Guide contains a +:ref:`language-reference`, but new Futhark programmers are probably +better served by reading `Parallel Programming in Futhark +`_ first. + +Documentation for the included basis library is also `available online +`_. + +The particularly interested reader may also want to peruse the +`publications `_, or +the `development blog `_. + +.. toctree:: + :caption: Table of Contents + :maxdepth: 2 + :numbered: + + installation.rst + usage.rst + language-reference.rst + package-management.rst + c-porting-guide.rst + versus-other-languages.rst + hacking.rst + binary-data-format.rst + +.. toctree:: + :caption: Manual Pages + :maxdepth: 1 + + man/futhark.rst + man/futhark-bench.rst + man/futhark-c.rst + man/futhark-csharp.rst + man/futhark-csopencl.rst + man/futhark-cuda.rst + man/futhark-dataset.rst + man/futhark-doc.rst + man/futhark-opencl.rst + man/futhark-pkg.rst + man/futhark-pyopencl.rst + man/futhark-python.rst + man/futhark-repl.rst + man/futhark-run.rst + man/futhark-test.rst diff --git a/docs/man/futhark-bench.rst b/docs/man/futhark-bench.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-bench.rst @@ -0,0 +1,131 @@ +.. role:: ref(emphasis) + +.. _futhark-bench(1): + +============= +futhark-bench +============= + +SYNOPSIS +======== + +futhark bench [options...] programs... + +DESCRIPTION +=========== + +This tool is the recommended way to benchmark Futhark programs. +Programs are compiled using the specified backend (``futhark c`` by +default), then run a number of times for each test case, and the +average runtime printed on standard output. Refer to +:ref:`futhark-test(1)` for information on how to format test data. A +program will be ignored if it contains no data sets - it will not even +be compiled. + +If compilation of a program fails, then ``futhark bench`` will abort +immediately. If execution of a test set fails, an error message will +be printed and benchmarking will continue (and ``--json`` will write +the file), but a non-zero exit code will be returned at the end. + +OPTIONS +======= + +--backend=name + + The backend used when compiling Futhark programs (without leading + ``futhark``, e.g. just ``opencl``). + +--entry-point=name + + Only run entry points with this name. + +--exclude-case=TAG + + Do not run test cases that contain the given tag. Cases marked with + "nobench" or "disable" are ignored by default. + +--futhark=program + + The program used to perform operations (eg. compilation). Defaults + to the binary running ``futhark bench`` itself. + +--ignore-files=REGEX + + Ignore files whose path match the given regular expression. + +--json=file + + Write raw results in JSON format to the specified file. + +--no-tuning + + Do not look for tuning files. + +--pass-option=opt + + Pass an option to benchmark programs that are being run. For + example, we might want to run OpenCL programs on a specific device:: + + futhark bench prog.fut --backend=opencl --pass-option=-dHawaii + +--runner=program + + If set to a non-empty string, compiled programs are not run + directly, but instead the indicated *program* is run with its first + argument being the path to the compiled Futhark program. This is + useful for compilation targets that cannot be executed directly (as + with :ref:`futhark-csharp(1)`), or when you wish to run the program + on a remote machine. + +--runs=count + + The number of runs per data set. + +--skip-compilation + + Do not run the compiler, and instead assume that each benchmark + program has already been compiled. Use with caution. + +--timeout=seconds + + If the runtime for a dataset exceeds this integral number of + seconds, it is aborted. Note that the time is allotted not *per + run*, but for *all runs* for a dataset. A twenty second limit for + ten runs thus means each run has only two seconds (minus + initialisation overhead). + + A negative timeout means to wait indefinitely. + +--tuning=EXTENSION + + For each program being run, look for a tuning file with this + extension, which is suffixed to the name of the program. For + example, given ``--tuning=tuning`` (the default), the program + ``foo.fut`` will be passed the tuning file ``foo.fut.tuning`` if it + exists. + +EXAMPLES +======== + +The following program benchmarks how quickly we can sum arrays of +different sizes:: + + -- How quickly can we reduce arrays? + -- + -- == + -- nobench input { 0 } + -- output { 0 } + -- input { 100 } + -- output { 4950 } + -- compiled input { 100000 } + -- output { 704982704 } + -- compiled input { 100000000 } + -- output { 887459712 } + + let main(n: i32): i32 = + reduce (+) 0 (iota n) + +SEE ALSO +======== + +:ref:`futhark-c(1)`, :ref:`futhark-test(1)` diff --git a/docs/man/futhark-c.rst b/docs/man/futhark-c.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-c.rst @@ -0,0 +1,60 @@ +.. role:: ref(emphasis) + +.. _futhark-c(1): + +========= +futhark-c +========= + +SYNOPSIS +======== + +futhark c [options...] infile + +DESCRIPTION +=========== + +``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``. + +The resulting program will read the arguments to the entry point +(``main`` by default) from standard input and print its return value +on standard output. The arguments are read and printed in Futhark +syntax. + +OPTIONS +======= + +-h + Print help text to standard output and exit. + +--library + Generate a library instead of an executable. Appends ``.c``/``.h`` + to the name indicated by the ``-o`` option to determine output + file names. + +-o outfile + Where to write the result. If the source program is named + ``foo.fut``, this defaults to ``foo``. + +--safe + Ignore ``unsafe`` in program and perform safety checks unconditionally. + +-v verbose + Enable debugging output. If compilation fails due to a compiler + error, the result of the last successful compiler step will be + printed to standard error. + +-V + Print version information on standard output and exit. + +--Werror + Treat warnings as errors. + +SEE ALSO +======== + +:ref:`futhark-opencl(1)`, :ref:`futhark-cuda(1)`, :ref:`futhark-test(1)` diff --git a/docs/man/futhark-csharp.rst b/docs/man/futhark-csharp.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-csharp.rst @@ -0,0 +1,65 @@ +.. role:: ref(emphasis) + +.. _futhark-csharp(1): + +============== +futhark-csharp +============== + +SYNOPSIS +======== + +futhark csharp [options...] infile + +DESCRIPTION +=========== + +``futhark csharp`` translates a Futhark program to sequential C# code, +and either compiles that C# code with the Roslyn C# Compiler ``csc`` +to an executable binary program, or produces a ``.dll`` file that can +be linked with other code.. The standard Futhark optimisation +pipeline is used, and ``csc`` is invoked with ``-lib:$MONO_PATH``, +``-r:Mono.Options.dll``, and ``/unsafe``. + +The resulting program will read the arguments to the entry point +(``main`` by default) from standard input and print its return value +on standard output. The arguments are read and printed in Futhark +syntax. + +OPTIONS +======= + +-h + Print help text to standard output and exit. + +--library + Generate a library instead of an executable. Appends ``.dll`` + to the name indicated by the ``-o`` option to determine output + file names. + +-o outfile + Where to write the result. If the source program is named + 'foo.fut', this defaults to 'foo'. + +--safe + Ignore ``unsafe`` in program and perform safety checks unconditionally. + +-v verbose + Enable debugging output. If compilation fails due to a compiler + error, the result of the last successful compiler step will be + printed to standard error. + +-V + Print version information on standard output and exit. + +REQUIREMENTS +============ +``futhark csharp`` uses the Mono implementation of the .NET framework. +To compile and execute the compiled binaries/libraries, you must have the ``MONO_PATH`` environment variable defined. ``MONO_PATH`` must be set to a directory containing the ``Mono.Options`` dll. + +Mono.Options is available on https://www.nuget.org/packages/Mono.Options/5.3.0.1 + +SEE ALSO +======== + +:ref:`futhark-csopencl(1)` diff --git a/docs/man/futhark-csopencl.rst b/docs/man/futhark-csopencl.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-csopencl.rst @@ -0,0 +1,69 @@ +.. role:: ref(emphasis) + +.. _futhark-csopencl(1): + +================ +futhark-csopencl +================ + +SYNOPSIS +======== + +futhark csopencl [options...] infile + +DESCRIPTION +=========== + + +``futhark csopencl`` translates a Futhark program to C# code invoking +OpenCL kernels, and either compiles that C# code with the Roslyn C# Compiler ``csc`` +to an executable binary program, or produces a ``.dll`` file that can be linked with +other code.. The standard Futhark optimisation pipeline is used, and +``csc`` is invoked with ``-lib:$MONO_PATH``, ``-r:Cloo.clSharp.dll``, +``-r:Mono.Options.dll``, and ``/unsafe``. + +The resulting program will otherwise behave exactly as +one compiled with ``futhark csharp``. + +OPTIONS +======= + +-h + Print help text to standard output and exit. + +--library + Generate a library instead of an executable. Appends ``.dll`` + to the name indicated by the ``-o`` option to determine output + file names. + +-o outfile + Where to write the result. If the source program is named + ``foo.fut``, this defaults to ``foo``. + +--safe + Ignore ``unsafe`` in program and perform safety checks unconditionally. + +-v verbose + Enable debugging output. If compilation fails due to a compiler + error, the result of the last successful compiler step will be + printed to standard error. + +-V + Print version information on standard output and exit. + +--Werror + Treat warnings as errors. + +REQUIREMENTS +============ +``futhark csopencl`` uses the Mono implementation of the .NET framework. +To compile and execute the compiled binaries/libraries, you must have the ``MONO_PATH`` environment variable defined. ``MONO_PATH`` must be set to a directory containing the ``Mono.Options`` and ``Cloo.clSharp`` dll's. + +Mono.Options is available on https://www.nuget.org/packages/Mono.Options/5.3.0.1 + +Cloo.clSharp is available on https://www.nuget.org/packages/Cloo.clSharp/ + +SEE ALSO +======== + +:ref:`futhark-test(1)`, :ref:`futhark-csharp(1)` diff --git a/docs/man/futhark-cuda.rst b/docs/man/futhark-cuda.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-cuda.rst @@ -0,0 +1,61 @@ +.. role:: ref(emphasis) + +.. _futhark-cuda(1): + +============== +futhark-cuda +============== + +SYNOPSIS +======== + +futhark cuda [options...] infile + +DESCRIPTION +=========== + + +``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 ``-O3``, ``-lm``, and ``-std=c99``. The +resulting program will otherwise behave exactly as one compiled with +``futhark c``. + +The generated programs use the NVRTC API for run-time compilation, +which must consequently be available. + +OPTIONS +======= + +-h + Print help text to standard output and exit. + +--library + Generate a library instead of an executable. Appends ``.c``/``.h`` + to the name indicated by the ``-o`` option to determine output + file names. + +-o outfile + Where to write the result. If the source program is named + ``foo.fut``, this defaults to ``foo``. + +--safe + Ignore ``unsafe`` in program and perform safety checks unconditionally. + +-v verbose + Enable debugging output. If compilation fails due to a compiler + error, the result of the last successful compiler step will be + printed to standard error. + +-V + Print version information on standard output and exit. + +--Werror + Treat warnings as errors. + +SEE ALSO +======== + +:ref:`futhark-opencl(1)` diff --git a/docs/man/futhark-dataset.rst b/docs/man/futhark-dataset.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-dataset.rst @@ -0,0 +1,87 @@ +.. role:: ref(emphasis) + +.. _futhark-dataset(1): + +=============== +futhark-dataset +=============== + +SYNOPSIS +======== + +futhark dataset [options...] + +DESCRIPTION +=========== + +Generate random values in Futhark syntax, which can be useful when +generating input datasets for program testing. All Futhark primitive +types are supported. Tuples are not supported. Arrays of specific +(non-random) sizes can be generated. You can specify maximum and +minimum bounds for values, as well as the random seed used when +generating the data. The generated values are written to standard +output. + +If no ``-g``/``--generate`` options are passed, values are read from +standard input, and printed to standard output in the indicated +format. The input format (whether textual or binary) is automatically +detected. + +OPTIONS +======= + +-b, --binary + Output data in binary Futhark format (must precede --generate). + +-g type, --generate type + Generate a value of the indicated type, e.g. ``-g i32`` or ``-g [10]f32``. + + The type may also be a value, in which case that literal value is + generated. + +-s int + Set the seed used for the RNG. Zero by default. + +--T-bounds= + Set inclusive lower and upper bounds on generated values of type + ``T``. ``T`` is any primitive type, e.g. ``i32`` or ``f32``. The + bounds apply to any following uses of the ``-g`` option. + +You can alter the output format using the following flags. To use them, add them +before data generation (--generate): + +--text + Output data in text format (must precede --generate). Default. + +-t, --type + Output the types of values (textually) instead of the values + themselves. Mostly useful when reading values on stdin. + +EXAMPLES +======== + +Generate a 4 by 2 integer matrix:: + + futhark dataset -g [4][2]i32 + +Generate an array of floating-point numbers and an array of indices into that array:: + + futhark dataset -g [10]f32 --i32-bounds=0:9 -g [100]i32 + +To generate binary data, the ``--binary`` must come before the ``--generate``:: + + futhark dataset --binary --generate=[42]i32 + +Create a binary data file from a data file:: + + futhark dataset --binary < any_data > binary_data + +Determine the types of values contained in a data file:: + + futhark dataset -t < any_data + + +SEE ALSO +======== + +:ref:`futhark-test(1)`, :ref:`futhark-bench(1)` diff --git a/docs/man/futhark-doc.rst b/docs/man/futhark-doc.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-doc.rst @@ -0,0 +1,66 @@ +.. role:: ref(emphasis) + +.. _futhark-doc(1): + +=========== +futhark-doc +=========== + +SYNOPSIS +======== + +futhark doc [options...] dir + +DESCRIPTION +=========== + +``futhark doc`` generates HTML-formatted documentation from Futhark +code. One HTML file will be created for each ``.fut`` file in the +given directory, as well as any file reachable through ``import`` +expressions. The given Futhark code will be considered as one +cohesive whole, and must be type-correct. + +Futhark definitions may be documented by prefixing them with a block +of line comments starting with :literal:`-- |` (see example below). +Simple Markdown syntax is supported within these comments. A link to +another identifier is possible with the notation +:literal:`\`name\`@namespace`, where ``namespace`` must be either +``term``, ``type``, or ``mtype`` (module names are in the ``term`` +namespace). A file may contain a leading documentation comment, which +will be considered the file *abstract*. + +``futhark doc`` will ignore any file whose documentation comment +consists solely of the word "ignore". This is useful for files that +contain tests, or are otherwise not relevant to the reader of the +documentation. + +OPTIONS +======= + +-h + Print help text to standard output and exit. + +-o outdir + The name of the directory that will contain the generated + documentation. This option is mandatory. + +-v, --verbose + Print status messages to stderr while running. + +-V + Print version information on standard output and exit. + +EXAMPLES +======== + +.. code-block:: futhark + + -- | Gratuitous re-implementation of `map`@term. + -- + -- Does exactly the same. + let mymap = ... + +SEE ALSO +======== + +:ref:`futhark-test(1)`, :ref:`futhark-bench(1)` diff --git a/docs/man/futhark-opencl.rst b/docs/man/futhark-opencl.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-opencl.rst @@ -0,0 +1,58 @@ +.. role:: ref(emphasis) + +.. _futhark-opencl(1): + +============== +futhark-opencl +============== + +SYNOPSIS +======== + +futhark opencl [options...] infile + +DESCRIPTION +=========== + + +``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 ``-O3``, ``-lm``, and +``-std=c99``. The resulting program will otherwise behave exactly as +one compiled with ``futhark c``. + +OPTIONS +======= + +-h + Print help text to standard output and exit. + +--library + Generate a library instead of an executable. Appends ``.c``/``.h`` + to the name indicated by the ``-o`` option to determine output + file names. + +-o outfile + Where to write the result. If the source program is named + ``foo.fut``, this defaults to ``foo``. + +--safe + Ignore ``unsafe`` in program and perform safety checks unconditionally. + +-v verbose + Enable debugging output. If compilation fails due to a compiler + error, the result of the last successful compiler step will be + printed to standard error. + +-V + Print version information on standard output and exit. + +--Werror + Treat warnings as errors. + +SEE ALSO +======== + +:ref:`futhark-test(1)`, :ref:`futhark-cuda(1)`, :ref:`futhark-c(1)` diff --git a/docs/man/futhark-pkg.rst b/docs/man/futhark-pkg.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-pkg.rst @@ -0,0 +1,167 @@ +.. role:: ref(emphasis) + +.. _futhark-pkg(1): + +=========== +futhark-pkg +=========== + +SYNOPSIS +======== + +futhark pkg add PKGPATH [X.Y.Z] + +futhark pkg check + +futhark pkg init PKGPATH + +futhark pkg fmt + +futhark pkg remove PKGPATH + +futhark pkg sync + +futhark pkg upgrade + +futhark pkg versions + +DESCRIPTION +=========== + +This tool is used to modify the package manifest (``futhark.pkg``) and +download the required packages it describes. ``futhark pkg`` is not a +build system; you will still need to compile your Futhark code with +the usual compilers. The only purpose of ``futhark pkg`` is to +download code (and perform other package management utility tasks). +This manpage is not a general introduction to package management in +Futhark; see the User's Guide for that. + +The ``futhark pkg`` subcommands will modify only two locations in the +file system (relative to the current working directory): the +``futhark.pkg`` file, and the contents of ``lib/``. When modifying +``lib/``, ``futhark pkg`` constructs the new version in ``lib~new/`` +and backs up the old version in ``lib~old``. If ``futhark pkg`` +should fail for any reason, you can recover the old state by moving +``lib~old`` back. These temporary directories are erased if +``futhark pkg`` finishes without errors. + +The ``futhark pkg sync`` and ``futhark pkg init`` subcommands are +the only ones that actually modifies ``lib/``; the others modify only +``futhark.pkg`` and require you to manually run ``futhark pkg sync`` +afterwards. + +Most commands take a ``-v``/``--verbose`` option that makes +``futhark pkg`` write running diagnostics to stderr. + +COMMANDS +======== + +futhark pkg add PKGPATH [X.Y.Z] +------------------------------- + +Add the specified package of the given minimum version as a +requirement to ``futhark.pkg``. If no version is provided, the newest +one is used. If the package is already required in ``futhark.pkg``, +the new version requirement will replace the old one. + +Note that adding a package does not automatically download it. Run +``futhark pkg sync`` to do that. + +futhark pkg check +----------------- + +Verify that the ``futhark.pkg`` is valid, that all required packages +are available in the indicated versions. This command does not check +that these versions contain well-formed code. If a package path is +defined in ``futhark.pkg``, also checks that ``.fut`` files are +located at the expected location in the file system. + +futhark pkg init PKGPATH +------------------------ + +Create a new ``futhark.pkg`` defining a package with the given package +path, and initially no requirements. + +futhark pkg fmt +--------------- + +Reformat the ``futhark.pkg`` file, while retaining any comments. + +futhark pkg remove PKGPATH +-------------------------- + +Remove a package from ``futhark.pkg``. Does *not* remove it from the +``lib/`` directory. + +futhark pkg sync +---------------- + +Populate the ``lib/`` directory with the packages listed in +``futhark.pkg``. **Warning**: this will delete everything in ``lib/`` +that does not relate to a file listed in ``futhark.pkg``, as well as +any local modifications. + +futhark pkg upgrade +------------------- + +Upgrade all package requirements in ``futhark.pkg`` to the newest +available versions. + +futhark pkg versions PKGPATH +---------------------------- + +Print all available versions for the given package path. + +COMMIT VERSIONS +=============== + +It is possible to use ``futhark pkg`` with packages that have not yet +made proper releases. This is done via pseudoversions of the form +``0.0.0-yyyymmddhhmmss+commitid``. The timestamp is not verified +against the actual commit. The timestamp ensures that newer commits +take precedence if multiple packages depend on a commit version for +the same package. If ``futhark pkg add`` is given a package with no +releases, the most recent commit will be used. In this case, the +timestamp is merely set to the current time. + +Commit versions are awkward and fragile, and should not be relied +upon. Issue proper releases (even experimental 0.x version) as soon +as feasible. Released versions also always take precedence over +commit versions, since any version number will be greater than 0.0.0. + +EXAMPLES +======== + +Create a new package that will be hosted at +``https://github.com/sturluson/edda``:: + + futhark pkg init github.com/sturluson/edda + +Add a package dependency:: + + futhark pkg add github.com/sturluson/hattatal + +Download the dependencies:: + + futhark pkg sync + +And then you're ready to start hacking! (Except that these packages +do not actually exist.) + +BUGS +==== + +Since the ``lib/`` directory is populated with transitive dependencies +as well, it is possible for a package to depend unwittingly on one of +the dependencies of its dependencies, without the ``futhark.pkg`` file +reflecting this. + +There is no caching of zipballs and version lists between invocations, +so the network traffic can be rather heavy. + +Only GitHub and GitLab are supported as code hosting sites. + +SEE ALSO +======== + +:ref:`futhark-test(1)`, :ref:`futhark-doc(1)` diff --git a/docs/man/futhark-pyopencl.rst b/docs/man/futhark-pyopencl.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-pyopencl.rst @@ -0,0 +1,69 @@ +.. role:: ref(emphasis) + +.. _futhark-pyopencl(1): + +================ +futhark-pyopencl +================ + +SYNOPSIS +======== + +futhark pyopencl [options...] infile + +DESCRIPTION +=========== + +``futhark pyopencl`` translates a Futhark program to Python code +invoking OpenCL kernels. By default, the program uses the first +device of the first OpenCL platform - this can be changed by passing +``-p`` and ``-d`` options to the generated program (not to +``futhark pyopencl`` itself). + +The resulting program will otherwise behave exactly as one compiled +with ``futhark py``. While the sequential host-level code is pure +Python and just as slow as in ``futhark py``, parallel sections will +have been compiled to OpenCL, and runs just as fast as when using +``futhark opencl``. The kernel launch overhead is significantly +higher, however, so a good rule of thumb when using +``futhark pyopencl`` is to aim for having fewer but longer-lasting +parallel sections. + +The generated code requires at least PyOpenCL version 2015.2. + +OPTIONS +======= + +-h + Print help text to standard output and exit. + +--library + Instead of compiling to an executable program, generate a Python + module that can be imported by other Python code. The module will + contain a class of the same name as the Futhark source file with + ``.fut`` removed. Objects of the class define one method per entry + point in the Futhark program, with matching parameters and return + value. + +-o outfile + Where to write the resulting binary. By default, if the source + program is named 'foo.fut', the binary will be named 'foo'. + +--safe + Ignore ``unsafe`` in program and perform safety checks unconditionally. + +-v verbose + Enable debugging output. If compilation fails due to a compiler + error, the result of the last successful compiler step will be + printed to standard error. + +-V + Print version information on standard output and exit. + +--Werror + Treat warnings as errors. + +SEE ALSO +======== + +:ref:`futhark-py(1)`, :ref:`futhark-opencl(1)` diff --git a/docs/man/futhark-python.rst b/docs/man/futhark-python.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-python.rst @@ -0,0 +1,63 @@ +.. role:: ref(emphasis) + +.. _futhark-py(1): + +============== +futhark-python +============== + +SYNOPSIS +======== + +futhark python [options...] infile + +DESCRIPTION +=========== + +``futhark python`` translates a Futhark program to sequential Python +code. + +The resulting program will read the arguments to the ``main`` function +from standard input and print its return value on standard output. +The arguments are read and printed in Futhark syntax. + +The generated code is very slow, likely too slow to be useful. It is +more interesting to use this command's big brother, +:ref:`futhark-pyopencl(1)`. + +OPTIONS +======= + +-h + Print help text to standard output and exit. + +--library + Instead of compiling to an executable program, generate a Python + module that can be imported by other Python code. The module will + contain a class of the same name as the Futhark source file with + ``.fut`` removed. Objects of the class define one method per entry + point in the Futhark program, with matching parameters and return + value. + +-o outfile + Where to write the resulting binary. By default, if the source + program is named 'foo.fut', the binary will be named 'foo'. + +--safe + Ignore ``unsafe`` in program and perform safety checks unconditionally. + +-v verbose + Enable debugging output. If compilation fails due to a compiler + error, the result of the last successful compiler step will be + printed to standard error. + +-V + Print version information on standard output and exit. + +--Werror + Treat warnings as errors. + +SEE ALSO +======== + +:ref:`futhark-pyopencl(1)` diff --git a/docs/man/futhark-repl.rst b/docs/man/futhark-repl.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-repl.rst @@ -0,0 +1,42 @@ +.. role:: ref(emphasis) + +.. _futhark-repl(1): + +============ +futhark-repl +============ + +SYNOPSIS +======== + +futhark repl + +DESCRIPTION +=========== + +Start an interactive Futhark session. This will let you interactively +enter expressions and declarations which are then immediately +interpreted. If the entered line can be either a declaration or an +expression, it is assumed to be a declaration. + +Futhark source files can be loaded using the ``:load`` command. This +will erase any interactively entered definitions. Use the ``:help`` +command to see a list of commands. All commands are prefixed with a +colon. + +``futhark-repl`` uses the Futhark interpreter, which grants access to +certain special functions. See :ref:`futhark-run(1)` for a description. + +OPTIONS +======= + +-h + Print help text to standard output and exit. + +-V + Print version information on standard output and exit. + +SEE ALSO +======== + +:ref:`futhark-run(1)`, :ref:`futhark-test(1)` diff --git a/docs/man/futhark-run.rst b/docs/man/futhark-run.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-run.rst @@ -0,0 +1,53 @@ +.. role:: ref(emphasis) + +.. _futhark-run(1): + +=========== +futhark-run +=========== + +SYNOPSIS +======== + +futhark run [program] + +DESCRIPTION +=========== + +Execute the given program by evaluating the ``main`` function with +arguments read from standard input, and write the results on standard +output. + +``futhark-run`` is very slow, and in practice only useful for testing, +teaching, and experimenting with the language. Certain special +debugging functions are available in ``futhark-run``: + +``trace 'a : a -> a`` + Semantically identity, but prints the value on standard output. + +``break 'a : a -> a`` + Semantically identity, but interrupts execution at the calling + point, such that the environment can be inspected. Continue + execution by entering an empty input line. Breakpoints are only + respected when starting a program from the prompt, not when + passing a program on the command line. + +OPTIONS +======= + +-e NAME + Run the given entry point instead of ``main``. + +-h + Print help text to standard output and exit. + +-V + Print version information on standard output and exit. + +-w, --no-warnings + Disable interpreter warnings. + +SEE ALSO +======== + +:ref:`futhark-repl(1)`, :ref:`futhark-test(1)` diff --git a/docs/man/futhark-test.rst b/docs/man/futhark-test.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-test.rst @@ -0,0 +1,224 @@ +.. role:: ref(emphasis) + +.. _futhark-test(1): + +============ +futhark-test +============ + +SYNOPSIS +======== + +futhark test [options...] infiles... + +DESCRIPTION +=========== + +This tool is used to test Futhark programs based on input/output +datasets. If a directory is given, all contained files with a +``.fut`` extension are considered. + +A Futhark test program is an ordinary Futhark program, with at least +one test block describing input/output test cases and possibly other +options. A test block consists of commented-out text with the +following overall format:: + + description + == + cases... + +The ``description`` is an arbitrary (and possibly multiline) +human-readable explanation of the test program. It is separated from +the test cases by a line containing just ``==``. Any comment starting +at the beginning of the line, and containing a line consisting of just +``==``, will be considered a test block. The format of a test case is +as follows:: + + [tags { tags... }] + [entry: names...] + [compiled|nobench|random] input ({ values... } | @ filename) + output { values... } | auto output | error: regex + +If ``compiled`` is present before the ``input`` keyword, this test +case will never be passed to the interpreter. This is useful for test +cases that are annoyingly slow to interpret. The ``nobench`` keyword +is for data sets that are too small to be worth benchmarking, and only +has meaning to :ref:`futhark-bench(1)`. + +If ``input`` is preceded by ``random``, the text between the curly +braces must consist of a sequence of Futhark types, including sizes in +the case of arrays. When ``futhark test`` is run, a file located in a +``data/`` subdirectory, containing values of the indicated types and +shapes is, automatically constructed with ``futhark-dataset``. Apart +from sizes, integer constants (without any type suffix) are also +permitted. These become ``ì32`` values. + +If ``input`` is followed by an ``@`` and a file name (which must not +contain any whitespace) instead of curly braces, values will be read +from the indicated file. This is recommended for large data sets. +This notation cannot be used with ``random`` input. + +After the ``input`` block, the expected result of the test case is +written as either ``output`` followed by another block of values, or +an expected run-time error, in which a regular expression can be used +to specify the exact error message expected. If no regular expression +is given, any error message is accepted. If neither ``output`` nor +``error`` is given, the program will be expected to execute +succesfully, but its output will not be validated. + +If ``output`` is preceded by ``auto`` (as in ``auto output``), the +expected values are automatically generated by compiling the program +with ``futhark-c`` and recording its result for the given input (which +must not fail). This is usually only useful for testing or +benchmarking alternative compilers, and not for testing the +correctness of Futhark programs. + +Alternatively, instead of input-output pairs, the test cases can +simply be a description of an expected compile time type error:: + + error: regex + +This is used to test the type checker. + +By default, both the interpreter and compiler is run on all test cases +(except those that have specified ``compiled``), although this can be +changed with command-line options to ``futhark test``. + +Tuple syntax is not supported when specifying input and output values. +Instead, you can write an N-tuple as its constituent N values. Beware +of syntax errors in the values - the errors reported by +``futhark test`` are very poor. + +An optional tags specification is permitted in the first test block. +This section can contain arbitrary tags that classify the benchmark:: + + tags { names... } + +Tag are sequences of alphanumeric characters, dashes, and underscores, +with each tag seperated by whitespace. Any program with the +``disable`` tag is ignored by ``futhark test``. + +Another optional directive is ``entry``, which specifies the entry +point to be used for testing. This is useful for writing programs +that test libraries with multiple entry points. Multiple entry points +can be specified on the same line by separating them with space, and +they will all be tested with the same input/output pairs. The +``entry`` directive affects subsequent input-output pairs in the same +comment block, and may only be present immediately preceding these +input-output pairs. If no ``entry`` is given, ``main`` is assumed. +See below for an example. + +For many usage examples, see the ``tests`` directory in the +Futhark source directory. A simple example can be found in +``EXAMPLES`` below. + +OPTIONS +======= + +--backend=program + + The backend used when compiling Futhark programs (without leading + ``futhark``, e.g. just ``opencl``). + +-c + Only run compiled code - do not run any interpreters. + +-C + Compile the programs, but do not run them. + +--exclude=tag + + Do not run test cases that contain the given tag. Cases marked with + "disable" are ignored by default. + +-i + Only interpret - do not run any compilers. + +-t + Type-check the programs, but do not run them. + +--futhark=program + + The program used to perform operations (eg. compilation). Defaults + to the binary running ``futhark test`` itself. + +--no-terminal + Print each result on a line by itself, without line buffering. + +--no-tuning + + Do not look for tuning files. + +--pass-option=opt + + Pass an option to benchmark programs that are being run. For + example, we might want to run OpenCL programs on a specific device:: + + futhark-bench prog.fut --backend=opencl --pass-option=-dHawaii + +--runner=program + + If set to a non-empty string, compiled programs are not run + directly, but instead the indicated *program* is run with its first + argument being the path to the compiled Futhark program. This is + useful for compilation targets that cannot be executed directly (as + with :ref:`futhark-csharp(1)`), or when you wish to run the program + on a remote machine. + +--tuning=EXTENSION + + For each program being run, look for a tuning file with this + extension, which is suffixed to the name of the program. For + example, given ``--tuning=tuning`` (the default), the program + ``foo.fut`` will be passed the tuning file ``foo.fut.tuning`` if it + exists. + +EXAMPLES +======== + +The following program tests simple indexing and bounds checking:: + + -- Test simple indexing of an array. + -- == + -- tags { firsttag secondtag } + -- input { [4,3,2,1] 1 } + -- output { 3 } + -- input { [4,3,2,1] 5 } + -- error: Assertion.*failed + + let main (a: []i32) (i: i32): i32 = + a[i] + +The following program contains two entry points, both of which are +tested:: + + let add(x: i32, y: i32): i32 = x + y + + -- Test the add1 function. + -- == + -- entry: add1 + -- input { 1 } output { 2 } + + entry add1 (x: i32): i32 = add x 1 + + -- Test the sub1 function. + -- == + -- entry: sub1 + -- input { 1 } output { 0 } + + entry sub1 (x: i32): i32 = add x (-1) + +The following program contains an entry point that is tested with +randomly generated data:: + + -- == + -- random input { [100]i32 [100]i32 } auto output + -- random input { [1000]i32 [1000]i32 } auto output + + let main xs ys = i32.product (map2 (*) xs ys) + + +SEE ALSO +======== + +:ref:`futhark-bench(1)`, :ref:`futhark-repl(1)` diff --git a/docs/man/futhark.rst b/docs/man/futhark.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark.rst @@ -0,0 +1,57 @@ +.. role:: ref(emphasis) + +.. _futhark(1): + +======= +futhark +======= + +SYNOPSIS +======== + +futhark options... + +DESCRIPTION +=========== + +Futhark is a data-parallel functional array language. Through various +subcommands, the ``futhark`` tool provides facilities for compiling, +developing, or analysing Futhark programs. Most subcommands are +documented in their own manpage. For example, ``futhark opencl`` is +documented as :ref:`futhark-opencl(1)`. The remaining subcommands are +documented in this page. + +COMMANDS +======== + +futhark check PROGRAM +--------------------- + +Check whether a Futhark program type checks. + +futhark datacmp FILE_A FILE_B +----------------------------- + +Check whether the two files contain the same Futhark values. The +files must be formatted using the general Futhark data format that is +used by all other executable and tools (such as +:ref:`futhark-dataset(1)`). All discrepancies will be reported. This +is in contrast to :ref:`futhark-test(1)`, which only reports the first +one. + +futhark dev options... PROGRAM +------------------------------ + +A Futhark compiler development command, intentionally undocumented and +intended for use in developing the Futhark compiler, not for +programmers writing in Futhark. + +futhark imports PROGRAM +----------------------- + +Print all non-builtin imported Futhark files to stdout, one per line. + +SEE ALSO +======== + +:ref:`futhark-opencl(1)`, :ref:`futhark-c(1)`, :ref:`futhark-py(1)`, :ref:`futhark-pyopencl(1)`, :ref:`futhark-dataset(1)`, :ref:`futhark-doc(1)`, :ref:`futhark-test(1)`, :ref:`futhark-bench(1)`, :ref:`futhark-run(1)`, :ref:`futhark-repl(1)` diff --git a/futhark.cabal b/futhark.cabal --- a/futhark.cabal +++ b/futhark.cabal @@ -1,13 +1,13 @@ cabal-version: 1.12 --- This file has been generated from package.yaml by hpack version 0.31.1. +-- This file has been generated from package.yaml by hpack version 0.31.2. -- -- see: https://github.com/sol/hpack -- --- hash: 1626fa99a399d6063ac04d34f5c73c94719161b2eab8876c29218c389c97d9e9 +-- hash: 38e40072e41929349f42259546b30abfbce169d4641549f5f1d5db10f60f2724 name: futhark -version: 0.11.2 +version: 0.12.1 synopsis: An optimising compiler for a functional, array-oriented language. description: Futhark is a small programming language designed to be compiled to efficient parallel code. It is a statically typed, data-parallel, @@ -25,7 +25,6 @@ license-file: LICENSE build-type: Simple extra-source-files: - rts/python/__init__.py rts/python/memory.py rts/python/opencl.py rts/python/panic.py @@ -55,6 +54,25 @@ futlib/prelude.fut futlib/soacs.fut futlib/zip.fut + docs/Makefile + docs/conf.py + docs/index.rst + docs/man/futhark-bench.rst + docs/man/futhark-c.rst + docs/man/futhark-csharp.rst + docs/man/futhark-csopencl.rst + docs/man/futhark-cuda.rst + docs/man/futhark-dataset.rst + docs/man/futhark-doc.rst + docs/man/futhark-opencl.rst + docs/man/futhark-pkg.rst + docs/man/futhark-pyopencl.rst + docs/man/futhark-python.rst + docs/man/futhark-repl.rst + docs/man/futhark-run.rst + docs/man/futhark-test.rst + docs/man/futhark.rst + package.yaml source-repository head type: git @@ -173,18 +191,17 @@ Futhark.Optimise.Simplify.Rule Futhark.Optimise.Simplify.Rules Futhark.Optimise.TileLoops - Futhark.Optimise.TileLoops.RegTiling3D Futhark.Optimise.Unstream Futhark.Pass Futhark.Pass.ExpandAllocations Futhark.Pass.ExplicitAllocations Futhark.Pass.ExtractKernels Futhark.Pass.ExtractKernels.BlockedKernel + Futhark.Pass.ExtractKernels.DistributeNests Futhark.Pass.ExtractKernels.Distribution Futhark.Pass.ExtractKernels.Interchange Futhark.Pass.ExtractKernels.Intragroup Futhark.Pass.ExtractKernels.ISRWIM - Futhark.Pass.ExtractKernels.Kernelise Futhark.Pass.FirstOrderTransform Futhark.Pass.KernelBabysitting Futhark.Pass.ResolveAssertions @@ -252,6 +269,7 @@ Language.Futhark.Syntax Language.Futhark.Traversals Language.Futhark.TypeChecker + Language.Futhark.TypeChecker.Modules Language.Futhark.TypeChecker.Monad Language.Futhark.TypeChecker.Terms Language.Futhark.TypeChecker.Types @@ -272,7 +290,7 @@ , binary >=0.8.3 , blaze-html >=0.9.0.1 , bytestring >=0.10.8 - , containers >=0.5 + , containers >=0.6.2.1 , data-binary-ieee754 >=0.1 , directory >=1.3.0.0 , directory-tree >=0.12.1 diff --git a/futlib/math.fut b/futlib/math.fut --- a/futlib/math.fut +++ b/futlib/math.fut @@ -83,7 +83,7 @@ val &: t -> t -> t val |: t -> t -> t val ^: t -> t -> t - val ~: t -> t + val !: t -> t val <<: t -> t -> t val >>: t -> t -> t @@ -205,7 +205,7 @@ let (x: i8) & (y: i8) = intrinsics.and8 x y let (x: i8) | (y: i8) = intrinsics.or8 x y let (x: i8) ^ (y: i8) = intrinsics.xor8 x y - let ~ (x: i8) = intrinsics.complement8 x + let ! (x: i8) = intrinsics.complement8 x let (x: i8) << (y: i8) = intrinsics.shl8 x y let (x: i8) >> (y: i8) = intrinsics.ashr8 x y @@ -234,7 +234,7 @@ let (x: i8) > (y: i8) = intrinsics.slt8 y x let (x: i8) <= (y: i8) = intrinsics.sle8 x y let (x: i8) >= (y: i8) = intrinsics.sle8 y x - let (x: i8) != (y: i8) = ! (x == y) + let (x: i8) != (y: i8) = intrinsics.! (x == y) let sgn (x: i8) = intrinsics.ssignum8 x let abs (x: i8) = intrinsics.abs8 x @@ -249,7 +249,7 @@ let num_bits = 8 let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) let set_bit (bit: i32) (x: t) (b: i32) = - ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + ((x & i32 (intrinsics.!(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) let iota (n: i8) = 0i8..1i8..> (y: i16) = intrinsics.ashr16 x y @@ -304,7 +304,7 @@ let (x: i16) > (y: i16) = intrinsics.slt16 y x let (x: i16) <= (y: i16) = intrinsics.sle16 x y let (x: i16) >= (y: i16) = intrinsics.sle16 y x - let (x: i16) != (y: i16) = ! (x == y) + let (x: i16) != (y: i16) = intrinsics.! (x == y) let sgn (x: i16) = intrinsics.ssignum16 x let abs (x: i16) = intrinsics.abs16 x @@ -319,7 +319,7 @@ let num_bits = 16 let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) let set_bit (bit: i32) (x: t) (b: i32) = - ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + ((x & i32 (intrinsics.!(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) let iota (n: i16) = 0i16..1i16..> (y: i32) = intrinsics.ashr32 x y @@ -377,7 +377,7 @@ let (x: i32) > (y: i32) = intrinsics.slt32 y x let (x: i32) <= (y: i32) = intrinsics.sle32 x y let (x: i32) >= (y: i32) = intrinsics.sle32 y x - let (x: i32) != (y: i32) = ! (x == y) + let (x: i32) != (y: i32) = intrinsics.! (x == y) let sgn (x: i32) = intrinsics.ssignum32 x let abs (x: i32) = intrinsics.abs32 x @@ -392,7 +392,7 @@ let num_bits = 32 let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) let set_bit (bit: i32) (x: t) (b: i32) = - ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + ((x & i32 (intrinsics.!(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) let iota (n: i32) = 0..1..> (y: i64) = intrinsics.ashr64 x y @@ -450,7 +450,7 @@ let (x: i64) > (y: i64) = intrinsics.slt64 y x let (x: i64) <= (y: i64) = intrinsics.sle64 x y let (x: i64) >= (y: i64) = intrinsics.sle64 y x - let (x: i64) != (y: i64) = ! (x == y) + let (x: i64) != (y: i64) = intrinsics.! (x == y) let sgn (x: i64) = intrinsics.ssignum64 x let abs (x: i64) = intrinsics.abs64 x @@ -465,7 +465,7 @@ let num_bits = 64 let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) let set_bit (bit: i32) (x: t) (b: i32) = - ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | intrinsics.zext_i32_i64 (b intrinsics.<< bit)) + ((x & i32 (intrinsics.!(1 intrinsics.<< bit))) | intrinsics.zext_i32_i64 (b intrinsics.<< bit)) let iota (n: i64) = 0i64..1i64..> (y: u8) = unsign (intrinsics.ashr8 (sign x) (sign y)) @@ -523,7 +523,7 @@ let (x: u8) > (y: u8) = intrinsics.ult8 (sign y) (sign x) let (x: u8) <= (y: u8) = intrinsics.ule8 (sign x) (sign y) let (x: u8) >= (y: u8) = intrinsics.ule8 (sign y) (sign x) - let (x: u8) != (y: u8) = ! (x == y) + let (x: u8) != (y: u8) = intrinsics.! (x == y) let sgn (x: u8) = unsign (intrinsics.usignum8 (sign x)) let abs (x: u8) = x @@ -538,7 +538,7 @@ let num_bits = 8 let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) let set_bit (bit: i32) (x: t) (b: i32) = - ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + ((x & i32 (intrinsics.!(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) let iota (n: u8) = 0u8..1u8..> (y: u16) = unsign (intrinsics.ashr16 (sign x) (sign y)) @@ -596,7 +596,7 @@ let (x: u16) > (y: u16) = intrinsics.ult16 (sign y) (sign x) let (x: u16) <= (y: u16) = intrinsics.ule16 (sign x) (sign y) let (x: u16) >= (y: u16) = intrinsics.ule16 (sign y) (sign x) - let (x: u16) != (y: u16) = ! (x == y) + let (x: u16) != (y: u16) = intrinsics.! (x == y) let sgn (x: u16) = unsign (intrinsics.usignum16 (sign x)) let abs (x: u16) = x @@ -611,7 +611,7 @@ let num_bits = 16 let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) let set_bit (bit: i32) (x: t) (b: i32) = - ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + ((x & i32 (intrinsics.!(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) let iota (n: u16) = 0u16..1u16..> (y: u32) = unsign (intrinsics.ashr32 (sign x) (sign y)) @@ -669,7 +669,7 @@ let (x: u32) > (y: u32) = intrinsics.ult32 (sign y) (sign x) let (x: u32) <= (y: u32) = intrinsics.ule32 (sign x) (sign y) let (x: u32) >= (y: u32) = intrinsics.ule32 (sign y) (sign x) - let (x: u32) != (y: u32) = ! (x == y) + let (x: u32) != (y: u32) = intrinsics.! (x == y) let sgn (x: u32) = unsign (intrinsics.usignum32 (sign x)) let abs (x: u32) = x @@ -684,7 +684,7 @@ let num_bits = 32 let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) let set_bit (bit: i32) (x: t) (b: i32) = - ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + ((x & i32 (intrinsics.!(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) let iota (n: u32) = 0u32..1u32..> (y: u64) = unsign (intrinsics.ashr64 (sign x) (sign y)) @@ -742,7 +742,7 @@ let (x: u64) > (y: u64) = intrinsics.ult64 (sign y) (sign x) let (x: u64) <= (y: u64) = intrinsics.ule64 (sign x) (sign y) let (x: u64) >= (y: u64) = intrinsics.ule64 (sign y) (sign x) - let (x: u64) != (y: u64) = ! (x == y) + let (x: u64) != (y: u64) = intrinsics.! (x == y) let sgn (x: u64) = unsign (intrinsics.usignum64 (sign x)) let abs (x: u64) = x @@ -757,7 +757,7 @@ let num_bits = 64 let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) let set_bit (bit: i32) (x: t) (b: i32) = - ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + ((x & i32 (intrinsics.!(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) let iota (n: u64) = 0u64..1u64.. (y: f64) = intrinsics.lt64 y x let (x: f64) <= (y: f64) = intrinsics.le64 x y let (x: f64) >= (y: f64) = intrinsics.le64 y x - let (x: f64) != (y: f64) = ! (x == y) + let (x: f64) != (y: f64) = intrinsics.! (x == y) let negate (x: t) = -x let max (x: t) (y: t) = intrinsics.fmax64 x y @@ -849,8 +849,6 @@ let trunc (x: f64) : f64 = i64 (i64m.f64 x) - let even (x: f64) = i64m.f64 x % 2i64 i64m.== 0i64 - let round = intrinsics.round64 let to_bits (x: f64): u64 = u64m.i64 (intrinsics.to_bits64 x) @@ -917,7 +915,7 @@ let (x: f32) > (y: f32) = intrinsics.lt32 y x let (x: f32) <= (y: f32) = intrinsics.le32 x y let (x: f32) >= (y: f32) = intrinsics.le32 y x - let (x: f32) != (y: f32) = ! (x == y) + let (x: f32) != (y: f32) = intrinsics.! (x == y) let negate (x: t) = -x let max (x: t) (y: t) = intrinsics.fmax32 x y @@ -959,8 +957,6 @@ else if ix > x then i32 (i i32m.- 1i32) else x let trunc (x: f32) : f32 = i32 (i32m.f32 x) - - let even (x: f32) = i32m.f32 x % 2i32 i32m.== 0i32 let round = intrinsics.round32 diff --git a/futlib/soacs.fut b/futlib/soacs.fut --- a/futlib/soacs.fut +++ b/futlib/soacs.fut @@ -156,7 +156,7 @@ let (as', is) = intrinsics.partition (3, p', as) in (as'[0:is[0]], as'[is[0]:is[0]+is[1]], as'[is[0]+is[1]:n]) --- | `stream_red op f as` splits `as` into chunks, applies `f` to each +-- | `reduce_stream op f as` splits `as` into chunks, applies `f` to each -- of these in parallel, and uses `op` (which must be associative) to -- combine the per-chunk results into a final result. The `i32` -- passed to `f` is the size of the chunk. This SOAC is useful when @@ -171,38 +171,38 @@ -- **Work:** *O(n)* -- -- **Span:** *O(log(n))* -let stream_red 'a 'b (op: b -> b -> b) (f: i32 -> []a -> b) (as: []a): b = - intrinsics.stream_red (op, f, as) +let reduce_stream 'a 'b (op: b -> b -> b) (f: i32 -> []a -> b) (as: []a): b = + intrinsics.reduce_stream (op, f, as) --- | As `stream_red`@term, but the chunks do not necessarily +-- | As `reduce_stream`@term, but the chunks do not necessarily -- correspond to subsequences of the original array (they may be -- interleaved). -- -- **Work:** *O(n)* -- -- **Span:** *O(log(n))* -let stream_red_per 'a 'b (op: b -> b -> b) (f: i32 -> []a -> b) (as: []a): b = - intrinsics.stream_red_per (op, f, as) +let reduce_stream_per 'a 'b (op: b -> b -> b) (f: i32 -> []a -> b) (as: []a): b = + intrinsics.reduce_stream_per (op, f, as) --- | Similar to `stream_red`@term, except that each chunk must produce +-- | Similar to `reduce_stream`@term, except that each chunk must produce -- an array *of the same size*. The per-chunk results are -- concatenated. -- -- **Work:** *O(n)* -- -- **Span:** *O(1)* -let stream_map 'a 'b (f: i32 -> []a -> []b) (as: []a): *[]b = - intrinsics.stream_map (f, as) +let map_stream 'a 'b (f: i32 -> []a -> []b) (as: []a): *[]b = + intrinsics.map_stream (f, as) --- | Similar to `stream_map`@term, but the chunks do not necessarily +-- | Similar to `map_stream`@term, but the chunks do not necessarily -- correspond to subsequences of the original array (they may be -- interleaved). -- -- **Work:** *O(n)* -- -- **Span:** *O(1)* -let stream_map_per 'a 'b (f: i32 -> []a -> []b) (as: []a): *[]b = - intrinsics.stream_map_per (f, as) +let map_stream_per 'a 'b (f: i32 -> []a -> []b) (as: []a): *[]b = + intrinsics.map_stream_per (f, as) -- | Return `true` if the given function returns `true` for all -- elements in the array. diff --git a/package.yaml b/package.yaml new file mode 100644 --- /dev/null +++ b/package.yaml @@ -0,0 +1,118 @@ +name: futhark +version: "0.12.1" +synopsis: An optimising compiler for a functional, array-oriented language. +description: | + Futhark is a small programming language designed to be compiled to + efficient parallel code. It is a statically typed, data-parallel, + and purely functional array language in the ML family, and comes + with a heavily optimising ahead-of-time compiler that presently + generates GPU code via CUDA and OpenCL, although the language itself + is hardware-agnostic. + + For more information, see the website at https://futhark-lang.org +homepage: https://futhark-lang.org +maintainer: Troels Henriksen athas@sigkill.dk +license: ISC +github: diku-dk/futhark +category: Language + +ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists + +extra-source-files: + - rts/python/*.py + - rts/c/*.h + - rts/csharp/*.cs + - rts/futhark-doc/*.css + + - futlib/*.fut + + - docs/Makefile + - docs/conf.py + - docs/index.rst + - docs/man/*.rst + - package.yaml + +library: + dependencies: + - base >= 4 && < 5 + - array >= 0.4 + - binary >= 0.8.3 + - data-binary-ieee754 >= 0.1 + - vector >= 0.12 + - vector-binary-instances >= 0.2.2.0 + - containers >= 0.6.2.1 + - mtl >= 2.2.1 + - transformers >= 0.3 + - srcloc >= 0.4 + - language-c-quote >= 0.12 + - mainland-pretty >= 0.6.1 + - megaparsec >= 7.0.1 + - parser-combinators >= 1.0.0 + - regex-tdfa >= 1.2 + - filepath >= 1.4.1.1 + - dlist >= 0.6.0.1 + - bytestring >= 0.10.8 + - text >= 1.2.2.2 + - neat-interpolation >= 0.3 + - file-embed >= 0.0.9 + - directory >= 1.3.0.0 + - directory-tree >= 0.12.1 + - gitrev >= 1.2.0 + - parallel >= 3.2.1.0 + - blaze-html >= 0.9.0.1 + - template-haskell >= 2.11.1 + - process >= 1.4.3.0 + - markdown >= 0.1.16 + - zlib >= 0.6.1.2 + - versions >= 3.3.1 + - http-client >= 0.5.7.0 + - http-client-tls >= 0.3.5.1 + - http-conduit >= 2.2.4 + - process-extras >= 0.7.2 + - free >= 4.12.4 + - zip-archive >= 0.3.1.1 + - time >= 1.6.0.1 + - ansi-terminal >= 0.6.3.1 + - random + - temporary + - aeson >= 1.0.0.0 + - haskeline + - utf8-string >= 1 + - terminal-size >= 0.3 + + build-tools: + - alex + - happy + + source-dirs: src + + other-modules: + - Language.Futhark.Parser.Parser + - Language.Futhark.Parser.Lexer + - Paths_futhark + +executables: + futhark: &futhark + main: src/futhark.hs + dependencies: + - base + - futhark + - text + ghc-options: -threaded -rtsopts "-with-rtsopts=-N -qg" + +tests: + unit: + source-dirs: unittests + main: futhark_tests.hs + dependencies: + - base + - futhark + - containers + - mtl + - text + - QuickCheck >= 2.8 + - tasty + - tasty-quickcheck + - tasty-hunit + - megaparsec + - parser-combinators diff --git a/rts/csharp/reader.cs b/rts/csharp/reader.cs --- a/rts/csharp/reader.cs +++ b/rts/csharp/reader.cs @@ -762,7 +762,8 @@ shape[i] = (int) bin_shape; } - var elem_size = Marshal.SizeOf(typeof(T)); + // For whatever reason, Marshal.SizeOf is 4, so special-case that here. + var elem_size = typeof(T) == typeof(bool) ? 1 : Marshal.SizeOf(); var num_bytes = elem_count * elem_size; var tmp = new byte[num_bytes]; var data = new T[elem_count]; @@ -774,6 +775,10 @@ var bytes_read = b.Read(tmp, have_read, to_read); to_read -= bytes_read; have_read += bytes_read; + if (bytes_read == 0) { + Console.WriteLine("binary-input: EOF after {0} bytes (expected {1})", have_read, num_bytes); + Environment.Exit(1); + } } if (!BitConverter.IsLittleEndian && elem_size != 1) diff --git a/rts/python/__init__.py b/rts/python/__init__.py deleted file mode 100644 --- a/rts/python/__init__.py +++ /dev/null diff --git a/rts/python/values.py b/rts/python/values.py --- a/rts/python/values.py +++ b/rts/python/values.py @@ -542,7 +542,7 @@ # work on things that are insufficiently file-like, like a network # stream. bytes = f.get_chars(elem_count * FUTHARK_PRIMTYPES[expected_type]['size']) - arr = np.fromstring(bytes, dtype='<'+bin_fmt) + arr = np.fromstring(bytes, dtype=FUTHARK_PRIMTYPES[bin_type_enum]['numpy_type']) arr.shape = shape return arr diff --git a/src/Futhark/Analysis/AlgSimplify.hs b/src/Futhark/Analysis/AlgSimplify.hs --- a/src/Futhark/Analysis/AlgSimplify.hs +++ b/src/Futhark/Analysis/AlgSimplify.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE TypeSynonymInstances, FlexibleInstances, LambdaCase #-} +{-# LANGUAGE FlexibleInstances, LambdaCase #-} module Futhark.Analysis.AlgSimplify ( ScalExp , Error @@ -62,7 +62,7 @@ step :: AlgSimplifyM () step = do modify (1+) - exceeded <- pure (>) <*> get <*> asks maxSteps + exceeded <- gets (>) <*> asks maxSteps when exceeded stepsExceeded stepsExceeded :: AlgSimplifyM a @@ -345,7 +345,7 @@ case f of MaxMin ismin ts -> do let id_set = mconcat $ map freeIn ts - if S.member ii id_set + if ii `nameIn` id_set then return (Just (MaxMin ismin ts), fs) else do (mm, fs') <- findMinMaxFact ii (NProd fs tp) return (mm, f:fs') @@ -456,7 +456,7 @@ pickSymToElim rangesrep elsyms0 e_scal = -- ranges <- asks ranges -- e_scal <- fromNumSofP e0 - let ids0= S.toList $ freeIn e_scal + let ids0= namesToList $ freeIn e_scal ids1= filter (\s -> not (S.member s elsyms0)) ids0 ids2= filter (\s -> case M.lookup s rangesrep of Nothing -> False @@ -516,7 +516,7 @@ [] -> Val $ IntValue $ Int32Value 1 f:fs' -> foldl STimes f fs' let b_ids = freeIn fs_scal - return $ acc && not (idd `S.member` b_ids) + return $ acc && not (idd `nameIn` b_ids) _ -> badAlgSimplifyM "linearForm: ILLEGAL222!!!!" ) True b_terms diff --git a/src/Futhark/Analysis/CallGraph.hs b/src/Futhark/Analysis/CallGraph.hs --- a/src/Futhark/Analysis/CallGraph.hs +++ b/src/Futhark/Analysis/CallGraph.hs @@ -14,9 +14,9 @@ import Futhark.Representation.SOACS -type FunctionTable = M.Map Name FunDef +type FunctionTable = M.Map Name (FunDef SOACS) -buildFunctionTable :: Prog -> FunctionTable +buildFunctionTable :: Prog SOACS -> FunctionTable buildFunctionTable = foldl expand M.empty . progFunctions where expand ftab f = M.insert (funDefName f) f ftab @@ -27,7 +27,7 @@ -- | @buildCallGraph prog@ build the program's Call Graph. The representation -- is a hashtable that maps function names to a list of callee names. -buildCallGraph :: Prog -> CallGraph +buildCallGraph :: Prog SOACS -> CallGraph buildCallGraph prog = foldl' (buildCGfun ftable) M.empty entry_points where entry_points = map funDefName $ filter (isJust . funDefEntryPoint) $ progFunctions prog ftable = buildFunctionTable prog diff --git a/src/Futhark/Analysis/DataDependencies.hs b/src/Futhark/Analysis/DataDependencies.hs --- a/src/Futhark/Analysis/DataDependencies.hs +++ b/src/Futhark/Analysis/DataDependencies.hs @@ -8,7 +8,6 @@ where import qualified Data.Map.Strict as M -import qualified Data.Set as S import Futhark.Representation.AST @@ -31,8 +30,8 @@ cdeps = depsOf deps c comb (pe, tres, fres) = (patElemName pe, - S.unions $ [freeIn pe, cdeps, depsOf tdeps tres, depsOf fdeps fres] ++ - map (depsOfVar deps) (S.toList $ freeIn pe)) + mconcat $ [freeIn pe, cdeps, depsOf tdeps tres, depsOf fdeps fres] ++ + map (depsOfVar deps) (namesToList $ freeIn pe)) branchdeps = M.fromList $ map comb $ zip3 (patternElements pat) (bodyResult tb) @@ -41,15 +40,15 @@ grow deps (Let pat _ e) = let free = freeIn pat <> freeIn e - freeDeps = S.unions $ map (depsOfVar deps) $ S.toList free + freeDeps = mconcat $ map (depsOfVar deps) $ namesToList free in M.fromList [ (name, freeDeps) | name <- patternNames pat ] `M.union` deps depsOf :: Dependencies -> SubExp -> Names -depsOf _ (Constant _) = S.empty +depsOf _ (Constant _) = mempty depsOf deps (Var v) = depsOfVar deps v depsOfVar :: Dependencies -> VName -> Names -depsOfVar deps name = S.insert name $ M.findWithDefault S.empty name deps +depsOfVar deps name = oneName name <> M.findWithDefault mempty name deps findNecessaryForReturned :: (Param attr -> Bool) -> [(Param attr, SubExp)] -> M.Map VName Names @@ -61,10 +60,10 @@ | otherwise = iterateNecessary necessary where necessary = mconcat $ map dependencies returnedResultSubExps usedAfterLoopOrNecessary param = - usedAfterLoop param || paramName param `S.member` prev_necessary + usedAfterLoop param || paramName param `nameIn` prev_necessary returnedResultSubExps = map snd $ filter (usedAfterLoopOrNecessary . fst) merge_and_res dependencies (Constant _) = - S.empty + mempty dependencies (Var v) = - M.findWithDefault (S.singleton v) v allDependencies + M.findWithDefault (oneName v) v allDependencies diff --git a/src/Futhark/Analysis/HORepresentation/MapNest.hs b/src/Futhark/Analysis/HORepresentation/MapNest.hs --- a/src/Futhark/Analysis/HORepresentation/MapNest.hs +++ b/src/Futhark/Analysis/HORepresentation/MapNest.hs @@ -16,11 +16,9 @@ import Data.List import Data.Maybe import qualified Data.Map.Strict as M -import qualified Data.Set as S import qualified Futhark.Analysis.HORepresentation.SOAC as SOAC import Futhark.Analysis.HORepresentation.SOAC (SOAC) - import qualified Futhark.Representation.SOACS.SOAC as Futhark import Futhark.Transform.Substitute import Futhark.Representation.AST hiding (typeOf) @@ -104,7 +102,7 @@ | otherwise = Nothing boundUsedInBody = - mapMaybe isBound $ S.toList $ freeIn lam + mapMaybe isBound $ namesToList $ freeIn lam newParams <- mapM (newIdent' (++"_wasfree")) boundUsedInBody let subst = M.fromList $ zip (map identName boundUsedInBody) (map identName newParams) diff --git a/src/Futhark/Analysis/PrimExp.hs b/src/Futhark/Analysis/PrimExp.hs --- a/src/Futhark/Analysis/PrimExp.hs +++ b/src/Futhark/Analysis/PrimExp.hs @@ -80,7 +80,7 @@ FunExp h <$> traverse (traverse f) args <*> pure t instance FreeIn v => FreeIn (PrimExp v) where - freeIn = foldMap freeIn + freeIn' = foldMap freeIn' -- | Perform quick and dirty constant folding on the top level of a -- PrimExp. This is necessary because we want to consider @@ -103,6 +103,16 @@ constFoldPrimExp (BinOpExp bop (ValueExp x) (ValueExp y)) | Just z <- doBinOp bop x y = ValueExp z +constFoldPrimExp (BinOpExp LogAnd x y) + | oneIshExp x = y + | oneIshExp y = x + | zeroIshExp x = x + | zeroIshExp y = y +constFoldPrimExp (BinOpExp LogOr x y) + | oneIshExp x = x + | oneIshExp y = y + | zeroIshExp x = y + | zeroIshExp y = x constFoldPrimExp e = e -- The Num instance performs a little bit of magic: whenever an @@ -161,32 +171,38 @@ -- | Lifted logical conjunction. (.&&.) :: PrimExp v -> PrimExp v -> PrimExp v -x .&&. y = BinOpExp LogAnd x y +x .&&. y = constFoldPrimExp $ BinOpExp LogAnd x y -- | Lifted logical conjunction. (.||.) :: PrimExp v -> PrimExp v -> PrimExp v -x .||. y = BinOpExp LogOr x y +x .||. y = constFoldPrimExp $ BinOpExp LogOr x y -- | Lifted relational operators; assuming signed numbers in case of -- integers. (.<.), (.>.), (.<=.), (.>=.), (.==.) :: PrimExp v -> PrimExp v -> PrimExp v -x .<. y = CmpOpExp cmp x y where cmp = case primExpType x of +x .<. y = constFoldPrimExp $ + CmpOpExp cmp x y where cmp = case primExpType x of IntType t -> CmpSlt $ t `min` primExpIntType y FloatType t -> FCmpLt t _ -> CmpLlt -x .<=. y = CmpOpExp cmp x y where cmp = case primExpType x of +x .<=. y = constFoldPrimExp $ + CmpOpExp cmp x y where cmp = case primExpType x of IntType t -> CmpSle $ t `min` primExpIntType y FloatType t -> FCmpLe t _ -> CmpLle -x .==. y = CmpOpExp (CmpEq $ primExpType x `min` primExpType y) x y +x .==. y = constFoldPrimExp $ + CmpOpExp (CmpEq $ primExpType x `min` primExpType y) x y x .>. y = y .<. x x .>=. y = y .<=. x -- | Lifted bitwise operators. (.&.), (.|.), (.^.) :: PrimExp v -> PrimExp v -> PrimExp v -x .&. y = BinOpExp (And $ primExpIntType x `min` primExpIntType y) x y -x .|. y = BinOpExp (Or $ primExpIntType x `min` primExpIntType y) x y -x .^. y = BinOpExp (Xor $ primExpIntType x `min` primExpIntType y) x y +x .&. y = constFoldPrimExp $ + BinOpExp (And $ primExpIntType x `min` primExpIntType y) x y +x .|. y = constFoldPrimExp $ + BinOpExp (Or $ primExpIntType x `min` primExpIntType y) x y +x .^. y = constFoldPrimExp $ + BinOpExp (Xor $ primExpIntType x `min` primExpIntType y) x y infix 4 .==., .<., .>., .<=., .>=. infixr 3 .&&. diff --git a/src/Futhark/Analysis/Range.hs b/src/Futhark/Analysis/Range.hs --- a/src/Futhark/Analysis/Range.hs +++ b/src/Futhark/Analysis/Range.hs @@ -96,8 +96,7 @@ runRangeM :: RangeM a -> a runRangeM = flip runReader emptyRangeEnv -bindFunParams :: Typed attr => - [ParamT attr] -> RangeM a -> RangeM a +bindFunParams :: Typed attr => [Param attr] -> RangeM a -> RangeM a bindFunParams [] m = m bindFunParams (param:params) m = do @@ -108,8 +107,7 @@ where bindFunParam = M.insert (paramName param) unknownRange dims = arrayDims $ paramType param -bindPattern :: Typed attr => - PatternT (Range, attr) -> RangeM a -> RangeM a +bindPattern :: Typed attr => PatternT (Range, attr) -> RangeM a -> RangeM a bindPattern pat m = do ranges <- rangesRep local bindPatElems $ diff --git a/src/Futhark/Analysis/Rephrase.hs b/src/Futhark/Analysis/Rephrase.hs --- a/src/Futhark/Analysis/Rephrase.hs +++ b/src/Futhark/Analysis/Rephrase.hs @@ -61,7 +61,7 @@ rephrasePatElem rephraser (PatElem ident from) = PatElem ident <$> rephraser from -rephraseParam :: Monad m => (from -> m to) -> ParamT from -> m (ParamT to) +rephraseParam :: Monad m => (from -> m to) -> Param from -> m (Param to) rephraseParam rephraser (Param name from) = Param name <$> rephraser from diff --git a/src/Futhark/Analysis/ScalExp.hs b/src/Futhark/Analysis/ScalExp.hs --- a/src/Futhark/Analysis/ScalExp.hs +++ b/src/Futhark/Analysis/ScalExp.hs @@ -13,7 +13,6 @@ where import Data.List -import qualified Data.Set as S import Data.Maybe import Data.Monoid ((<>)) @@ -289,22 +288,22 @@ ] instance FreeIn ScalExp where - freeIn (Val _) = mempty - freeIn (Id i _) = S.singleton i - freeIn (SNeg e) = freeIn e - freeIn (SNot e) = freeIn e - freeIn (SAbs e) = freeIn e - freeIn (SSignum e) = freeIn e - freeIn (SPlus x y) = freeIn x <> freeIn y - freeIn (SMinus x y) = freeIn x <> freeIn y - freeIn (SPow x y) = freeIn x <> freeIn y - freeIn (STimes x y) = freeIn x <> freeIn y - freeIn (SDiv x y) = freeIn x <> freeIn y - freeIn (SMod x y) = freeIn x <> freeIn y - freeIn (SQuot x y) = freeIn x <> freeIn y - freeIn (SRem x y) = freeIn x <> freeIn y - freeIn (SLogOr x y) = freeIn x <> freeIn y - freeIn (SLogAnd x y) = freeIn x <> freeIn y - freeIn (RelExp LTH0 e) = freeIn e - freeIn (RelExp LEQ0 e) = freeIn e - freeIn (MaxMin _ es) = mconcat $ map freeIn es + freeIn' (Val _) = mempty + freeIn' (Id i _) = fvName i + freeIn' (SNeg e) = freeIn' e + freeIn' (SNot e) = freeIn' e + freeIn' (SAbs e) = freeIn' e + freeIn' (SSignum e) = freeIn' e + freeIn' (SPlus x y) = freeIn' x <> freeIn' y + freeIn' (SMinus x y) = freeIn' x <> freeIn' y + freeIn' (SPow x y) = freeIn' x <> freeIn' y + freeIn' (STimes x y) = freeIn' x <> freeIn' y + freeIn' (SDiv x y) = freeIn' x <> freeIn' y + freeIn' (SMod x y) = freeIn' x <> freeIn' y + freeIn' (SQuot x y) = freeIn' x <> freeIn' y + freeIn' (SRem x y) = freeIn' x <> freeIn' y + freeIn' (SLogOr x y) = freeIn' x <> freeIn' y + freeIn' (SLogAnd x y) = freeIn' x <> freeIn' y + freeIn' (RelExp LTH0 e) = freeIn' e + freeIn' (RelExp LEQ0 e) = freeIn' e + freeIn' (MaxMin _ es) = freeIn' es diff --git a/src/Futhark/Analysis/SymbolTable.hs b/src/Futhark/Analysis/SymbolTable.hs --- a/src/Futhark/Analysis/SymbolTable.hs +++ b/src/Futhark/Analysis/SymbolTable.hs @@ -3,7 +3,7 @@ {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE TypeFamilies #-} module Futhark.Analysis.SymbolTable - ( SymbolTable (bindings, loopDepth, availableAtClosestLoop) + ( SymbolTable (bindings, loopDepth, availableAtClosestLoop, simplifyMemory) , empty , fromScope , toScope @@ -51,6 +51,8 @@ -- * Misc , enclosingLoopVars , rangesRep + , hideIf + , hideCertified ) where @@ -67,7 +69,7 @@ import Prelude hiding (elem, lookup) import Futhark.Analysis.PrimExp.Convert -import Futhark.Representation.AST hiding (FParam, ParamT (..), lookupType) +import Futhark.Representation.AST hiding (FParam, lookupType) import qualified Futhark.Representation.AST as AST import Futhark.Analysis.ScalExp @@ -84,6 +86,10 @@ , availableAtClosestLoop :: Names -- ^ Which names are available just before the most enclosing -- loop? + , simplifyMemory :: Bool + -- ^ We are in a situation where we should + -- simplify/hoist/un-existentialise memory as much as possible - + -- typically, inside a kernel. } instance Semigroup (SymbolTable lore) where @@ -92,13 +98,14 @@ , bindings = bindings table1 <> bindings table2 , availableAtClosestLoop = availableAtClosestLoop table1 <> availableAtClosestLoop table2 + , simplifyMemory = simplifyMemory table1 || simplifyMemory table2 } instance Monoid (SymbolTable lore) where mempty = empty empty :: SymbolTable lore -empty = SymbolTable 0 M.empty mempty +empty = SymbolTable 0 M.empty mempty False fromScope :: Attributes lore => Scope lore -> SymbolTable lore fromScope = M.foldlWithKey' insertFreeVar' empty @@ -144,8 +151,8 @@ -> (FreeVarEntry fromlore -> Entry tolore) -> SymbolTable fromlore -> SymbolTable tolore -genCastSymbolTable loopVar letBound fParam lParam freeVar (SymbolTable depth entries loopfree) = - SymbolTable depth (M.map onEntry entries) loopfree +genCastSymbolTable loopVar letBound fParam lParam freeVar vtable = + vtable { bindings = M.map onEntry $ bindings vtable } where onEntry (LoopVar entry) = loopVar entry onEntry (LetBound entry) = letBound entry onEntry (FParam entry) = fParam entry @@ -154,7 +161,7 @@ deepen :: SymbolTable lore -> SymbolTable lore deepen vtable = vtable { loopDepth = loopDepth vtable + 1, - availableAtClosestLoop = S.fromList $ M.keys $ bindings vtable + availableAtClosestLoop = namesFromList $ M.keys $ bindings vtable } -- | Indexing a delayed array if possible. @@ -540,7 +547,7 @@ -> SymbolTable lore -> SymbolTable lore insertStm stm vtable = - flip (foldl' $ flip consume) stm_consumed $ + flip (foldl' $ flip consume) (namesToList stm_consumed) $ flip (foldl' addRevAliases) (patternElements $ stmPattern stm) $ insertEntries (zip names $ map LetBound $ bindingEntries stm vtable) vtable where names = patternNames $ stmPattern stm @@ -548,19 +555,19 @@ stm_consumed = expandAliases (Aliases.consumedInStm stm) vtable addRevAliases vtable' pe = vtable' { bindings = adjustSeveral update inedges $ bindings vtable' } - where inedges = expandAliases (Aliases.aliasesOf pe) vtable' + where inedges = namesToList $ expandAliases (Aliases.aliasesOf pe) vtable' update (LetBound entry) = LetBound entry - { letBoundAliases = patElemName pe `S.insert` letBoundAliases entry } + { letBoundAliases = oneName (patElemName pe) <> letBoundAliases entry } update (FParam entry) = FParam entry - { fparamAliases = patElemName pe `S.insert` fparamAliases entry } + { fparamAliases = oneName (patElemName pe) <> fparamAliases entry } update e = e expandAliases :: Names -> SymbolTable lore -> Names -expandAliases names vtable = names `S.union` aliasesOfAliases +expandAliases names vtable = names <> aliasesOfAliases where aliasesOfAliases = - mconcat . map (`lookupAliases` vtable) . S.toList $ names + mconcat . map (`lookupAliases` vtable) . namesToList $ names insertFParam :: Attributes lore => AST.FParam lore @@ -731,26 +738,28 @@ let candidates = freeIn e_scal sym0 = AS.pickSymToElim ranges elsyms0 e_scal case sym0 of - Just sy -> let trclsyms = foldl trClSymsInRange S.empty $ S.toList $ - candidates `S.difference` S.singleton sy - in if S.member sy trclsyms + Just sy -> let trclsyms = foldl trClSymsInRange mempty $ namesToList $ + candidates `namesSubtract` oneName sy + in if sy `nameIn` trclsyms then pickRefinedSym (S.insert sy elsyms0) e_scal else sym0 Nothing -> sym0 + -- computes the transitive closure of the symbols appearing -- in the ranges of a symbol - trClSymsInRange :: S.Set VName -> VName -> S.Set VName + trClSymsInRange :: Names -> VName -> Names trClSymsInRange cur_syms sym = - if S.member sym cur_syms then cur_syms + if sym `nameIn` cur_syms then cur_syms else case M.lookup sym ranges of - Just (_,lb,ub) -> let sym_bds = concatMap (S.toList . freeIn) (catMaybes [lb, ub]) + Just (_,lb,ub) -> let sym_bds = concatMap (namesToList . freeIn) (catMaybes [lb, ub]) in foldl trClSymsInRange - (S.insert sym cur_syms) - (S.toList $ S.fromList sym_bds) - Nothing -> S.insert sym cur_syms + (oneName sym <> cur_syms) + sym_bds + Nothing -> oneName sym <> cur_syms consume :: Attributes lore => VName -> SymbolTable lore -> SymbolTable lore -consume consumee vtable = foldl' consume' vtable $ expandAliases (S.singleton consumee) vtable +consume consumee vtable = foldl' consume' vtable $ namesToList $ + expandAliases (oneName consumee) vtable where consume' vtable' v | Just e <- lookup v vtable = insertEntry v (consume'' e) vtable' | otherwise = vtable' consume'' (FreeVar e) = FreeVar e { freeVarConsumed = True } @@ -794,3 +803,21 @@ isAtLeast :: VName -> Int -> SymbolTable lore -> SymbolTable lore isAtLeast name x = setLowerBound name $ fromIntegral x + +-- | Hide definitions of those entries that satisfy some predicate. +hideIf :: (Entry lore -> Bool) -> SymbolTable lore -> SymbolTable lore +hideIf hide vtable = vtable { bindings = M.map maybeHide $ bindings vtable } + where maybeHide entry + | hide entry = FreeVar FreeVarEntry { freeVarAttr = entryInfo entry + , freeVarStmDepth = bindingDepth entry + , freeVarRange = valueRange entry + , freeVarIndex = \_ _ -> Nothing + , freeVarConsumed = consumed entry + } + | otherwise = entry + +-- | Hide these definitions, if they are protected by certificates in +-- the set of names. +hideCertified :: Names -> SymbolTable lore -> SymbolTable lore +hideCertified to_hide = hideIf $ maybe False hide . entryStm + where hide = any (`nameIn` to_hide) . unCertificates . stmCerts diff --git a/src/Futhark/Analysis/Usage.hs b/src/Futhark/Analysis/Usage.hs --- a/src/Futhark/Analysis/Usage.hs +++ b/src/Futhark/Analysis/Usage.hs @@ -2,7 +2,6 @@ module Futhark.Analysis.Usage ( usageInStm ) where import Data.Foldable -import qualified Data.Set as S import Futhark.Representation.AST import Futhark.Representation.AST.Attributes.Aliases @@ -16,25 +15,25 @@ UT.usages (freeIn e)] where usageInPat = UT.usages (mconcat (map freeIn $ patternElements pat) - `S.difference` - S.fromList (patternNames pat)) + `namesSubtract` + namesFromList (patternNames pat)) usageInExpLore = UT.usages $ freeIn lore usageInExp :: Aliased lore => Exp lore -> UT.UsageTable usageInExp (Apply _ args _ _) = mconcat [ mconcat $ map UT.consumedUsage $ - S.toList $ subExpAliases arg + namesToList $ subExpAliases arg | (arg,d) <- args, d == Consume ] usageInExp (DoLoop _ merge _ _) = mconcat [ mconcat $ map UT.consumedUsage $ - S.toList $ subExpAliases se + namesToList $ subExpAliases se | (v,se) <- merge, unique $ paramDeclType v ] usageInExp (If _ tbranch fbranch _) = - fold $ map UT.consumedUsage $ S.toList $ + fold $ map UT.consumedUsage $ namesToList $ consumedInBody tbranch <> consumedInBody fbranch usageInExp (BasicOp (Update src _ _)) = UT.consumedUsage src usageInExp (Op op) = - mconcat $ map UT.consumedUsage (S.toList $ consumedInOp op) + mconcat $ map UT.consumedUsage (namesToList $ consumedInOp op) usageInExp _ = UT.empty diff --git a/src/Futhark/Analysis/UsageTable.hs b/src/Futhark/Analysis/UsageTable.hs --- a/src/Futhark/Analysis/UsageTable.hs +++ b/src/Futhark/Analysis/UsageTable.hs @@ -27,7 +27,6 @@ import qualified Data.Foldable as Foldable import Data.List (foldl') import qualified Data.Map.Strict as M -import qualified Data.Set as S import Prelude hiding (lookup) @@ -72,7 +71,8 @@ -- | Expand the usage table based on aliasing information. expand :: (VName -> Names) -> UsageTable -> UsageTable expand look (UsageTable m) = UsageTable $ foldl' grow m $ M.toList m - where grow m' (k, v) = foldl' (grow'' $ v `withoutU` presentU) m' $ look k + where grow m' (k, v) = foldl' (grow'' $ v `withoutU` presentU) m' $ + namesToList $ look k grow'' v m'' k = M.insertWith (<>) k v m'' keys :: UsageTable -> [VName] @@ -94,10 +94,10 @@ allConsumed :: UsageTable -> Names allConsumed (UsageTable m) = - S.fromList . map fst . filter (matches consumedU . snd) $ M.toList m + namesFromList . map fst . filter (matches consumedU . snd) $ M.toList m usages :: Names -> UsageTable -usages names = UsageTable $ M.fromList [ (name, presentU) | name <- S.toList names ] +usages names = UsageTable $ M.fromList [ (name, presentU) | name <- namesToList names ] usage :: VName -> Usages -> UsageTable usage name uses = UsageTable $ M.singleton name uses diff --git a/src/Futhark/Binder.hs b/src/Futhark/Binder.hs --- a/src/Futhark/Binder.hs +++ b/src/Futhark/Binder.hs @@ -5,7 +5,8 @@ module Futhark.Binder ( -- * A concrete @MonadBinder@ monad. BinderT - , runBinderT + , runBinderT, runBinderT_ + , runBinderT', runBinderT'_ , BinderOps (..) , bindableMkExpAttrB , bindableMkBodyB @@ -106,6 +107,21 @@ runBinderT (BinderT m) scope = do (x, (stms, _)) <- runStateT m (mempty, scope) return (x, stms) + +runBinderT_ :: MonadFreshNames m => + BinderT lore m a -> Scope lore -> m (Stms lore) +runBinderT_ m = fmap snd . runBinderT m + +runBinderT' :: (MonadFreshNames m, HasScope somelore m, SameScope somelore lore) => + BinderT lore m a + -> m (a, Stms lore) +runBinderT' m = do + scope <- askScope + runBinderT m $ castScope scope + +runBinderT'_ :: (MonadFreshNames m, HasScope somelore m, SameScope somelore lore) => + BinderT lore m a -> m (Stms lore) +runBinderT'_ = fmap snd . runBinderT' runBinder :: (MonadFreshNames m, HasScope somelore m, SameScope somelore lore) => diff --git a/src/Futhark/CLI/Bench.hs b/src/Futhark/CLI/Bench.hs --- a/src/Futhark/CLI/Bench.hs +++ b/src/Futhark/CLI/Bench.hs @@ -19,6 +19,7 @@ import System.Console.GetOpt import System.FilePath import System.Directory +import System.Environment import System.IO import System.IO.Temp import System.Timeout @@ -35,7 +36,7 @@ data BenchOptions = BenchOptions { optBackend :: String - , optFuthark :: String + , optFuthark :: Maybe String , optRunner :: String , optRuns :: Int , optExtraOptions :: [String] @@ -49,7 +50,7 @@ } initialBenchOptions :: BenchOptions -initialBenchOptions = BenchOptions "c" "futhark" "" 10 [] Nothing (-1) False +initialBenchOptions = BenchOptions "c" Nothing "" 10 [] Nothing (-1) False ["nobench", "disable"] [] Nothing (Just "tuning") -- | The name we use for compiled programs. @@ -133,6 +134,8 @@ else do putStr $ "Compiling " ++ program ++ "...\n" + futhark <- maybe getExecutablePath return $ optFuthark opts + ref_res <- runExceptT $ ensureReferenceOutput futhark "c" program cases case ref_res of Left err -> do @@ -154,7 +157,6 @@ _ -> return $ Left Skipped where hasRuns (InputOutputs _ runs) = not $ null runs - futhark = optFuthark opts runBenchmark :: BenchOptions -> (FilePath, [InputOutputs]) -> IO [BenchResult] runBenchmark opts (program, cases) = mapM forInputOutputs $ filter relevant cases @@ -327,9 +329,9 @@ "PROGRAM") "The compiler used (defaults to 'futhark-c')." , Option [] ["futhark"] - (ReqArg (\prog -> Right $ \config -> config { optFuthark = prog }) + (ReqArg (\prog -> Right $ \config -> config { optFuthark = Just prog }) "PROGRAM") - "The binary used for operations (defaults to 'futhark')." + "The binary used for operations (defaults to same binary as 'futhark bench')." , Option [] ["runner"] (ReqArg (\prog -> Right $ \config -> config { optRunner = prog }) "PROGRAM") "The program used to run the Futhark-generated programs (defaults to nothing)." diff --git a/src/Futhark/CLI/Dataset.hs b/src/Futhark/CLI/Dataset.hs --- a/src/Futhark/CLI/Dataset.hs +++ b/src/Futhark/CLI/Dataset.hs @@ -140,7 +140,7 @@ toValueType TERecord{} = Left "Cannot handle records yet." toValueType TEApply{} = Left "Cannot handle type applications yet." toValueType TEArrow{} = Left "Cannot generate functions." -toValueType TEEnum{} = Left "Cannot handle enums yet." +toValueType TESum{} = Left "Cannot handle sumtypes yet." toValueType (TEUnique t _) = toValueType t toValueType (TEArray t d _) = do d' <- constantDim d diff --git a/src/Futhark/CLI/Dev.hs b/src/Futhark/CLI/Dev.hs --- a/src/Futhark/CLI/Dev.hs +++ b/src/Futhark/CLI/Dev.hs @@ -376,7 +376,7 @@ prog <- runPipelineOnProgram (futharkConfig config) id file runPolyPasses config prog -runPolyPasses :: Config -> SOACS.Prog -> FutharkM () +runPolyPasses :: Config -> Prog SOACS -> FutharkM () runPolyPasses config prog = do prog' <- foldM (runPolyPass pipeline_config) (SOACS prog) (getFutharkPipeline config) case (prog', futharkAction config) of diff --git a/src/Futhark/CLI/Test.hs b/src/Futhark/CLI/Test.hs --- a/src/Futhark/CLI/Test.hs +++ b/src/Futhark/CLI/Test.hs @@ -19,6 +19,7 @@ import qualified Data.Text.IO as T import System.Console.ANSI import System.Process.ByteString (readProcessWithExitCode) +import System.Environment import System.Exit import System.FilePath import System.Console.GetOpt @@ -98,14 +99,15 @@ GpuPipeline -> check "--gpu" where check opt = do + futhark <- io $ maybe getExecutablePath return $ configFuthark programs (code, output, err) <- - io $ readProcessWithExitCode (configFuthark programs) ["dev", opt, "--metrics", program] "" + io $ readProcessWithExitCode futhark ["dev", opt, "--metrics", program] "" let output' = T.decodeUtf8 output case code of ExitSuccess | [(m, [])] <- reads $ T.unpack output' -> return m | otherwise -> throwError $ "Could not read metrics output:\n" <> output' - ExitFailure 127 -> throwError $ progNotFound $ T.pack $ configFuthark programs + ExitFailure 127 -> throwError $ progNotFound $ T.pack futhark ExitFailure _ -> throwError $ T.decodeUtf8 err testMetrics :: ProgConfig -> FilePath -> StructureTest -> TestM () @@ -122,7 +124,7 @@ Just actual_occurences | expected_occurences /= actual_occurences -> throwError $ name <> " should have occurred " <> T.pack (show expected_occurences) <> - " times, but occured " <> T.pack (show actual_occurences) <> " times." + " times, but occurred " <> T.pack (show actual_occurences) <> " times." _ -> return () testWarnings :: [WarningTest] -> SBS.ByteString -> TestM () @@ -134,7 +136,8 @@ | otherwise = return () runTestCase :: TestCase -> TestM () -runTestCase (TestCase mode program testcase progs) = +runTestCase (TestCase mode program testcase progs) = do + futhark <- io $ maybe getExecutablePath return $ configFuthark progs case testAction testcase of CompileTimeFailure expected_error -> @@ -180,7 +183,6 @@ unless (mode == Compile || mode == Compiled) $ context "Interpreting" $ accErrors_ $ map (runInterpretedEntry futhark program) ios - where futhark = configFuthark progs runInterpretedEntry :: String -> FilePath -> InputOutputs -> TestM() runInterpretedEntry futhark program (InputOutputs entry run_cases) = @@ -317,6 +319,19 @@ excludedTest config = any (`elem` configExclude config) . testTags . testCaseTest +-- | Exclude those test cases that have tags we do not wish to run. +excludeCases :: TestConfig -> TestCase -> TestCase +excludeCases config tcase = + tcase { testCaseTest = onTest $ testCaseTest tcase } + where onTest (ProgramTest desc tags action) = + ProgramTest desc tags $ onAction action + onAction (RunCases ios stest wtest) = + RunCases (map onIOs ios) stest wtest + onAction action = action + onIOs (InputOutputs entry runs) = + InputOutputs entry $ filter (not . any excluded . runTags) runs + excluded = (`elem` configExclude config) . T.pack + statusTable :: TestStatus -> String statusTable ts = buildTable rows 1 where rows = @@ -383,7 +398,7 @@ replicateM_ concurrency $ forkIO $ runTest testmvar reportmvar let (excluded, included) = partition (excludedTest config) all_tests - _ <- forkIO $ mapM_ (putMVar testmvar) included + _ <- forkIO $ mapM_ (putMVar testmvar . excludeCases config) included isTTY <- (&& not (configLineOutput config)) <$> hIsTerminalDevice stdout let report | isTTY = reportTable @@ -476,7 +491,7 @@ , configPrograms = ProgConfig { configBackend = "c" - , configFuthark = "futhark" + , configFuthark = Nothing , configRunner = "" , configExtraOptions = [] , configExtraCompilerOptions = [] @@ -487,7 +502,7 @@ data ProgConfig = ProgConfig { configBackend :: String - , configFuthark :: FilePath + , configFuthark :: Maybe FilePath , configRunner :: FilePath , configExtraCompilerOptions :: [String] , configTuning :: Maybe String @@ -505,7 +520,7 @@ setFuthark :: FilePath -> ProgConfig -> ProgConfig setFuthark futhark config = - config { configFuthark = futhark } + config { configFuthark = Just futhark } setRunner :: FilePath -> ProgConfig -> ProgConfig setRunner runner config = @@ -548,7 +563,7 @@ "Backend used for compilation (defaults to 'c')." , Option [] ["futhark"] (ReqArg (Right . changeProgConfig . setFuthark) "PROGRAM") - "Program to run for subcommands (defaults to 'futhark')." + "Program to run for subcommands (defaults to same binary as 'futhark test')." , Option [] ["runner"] (ReqArg (Right . changeProgConfig . setRunner) "PROGRAM") "The program used to run the Futhark-generated programs (defaults to nothing)." @@ -564,6 +579,9 @@ , Option [] ["pass-compiler-option"] (ReqArg (Right . changeProgConfig . addCompilerOption) "OPT") "Pass this option to the compiler (or typechecker if in -t mode)." + , Option [] ["no-tuning"] + (NoArg $ Right $ changeProgConfig $ \config -> config { configTuning = Nothing }) + "Do not load tuning files." ] main :: String -> [String] -> IO () diff --git a/src/Futhark/CodeGen/Backends/COpenCL.hs b/src/Futhark/CodeGen/Backends/COpenCL.hs --- a/src/Futhark/CodeGen/Backends/COpenCL.hs +++ b/src/Futhark/CodeGen/Backends/COpenCL.hs @@ -320,7 +320,7 @@ |] setKernelArg i (SharedMemoryKArg num_bytes) = do - num_bytes' <- GC.compileExp $ innerExp num_bytes + num_bytes' <- GC.compileExp $ unCount num_bytes GC.stm [C.cstm| OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->$id:name, $int:i, $exp:num_bytes', NULL)); |] diff --git a/src/Futhark/CodeGen/Backends/GenericC.hs b/src/Futhark/CodeGen/Backends/GenericC.hs --- a/src/Futhark/CodeGen/Backends/GenericC.hs +++ b/src/Futhark/CodeGen/Backends/GenericC.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE QuasiQuotes, GeneralizedNewtypeDeriving, TypeSynonymInstances, FlexibleInstances #-} +{-# LANGUAGE QuasiQuotes, GeneralizedNewtypeDeriving, FlexibleInstances #-} {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE TupleSections #-} {-# OPTIONS_GHC -fno-warn-orphans #-} diff --git a/src/Futhark/CodeGen/Backends/GenericCSharp.hs b/src/Futhark/CodeGen/Backends/GenericCSharp.hs --- a/src/Futhark/CodeGen/Backends/GenericCSharp.hs +++ b/src/Futhark/CodeGen/Backends/GenericCSharp.hs @@ -39,6 +39,7 @@ , CompilerEnv(..) , CompilerState(..) + , CompilerAcc , stm , stms , atInit diff --git a/src/Futhark/CodeGen/Backends/GenericCSharp/AST.hs b/src/Futhark/CodeGen/Backends/GenericCSharp/AST.hs --- a/src/Futhark/CodeGen/Backends/GenericCSharp/AST.hs +++ b/src/Futhark/CodeGen/Backends/GenericCSharp/AST.hs @@ -10,6 +10,7 @@ , CSUInt(..) , CSFloat(..) , CSIdx (..) + , ArgMemType(..) , CSArg (..) , CSStmt(..) , module Language.Futhark.Core diff --git a/src/Futhark/CodeGen/ImpCode.hs b/src/Futhark/CodeGen/ImpCode.hs --- a/src/Futhark/CodeGen/ImpCode.hs +++ b/src/Futhark/CodeGen/ImpCode.hs @@ -1,7 +1,5 @@ {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TypeSynonymInstances #-} {-# LANGUAGE TupleSections #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -- | Imperative intermediate language used as a stepping stone in code generation. -- -- This is a generic representation parametrised on an extensible @@ -38,7 +36,6 @@ , ArrayContents(..) -- * Typed enumerations - , Count (..) , Bytes , Elements , elements @@ -56,6 +53,7 @@ , module Language.Futhark.Core , module Futhark.Representation.Primitive , module Futhark.Analysis.PrimExp + , module Futhark.Representation.Kernels.Sizes ) where @@ -63,7 +61,6 @@ import Data.List import Data.Loc import Data.Traversable -import qualified Data.Set as S import Language.Futhark.Core import Futhark.Representation.Primitive @@ -71,9 +68,9 @@ (Space(..), SpaceId, ErrorMsg(..), ErrorMsgPart(..)) import Futhark.Representation.AST.Attributes.Names import Futhark.Representation.AST.Pretty () -import Futhark.Util.IntegralExp import Futhark.Analysis.PrimExp import Futhark.Util.Pretty hiding (space) +import Futhark.Representation.Kernels.Sizes (Count(..)) data Size = ConstSize Int64 | VarSize VName @@ -160,7 +157,7 @@ -- This is mostly used for constant arrays, but also for -- some bookkeeping data, like the synchronisation -- counts used to implement reduction. - | Allocate VName (Count Bytes) Space + | Allocate VName (Count Bytes Exp) Space -- ^ Memory space must match the corresponding -- 'DeclareMem'. | Free VName Space @@ -171,11 +168,11 @@ -- is the last reference. There is no guarantee that -- all memory blocks will be freed with this statement. -- Backends are free to ignore it entirely. - | Copy VName (Count Bytes) Space VName (Count Bytes) Space (Count Bytes) + | Copy VName (Count Bytes Exp) Space VName (Count Bytes Exp) Space (Count Bytes Exp) -- ^ Destination, offset in destination, destination -- space, source, offset in source, offset space, number -- of bytes. - | Write VName (Count Elements) PrimType Space Volatility Exp + | Write VName (Count Elements Exp) PrimType Space Volatility Exp | SetScalar VName Exp | SetMem VName VName Space -- ^ Must be in same space. @@ -210,7 +207,7 @@ data ExpLeaf = ScalarVar VName | SizeOf PrimType - | Index VName (Count Elements) PrimType Space Volatility + | Index VName (Count Elements Exp) PrimType Space Volatility deriving (Eq, Show) type Exp = PrimExp ExpLeaf @@ -220,32 +217,27 @@ | MemArg VName deriving (Show) --- | A wrapper around 'Imp.Exp' that maintains a unit as a phantom --- type. -newtype Count u = Count { innerExp :: Exp } - deriving (Eq, Show, Num, IntegralExp, FreeIn, Pretty) - -- | Phantom type for a count of elements. data Elements -- | Phantom type for a count of bytes. data Bytes -elements :: Exp -> Count Elements +elements :: Exp -> Count Elements Exp elements = Count -bytes :: Exp -> Count Bytes +bytes :: Exp -> Count Bytes Exp bytes = Count -- | Convert a count of elements into a count of bytes, given the -- per-element size. -withElemType :: Count Elements -> PrimType -> Count Bytes +withElemType :: Count Elements Exp -> PrimType -> Count Bytes Exp withElemType (Count e) t = bytes $ e * LeafExp (SizeOf t) (IntType Int32) -dimSizeToExp :: DimSize -> Count Elements +dimSizeToExp :: DimSize -> Count Elements Exp dimSizeToExp = elements . sizeToExp -memSizeToExp :: MemSize -> Count Bytes +memSizeToExp :: MemSize -> Count Bytes Exp memSizeToExp = bytes . sizeToExp sizeToExp :: Size -> Exp @@ -259,7 +251,7 @@ vi32 :: VName -> Exp vi32 = flip var $ IntType Int32 -index :: VName -> Count Elements -> PrimType -> Space -> Volatility -> Exp +index :: VName -> Count Elements Exp -> PrimType -> Space -> Volatility -> Exp index arr i t s vol = LeafExp (Index arr i t s vol) t -- Prettyprinting definitions. @@ -457,65 +449,65 @@ pure $ DebugPrint s v declaredIn :: Code a -> Names -declaredIn (DeclareMem name _) = S.singleton name -declaredIn (DeclareScalar name _) = S.singleton name -declaredIn (DeclareArray name _ _ _) = S.singleton name +declaredIn (DeclareMem name _) = oneName name +declaredIn (DeclareScalar name _) = oneName name +declaredIn (DeclareArray name _ _ _) = oneName name declaredIn (If _ t f) = declaredIn t <> declaredIn f declaredIn (x :>>: y) = declaredIn x <> declaredIn y -declaredIn (For i _ _ body) = S.singleton i <> declaredIn body +declaredIn (For i _ _ body) = oneName i <> declaredIn body declaredIn (While _ body) = declaredIn body declaredIn (Comment _ body) = declaredIn body declaredIn _ = mempty instance FreeIn a => FreeIn (Code a) where - freeIn (x :>>: y) = - freeIn x <> freeIn y `S.difference` declaredIn x - freeIn Skip = + freeIn' (x :>>: y) = + fvBind (declaredIn x) $ freeIn' x <> freeIn' y + freeIn' Skip = mempty - freeIn (For i _ bound body) = - i `S.delete` (freeIn bound <> freeIn body) - freeIn (While cond body) = - freeIn cond <> freeIn body - freeIn DeclareMem{} = + freeIn' (For i _ bound body) = + fvBind (oneName i) $ freeIn' bound <> freeIn' body + freeIn' (While cond body) = + freeIn' cond <> freeIn' body + freeIn' DeclareMem{} = mempty - freeIn DeclareScalar{} = + freeIn' DeclareScalar{} = mempty - freeIn DeclareArray{} = + freeIn' DeclareArray{} = mempty - freeIn (Allocate name size _) = - freeIn name <> freeIn size - freeIn (Free name _) = - freeIn name - freeIn (Copy dest x _ src y _ n) = - freeIn dest <> freeIn x <> freeIn src <> freeIn y <> freeIn n - freeIn (SetMem x y _) = - freeIn x <> freeIn y - freeIn (Write v i _ _ _ e) = - freeIn v <> freeIn i <> freeIn e - freeIn (SetScalar x y) = - freeIn x <> freeIn y - freeIn (Call dests _ args) = - freeIn dests <> freeIn args - freeIn (If cond t f) = - freeIn cond <> freeIn t <> freeIn f - freeIn (Assert e _ _) = - freeIn e - freeIn (Op op) = - freeIn op - freeIn (Comment _ code) = - freeIn code - freeIn (DebugPrint _ v) = - maybe mempty (freeIn . snd) v + freeIn' (Allocate name size _) = + freeIn' name <> freeIn' size + freeIn' (Free name _) = + freeIn' name + freeIn' (Copy dest x _ src y _ n) = + freeIn' dest <> freeIn' x <> freeIn' src <> freeIn' y <> freeIn' n + freeIn' (SetMem x y _) = + freeIn' x <> freeIn' y + freeIn' (Write v i _ _ _ e) = + freeIn' v <> freeIn' i <> freeIn' e + freeIn' (SetScalar x y) = + freeIn' x <> freeIn' y + freeIn' (Call dests _ args) = + freeIn' dests <> freeIn' args + freeIn' (If cond t f) = + freeIn' cond <> freeIn' t <> freeIn' f + freeIn' (Assert e _ _) = + freeIn' e + freeIn' (Op op) = + freeIn' op + freeIn' (Comment _ code) = + freeIn' code + freeIn' (DebugPrint _ v) = + maybe mempty (freeIn' . snd) v instance FreeIn ExpLeaf where - freeIn (Index v e _ _ _) = freeIn v <> freeIn e - freeIn (ScalarVar v) = freeIn v - freeIn (SizeOf _) = mempty + freeIn' (Index v e _ _ _) = freeIn' v <> freeIn' e + freeIn' (ScalarVar v) = freeIn' v + freeIn' (SizeOf _) = mempty instance FreeIn Arg where - freeIn (MemArg m) = freeIn m - freeIn (ExpArg e) = freeIn e + freeIn' (MemArg m) = freeIn' m + freeIn' (ExpArg e) = freeIn' e instance FreeIn Size where - freeIn (VarSize name) = S.singleton name - freeIn (ConstSize _) = mempty + freeIn' (VarSize name) = fvName name + freeIn' (ConstSize _) = mempty diff --git a/src/Futhark/CodeGen/ImpCode/Kernels.hs b/src/Futhark/CodeGen/ImpCode/Kernels.hs --- a/src/Futhark/CodeGen/ImpCode/Kernels.hs +++ b/src/Futhark/CodeGen/ImpCode/Kernels.hs @@ -69,7 +69,7 @@ deriving (Show) -- ^ In-kernel name and per-workgroup size in bytes. -type LocalMemoryUse = (VName, Either (Count Bytes) KernelConstExp) +type LocalMemoryUse = (VName, Either (Count Bytes Exp) KernelConstExp) data KernelUse = ScalarUse VName PrimType | MemoryUse VName @@ -85,7 +85,7 @@ sameKernel _ _ = False -- | Get an atomic operator corresponding to a binary operator. -atomicBinOp :: BinOp -> Maybe (VName -> VName -> Count Elements -> Exp -> AtomicOp) +atomicBinOp :: BinOp -> Maybe (VName -> VName -> Count Elements Imp.Exp -> Exp -> AtomicOp) atomicBinOp = flip lookup [ (Add Int32, AtomicAdd) , (SMax Int32, AtomicSMax) , (SMin Int32, AtomicSMin) @@ -121,17 +121,18 @@ ppr c instance FreeIn HostOp where - freeIn (CallKernel c) = freeIn c - freeIn (CmpSizeLe dest _ _ x) = - freeIn dest <> freeIn x - freeIn (GetSizeMax dest _) = - freeIn dest - freeIn (GetSize dest _ _) = - freeIn dest + freeIn' (CallKernel c) = + freeIn' c + freeIn' (CmpSizeLe dest _ _ x) = + freeIn' dest <> freeIn' x + freeIn' (GetSizeMax dest _) = + freeIn' dest + freeIn' (GetSize dest _ _) = + freeIn' dest instance FreeIn Kernel where - freeIn kernel = freeIn (kernelBody kernel) <> - freeIn [kernelNumGroups kernel, kernelGroupSize kernel] + freeIn' kernel = freeIn' (kernelBody kernel) <> + freeIn' [kernelNumGroups kernel, kernelGroupSize kernel] instance Pretty Kernel where ppr kernel = @@ -152,35 +153,35 @@ | GlobalBarrier | MemFenceLocal | MemFenceGlobal - | PrivateAlloc VName (Count Bytes) - | LocalAlloc VName (Either (Count Bytes) KernelConstExp) + | PrivateAlloc VName (Count Bytes Imp.Exp) + | LocalAlloc VName (Either (Count Bytes Imp.Exp) KernelConstExp) deriving (Show) -- Atomic operations return the value stored before the update. -- This value is stored in the first VName. -data AtomicOp = AtomicAdd VName VName (Count Elements) Exp - | AtomicSMax VName VName (Count Elements) Exp - | AtomicSMin VName VName (Count Elements) Exp - | AtomicUMax VName VName (Count Elements) Exp - | AtomicUMin VName VName (Count Elements) Exp - | AtomicAnd VName VName (Count Elements) Exp - | AtomicOr VName VName (Count Elements) Exp - | AtomicXor VName VName (Count Elements) Exp - | AtomicCmpXchg VName VName (Count Elements) Exp Exp - | AtomicXchg VName VName (Count Elements) Exp +data AtomicOp = AtomicAdd VName VName (Count Elements Imp.Exp) Exp + | AtomicSMax VName VName (Count Elements Imp.Exp) Exp + | AtomicSMin VName VName (Count Elements Imp.Exp) Exp + | AtomicUMax VName VName (Count Elements Imp.Exp) Exp + | AtomicUMin VName VName (Count Elements Imp.Exp) Exp + | AtomicAnd VName VName (Count Elements Imp.Exp) Exp + | AtomicOr VName VName (Count Elements Imp.Exp) Exp + | AtomicXor VName VName (Count Elements Imp.Exp) Exp + | AtomicCmpXchg VName VName (Count Elements Imp.Exp) Exp Exp + | AtomicXchg VName VName (Count Elements Imp.Exp) Exp deriving (Show) instance FreeIn AtomicOp where - freeIn (AtomicAdd _ arr i x) = freeIn arr <> freeIn i <> freeIn x - freeIn (AtomicSMax _ arr i x) = freeIn arr <> freeIn i <> freeIn x - freeIn (AtomicSMin _ arr i x) = freeIn arr <> freeIn i <> freeIn x - freeIn (AtomicUMax _ arr i x) = freeIn arr <> freeIn i <> freeIn x - freeIn (AtomicUMin _ arr i x) = freeIn arr <> freeIn i <> freeIn x - freeIn (AtomicAnd _ arr i x) = freeIn arr <> freeIn i <> freeIn x - freeIn (AtomicOr _ arr i x) = freeIn arr <> freeIn i <> freeIn x - freeIn (AtomicXor _ arr i x) = freeIn arr <> freeIn i <> freeIn x - freeIn (AtomicCmpXchg _ arr i x y) = freeIn arr <> freeIn i <> freeIn x <> freeIn y - freeIn (AtomicXchg _ arr i x) = freeIn arr <> freeIn i <> freeIn x + freeIn' (AtomicAdd _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicSMax _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicSMin _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicUMax _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicUMin _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicAnd _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicOr _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicXor _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicCmpXchg _ arr i x y) = freeIn' arr <> freeIn' i <> freeIn' x <> freeIn' y + freeIn' (AtomicXchg _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x instance Pretty KernelOp where ppr (GetGroupId dest i) = @@ -247,8 +248,8 @@ parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) instance FreeIn KernelOp where - freeIn (Atomic _ op) = freeIn op - freeIn _ = mempty + freeIn' (Atomic _ op) = freeIn' op + freeIn' _ = mempty brace :: Doc -> Doc brace body = text " {" indent 2 body text "}" diff --git a/src/Futhark/CodeGen/ImpCode/OpenCL.hs b/src/Futhark/CodeGen/ImpCode/OpenCL.hs --- a/src/Futhark/CodeGen/ImpCode/OpenCL.hs +++ b/src/Futhark/CodeGen/ImpCode/OpenCL.hs @@ -54,7 +54,7 @@ -- ^ Pass the value of this scalar expression as argument. | MemKArg VName -- ^ Pass this pointer as argument. - | SharedMemoryKArg (Count Bytes) + | SharedMemoryKArg (Count Bytes Exp) -- ^ Create this much local memory per workgroup. deriving (Show) diff --git a/src/Futhark/CodeGen/ImpCode/Sequential.hs b/src/Futhark/CodeGen/ImpCode/Sequential.hs --- a/src/Futhark/CodeGen/ImpCode/Sequential.hs +++ b/src/Futhark/CodeGen/ImpCode/Sequential.hs @@ -31,4 +31,4 @@ ppr _ = empty instance FreeIn Sequential where - freeIn _ = mempty + freeIn' _ = mempty diff --git a/src/Futhark/CodeGen/ImpGen.hs b/src/Futhark/CodeGen/ImpGen.hs --- a/src/Futhark/CodeGen/ImpGen.hs +++ b/src/Futhark/CodeGen/ImpGen.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts, LambdaCase, TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses #-} +{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts, LambdaCase, FlexibleInstances, MultiParamTypeClasses #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE ConstraintKinds #-} module Futhark.CodeGen.ImpGen @@ -122,11 +122,11 @@ type CopyCompiler lore op = PrimType -> MemLocation -> MemLocation - -> Count Elements -- ^ Number of row elements of the source. + -> Count Elements Imp.Exp -- ^ Number of row elements of the source. -> ImpM lore op () -- | An alternate way of compiling an allocation. -type AllocCompiler lore op = VName -> Count Bytes -> ImpM lore op () +type AllocCompiler lore op = VName -> Count Bytes Imp.Exp -> ImpM lore op () data Operations lore op = Operations { opsExpCompiler :: ExpCompiler lore op , opsOpCompiler :: OpCompiler lore op @@ -265,11 +265,11 @@ -> Either InternalError (a, State lore op, Imp.Code op) runImpM (ImpM m) ops space fname = runRWST m $ newEnv ops space fname -subImpM_ :: Operations lore' op' -> ImpM lore' op' a +subImpM_ :: Operations lore op' -> ImpM lore op' a -> ImpM lore op (Imp.Code op') subImpM_ ops m = snd <$> subImpM ops m -subImpM :: Operations lore' op' -> ImpM lore' op' a +subImpM :: Operations lore op' -> ImpM lore op' a -> ImpM lore op (a, Imp.Code op') subImpM ops (ImpM m) = do env <- ask @@ -280,15 +280,12 @@ , envOpCompiler = opsOpCompiler ops , envAllocCompilers = opsAllocCompilers ops } - s { stateVTable = M.map scrubExps $ stateVTable s + s { stateVTable = stateVTable s , stateFunctions = mempty } of Left err -> throwError err Right (x, s', code) -> do putNameSource $ stateNameSource s' return (x, code) - where scrubExps (ArrayVar _ entry) = ArrayVar Nothing entry - scrubExps (MemVar _ entry) = MemVar Nothing entry - scrubExps (ScalarVar _ entry) = ScalarVar Nothing entry -- | Execute a code generation action, returning the code that was -- emitted. @@ -536,8 +533,8 @@ e_code <- collect $ compileExp pat e (live_after, bs_code) <- collect' $ compileStms' (patternAllocs pat <> allocs) bs - let dies_here v = not (v `S.member` live_after) && - v `S.member` freeIn e_code + let dies_here v = not (v `nameIn` live_after) && + v `nameIn` freeIn e_code to_free = S.filter (dies_here . fst) allocs emit e_code @@ -701,7 +698,7 @@ let srcloc = entryArrayLocation yentry rows = case drop i $ entryArrayShape yentry of [] -> error $ "defCompileBasicOp Concat: empty array shape for " ++ pretty y - r:_ -> innerExp $ Imp.dimSizeToExp r + r:_ -> unCount $ Imp.dimSizeToExp r copy (elemType $ patElemType pe) destloc srcloc $ arrayOuterSize yentry emit $ Imp.SetScalar offs_glb $ Imp.var offs_glb int32 + rows @@ -943,13 +940,13 @@ return $ ScalarDestination name fullyIndexArray :: VName -> [Imp.Exp] - -> ImpM lore op (VName, Imp.Space, Count Elements) + -> ImpM lore op (VName, Imp.Space, Count Elements Imp.Exp) fullyIndexArray name indices = do arr <- lookupArray name fullyIndexArray' (entryArrayLocation arr) indices fullyIndexArray' :: MemLocation -> [Imp.Exp] - -> ImpM lore op (VName, Imp.Space, Count Elements) + -> ImpM lore op (VName, Imp.Space, Count Elements Imp.Exp) fullyIndexArray' (MemLocation mem _ ixfun) indices = do space <- entryMemSpace <$> lookupMemory mem return (mem, space, @@ -976,10 +973,10 @@ strideArray (MemLocation mem shape ixfun) stride = MemLocation mem shape $ IxFun.strideIndex ixfun stride -arrayOuterSize :: ArrayEntry -> Count Elements +arrayOuterSize :: ArrayEntry -> Count Elements Imp.Exp arrayOuterSize = arrayDimSize 0 -arrayDimSize :: Int -> ArrayEntry -> Count Elements +arrayDimSize :: Int -> ArrayEntry -> Count Elements Imp.Exp arrayDimSize i = product . map Imp.dimSizeToExp . take 1 . drop i . entryArrayShape @@ -1020,7 +1017,7 @@ let ivars = map Imp.vi32 is destidx = IxFun.index destIxFun ivars srcidx = IxFun.index srcIxFun ivars - bounds = map innerExp $ n : drop 1 (map Imp.dimSizeToExp srcshape) + bounds = map unCount $ n : drop 1 (map Imp.dimSizeToExp srcshape) srcspace <- entryMemSpace <$> lookupMemory srcmem destspace <- entryMemSpace <$> lookupMemory destmem vol <- asks envVolatility @@ -1108,12 +1105,11 @@ (_, ScalarVar _ (ScalarEntry _)) | not $ null src_is -> compilerBugS $ - unwords ["copyDWIMDest: prim-typed source", pretty src, "with nonzero indices."] - + unwords ["copyDWIMDest: prim-typed source", pretty src, "with nonzero indices", pretty src_is] (ScalarDestination name, _) | not $ null dest_is -> compilerBugS $ - unwords ["copyDWIMDest: prim-typed target", pretty name, "with nonzero indices."] + unwords ["copyDWIMDest: prim-typed target", pretty name, "with nonzero indices", pretty dest_is] (ScalarDestination name, ScalarVar _ (ScalarEntry pt)) -> emit $ Imp.SetScalar name $ Imp.var src pt @@ -1183,7 +1179,7 @@ -- | The number of bytes needed to represent the array in a -- straightforward contiguous format. -typeSize :: Type -> Count Bytes +typeSize :: Type -> Count Bytes Imp.Exp typeSize t = Imp.bytes $ Imp.LeafExp (Imp.SizeOf $ elemType t) int32 * product (map (toExp' int32) (arrayDims t)) @@ -1227,14 +1223,14 @@ addVar name' $ MemVar Nothing $ MemEntry space return name' -sAlloc_ :: VName -> Count Bytes -> Space -> ImpM lore op () +sAlloc_ :: VName -> Count Bytes Imp.Exp -> Space -> ImpM lore op () sAlloc_ name' size' space = do allocator <- asks $ M.lookup space . envAllocCompilers case allocator of Nothing -> emit $ Imp.Allocate name' size' space Just allocator' -> allocator' name' size' -sAlloc :: String -> Count Bytes -> Space -> ImpM lore op VName +sAlloc :: String -> Count Bytes Imp.Exp -> Space -> ImpM lore op VName sAlloc name size space = do name' <- sDeclareMem name space sAlloc_ name' size space diff --git a/src/Futhark/CodeGen/ImpGen/Kernels.hs b/src/Futhark/CodeGen/ImpGen/Kernels.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels.hs @@ -33,11 +33,11 @@ callKernelOperations :: Operations ExplicitMemory Imp.HostOp callKernelOperations = Operations { opsExpCompiler = expCompiler - , opsCopyCompiler = callKernelCopy - , opsOpCompiler = opCompiler - , opsStmsCompiler = defCompileStms - , opsAllocCompilers = mempty - } + , opsCopyCompiler = callKernelCopy + , opsOpCompiler = opCompiler + , opsStmsCompiler = defCompileStms + , opsAllocCompilers = mempty + } compileProg :: MonadFreshNames m => Prog ExplicitMemory -> m (Either InternalError Imp.Program) compileProg prog = @@ -59,8 +59,8 @@ =<< toExp x opCompiler (Pattern _ [pe]) (Inner (GetSizeMax size_class)) = sOp $ Imp.GetSizeMax (patElemName pe) size_class -opCompiler dest (Inner (HostOp kernel)) = - kernelCompiler dest kernel +opCompiler dest (Inner (SegOp op)) = + segOpCompiler dest op opCompiler pat e = compilerBugS $ "opCompiler: Invalid pattern\n " ++ pretty pat ++ "\nfor expression\n " ++ pretty e @@ -71,46 +71,17 @@ where f (name, x) = (keyWithEntryPoint fname name, x) sizeClassWithEntryPoint _ size_class = size_class -kernelCompiler :: Pattern ExplicitMemory -> Kernel InKernel - -> CallKernelGen () - -kernelCompiler pat (Kernel desc space _ kernel_body) = do - (constants, init_constants) <- kernelInitialisationSetSpace space $ return () - - forM_ (kernelHints desc) $ \(s,v) -> do - ty <- case v of - Constant pv -> return $ Prim $ primValueType pv - Var vn -> lookupType vn - unless (primType ty) $ fail $ concat [ "debugKernelHint '", s, "'" - , " in kernel '", kernelName desc, "'" - , " did not have primType value." ] - - emit $ Imp.DebugPrint s $ Just (elemType ty, toExp' (elemType ty) v) - - let virt_groups = toExp' int32 (spaceNumVirtGroups space) - sKernel constants (kernelName desc) $ do - init_constants - virtualiseGroups constants virt_groups $ \group_id -> do - let flat_id = - if kernelGroupIdVar constants /= group_id - then Imp.vi32 group_id * kernelGroupSize constants + kernelLocalThreadId constants - else kernelGlobalThreadId constants - setSpaceIndices flat_id space - compileKernelStms constants (kernelBodyStms kernel_body) $ - zipWithM_ (compileKernelResult constants) (patternElements pat) $ - kernelBodyResult kernel_body - -kernelCompiler pat (SegMap space _ body) = - compileSegMap pat space body - -kernelCompiler pat (SegRed space reds _ body) = - compileSegRed pat space reds body - -kernelCompiler pat (SegScan space red_op nes _ kbody) = - compileSegScan pat space red_op nes kbody - -kernelCompiler pat (SegGenRed space ops _ body) = - compileSegGenRed pat space ops body +segOpCompiler :: Pattern ExplicitMemory -> SegOp ExplicitMemory -> CallKernelGen () +segOpCompiler pat (SegMap lvl space _ kbody) = + compileSegMap pat lvl space kbody +segOpCompiler pat (SegRed lvl@SegThread{} space reds _ kbody) = + compileSegRed pat lvl space reds kbody +segOpCompiler pat (SegScan lvl@SegThread{} space scan_op nes _ kbody) = + compileSegScan pat lvl space scan_op nes kbody +segOpCompiler pat (SegGenRed (SegThread num_groups group_size _) space ops _ kbody) = + compileSegGenRed pat num_groups group_size space ops kbody +segOpCompiler pat segop = + compilerBugS $ "segOpCompiler: unexpected " ++ pretty (segLevel segop) ++ " for rhs of pattern " ++ pretty pat expCompiler :: ExpCompiler ExplicitMemory Imp.HostOp diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs b/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs @@ -3,29 +3,25 @@ {-# LANGUAGE TypeFamilies #-} module Futhark.CodeGen.ImpGen.Kernels.Base ( KernelConstants (..) - , inKernelOperations , keyWithEntryPoint , CallKernelGen , InKernelGen , computeThreadChunkSize - , simpleKernelConstants - , kernelInitialisation - , kernelInitialisationSimple - , kernelInitialisationSetSpace - , setSpaceIndices - , makeAllMemoryGlobal - , allThreads - , compileKernelStms , groupReduce , groupScan , isActive - , sKernel + , sKernelThread + , sKernelGroup + , sKernelSimple , sReplicate , sIota , sCopy - , compileKernelResult + , compileThreadResult + , compileGroupResult , virtualiseGroups + , getSize + , atomicUpdate , atomicUpdateLocking , Locking(..) @@ -33,13 +29,12 @@ ) where -import Control.Arrow ((&&&)) import Control.Monad.Except import Control.Monad.Reader import Data.Maybe import qualified Data.Map.Strict as M -import qualified Data.Set as S import Data.List +import Data.Loc import Prelude hiding (quot, rem) @@ -50,268 +45,257 @@ import qualified Futhark.CodeGen.ImpCode.Kernels as Imp import Futhark.CodeGen.ImpCode.Kernels (elements) import Futhark.CodeGen.ImpGen -import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem, IntegralExp) -import Futhark.Util (splitAt3, maybeNth, takeLast) +import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem) +import Futhark.Util (chunks, maybeNth) type CallKernelGen = ImpM ExplicitMemory Imp.HostOp -type InKernelGen = ImpM InKernel Imp.KernelOp +type InKernelGen = ImpM ExplicitMemory Imp.KernelOp data KernelConstants = KernelConstants - { kernelOuterVTable :: VTable ExplicitMemory -- XXX - , kernelGlobalThreadId :: Imp.Exp + { kernelGlobalThreadId :: Imp.Exp , kernelLocalThreadId :: Imp.Exp , kernelGroupId :: Imp.Exp , kernelGlobalThreadIdVar :: VName , kernelLocalThreadIdVar :: VName , kernelGroupIdVar :: VName - , kernelGroupSize :: Imp.Exp , kernelNumGroups :: Imp.Exp + , kernelGroupSize :: Imp.Exp , kernelNumThreads :: Imp.Exp , kernelWaveSize :: Imp.Exp - , kernelDimensions :: [(VName, Imp.Exp)] , kernelThreadActive :: Imp.Exp - , kernelStreamed :: [(VName, Imp.DimSize)] - -- ^ Chunk sizes and their maximum size. Hint - -- for unrolling. } -inKernelOperations :: KernelConstants -> Operations InKernel Imp.KernelOp -inKernelOperations constants = - (defaultOperations $ compileInKernelOp constants) - { opsCopyCompiler = inKernelCopy - , opsExpCompiler = inKernelExpCompiler - , opsStmsCompiler = \_ -> compileKernelStms constants - , opsAllocCompilers = - M.fromList [ (Space "local", allocLocal) - , (Space "private", allocPrivate) ] - } - where allocLocal :: AllocCompiler InKernel Imp.KernelOp - allocLocal mem size = do - size' <- localMemSize (kernelOuterVTable constants) size - sOp $ Imp.LocalAlloc mem size' - allocPrivate mem size = - sOp $ Imp.PrivateAlloc mem size - keyWithEntryPoint :: Name -> Name -> Name keyWithEntryPoint fname key = nameFromString $ nameToString fname ++ "." ++ nameToString key --- | We have no bulk copy operation (e.g. memmove) inside kernels, so --- turn any copy into a loop. -inKernelCopy :: CopyCompiler InKernel Imp.KernelOp -inKernelCopy = copyElementWise +noAssert :: MonadError InternalError m => [SrcLoc] -> m a +noAssert locs = + compilerLimitationS $ + unlines [ "Cannot compile assertion at " ++ + intercalate " -> " (reverse $ map locStr locs) ++ + " inside parallel kernel." + , "As a workaround, surround the expression with 'unsafe'."] -compileInKernelOp :: KernelConstants -> Pattern InKernel -> Op InKernel - -> InKernelGen () -compileInKernelOp _ (Pattern _ [mem]) (Alloc size (Space "private")) = do + +allocLocal, allocPrivate :: AllocCompiler ExplicitMemory Imp.KernelOp +allocLocal mem size = do + vtable <- getVTable + size' <- localMemSize vtable size + sOp $ Imp.LocalAlloc mem size' +allocPrivate mem size = + sOp $ Imp.PrivateAlloc mem size + +kernelAlloc :: KernelConstants + -> Pattern ExplicitMemory + -> SubExp -> Space + -> ImpM ExplicitMemory Imp.KernelOp () +kernelAlloc _ (Pattern _ [_]) _ (Space space) + | space `M.member` allScalarMemory = + return () -- Handled by the declaration of the memory block, + -- which is then translated to an actual scalar + -- variable during C code generation. +kernelAlloc _ (Pattern _ [mem]) size (Space "private") = do size' <- toExp size - sOp $ Imp.PrivateAlloc (patElemName mem) $ Imp.bytes size' -compileInKernelOp constants (Pattern _ [mem]) (Alloc size (Space "local")) = do - size' <- localMemSize (kernelOuterVTable constants) . Imp.bytes =<< - toExp size - sOp $ Imp.LocalAlloc (patElemName mem) size' -compileInKernelOp _ (Pattern _ [mem]) Alloc{} = + allocPrivate (patElemName mem) $ Imp.bytes size' +kernelAlloc _ (Pattern _ [mem]) size (Space "local") = do + size' <- toExp size + allocLocal (patElemName mem) $ Imp.bytes size' +kernelAlloc _ (Pattern _ [mem]) _ _ = compilerLimitationS $ "Cannot allocate memory block " ++ pretty mem ++ " in kernel." -compileInKernelOp _ dest Alloc{} = +kernelAlloc _ dest _ _ = compilerBugS $ "Invalid target for in-kernel allocation: " ++ show dest -compileInKernelOp constants pat (Inner op) = - compileKernelExp constants pat op -inKernelExpCompiler :: ExpCompiler InKernel Imp.KernelOp -inKernelExpCompiler _ (BasicOp (Assert _ _ (loc, locs))) = - compilerLimitationS $ - unlines [ "Cannot compile assertion at " ++ - intercalate " -> " (reverse $ map locStr $ loc:locs) ++ - " inside parallel kernel." - , "As a workaround, surround the expression with 'unsafe'."] --- The static arrays stuff does not work inside kernels. -inKernelExpCompiler (Pattern _ [dest]) (BasicOp (ArrayLit es _)) = +splitSpace :: (ToExp w, ToExp i, ToExp elems_per_thread) => + Pattern ExplicitMemory -> SplitOrdering -> w -> i -> elems_per_thread + -> ImpM lore op () +splitSpace (Pattern [] [size]) o w i elems_per_thread = do + num_elements <- Imp.elements <$> toExp w + i' <- toExp i + elems_per_thread' <- Imp.elements <$> toExp elems_per_thread + computeThreadChunkSize o i' elems_per_thread' num_elements (patElemName size) +splitSpace pat _ _ _ _ = + compilerBugS $ "Invalid target for splitSpace: " ++ pretty pat + +compileThreadExp :: ExpCompiler ExplicitMemory Imp.KernelOp +compileThreadExp _ (BasicOp (Assert _ _ (loc, locs))) = noAssert $ loc:locs +compileThreadExp (Pattern _ [dest]) (BasicOp (ArrayLit es _)) = forM_ (zip [0..] es) $ \(i,e) -> copyDWIM (patElemName dest) [fromIntegral (i::Int32)] e [] -inKernelExpCompiler dest e = +compileThreadExp dest e = defCompileExp dest e -compileKernelExp :: KernelConstants -> Pattern InKernel -> KernelExp InKernel - -> InKernelGen () +-- | Assign iterations of a for-loop to threads in the workgroup. The +-- passed-in function is invoked with the (symbolic) iteration. For +-- multidimensional loops, use 'groupCoverSpace'. +groupLoop :: KernelConstants -> Imp.Exp + -> (Imp.Exp -> InKernelGen ()) -> InKernelGen () +groupLoop constants n f = do + i <- newVName "i" -compileKernelExp _ pat (Barrier ses) = do - forM_ (zip (patternNames pat) ses) $ \(d, se) -> - copyDWIM d [] se [] + -- Compute how many elements this thread is responsible for. + -- Formula: (n - ltid) / group_size (rounded up). + let ltid = kernelLocalThreadId constants + elems_for_this = (n - ltid) `quotRoundingUp` kernelGroupSize constants + + sFor i Int32 elems_for_this $ f $ Imp.vi32 i + +-- | Iterate collectively though a multidimensional space, such that +-- all threads in the group participate. The passed-in function is +-- invoked with a (symbolic) point in the index space. +groupCoverSpace :: KernelConstants -> [Imp.Exp] + -> ([Imp.Exp] -> InKernelGen ()) -> InKernelGen () +groupCoverSpace constants ds f = + groupLoop constants (product ds) $ \i -> do + let is = unflattenIndex ds $ + i * kernelGroupSize constants + + kernelLocalThreadId constants + f is + +groupCopy :: KernelConstants -> VName -> [Imp.Exp] -> SubExp -> [Imp.Exp] -> InKernelGen () +groupCopy constants to to_is from from_is = do + ds <- mapM toExp . arrayDims =<< subExpType from + groupCoverSpace constants ds $ \is -> + copyDWIM to (to_is++ is) from (from_is ++ is) + +compileGroupExp :: KernelConstants -> ExpCompiler ExplicitMemory Imp.KernelOp +compileGroupExp _ _ (BasicOp (Assert _ _ (loc, locs))) = noAssert $ loc:locs +-- The static arrays stuff does not work inside kernels. +compileGroupExp _ (Pattern _ [dest]) (BasicOp (ArrayLit es _)) = + forM_ (zip [0..] es) $ \(i,e) -> + copyDWIM (patElemName dest) [fromIntegral (i::Int32)] e [] +compileGroupExp constants (Pattern _ [dest]) (BasicOp (Copy arr)) = do + groupCopy constants (patElemName dest) [] (Var arr) [] sOp Imp.LocalBarrier +compileGroupExp constants (Pattern _ [dest]) (BasicOp (Manifest _ arr)) = do + groupCopy constants (patElemName dest) [] (Var arr) [] + sOp Imp.LocalBarrier +compileGroupExp constants (Pattern _ [dest]) (BasicOp (Replicate ds se)) = do + ds' <- mapM toExp $ shapeDims ds + groupCoverSpace constants ds' $ \is -> + copyDWIM (patElemName dest) is se (drop (shapeRank ds) is) +compileGroupExp constants (Pattern _ [dest]) (BasicOp (Iota n e s _)) = do + n' <- toExp n + e' <- toExp e + s' <- toExp s + groupLoop constants n' $ \i' -> do + x <- dPrimV "x" $ e' + i' * s' + copyDWIM (patElemName dest) [i'] (Var x) [] -compileKernelExp _ (Pattern [] [size]) (SplitSpace o w i elems_per_thread) = do - num_elements <- Imp.elements <$> toExp w - i' <- toExp i - elems_per_thread' <- Imp.elements <$> toExp elems_per_thread - computeThreadChunkSize o i' elems_per_thread' num_elements (patElemName size) +compileGroupExp _ dest e = + defCompileExp dest e -compileKernelExp constants pat (Combine (CombineSpace scatter cspace) _ aspace body) = do - -- First we compute how many times we have to iterate to cover - -- cspace with our group size. It is a fairly common case that - -- we statically know that this requires 1 iteration, so we - -- could detect it and not generate a loop in that case. - -- However, it seems to have no impact on performance (an extra - -- conditional jump), so for simplicity we just always generate - -- the loop. - let cspace_dims = map (streamBounded . snd) cspace - num_iters - | cspace_dims == [kernelGroupSize constants] = 1 - | otherwise = product cspace_dims `quotRoundingUp` - kernelGroupSize constants +sanityCheckLevel :: SegLevel -> InKernelGen () +sanityCheckLevel SegThread{} = return () +sanityCheckLevel SegThreadScalar{} = return () +sanityCheckLevel SegGroup{} = + fail "compileGroupOp: unexpected group-level SegOp." - iter <- newVName "comb_iter" +compileGroupSpace :: KernelConstants -> SegLevel -> SegSpace -> InKernelGen () +compileGroupSpace constants lvl space = do + sanityCheckLevel lvl - sFor iter Int32 num_iters $ do - mapM_ ((`dPrim_` int32) . fst) cspace - -- Compute the *flat* array index. - cid <- dPrimV "flat_comb_id" $ - Imp.var iter int32 * kernelGroupSize constants + - kernelLocalThreadId constants + let (ltids, dims) = unzip $ unSegSpace space + dims' <- mapM toExp dims + zipWithM_ dPrimV_ ltids $ unflattenIndex dims' $ kernelLocalThreadId constants - -- Turn it into a nested array index. - zipWithM_ (<--) (map fst cspace) $ unflattenIndex cspace_dims (Imp.var cid int32) + dPrimV_ (segFlat space) $ kernelLocalThreadId constants - -- Construct the body. This is mostly about the book-keeping - -- for the scatter-like part. - let (scatter_ws, scatter_ns, _scatter_vs) = unzip3 scatter - scatter_ws_repl = concat $ zipWith replicate scatter_ns scatter_ws - (scatter_pes, normal_pes) = - splitAt (sum scatter_ns) $ patternElements pat - (res_is, res_vs, res_normal) = - splitAt3 (sum scatter_ns) (sum scatter_ns) $ bodyResult body +compileGroupOp :: KernelConstants -> OpCompiler ExplicitMemory Imp.KernelOp - -- Execute the body if we are within bounds. - sWhen (isActive cspace .&&. isActive aspace) $ allThreads constants $ - compileStms (freeIn $ bodyResult body) (bodyStms body) $ do +compileGroupOp constants pat (Alloc size space) = + kernelAlloc constants pat size space - forM_ (zip4 scatter_ws_repl res_is res_vs scatter_pes) $ - \(w, res_i, res_v, scatter_pe) -> do - let res_i' = toExp' int32 res_i - w' = toExp' int32 w - -- We have to check that 'res_i' is in-bounds wrt. an array of size 'w'. - in_bounds = 0 .<=. res_i' .&&. res_i' .<. w' - sWhen in_bounds $ copyDWIM (patElemName scatter_pe) [res_i'] res_v [] +compileGroupOp _ pat (Inner (SplitSpace o w i elems_per_thread)) = + splitSpace pat o w i elems_per_thread - forM_ (zip normal_pes res_normal) $ \(pe, res) -> - copyDWIM (patElemName pe) local_index res [] +compileGroupOp constants pat (Inner (SegOp (SegMap lvl space _ body))) = do + void $ compileGroupSpace constants lvl space - sOp Imp.LocalBarrier + sWhen (isActive $ unSegSpace space) $ + compileStms mempty (kernelBodyStms body) $ + zipWithM_ (compileThreadResult space constants) (patternElements pat) $ + kernelBodyResult body - where streamBounded (Var v) - | Just x <- lookup v $ kernelStreamed constants = - Imp.sizeToExp x - streamBounded se = toExp' int32 se + sOp Imp.LocalBarrier - local_index = map (toExp' int32 . Var . fst) cspace +compileGroupOp constants pat (Inner (SegOp (SegScan lvl space scan_op _ _ body))) = do + compileGroupSpace constants lvl space + let (ltids, dims) = unzip $ unSegSpace space + dims' <- mapM toExp dims -compileKernelExp constants (Pattern _ dests) (GroupReduce w lam input) = do - w' <- toExp w - groupReduce constants w' lam $ map snd input + sWhen (isActive $ unSegSpace space) $ + compileStms mempty (kernelBodyStms body) $ + forM_ (zip (patternNames pat) $ kernelBodyResult body) $ \(dest, res) -> + copyDWIM dest + (map (`Imp.var` int32) ltids) + (kernelResultSubExp res) [] sOp Imp.LocalBarrier - -- The final result will be stored in element 0 of the local memory array. - forM_ (zip dests input) $ \(dest, (_, arr)) -> - copyDWIM (patElemName dest) [] (Var arr) [0] + let segment_size = last dims' + crossesSegment from to = (to-from) .>. (to `rem` segment_size) + groupScan constants (Just crossesSegment) (product dims') scan_op $ + patternNames pat -compileKernelExp constants _ (GroupScan w lam input) = do - w' <- toExp w - groupScan constants Nothing w' lam $ map snd input +compileGroupOp constants pat (Inner (SegOp (SegRed lvl space ops _ body))) = do + compileGroupSpace constants lvl space -compileKernelExp constants (Pattern _ final) (GroupStream w maxchunk lam accs _arrs) = do - let GroupStreamLambda block_size block_offset acc_params arr_params body = lam - block_offset' = Imp.var block_offset int32 - w' <- toExp w - max_block_size <- toExp maxchunk + let (ltids, dims) = unzip $ unSegSpace space + (red_pes, map_pes) = + splitAt (segRedResults ops) $ patternElements pat - dLParams (acc_params++arr_params) - forM_ (zip acc_params accs) $ \(p, acc) -> - copyDWIM (paramName p) [] acc [] - dPrim_ block_size int32 + dims' <- mapM toExp dims - -- If the GroupStream is morally just a do-loop, generate simpler code. - case mapM isSimpleThreadInSpace $ stmsToList $ bodyStms body of - Just stms' | ValueExp x <- max_block_size, oneIsh x -> do - let body' = body { bodyStms = stmsFromList stms' } - body'' = allThreads constants $ - compileLoopBody acc_params body' - block_size <-- 1 + let mkTempArr t = + sAllocArray "red_arr" (elemType t) (Shape dims <> arrayShape t) $ Space "local" + tmp_arrs <- mapM mkTempArr $ concatMap (lambdaReturnType . segRedLambda) ops + let tmps_for_ops = chunks (map (length . segRedNeutral) ops) tmp_arrs - -- Check if loop is candidate for unrolling. - let loop = - case w of - Var w_var | Just w_bound <- lookup w_var $ kernelStreamed constants, - w_bound /= Imp.ConstSize 1 -> - -- Candidate for unrolling, so generate two loops. - sIf (w' .==. Imp.sizeToExp w_bound) - (sFor block_offset Int32 (Imp.sizeToExp w_bound) body'') - (sFor block_offset Int32 w' body'') - _ -> sFor block_offset Int32 w' body'' + sWhen (isActive $ unSegSpace space) $ + compileStms mempty (kernelBodyStms body) $ do + let (red_res, map_res) = + splitAt (segRedResults ops) $ kernelBodyResult body + forM_ (zip tmp_arrs red_res) $ \(dest, res) -> + copyDWIM dest (map (`Imp.var` int32) ltids) (kernelResultSubExp res) [] + zipWithM_ (compileThreadResult space constants) map_pes map_res - if kernelThreadActive constants == Imp.ValueExp (BoolValue True) - then loop - else sWhen (kernelThreadActive constants) loop + sOp Imp.LocalBarrier - _ -> do - dPrim_ block_offset int32 - let body' = streaming constants block_size maxchunk $ - compileBody' acc_params body + case dims' of + [dim'] -> do + forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> + groupReduce constants dim' (segRedLambda op) tmps - block_offset <-- 0 + sOp Imp.LocalBarrier - let not_at_end = block_offset' .<. w' - set_block_size = - sIf (w' - block_offset' .<. max_block_size) - (block_size <-- (w' - block_offset')) - (block_size <-- max_block_size) - increase_offset = - block_offset <-- block_offset' + max_block_size + forM_ (zip red_pes tmp_arrs) $ \(pe, arr) -> + copyDWIM (patElemName pe) [] (Var arr) [0] - -- Three cases to consider for simpler generated code based - -- on max block size: (0) if full input size, do not - -- generate a loop; (1) if one, generate for-loop (2) - -- otherwise, generate chunked while-loop. - if max_block_size == w' then - (block_size <-- w') >> body' - else if max_block_size == Imp.ValueExp (value (1::Int32)) then do - block_size <-- w' - sFor block_offset Int32 w' body' - else - sWhile not_at_end $ - set_block_size >> body' >> increase_offset + _ -> do + let segment_size = last dims' + crossesSegment from to = (to-from) .>. (to `rem` segment_size) - forM_ (zip final acc_params) $ \(pe, p) -> - copyDWIM (patElemName pe) [] (Var $ paramName p) [] + forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> + groupScan constants (Just crossesSegment) (product dims') (segRedLambda op) tmps - where isSimpleThreadInSpace (Let _ _ Op{}) = Nothing - isSimpleThreadInSpace bnd = Just bnd + sOp Imp.LocalBarrier -compileKernelExp _ _ (GroupGenReduce w arrs op bucket values locks) = do - -- Check if bucket is in-bounds - bucket' <- mapM toExp bucket - w' <- mapM toExp w - num_locks <- toExp' int32 . arraySize 0 <$> lookupType locks - let locking = Locking locks 0 1 0 $ (`rem` num_locks) . sum - values_params = takeLast (length values) $ lambdaParams op + let segment_is = map Imp.vi32 $ init ltids + forM_ (zip red_pes tmp_arrs) $ \(pe, arr) -> + copyDWIM (patElemName pe) segment_is (Var arr) (segment_is ++ [last dims'-1]) - sWhen (indexInBounds bucket' w') $ do - forM_ (zip values_params values) $ \(p, v) -> - copyDWIM (paramName p) [] v [] - atomicUpdate DefaultSpace arrs bucket' op locking - where indexInBounds inds bounds = - foldl1 (.&&.) $ zipWith checkBound inds bounds - where checkBound ind bound = 0 .<=. ind .&&. ind .<. bound -compileKernelExp _ dest e = - compilerBugS $ unlines ["Invalid target", " " ++ show dest, - "for kernel expression", " " ++ pretty e] +compileGroupOp _ pat _ = + compilerBugS $ "compileGroupOp: cannot compile rhs of binding " ++ pretty pat -streaming :: KernelConstants -> VName -> SubExp -> InKernelGen () -> InKernelGen () -streaming constants chunksize bound m = do - bound' <- subExpToDimSize bound - let constants' = - constants { kernelStreamed = (chunksize, bound') : kernelStreamed constants } - emit =<< subImpM_ (inKernelOperations constants') m +compileThreadOp :: KernelConstants -> OpCompiler ExplicitMemory Imp.KernelOp +compileThreadOp constants pat (Alloc size space) = + kernelAlloc constants pat size space +compileThreadOp _ pat (Inner (SplitSpace o w i elems_per_thread)) = + splitSpace pat o w i elems_per_thread +compileThreadOp _ pat _ = + compilerBugS $ "compileThreadOp: cannot compile rhs of binding " ++ pretty pat -- | Locking strategy used for an atomic update. data Locking = @@ -504,14 +488,14 @@ a -> [VName] -> CallKernelGen [Imp.KernelUse] computeKernelUses kernel_body bound_in_kernel = do - let actually_free = freeIn kernel_body `S.difference` S.fromList bound_in_kernel + 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 readsFromSet :: Names -> CallKernelGen [Imp.KernelUse] readsFromSet free = fmap catMaybes $ - forM (S.toList free) $ \var -> do + forM (namesToList free) $ \var -> do t <- lookupType var vtable <- getVTable case t of @@ -524,9 +508,9 @@ Nothing | bt == Cert -> return Nothing | otherwise -> return $ Just $ Imp.ScalarUse var bt -localMemSize :: VTable ExplicitMemory -> Imp.Count Imp.Bytes - -> ImpM lore op (Either (Imp.Count Imp.Bytes) Imp.KernelConstExp) -localMemSize vtable e = isConstExp vtable (Imp.innerExp e) >>= \case +localMemSize :: VTable ExplicitMemory -> Imp.Count Imp.Bytes Imp.Exp + -> ImpM lore op (Either (Imp.Count Imp.Bytes Imp.Exp) Imp.KernelConstExp) +localMemSize vtable e = isConstExp vtable (Imp.unCount e) >>= \case Just e' | isStaticExp e' -> return $ Right e' _ -> return $ Left e @@ -563,61 +547,54 @@ computeThreadChunkSize :: SplitOrdering -> Imp.Exp - -> Imp.Count Imp.Elements - -> Imp.Count Imp.Elements + -> Imp.Count Imp.Elements Imp.Exp + -> Imp.Count Imp.Elements Imp.Exp -> VName -> ImpM lore op () computeThreadChunkSize (SplitStrided stride) thread_index elements_per_thread num_elements chunk_var = do stride' <- toExp stride chunk_var <-- Imp.BinOpExp (SMin Int32) - (Imp.innerExp elements_per_thread) - ((Imp.innerExp num_elements - thread_index) `quotRoundingUp` stride') + (Imp.unCount elements_per_thread) + ((Imp.unCount num_elements - thread_index) `quotRoundingUp` stride') computeThreadChunkSize SplitContiguous thread_index elements_per_thread num_elements chunk_var = do starting_point <- dPrimV "starting_point" $ - thread_index * Imp.innerExp elements_per_thread + thread_index * Imp.unCount elements_per_thread remaining_elements <- dPrimV "remaining_elements" $ - Imp.innerExp num_elements - Imp.var starting_point int32 + Imp.unCount num_elements - Imp.var starting_point int32 let no_remaining_elements = Imp.var remaining_elements int32 .<=. 0 - beyond_bounds = Imp.innerExp num_elements .<=. Imp.var starting_point int32 + beyond_bounds = Imp.unCount num_elements .<=. Imp.var starting_point int32 sIf (no_remaining_elements .||. beyond_bounds) (chunk_var <-- 0) (sIf is_last_thread - (chunk_var <-- Imp.innerExp last_thread_elements) - (chunk_var <-- Imp.innerExp elements_per_thread)) + (chunk_var <-- Imp.unCount last_thread_elements) + (chunk_var <-- Imp.unCount elements_per_thread)) where last_thread_elements = num_elements - Imp.elements thread_index * elements_per_thread is_last_thread = - Imp.innerExp num_elements .<. - (thread_index + 1) * Imp.innerExp elements_per_thread + Imp.unCount num_elements .<. + (thread_index + 1) * Imp.unCount elements_per_thread -kernelInitialisationSimple :: Imp.Exp -> Imp.Exp - -> Maybe (VName, VName, VName) - -> CallKernelGen (KernelConstants, ImpM InKernel Imp.KernelOp ()) -kernelInitialisationSimple num_groups group_size names = do - (global_tid, local_tid, group_id) <- - case names of Nothing -> - (,,) - <$> newVName "global_tid" - <*> newVName "local_tid" - <*> newVName "group_id" - Just (global_tid, local_tid, group_id) -> - return (global_tid, local_tid, group_id) +kernelInitialisationSimple :: Count NumGroups Imp.Exp -> Count GroupSize Imp.Exp + -> CallKernelGen (KernelConstants, InKernelGen ()) +kernelInitialisationSimple (Count num_groups) (Count group_size) = do + global_tid <- newVName "global_tid" + local_tid <- newVName "local_tid" + group_id <- newVName "group_tid" wave_size <- newVName "wave_size" inner_group_size <- newVName "group_size" - vtable <- getVTable let constants = - KernelConstants vtable + KernelConstants (Imp.var global_tid int32) (Imp.var local_tid int32) (Imp.var group_id int32) global_tid local_tid group_id - group_size num_groups (group_size*num_groups) - (Imp.var wave_size int32) [] - true mempty + num_groups group_size (group_size*num_groups) + (Imp.var wave_size int32) + true let set_constants = do dPrim_ global_tid int32 @@ -634,61 +611,6 @@ return (constants, set_constants) -kernelInitialisationSetSpace :: KernelSpace -> InKernelGen () - -> CallKernelGen (KernelConstants, ImpM InKernel Imp.KernelOp ()) -kernelInitialisationSetSpace space set_space = do - group_size <- toExp $ spaceGroupSize space - num_groups <- toExp $ spaceNumGroups space - - (constants, set_constants) <- - kernelInitialisationSimple num_groups group_size $ - Just (spaceGlobalId space, spaceLocalId space, spaceGroupId space) - - let set_constants' = do - set_constants - case spaceStructure space of - FlatThreadSpace is_and_dims -> - mapM_ ((`dPrim_` int32) . fst) is_and_dims - NestedThreadSpace is_and_dims -> do - let (gtids, _, ltids, _) = unzip4 is_and_dims - mapM_ (`dPrim_` int32) $ gtids ++ ltids - set_space - - let (space_is, space_dims) = unzip $ spaceDimensions space - space_dims' <- mapM toExp space_dims - - return (constants { kernelThreadActive = - if null $ spaceDimensions space - then true - else isActive $ spaceDimensions space - , kernelDimensions = - zip space_is space_dims' - }, - set_constants') - -kernelInitialisation :: KernelSpace - -> CallKernelGen (KernelConstants, ImpM InKernel Imp.KernelOp ()) -kernelInitialisation space = - kernelInitialisationSetSpace space $ - setSpaceIndices (Imp.var (spaceGlobalId space) int32) space - -setSpaceIndices :: Imp.Exp -> KernelSpace -> InKernelGen () -setSpaceIndices gtid space = - case spaceStructure space of - FlatThreadSpace is_and_dims -> - flatSpaceWith gtid is_and_dims - NestedThreadSpace is_and_dims -> do - let (gtids, gdims, ltids, ldims) = unzip4 is_and_dims - gdims' <- mapM toExp gdims - ldims' <- mapM toExp ldims - let (gtid_es, ltid_es) = unzip $ unflattenNestedIndex gdims' ldims' gtid - zipWithM_ (<--) gtids gtid_es - zipWithM_ (<--) ltids ltid_es - where flatSpaceWith base is_and_dims = do - let (is, dims) = unzip is_and_dims - dims' <- mapM toExp dims - zipWithM_ (<--) is $ unflattenIndex dims' base - isActive :: [(VName, SubExp)] -> Imp.Exp isActive limit = case actives of [] -> Imp.ValueExp $ BoolValue True @@ -697,19 +619,6 @@ actives = zipWith active is $ map (toExp' Bool) ws active i = (Imp.var i int32 .<.) -unflattenNestedIndex :: IntegralExp num => [num] -> [num] -> num -> [(num,num)] -unflattenNestedIndex global_dims group_dims global_id = - zip global_is local_is - where num_groups_dims = zipWith quotRoundingUp global_dims group_dims - group_size = product group_dims - group_id = global_id `Futhark.Util.IntegralExp.quot` group_size - local_id = global_id `Futhark.Util.IntegralExp.rem` group_size - - group_is = unflattenIndex num_groups_dims group_id - local_is = unflattenIndex group_dims local_id - global_is = zipWith (+) local_is $ zipWith (*) group_is group_dims - - -- | Change every memory block to be in the global address space, -- except those who are in the local memory space. This only affects -- generated code - we still need to make sure that the memory is @@ -725,13 +634,6 @@ globalMemory entry = entry -allThreads :: KernelConstants -> InKernelGen () -> InKernelGen () -allThreads constants = emit <=< subImpM_ (inKernelOperations constants') - where constants' = - constants { kernelThreadActive = Imp.ValueExp (BoolValue True) } - - - writeParamToLocalMemory :: Typed (MemBound u) => Imp.Exp -> (VName, t) -> Param (MemBound u) -> ImpM lore op () @@ -852,9 +754,9 @@ groupScan :: KernelConstants -> Maybe (Imp.Exp -> Imp.Exp -> Imp.Exp) -> Imp.Exp - -> Lambda InKernel + -> Lambda ExplicitMemory -> [VName] - -> ImpM InKernel Imp.KernelOp () + -> ImpM ExplicitMemory Imp.KernelOp () groupScan constants seg_flag w lam arrs = do when (any (not . primType . paramType) $ lambdaParams lam) $ compilerLimitationS "Cannot compile parallel scans with array element type." @@ -940,7 +842,7 @@ -> Imp.Exp -> Imp.Exp -> [(VName, t)] - -> Lambda InKernel + -> Lambda ExplicitMemory -> InKernelGen () inBlockScan seg_flag lockstep_width block_size active ltid acc_local_mem scan_lam = everythingVolatile $ do skip_threads <- dPrim "skip_threads" int32 @@ -987,40 +889,13 @@ where block_id = ltid `quot` block_size in_block_id = ltid - block_id * block_size -compileKernelStms :: KernelConstants -> Stms InKernel - -> InKernelGen a - -> InKernelGen a -compileKernelStms constants ungrouped_bnds m = - compileGroupedKernelStms' $ groupStmsByGuard constants $ stmsToList ungrouped_bnds - where compileGroupedKernelStms' [] = m - compileGroupedKernelStms' ((g, bnds):rest_bnds) = do - dScopes (map ((Just . stmExp) &&& (castScope . scopeOf)) bnds) - protect g $ mapM_ compileKernelStm bnds - compileGroupedKernelStms' rest_bnds - - protect Nothing body_m = - body_m - protect (Just (Imp.ValueExp (BoolValue True))) body_m = - body_m - protect (Just g) body_m = - sWhen g $ allThreads constants body_m - - compileKernelStm (Let pat _ e) = compileExp pat e - -groupStmsByGuard :: KernelConstants - -> [Stm InKernel] - -> [(Maybe Imp.Exp, [Stm InKernel])] -groupStmsByGuard constants bnds = - map collapse $ groupBy sameGuard $ zip (map bindingGuard bnds) bnds - where bindingGuard (Let _ _ Op{}) = Nothing - bindingGuard _ = Just $ kernelThreadActive constants - - sameGuard (g1, _) (g2, _) = g1 == g2 - - collapse [] = - (Nothing, []) - collapse l@((g,_):_) = - (g, map snd l) +getSize :: String -> Imp.SizeClass -> CallKernelGen VName +getSize desc sclass = do + size <- dPrim desc int32 + fname <- asks envFunction + let size_key = keyWithEntryPoint fname $ nameFromString $ pretty size + sOp $ Imp.GetSize size size_key sclass + return size computeMapKernelGroups :: Imp.Exp -> CallKernelGen (Imp.Exp, Imp.Exp) computeMapKernelGroups kernel_size = do @@ -1030,15 +905,15 @@ group_size_key = keyWithEntryPoint fname $ nameFromString $ pretty group_size sOp $ Imp.GetSize group_size group_size_key Imp.SizeGroup num_groups <- dPrimV "num_groups" $ kernel_size `quotRoundingUp` Imp.ConvOpExp (SExt Int32 Int32) group_size_var - return (Imp.var group_size int32, Imp.var num_groups int32) + return (Imp.var num_groups int32, Imp.var group_size int32) simpleKernelConstants :: Imp.Exp -> String - -> CallKernelGen (KernelConstants, ImpM InKernel Imp.KernelOp ()) + -> CallKernelGen (KernelConstants, InKernelGen ()) simpleKernelConstants kernel_size desc = do thread_gtid <- newVName $ desc ++ "_gtid" thread_ltid <- newVName $ desc ++ "_ltid" group_id <- newVName $ desc ++ "_gid" - (group_size, num_groups) <- computeMapKernelGroups kernel_size + (num_groups, group_size) <- computeMapKernelGroups kernel_size let set_constants = do dPrim_ thread_gtid int32 dPrim_ thread_ltid int32 @@ -1048,12 +923,11 @@ sOp (Imp.GetGroupId group_id 0) - vtable <- getVTable - return (KernelConstants vtable + return (KernelConstants (Imp.var thread_gtid int32) (Imp.var thread_ltid int32) (Imp.var group_id int32) thread_gtid thread_ltid group_id - group_size num_groups (group_size*num_groups) 0 - [] (Imp.var thread_gtid int32 .<. kernel_size) mempty, + num_groups group_size (group_size*num_groups) 0 + (Imp.var thread_gtid int32 .<. kernel_size), set_constants) @@ -1065,13 +939,13 @@ -- *physical* threads (hardware parallelism), not the amount of -- application parallelism. virtualiseGroups :: KernelConstants + -> SegVirt -> Imp.Exp -> (VName -> InKernelGen ()) -> InKernelGen () -virtualiseGroups constants required_groups m - | kernelNumGroups constants == required_groups = - m $ kernelGroupIdVar constants - | otherwise = do +virtualiseGroups constants SegNoVirt _ m = + m $ kernelGroupIdVar constants +virtualiseGroups constants SegVirt required_groups m = do phys_group_id <- dPrim "phys_group_id" int32 sOp $ Imp.GetGroupId phys_group_id 0 let iterations = (required_groups - Imp.vi32 phys_group_id) `quotRoundingUp` @@ -1080,21 +954,79 @@ sFor i Int32 iterations $ m =<< dPrimV "virt_group_id" (Imp.vi32 phys_group_id + Imp.vi32 i * kernelNumGroups constants) -sKernel :: KernelConstants -> String -> ImpM InKernel Imp.KernelOp a -> CallKernelGen () -sKernel constants name m = do - body <- makeAllMemoryGlobal $ subImpM_ (inKernelOperations constants) m - uses <- computeKernelUses body mempty +sKernelThread, sKernelGroup :: String + -> Count NumGroups Imp.Exp -> Count GroupSize Imp.Exp + -> VName + -> (KernelConstants -> InKernelGen ()) + -> CallKernelGen () +(sKernelThread, sKernelGroup) = (sKernel' threadOperations kernelGlobalThreadId, + sKernel' groupOperations kernelGroupId) + where sKernel' ops flatf name num_groups group_size v f = do + (constants, set_constants) <- kernelInitialisationSimple num_groups group_size + let name' = nameFromString $ name ++ "_" ++ show (baseTag v) + sKernel (ops constants) constants name' $ do + set_constants + dPrimV_ v $ flatf constants + f constants +sKernel :: Operations ExplicitMemory Imp.KernelOp + -> KernelConstants -> Name -> ImpM ExplicitMemory Imp.KernelOp a -> CallKernelGen () +sKernel ops constants name m = do + body <- makeAllMemoryGlobal $ subImpM_ ops m + uses <- computeKernelUses body mempty emit $ Imp.Op $ Imp.CallKernel Imp.Kernel { Imp.kernelBody = body , Imp.kernelUses = uses , Imp.kernelNumGroups = [kernelNumGroups constants] , Imp.kernelGroupSize = [kernelGroupSize constants] - , Imp.kernelName = - nameFromString $ name ++ "_" ++ show tag + , Imp.kernelName = name } - where tag = baseTag $ kernelGlobalThreadIdVar constants +-- | A kernel with the given number of threads, running per-thread code. +sKernelSimple :: String -> Imp.Exp + -> (KernelConstants -> InKernelGen ()) + -> CallKernelGen () +sKernelSimple name kernel_size f = do + (constants, init_constants) <- simpleKernelConstants kernel_size name + let name' = nameFromString $ name ++ "_" ++ + show (baseTag $ kernelGlobalThreadIdVar constants) + sKernel (threadOperations constants) constants name' $ do + init_constants + f constants + +copyInGroup :: CopyCompiler ExplicitMemory Imp.KernelOp +copyInGroup pt destloc srcloc n = do + dest_space <- entryMemSpace <$> lookupMemory (memLocationName destloc) + src_space <- entryMemSpace <$> lookupMemory (memLocationName srcloc) + + if isScalarMem dest_space && isScalarMem src_space + then memLocationName destloc <-- Imp.var (memLocationName srcloc) pt + else copyElementWise pt destloc srcloc n + + where isScalarMem (Space space) = space `M.member` allScalarMemory + isScalarMem DefaultSpace = False + +threadOperations, groupOperations :: KernelConstants + -> Operations ExplicitMemory Imp.KernelOp +threadOperations constants = + (defaultOperations $ compileThreadOp constants) + { opsCopyCompiler = copyElementWise + , opsExpCompiler = compileThreadExp + , opsStmsCompiler = \_ -> defCompileStms mempty + , opsAllocCompilers = + M.fromList [ (Space "local", allocLocal) + , (Space "private", allocPrivate) ] + } +groupOperations constants = + (defaultOperations $ compileGroupOp constants) + { opsCopyCompiler = copyInGroup + , opsExpCompiler = compileGroupExp constants + , opsStmsCompiler = \_ -> defCompileStms mempty + , opsAllocCompilers = + M.fromList [ (Space "local", allocLocal) + , (Space "private", allocPrivate) ] + } + -- | Perform a Replicate with a kernel. sReplicate :: VName -> Shape -> SubExp -> CallKernelGen () @@ -1106,8 +1038,10 @@ simpleKernelConstants (product dims) "replicate" let is' = unflattenIndex dims $ kernelGlobalThreadId constants + name = nameFromString $ "replicate_" ++ + show (baseTag $ kernelGlobalThreadIdVar constants) - sKernel constants "replicate" $ do + sKernel (threadOperations constants) constants name $ do set_constants sWhen (kernelThreadActive constants) $ copyDWIM arr is' se $ drop (length ds) is' @@ -1119,7 +1053,10 @@ destloc <- entryArrayLocation <$> lookupArray arr (constants, set_constants) <- simpleKernelConstants n "iota" - sKernel constants "iota" $ do + let name = nameFromString $ "iota_" ++ + show (baseTag $ kernelGlobalThreadIdVar constants) + + sKernel (threadOperations constants) constants name $ do set_constants let gtid = kernelGlobalThreadId constants sWhen (kernelThreadActive constants) $ do @@ -1132,7 +1069,7 @@ sCopy :: PrimType -> MemLocation -> MemLocation - -> Imp.Count Imp.Elements + -> Imp.Count Imp.Elements Imp.Exp -> CallKernelGen () sCopy bt destloc@(MemLocation destmem _ _) @@ -1141,12 +1078,15 @@ -- Note that the shape of the destination and the source are -- necessarily the same. let shape = map Imp.sizeToExp srcshape - shape_se = map (Imp.innerExp . dimSizeToExp) srcshape - kernel_size = Imp.innerExp n * product (drop 1 shape) + shape_se = map (Imp.unCount . dimSizeToExp) srcshape + kernel_size = Imp.unCount n * product (drop 1 shape) (constants, set_constants) <- simpleKernelConstants kernel_size "copy" - sKernel constants "copy" $ do + let name = nameFromString $ "copy_" ++ + show (baseTag $ kernelGlobalThreadIdVar constants) + + sKernel (threadOperations constants) constants name $ do set_constants let gtid = kernelGlobalThreadId constants @@ -1160,66 +1100,85 @@ Imp.Write destmem destidx bt destspace Imp.Nonvolatile $ Imp.index srcmem srcidx bt srcspace Imp.Nonvolatile +compileGroupResult :: SegSpace + -> KernelConstants -> PatElem ExplicitMemory -> KernelResult + -> InKernelGen () -compileKernelResult :: KernelConstants -> PatElem InKernel -> KernelResult - -> InKernelGen () +compileGroupResult _ constants pe (TileReturns [(w,per_group_elems)] what) = do + dest_loc <- entryArrayLocation <$> lookupArray (patElemName pe) + let dest_loc_offset = offsetArray dest_loc offset + dest' = arrayDestination dest_loc_offset + n <- toExp . arraySize 0 =<< lookupType what -compileKernelResult constants pe (GroupsReturn what) = do - i <- newVName "i" + -- Avoid loop for the common case where each thread is statically + -- known to write at most one element. + if toExp' int32 per_group_elems == kernelGroupSize constants + then sWhen (offset + ltid .<. toExp' int32 w) $ + copyDWIMDest dest' [ltid] (Var what) [ltid] + else do + i <- newVName "i" + sFor i Int32 (n `quotRoundingUp` kernelGroupSize constants) $ do + j <- fmap Imp.vi32 $ dPrimV "j" $ + kernelGroupSize constants * Imp.vi32 i + ltid + sWhen (j .<. n) $ copyDWIMDest dest' [j] (Var what) [j] + where ltid = kernelLocalThreadId constants + offset = toExp' int32 per_group_elems * kernelGroupId constants +compileGroupResult space constants pe (TileReturns dims what) = do + let gids = map fst $ unSegSpace space + out_tile_sizes = map (toExp' int32 . snd) dims + local_is = unflattenIndex out_tile_sizes $ kernelLocalThreadId constants + group_is = zipWith (*) (map Imp.vi32 gids) out_tile_sizes + is_for_thread <- mapM (dPrimV "thread_out_index") $ zipWith (+) group_is local_is + + sWhen (isActive $ zip is_for_thread $ map fst dims) $ + copyDWIM (patElemName pe) (map Imp.vi32 is_for_thread) (Var what) local_is + +compileGroupResult space constants pe (Returns what) = do in_local_memory <- arrayInLocalMemory what - let me = kernelLocalThreadId constants + let gids = map (Imp.vi32 . fst) $ unSegSpace space - if not in_local_memory then do - who' <- toExp $ intConst Int32 0 - sWhen (me .==. who') $ - copyDWIM (patElemName pe) [kernelGroupId constants] what [] - else do + if not in_local_memory then + sWhen (kernelLocalThreadId constants .==. 0) $ + copyDWIM (patElemName pe) gids what [] + else -- If the result of the group is an array in local memory, we -- store it by collective copying among all the threads of the -- group. TODO: also do this if the array is in global memory -- (but this is a bit more tricky, synchronisation-wise). - -- - -- We do the reads/writes multidimensionally, but the loop is - -- single-dimensional. - ws <- mapM toExp . arrayDims =<< subExpType what - -- Compute how many elements this thread is responsible for. - -- Formula: (w - ltid) / group_size (rounded up). - let w = product ws - ltid = kernelLocalThreadId constants - group_size = kernelGroupSize constants - to_write = (w - ltid) `quotRoundingUp` group_size - is = unflattenIndex ws $ Imp.vi32 i * group_size + ltid + groupCopy constants (patElemName pe) gids what [] - sFor i Int32 to_write $ - copyDWIM (patElemName pe) (kernelGroupId constants : is) what is +compileGroupResult _ _ _ WriteReturns{} = + compilerLimitationS "compileGroupResult: WriteReturns not handled yet." -compileKernelResult constants pe (ThreadsReturn what) = do - let is = map (Imp.vi32 . fst) $ kernelDimensions constants - sWhen (kernelThreadActive constants) $ copyDWIM (patElemName pe) is what [] +compileGroupResult _ _ _ ConcatReturns{} = + compilerLimitationS "compileGroupResult: ConcatReturns not handled yet." -compileKernelResult constants pe (ConcatReturns SplitContiguous _ per_thread_elems moffset what) = do +compileThreadResult :: SegSpace + -> KernelConstants -> PatElem ExplicitMemory -> KernelResult + -> InKernelGen () + +compileThreadResult space _ pe (Returns what) = do + let is = map (Imp.vi32 . fst) $ unSegSpace space + copyDWIM (patElemName pe) is what [] + +compileThreadResult _ constants pe (ConcatReturns SplitContiguous _ per_thread_elems what) = do dest_loc <- entryArrayLocation <$> lookupArray (patElemName pe) let dest_loc_offset = offsetArray dest_loc offset dest' = arrayDestination dest_loc_offset copyDWIMDest dest' [] (Var what) [] - where offset = case moffset of - Nothing -> toExp' int32 per_thread_elems * - kernelGlobalThreadId constants - Just se -> toExp' int32 se + where offset = toExp' int32 per_thread_elems * kernelGlobalThreadId constants -compileKernelResult constants pe (ConcatReturns (SplitStrided stride) _ _ moffset what) = do +compileThreadResult _ constants pe (ConcatReturns (SplitStrided stride) _ _ what) = do dest_loc <- entryArrayLocation <$> lookupArray (patElemName pe) let dest_loc' = strideArray (offsetArray dest_loc offset) $ toExp' int32 stride dest' = arrayDestination dest_loc' copyDWIMDest dest' [] (Var what) [] - where offset = case moffset of - Nothing -> kernelGlobalThreadId constants - Just se -> toExp' int32 se + where offset = kernelGlobalThreadId constants -compileKernelResult constants pe (WriteReturn rws _arr dests) = do +compileThreadResult _ constants pe (WriteReturns rws _arr dests) = do rws' <- mapM toExp rws forM_ dests $ \(is, e) -> do is' <- mapM toExp is @@ -1227,6 +1186,9 @@ write = foldl (.&&.) (kernelThreadActive constants) $ zipWith condInBounds is' rws' sWhen write $ copyDWIM (patElemName pe) (map (toExp' int32) is) e [] + +compileThreadResult _ _ _ TileReturns{} = + compilerBugS "compileThreadResult: TileReturns unhandled." arrayInLocalMemory :: SubExp -> InKernelGen Bool arrayInLocalMemory (Var name) = do diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegGenRed.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegGenRed.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegGenRed.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegGenRed.hs @@ -68,18 +68,18 @@ } data SegGenRedSlug = SegGenRedSlug - { slugOp :: GenReduceOp InKernel + { slugOp :: GenReduceOp ExplicitMemory , slugNumSubhistos :: VName , slugSubhistos :: [SubhistosInfo] } -- | Figure out how much memory is needed per histogram, and compute -- some other auxiliary information. -computeHistoUsage :: KernelSpace - -> GenReduceOp InKernel - -> CallKernelGen (Imp.Count Imp.Bytes, SegGenRedSlug) +computeHistoUsage :: SegSpace + -> GenReduceOp ExplicitMemory + -> CallKernelGen (Imp.Count Imp.Bytes Imp.Exp, SegGenRedSlug) computeHistoUsage space op = do - let segment_dims = init $ spaceDimensions space + let segment_dims = init $ unSegSpace space num_segments = length segment_dims op_h <- fmap (sum . map typeSize) $ mapM lookupType $ genReduceDest op @@ -113,7 +113,7 @@ let subhistos_mem_size = Imp.bytes $ - Imp.innerExp (Imp.elements num_elems `Imp.withElemType` elemType dest_t) + Imp.unCount (Imp.elements num_elems `Imp.withElemType` elemType dest_t) sAlloc_ subhistos_mem subhistos_mem_size $ Space "device" sReplicate subhistos (Shape (map snd segment_dims ++ @@ -129,20 +129,20 @@ return (op_h, SegGenRedSlug op num_subhistos subhisto_infos) -localMemLockArray :: KernelSpace -> Type -localMemLockArray space = Array int32 (Shape [spaceGroupSize space]) NoUniqueness +localMemLockArray :: Count GroupSize SubExp -> Type +localMemLockArray (Count group_size) = Array int32 (Shape [group_size]) NoUniqueness -- | How many bytes will be spent on lock arrays if we use a local -- memory implementation? -localMemLockUsage :: KernelSpace -> [SegGenRedSlug] -> Imp.Count Imp.Bytes -localMemLockUsage space slugs = +localMemLockUsage :: Count GroupSize SubExp -> [SegGenRedSlug] -> Imp.Count Imp.Bytes Imp.Exp +localMemLockUsage group_size slugs = if any (isRight . atomicUpdateLocking . genReduceOp . slugOp) slugs - then typeSize $ localMemLockArray space + then typeSize $ localMemLockArray group_size else 0 prepareAtomicUpdateGlobal :: Maybe Locking -> [VName] -> SegGenRedSlug -> CallKernelGen (Maybe Locking, - [Imp.Exp] -> ImpM InKernel Imp.KernelOp ()) + [Imp.Exp] -> InKernelGen ()) prepareAtomicUpdateGlobal l dests slug = -- We need a separate lock array if the operators are not all of a -- particularly simple form that permits pure atomic operations. @@ -172,7 +172,7 @@ -> CallKernelGen [(VName, [VName], - [Imp.Exp] -> ImpM InKernel Imp.KernelOp ())] + [Imp.Exp] -> InKernelGen ())] prepareIntermediateArraysGlobal num_threads = fmap snd . mapAccumLM onOp Nothing where onOp l slug@(SegGenRedSlug op num_subhistos subhisto_info) = do @@ -217,26 +217,26 @@ return (l', (num_subhistos, dests, do_op)) genRedKernelGlobal :: [PatElem ExplicitMemory] - -> KernelSpace - -> [SegGenRedSlug] - -> KernelBody InKernel - -> CallKernelGen () -genRedKernelGlobal map_pes space slugs kbody = do - (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return () - let constants = base_constants { kernelThreadActive = true } - (space_is, space_sizes) = unzip $ spaceDimensions space + -> Count NumGroups SubExp -> Count GroupSize SubExp + -> SegSpace + -> [SegGenRedSlug] + -> KernelBody ExplicitMemory + -> CallKernelGen () +genRedKernelGlobal map_pes num_groups group_size space slugs kbody = do + num_groups' <- traverse toExp num_groups + group_size' <- traverse toExp group_size + let (space_is, space_sizes) = unzip $ unSegSpace space space_sizes_64 = map (i32Toi64 . toExp' int32) space_sizes total_w_64 = product space_sizes_64 + num_threads = unCount num_groups' * unCount group_size' - histograms <- prepareIntermediateArraysGlobal (kernelNumThreads constants) slugs + histograms <- prepareIntermediateArraysGlobal num_threads slugs elems_per_thread_64 <- dPrimV "elems_per_thread_64" $ total_w_64 `quotRoundingUp` - ConvOpExp (SExt Int32 Int64) (kernelNumThreads constants) - - sKernel constants "seggenred_global" $ allThreads constants $ do - init_constants + ConvOpExp (SExt Int32 Int64) num_threads + sKernelThread "seggenred_global" num_groups' group_size' (segFlat space) $ \constants -> do -- Compute subhistogram index for each thread, per histogram. subhisto_inds <- forM histograms $ \(num_histograms, _, _) -> dPrimV "subhisto_ind" $ @@ -276,7 +276,7 @@ sComment "save map-out results" $ forM_ (zip map_pes map_res) $ \(pe, res) -> copyDWIM (patElemName pe) - (map ((`Imp.var` int32) . fst) $ kernelDimensions constants) + (map ((`Imp.var` int32) . fst) $ unSegSpace space) (kernelResultSubExp res) [] let (buckets, vs) = splitAt (length slugs) red_res @@ -301,19 +301,19 @@ copyDWIM (paramName p) [] (kernelResultSubExp res) is do_op (bucket_is ++ is) -prepareIntermediateArraysLocal :: KernelSpace -> KernelConstants +prepareIntermediateArraysLocal :: Count NumGroups SubExp -> Count GroupSize SubExp -> VName -> [SegGenRedSlug] -> CallKernelGen [([VName], + KernelConstants -> InKernelGen ([VName], - [Imp.Exp] -> ImpM InKernel Imp.KernelOp ()))] -prepareIntermediateArraysLocal space constants num_subhistos_per_group = + [Imp.Exp] -> InKernelGen ()))] +prepareIntermediateArraysLocal num_groups group_size num_subhistos_per_group = fmap snd . mapAccumLM onOp Nothing where onOp l (SegGenRedSlug op num_subhistos subhisto_info) = do - num_subhistos <-- - toExp' int32 (spaceNumGroups space) + num_subhistos <-- toExp' int32 (unCount num_groups) emit $ Imp.DebugPrint "Number of subhistograms in global memory" $ Just (int32, Imp.vi32 num_subhistos) @@ -325,20 +325,20 @@ -- operators need locking. (l', mk_op) <- case (l, atomicUpdateLocking $ genReduceOp op) of - (_, Left f) -> return (l, return f) - (Just l', Right f) -> return (l, return $ f l') + (_, Left f) -> return (l, const $ return f) + (Just l', Right f) -> return (l, const $ return $ f l') (Nothing, Right f) -> do locks <- newVName "locks" - num_locks <- toExp $ spaceGroupSize space + num_locks <- toExp $ unCount group_size let dims = map (toExp' int32) $ Var num_subhistos_per_group : shapeDims (genReduceShape op) ++ [genReduceWidth op] l' = Locking locks 0 1 0 ((`rem` num_locks) . flattenIndex dims) - locks_t = localMemLockArray space + locks_t = localMemLockArray group_size - mk_op = do + mk_op constants = do locks_mem <- sAlloc "locks_mem" (typeSize locks_t) $ Space "local" dArray locks int32 (arrayShape locks_t) $ ArrayIn locks_mem $ IxFun.iota $ @@ -353,7 +353,7 @@ -- Initialise local-memory sub-histograms. These are -- represented as two-dimensional arrays. - let init_local_subhistos = do + let init_local_subhistos constants = do local_subhistos <- forM (genReduceDest op) $ \dest -> do dest_t <- lookupType dest @@ -363,7 +363,7 @@ sAllocArray "subhistogram_local" (elemType dest_t) sub_local_shape (Space "local") - do_op <- mk_op + do_op <- mk_op constants return (local_subhistos, do_op (Space "local") local_subhistos) @@ -376,34 +376,34 @@ genRedKernelLocal :: VName -> [PatElem ExplicitMemory] - -> KernelSpace + -> Count NumGroups SubExp -> Count GroupSize SubExp + -> SegSpace -> [SegGenRedSlug] - -> KernelBody InKernel + -> KernelBody ExplicitMemory -> CallKernelGen () -genRedKernelLocal num_subhistos_per_group_var map_pes space slugs kbody = do - (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return () - let (space_is, space_sizes) = unzip $ spaceDimensions space +genRedKernelLocal num_subhistos_per_group_var map_pes num_groups group_size space slugs kbody = do + num_groups' <- traverse toExp num_groups + group_size' <- traverse toExp group_size + let num_threads = unCount num_groups' * unCount group_size' + (space_is, space_sizes) = unzip $ unSegSpace space segment_dims = init space_sizes num_segments = length segment_dims - constants = base_constants { kernelThreadActive = true } space_sizes_64 = map (i32Toi64 . toExp' int32) space_sizes total_w_64 = product space_sizes_64 num_subhistos_per_group = Imp.var num_subhistos_per_group_var int32 emit $ Imp.DebugPrint "Number of local subhistograms per group" $ Just (int32, num_subhistos_per_group) - init_histograms <- prepareIntermediateArraysLocal space constants num_subhistos_per_group_var slugs + init_histograms <- prepareIntermediateArraysLocal num_groups group_size num_subhistos_per_group_var slugs elems_per_thread_64 <- dPrimV "elems_per_thread_64" $ total_w_64 `quotRoundingUp` - ConvOpExp (SExt Int32 Int64) (kernelNumThreads constants) - - sKernel constants "seggenred_local" $ allThreads constants $ do - init_constants + ConvOpExp (SExt Int32 Int64) num_threads + sKernelThread "seggenred_local" num_groups' group_size' (segFlat space) $ \constants -> do histograms <- forM init_histograms $ \(glob_subhistos, init_local_subhistos) -> do - (local_subhistos, do_op) <- init_local_subhistos + (local_subhistos, do_op) <- init_local_subhistos constants return (zip glob_subhistos local_subhistos, do_op) -- Find index of local subhistograms updated by this thread. We @@ -496,7 +496,7 @@ sComment "save map-out results" $ forM_ (zip map_pes map_res) $ \(pe, se) -> copyDWIM (patElemName pe) - (map ((`Imp.var` int32) . fst) $ kernelDimensions constants) se [] + (map (`Imp.var` int32) space_is) se [] forM_ (zip4 (map slugOp slugs) histograms buckets (perOp vs)) $ \(GenReduceOp dest_w _ _ shape lam, @@ -585,22 +585,25 @@ -- well as collapsing the subhistograms produced (which are always in -- global memory, but their number may vary). compileSegGenRed :: Pattern ExplicitMemory - -> KernelSpace - -> [GenReduceOp InKernel] - -> KernelBody InKernel + -> Count NumGroups SubExp -> Count GroupSize SubExp + -> SegSpace + -> [GenReduceOp ExplicitMemory] + -> KernelBody ExplicitMemory -> CallKernelGen () -compileSegGenRed (Pattern _ pes) space ops kbody = do +compileSegGenRed (Pattern _ pes) num_groups group_size space ops kbody = do + group_size' <- traverse toExp group_size + let num_red_res = length ops + sum (map (length . genReduceNeutral) ops) (all_red_pes, map_pes) = splitAt num_red_res pes let t = 8 * 4 - g <- toExp $ spaceGroupSize space + g = unCount group_size' lmax <- dPrim "lmax" int32 sOp $ Imp.GetSizeMax lmax Imp.SizeLocalMemory (op_hs, slugs) <- unzip <$> mapM (computeHistoUsage space) ops h <- fmap (`Imp.var` int32) $ - dPrimV "h" $ Imp.innerExp $ sum op_hs + dPrimV "h" $ Imp.unCount $ sum op_hs coop <- fmap (`Imp.var` int32) $ dPrimV "coop" $ h `quotRoundingUp` t @@ -613,10 +616,11 @@ emit $ Imp.DebugPrint "Memory per set of subhistograms" $ Just (int32, h) emit $ Imp.DebugPrint "Desired group size" $ Just (int32, g) - sIf (Imp.innerExp (localMemLockUsage space slugs) + h * Imp.vi32 lh .<=. Imp.vi32 lmax + sIf (Imp.unCount (localMemLockUsage group_size slugs) + h * Imp.vi32 lh + .<=. Imp.vi32 lmax .&&. coop .<=. g) - (genRedKernelLocal lh map_pes space slugs kbody) - (genRedKernelGlobal map_pes space slugs kbody) + (genRedKernelLocal lh map_pes num_groups group_size space slugs kbody) + (genRedKernelGlobal map_pes num_groups group_size space slugs kbody) let pes_per_op = chunks (map (length . genReduceDest) ops) all_red_pes @@ -645,21 +649,21 @@ subhistogram_id <- newVName "subhistogram_id" vector_ids <- mapM (const $ newVName "vector_id") $ shapeDims $ genReduceShape op - gtid <- newVName $ baseString $ spaceGlobalId space - let segred_space = - space { spaceStructure = - FlatThreadSpace $ - segment_dims ++ - [(bucket_id, num_buckets)] ++ - zip vector_ids (shapeDims $ genReduceShape op) ++ - [(subhistogram_id, Var num_histos)] - , spaceGlobalId = gtid - } + flat_gtid <- newVName "flat_gtid" + + let lvl = SegThread num_groups group_size SegVirt + segred_space = + SegSpace flat_gtid $ + segment_dims ++ + [(bucket_id, num_buckets)] ++ + zip vector_ids (shapeDims $ genReduceShape op) ++ + [(subhistogram_id, Var num_histos)] + let segred_op = SegRedOp Commutative (genReduceOp op) (genReduceNeutral op) mempty - compileSegRed' (Pattern [] red_pes) segred_space [segred_op] $ \_ red_cont -> + compileSegRed' (Pattern [] red_pes) lvl segred_space [segred_op] $ \_ red_cont -> red_cont $ flip map subhistos $ \subhisto -> (Var subhisto, map (`Imp.var` int32) $ map fst segment_dims ++ [subhistogram_id, bucket_id] ++ vector_ids) - where segment_dims = init $ spaceDimensions space + where segment_dims = init $ unSegSpace space diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs @@ -1,7 +1,8 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleContexts #-} module Futhark.CodeGen.ImpGen.Kernels.SegMap - ( compileSegMap ) where + ( compileSegMap ) +where import Control.Monad.Except @@ -9,18 +10,52 @@ import Futhark.Representation.ExplicitMemory import Futhark.CodeGen.ImpGen.Kernels.Base +import Futhark.CodeGen.ImpGen +import qualified Futhark.CodeGen.ImpCode.Kernels as Imp +import Futhark.Util.IntegralExp (quotRoundingUp) --- | Compile 'SegMap' instance to host-level code with calls to --- various kernels. +-- | Compile 'SegMap' instance code. compileSegMap :: Pattern ExplicitMemory - -> KernelSpace - -> KernelBody InKernel + -> SegLevel + -> SegSpace + -> KernelBody ExplicitMemory -> CallKernelGen () -compileSegMap pat space kbody = do - (constants, init_constants) <- kernelInitialisation space - sKernel constants "segmap" $ do - init_constants - compileKernelStms constants (kernelBodyStms kbody) $ - zipWithM_ (compileKernelResult constants) (patternElements pat) $ - kernelBodyResult kbody +compileSegMap _ SegThreadScalar{} _ _ = + fail "compileSegMap: SegThreadScalar cannot be compiled at top level." + +compileSegMap pat lvl space kbody = do + let (is, dims) = unzip $ unSegSpace space + dims' <- mapM toExp dims + + num_groups' <- traverse toExp $ segNumGroups lvl + group_size' <- traverse toExp $ segGroupSize lvl + + case lvl of + SegThreadScalar{} -> + fail "compileSegMap: SegThreadScalar cannot be compiled at top level." + + SegThread{} -> + sKernelThread "segmap" num_groups' group_size' (segFlat space) $ \constants -> do + let virt_num_groups = product dims' `quotRoundingUp` unCount group_size' + virtualiseGroups constants (segVirt lvl) virt_num_groups $ \group_id -> do + let global_tid = Imp.vi32 group_id * unCount group_size' + + kernelLocalThreadId constants + + zipWithM_ dPrimV_ is $ unflattenIndex dims' global_tid + + sWhen (isActive $ unSegSpace space) $ + compileStms mempty (kernelBodyStms kbody) $ + zipWithM_ (compileThreadResult space constants) (patternElements pat) $ + kernelBodyResult kbody + + SegGroup{} -> + sKernelGroup "segmap_intragroup" num_groups' group_size' (segFlat space) $ \constants -> do + let virt_num_groups = product dims' + virtualiseGroups constants (segVirt lvl) virt_num_groups $ \group_id -> do + + zipWithM_ dPrimV_ is $ unflattenIndex dims' $ Imp.vi32 group_id + + compileStms mempty (kernelBodyStms kbody) $ + zipWithM_ (compileGroupResult space constants) (patternElements pat) $ + kernelBodyResult kbody diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs @@ -44,6 +44,7 @@ module Futhark.CodeGen.ImpGen.Kernels.SegRed ( compileSegRed , compileSegRed' + , DoSegBody ) where @@ -74,41 +75,42 @@ -- | Compile 'SegRed' instance to host-level code with calls to -- various kernels. compileSegRed :: Pattern ExplicitMemory - -> KernelSpace - -> [SegRedOp InKernel] - -> KernelBody InKernel + -> SegLevel -> SegSpace + -> [SegRedOp ExplicitMemory] + -> KernelBody ExplicitMemory -> CallKernelGen () -compileSegRed pat space reds body = - compileSegRed' pat space reds $ \constants red_cont -> - compileKernelStms constants (kernelBodyStms body) $ do +compileSegRed pat lvl space reds body = + compileSegRed' pat lvl space reds $ \constants red_cont -> + compileStms mempty (kernelBodyStms body) $ do let (red_res, map_res) = splitAt (segRedResults reds) $ kernelBodyResult body sComment "save map-out results" $ do let map_arrs = drop (segRedResults reds) $ patternElements pat - zipWithM_ (compileKernelResult constants) map_arrs map_res + zipWithM_ (compileThreadResult space constants) map_arrs map_res red_cont $ zip (map kernelResultSubExp red_res) $ repeat [] -- | Like 'compileSegRed', but where the body is a monadic action. compileSegRed' :: Pattern ExplicitMemory - -> KernelSpace - -> [SegRedOp InKernel] + -> SegLevel -> SegSpace + -> [SegRedOp ExplicitMemory] -> DoSegBody -> CallKernelGen () -compileSegRed' pat space reds body +compileSegRed' pat lvl space reds body | genericLength reds > maxNumOps = compilerLimitationS $ "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported." - | [(_, Constant (IntValue (Int32Value 1))), _] <- spaceDimensions space = - nonsegmentedReduction pat space reds body + | [(_, Constant (IntValue (Int32Value 1))), _] <- unSegSpace space = + nonsegmentedReduction pat num_groups group_size space reds body | otherwise = do - segment_size <- - toExp $ last $ map snd $ spaceDimensions space - group_size <- toExp $ spaceGroupSize space - let use_small_segments = segment_size * 2 .<. group_size + group_size' <- toExp $ unCount group_size + segment_size <- toExp $ last $ segSpaceDims space + let use_small_segments = segment_size * 2 .<. group_size' sIf use_small_segments - (smallSegmentsReduction pat space reds body) - (largeSegmentsReduction pat space reds body) + (smallSegmentsReduction pat num_groups group_size space reds body) + (largeSegmentsReduction pat num_groups group_size space reds body) + where num_groups = segNumGroups lvl + group_size = segGroupSize lvl -- | Prepare intermediate arrays for the reduction. Prim-typed -- arguments go in local memory (so we need to do the allocation of @@ -116,19 +118,21 @@ -- global memory. Allocations for the former have already been -- performed. This policy is baked into how the allocations are done -- in ExplicitAllocations. -intermediateArrays :: KernelSpace -> SegRedOp InKernel -> InKernelGen [VName] -intermediateArrays space (SegRedOp _ red_op nes _) = do +intermediateArrays :: Count GroupSize SubExp -> SubExp + -> SegRedOp ExplicitMemory + -> InKernelGen [VName] +intermediateArrays (Count group_size) num_threads (SegRedOp _ red_op nes _) = do let red_op_params = lambdaParams red_op (red_acc_params, _) = splitAt (length nes) red_op_params forM red_acc_params $ \p -> case paramAttr p of MemArray pt shape _ (ArrayIn mem _) -> do - let shape' = Shape [spaceNumThreads space] <> shape + let shape' = Shape [num_threads] <> shape sArray "red_arr" pt shape' $ ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape' _ -> do let pt = elemType $ paramType p - shape = Shape [spaceGroupSize space] + shape = Shape [group_size] sAllocArray "red_arr" pt shape $ Space "local" -- | Arrays for storing group results. @@ -137,8 +141,10 @@ -- because they are also used for keeping vectorised accumulators for -- first-stage reduction, if necessary. When actually storing group -- results, the first index is set to 0. -groupResultArrays :: SubExp -> SubExp -> [SegRedOp InKernel] -> CallKernelGen [[VName]] -groupResultArrays virt_num_groups group_size reds = +groupResultArrays :: Count NumGroups SubExp -> Count GroupSize SubExp + -> [SegRedOp ExplicitMemory] + -> CallKernelGen [[VName]] +groupResultArrays (Count virt_num_groups) (Count group_size) reds = forM reds $ \(SegRedOp _ lam _ shape) -> forM (lambdaReturnType lam) $ \t -> do let pt = elemType t @@ -149,47 +155,43 @@ sAllocArrayPerm "group_res_arr" pt full_shape (Space "device") perm nonsegmentedReduction :: Pattern ExplicitMemory - -> KernelSpace - -> [SegRedOp InKernel] + -> Count NumGroups SubExp -> Count GroupSize SubExp -> SegSpace + -> [SegRedOp ExplicitMemory] -> DoSegBody -> CallKernelGen () -nonsegmentedReduction segred_pat space reds body = do - (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return () - let constants = base_constants { kernelThreadActive = true } - global_tid = kernelGlobalThreadId constants - (_, w) = last $ spaceDimensions space +nonsegmentedReduction segred_pat num_groups group_size space reds body = do + let (gtids, dims) = unzip $ unSegSpace space + dims' <- mapM toExp dims + num_groups' <- traverse toExp num_groups + group_size' <- traverse toExp group_size + + let global_tid = Imp.vi32 $ segFlat space + w = last dims' + counter <- sStaticArray "counter" (Space "device") int32 $ Imp.ArrayValues $ replicate (fromIntegral maxNumOps) $ IntValue $ Int32Value 0 - -- The group-result arrays have an extra dimension (of size - -- groupsize) because they are also used for keeping vectorised - -- accumulators for first-stage reduction, if necessary. When - -- actually storing group results, the first index is set to 0. - reds_group_res_arrs <- groupResultArrays (spaceNumGroups space) (spaceGroupSize space) reds - - num_threads <- dPrimV "num_threads" $ kernelNumThreads constants + reds_group_res_arrs <- groupResultArrays num_groups group_size reds - sKernel constants "segred_nonseg" $ allThreads constants $ do - init_constants + num_threads <- dPrimV "num_threads" $ unCount num_groups' * unCount group_size' + sKernelThread "segred_nonseg" num_groups' group_size' (segFlat space) $ \constants -> do sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local" - reds_arrs <- mapM (intermediateArrays space) reds + reds_arrs <- mapM (intermediateArrays group_size (Var num_threads)) reds -- Since this is the nonsegmented case, all outer segment IDs must -- necessarily be 0. - let gtids = map fst $ spaceDimensions space - forM_ (init gtids) $ \v -> - v <-- 0 + forM_ gtids $ \v -> dPrimV_ v 0 - num_elements <- Imp.elements <$> toExp w + let num_elements = Imp.elements w let elems_per_thread = num_elements `quotRoundingUp` Imp.elements (kernelNumThreads constants) slugs <- mapM (segRedOpSlug (kernelLocalThreadId constants) (kernelGroupId constants)) $ zip3 reds reds_arrs reds_group_res_arrs reds_op_renamed <- - reductionStageOne constants num_elements + reductionStageOne constants (zip gtids dims') num_elements global_tid elems_per_thread num_threads slugs body @@ -206,24 +208,26 @@ sync_arr group_res_arrs red_arrs smallSegmentsReduction :: Pattern ExplicitMemory - -> KernelSpace - -> [SegRedOp InKernel] + -> Count NumGroups SubExp -> Count GroupSize SubExp + -> SegSpace + -> [SegRedOp ExplicitMemory] -> DoSegBody -> CallKernelGen () -smallSegmentsReduction (Pattern _ segred_pes) space reds body = do - (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return () - let constants = base_constants { kernelThreadActive = true } - - let (gtids, dims) = unzip $ spaceDimensions space +smallSegmentsReduction (Pattern _ segred_pes) num_groups group_size space reds body = do + let (gtids, dims) = unzip $ unSegSpace space dims' <- mapM toExp dims let segment_size = last dims' -- Careful to avoid division by zero now. segment_size_nonzero_v <- dPrimV "segment_size_nonzero" $ BinOpExp (SMax Int32) 1 segment_size + + num_groups' <- traverse toExp num_groups + group_size' <- traverse toExp group_size + num_threads <- dPrimV "num_threads" $ unCount num_groups' * unCount group_size' let segment_size_nonzero = Imp.var segment_size_nonzero_v int32 num_segments = product $ init dims' - segments_per_group = kernelGroupSize constants `quot` segment_size_nonzero + segments_per_group = unCount group_size' `quot` segment_size_nonzero required_groups = num_segments `quotRoundingUp` segments_per_group emit $ Imp.DebugPrint "\n# SegRed-small" Nothing @@ -232,15 +236,14 @@ emit $ Imp.DebugPrint "segments_per_group" $ Just (int32, segments_per_group) emit $ Imp.DebugPrint "required_groups" $ Just (int32, required_groups) - sKernel constants "segred_small" $ allThreads constants $ do - init_constants + sKernelThread "segred_small" num_groups' group_size' (segFlat space) $ \constants -> do - reds_arrs <- mapM (intermediateArrays space) reds + reds_arrs <- mapM (intermediateArrays group_size (Var num_threads)) reds -- We probably do not have enough actual workgroups to cover the -- entire iteration space. Some groups thus have to perform double -- duty; we put an outer loop to accomplish this. - virtualiseGroups constants required_groups $ \group_id_var' -> do + virtualiseGroups constants SegVirt required_groups $ \group_id_var' -> do let group_id' = Imp.vi32 group_id_var' -- Compute the 'n' input indices. The outer 'n-1' correspond to -- the segment ID, and are computed from the group id. The inner @@ -249,8 +252,8 @@ segment_index = (ltid `quot` segment_size_nonzero) + (group_id' * segments_per_group) index_within_segment = ltid `rem` segment_size - zipWithM_ (<--) (init gtids) $ unflattenIndex (init dims') segment_index - last gtids <-- index_within_segment + zipWithM_ dPrimV_ (init gtids) $ unflattenIndex (init dims') segment_index + dPrimV_ (last gtids) index_within_segment let out_of_bounds = forM_ (zip reds reds_arrs) $ \(SegRedOp _ _ nes _, red_arrs) -> @@ -294,38 +297,41 @@ sOp Imp.LocalBarrier largeSegmentsReduction :: Pattern ExplicitMemory - -> KernelSpace - -> [SegRedOp InKernel] + -> Count NumGroups SubExp -> Count GroupSize SubExp + -> SegSpace + -> [SegRedOp ExplicitMemory] -> DoSegBody -> CallKernelGen () -largeSegmentsReduction segred_pat space reds body = do - (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return () - let (gtids, dims) = unzip $ spaceDimensions space +largeSegmentsReduction segred_pat num_groups group_size space reds body = do + let (gtids, dims) = unzip $ unSegSpace space dims' <- mapM toExp dims let segment_size = last dims' num_segments = product $ init dims' + num_groups' <- traverse toExp num_groups + group_size' <- traverse toExp group_size + let (groups_per_segment, elems_per_thread) = groupsPerSegmentAndElementsPerThread segment_size num_segments - (kernelNumGroups base_constants) (kernelGroupSize base_constants) - virt_num_groups <- dPrimV "virt_num_groups" $ + num_groups' group_size' + virt_num_groups <- dPrimV "vit_num_groups" $ groups_per_segment * num_segments - threads_per_segment <- dPrimV "thread_per_segment" $ - groups_per_segment * kernelGroupSize base_constants + num_threads <- dPrimV "num_threads" $ unCount num_groups' * unCount group_size' - let constants = base_constants { kernelThreadActive = true } + threads_per_segment <- dPrimV "thread_per_segment" $ + groups_per_segment * unCount group_size' emit $ Imp.DebugPrint "\n# SegRed-large" Nothing emit $ Imp.DebugPrint "num_segments" $ Just (int32, num_segments) emit $ Imp.DebugPrint "segment_size" $ Just (int32, segment_size) emit $ Imp.DebugPrint "virt_num_groups" $ Just (int32, Imp.vi32 virt_num_groups) - emit $ Imp.DebugPrint "num_groups" $ Just (int32, kernelNumGroups constants) - emit $ Imp.DebugPrint "group_size" $ Just (int32, kernelGroupSize constants) - emit $ Imp.DebugPrint "elems_per_thread" $ Just (int32, Imp.innerExp elems_per_thread) + emit $ Imp.DebugPrint "num_groups" $ Just (int32, Imp.unCount num_groups') + emit $ Imp.DebugPrint "group_size" $ Just (int32, Imp.unCount group_size') + emit $ Imp.DebugPrint "elems_per_thread" $ Just (int32, Imp.unCount elems_per_thread) emit $ Imp.DebugPrint "groups_per_segment" $ Just (int32, groups_per_segment) - reds_group_res_arrs <- groupResultArrays (Var virt_num_groups) (spaceGroupSize space) reds + reds_group_res_arrs <- groupResultArrays (Count (Var virt_num_groups)) group_size reds -- In principle we should have a counter for every segment. Since -- the number of segments is a dynamic quantity, we would have to @@ -342,35 +348,33 @@ sStaticArray "counter" (Space "device") int32 $ Imp.ArrayZeros num_counters - sKernel constants "segred_large" $ allThreads constants $ do - init_constants + sKernelThread "segred_large" num_groups' group_size' (segFlat space) $ \constants -> do - reds_arrs <- mapM (intermediateArrays space) reds + reds_arrs <- mapM (intermediateArrays group_size (Var num_threads)) reds sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local" -- We probably do not have enough actual workgroups to cover the -- entire iteration space. Some groups thus have to perform double -- duty; we put an outer loop to accomplish this. - virtualiseGroups constants (Imp.vi32 virt_num_groups) $ \group_id_var -> do - + virtualiseGroups constants SegVirt (Imp.vi32 virt_num_groups) $ \group_id_var -> do let segment_gtids = init gtids group_id = Imp.vi32 group_id_var - group_size = kernelGroupSize constants flat_segment_id = group_id `quot` groups_per_segment local_tid = kernelLocalThreadId constants - global_tid = (group_id * group_size + local_tid) - `rem` (group_size * groups_per_segment) + global_tid = (group_id * unCount group_size' + local_tid) + `rem` (unCount group_size' * groups_per_segment) w = last dims first_group_for_segment = flat_segment_id * groups_per_segment - zipWithM_ (<--) segment_gtids $ unflattenIndex (init dims') flat_segment_id + zipWithM_ dPrimV_ segment_gtids $ unflattenIndex (init dims') flat_segment_id + dPrim_ (last gtids) int32 num_elements <- Imp.elements <$> toExp w slugs <- mapM (segRedOpSlug local_tid group_id) $ zip3 reds reds_arrs reds_group_res_arrs reds_op_renamed <- - reductionStageOne constants num_elements + reductionStageOne constants (zip gtids dims') num_elements global_tid elems_per_thread threads_per_segment slugs body @@ -399,20 +403,22 @@ sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment --- Careful to avoid division by zero here. -groupsPerSegmentAndElementsPerThread :: Imp.Exp -> Imp.Exp -> Imp.Exp -> Imp.Exp - -> (Imp.Exp, Imp.Count Imp.Elements) +-- Careful to avoid division by zero here. We have at least one group +-- per segment. +groupsPerSegmentAndElementsPerThread :: Imp.Exp -> Imp.Exp + -> Count NumGroups Imp.Exp -> Count GroupSize Imp.Exp + -> (Imp.Exp, Imp.Count Imp.Elements Imp.Exp) groupsPerSegmentAndElementsPerThread segment_size num_segments num_groups_hint group_size = let groups_per_segment = - num_groups_hint `quotRoundingUp` BinOpExp (SMax Int32) 1 num_segments + unCount num_groups_hint `quotRoundingUp` BinOpExp (SMax Int32) 1 num_segments elements_per_thread = - segment_size `quotRoundingUp` (group_size * groups_per_segment) + segment_size `quotRoundingUp` (unCount group_size * groups_per_segment) in (groups_per_segment, Imp.elements elements_per_thread) -- | A SegRedOp with auxiliary information. data SegRedOpSlug = SegRedOpSlug - { slugOp :: SegRedOp InKernel + { slugOp :: SegRedOp ExplicitMemory , slugArrs :: [VName] -- ^ The arrays used for computing the intra-group reduction -- (either local or global memory). @@ -420,10 +426,10 @@ -- ^ Places to store accumulator in stage 1 reduction. } -slugBody :: SegRedOpSlug -> Body InKernel +slugBody :: SegRedOpSlug -> Body ExplicitMemory slugBody = lambdaBody . segRedLambda . slugOp -slugParams :: SegRedOpSlug -> [LParam InKernel] +slugParams :: SegRedOpSlug -> [LParam ExplicitMemory] slugParams = lambdaParams . segRedLambda . slugOp slugNeutral :: SegRedOpSlug -> [SubExp] @@ -435,11 +441,11 @@ slugsComm :: [SegRedOpSlug] -> Commutativity slugsComm = mconcat . map (segRedComm . slugOp) -accParams, nextParams :: SegRedOpSlug -> [LParam InKernel] +accParams, nextParams :: SegRedOpSlug -> [LParam ExplicitMemory] accParams slug = take (length (slugNeutral slug)) $ slugParams slug nextParams slug = drop (length (slugNeutral slug)) $ slugParams slug -segRedOpSlug :: Imp.Exp -> Imp.Exp -> (SegRedOp InKernel, [VName], [VName]) -> InKernelGen SegRedOpSlug +segRedOpSlug :: Imp.Exp -> Imp.Exp -> (SegRedOp ExplicitMemory, [VName], [VName]) -> InKernelGen SegRedOpSlug segRedOpSlug local_tid group_id (op, group_res_arrs, param_arrs) = SegRedOpSlug op group_res_arrs <$> zipWithM mkAcc (lambdaParams (segRedLambda op)) param_arrs @@ -452,16 +458,16 @@ return (param_arr, [local_tid, group_id]) reductionStageZero :: KernelConstants - -> Imp.Count Imp.Elements + -> [(VName, Imp.Exp)] + -> Imp.Count Imp.Elements Imp.Exp -> Imp.Exp - -> Imp.Count Imp.Elements + -> Imp.Count Imp.Elements Imp.Exp -> VName -> [SegRedOpSlug] -> DoSegBody - -> InKernelGen ([Lambda InKernel], InKernelGen ()) -reductionStageZero constants num_elements global_tid elems_per_thread threads_per_segment slugs body = do - - let (gtids, _dims) = unzip $ kernelDimensions constants + -> InKernelGen ([Lambda ExplicitMemory], InKernelGen ()) +reductionStageZero constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body = do + let (gtids, _dims) = unzip ispace gtid = last gtids local_tid = kernelLocalThreadId constants @@ -512,8 +518,8 @@ (bound, check_bounds) = case comm of Commutative -> (Imp.var chunk_size int32, id) - Noncommutative -> (Imp.innerExp elems_per_thread, - sWhen (Imp.var gtid int32 .<. Imp.innerExp num_elements)) + Noncommutative -> (Imp.unCount elems_per_thread, + sWhen (Imp.var gtid int32 .<. Imp.unCount num_elements)) sFor i Int32 bound $ do gtid <-- @@ -524,7 +530,7 @@ Noncommutative -> let index_in_segment = global_tid `quot` kernelGroupSize constants in local_tid + - (index_in_segment * Imp.innerExp elems_per_thread + Imp.var i int32) * + (index_in_segment * Imp.unCount elems_per_thread + Imp.var i int32) * kernelGroupSize constants check_bounds $ sComment "apply map function" $ @@ -563,16 +569,17 @@ return (slugs_op_renamed, doTheReduction) reductionStageOne :: KernelConstants - -> Imp.Count Imp.Elements + -> [(VName, Imp.Exp)] + -> Imp.Count Imp.Elements Imp.Exp -> Imp.Exp - -> Imp.Count Imp.Elements + -> Imp.Count Imp.Elements Imp.Exp -> VName -> [SegRedOpSlug] -> DoSegBody - -> InKernelGen [Lambda InKernel] -reductionStageOne constants num_elements global_tid elems_per_thread threads_per_segment slugs body = do + -> InKernelGen [Lambda ExplicitMemory] +reductionStageOne constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body = do (slugs_op_renamed, doTheReduction) <- - reductionStageZero constants num_elements global_tid elems_per_thread threads_per_segment slugs body + reductionStageZero constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body case slugsComm slugs of Noncommutative -> @@ -591,8 +598,8 @@ -> Imp.Exp -> Imp.Exp -> SegRedOpSlug - -> [LParam InKernel] - -> Lambda InKernel -> [SubExp] + -> [LParam ExplicitMemory] + -> Lambda ExplicitMemory -> [SubExp] -> Imp.Exp -> VName -> Imp.Exp -> VName -> [VName] -> [VName] -> InKernelGen () reductionStageTwo constants segred_pes @@ -618,6 +625,7 @@ sWrite sync_arr [0] $ Imp.var old_counter int32 .==. groups_per_segment - 1 sOp Imp.LocalBarrier + sOp Imp.GlobalBarrier is_last_group <- dPrim "is_last_group" Bool copyDWIM is_last_group [] (Var sync_arr) [0] diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs @@ -19,9 +19,9 @@ import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem) -makeLocalArrays :: SubExp -> SubExp -> [SubExp] -> Lambda InKernel +makeLocalArrays :: Count GroupSize SubExp -> SubExp -> [SubExp] -> Lambda ExplicitMemory -> InKernelGen [VName] -makeLocalArrays group_size num_threads nes scan_op = do +makeLocalArrays (Count group_size) num_threads nes scan_op = do let (scan_x_params, _scan_y_params) = splitAt (length nes) $ lambdaParams scan_op forM scan_x_params $ \p -> @@ -39,18 +39,21 @@ -- | Produce partially scanned intervals; one per workgroup. scanStage1 :: Pattern ExplicitMemory - -> KernelSpace - -> Lambda InKernel -> [SubExp] - -> KernelBody InKernel + -> Count NumGroups SubExp -> Count GroupSize SubExp -> SegSpace + -> Lambda ExplicitMemory -> [SubExp] + -> KernelBody ExplicitMemory -> CallKernelGen (Imp.Exp, CrossesSegment) -scanStage1 (Pattern _ pes) space scan_op nes kbody = do - (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return () - let (gtids, dims) = unzip $ spaceDimensions space +scanStage1 (Pattern _ pes) num_groups group_size space scan_op nes kbody = do + num_groups' <- traverse toExp num_groups + group_size' <- traverse toExp group_size + num_threads <- dPrimV "num_threads" $ + unCount num_groups' * unCount group_size' + + let (gtids, dims) = unzip $ unSegSpace space dims' <- mapM toExp dims - let constants = base_constants { kernelThreadActive = true } - num_elements = product dims' - elems_per_thread = num_elements `quotRoundingUp` kernelNumThreads constants - elems_per_group = kernelGroupSize constants * elems_per_thread + let num_elements = product dims' + elems_per_thread = num_elements `quotRoundingUp` Imp.vi32 num_threads + elems_per_group = unCount group_size' * elems_per_thread -- Squirrel away a copy of the operator with unique names that we -- can pass to groupScan. @@ -62,12 +65,9 @@ (to-from) .>. (to `rem` segment_size) _ -> Nothing - sKernel constants "scan_stage1" $ allThreads constants $ do - init_constants - + sKernelThread "scan_stage1" num_groups' group_size' (segFlat space) $ \constants -> do local_arrs <- - makeLocalArrays (spaceGroupSize space) (spaceNumThreads space) - nes scan_op + makeLocalArrays group_size (Var num_threads) nes scan_op -- The variables from scan_op will be used for the carry and such -- in the big chunking loop. @@ -86,7 +86,7 @@ flat_idx <- dPrimV "flat_idx" $ Imp.var chunk_offset int32 + kernelLocalThreadId constants -- Construct segment indices. - zipWithM_ (<--) gtids $ unflattenIndex dims' $ Imp.var flat_idx int32 + zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ Imp.var flat_idx int32 let in_bounds = foldl1 (.&&.) $ zipWith (.<.) (map (`Imp.var` int32) gtids) dims' @@ -148,26 +148,23 @@ return (elems_per_group, crossesSegment) scanStage2 :: Pattern ExplicitMemory - -> Imp.Exp -> CrossesSegment -> KernelSpace - -> Lambda InKernel -> [SubExp] + -> Imp.Exp -> Count NumGroups SubExp -> CrossesSegment -> SegSpace + -> Lambda ExplicitMemory -> [SubExp] -> CallKernelGen () -scanStage2 (Pattern _ pes) elems_per_group crossesSegment space scan_op nes = do - -- A single group, with one thread for each group in stage 1. - group_size <- toExp $ spaceNumGroups space - (constants, init_constants) <- - kernelInitialisationSimple 1 group_size Nothing +scanStage2 (Pattern _ pes) elems_per_group num_groups crossesSegment space scan_op nes = do + -- Our group size is the number of groups for the stage 1 kernel. + let group_size = Count $ unCount num_groups + group_size' <- traverse toExp group_size - let (gtids, dims) = unzip $ spaceDimensions space + let (gtids, dims) = unzip $ unSegSpace space dims' <- mapM toExp dims let crossesSegment' = do f <- crossesSegment Just $ \from to -> f ((from + 1) * elems_per_group - 1) ((to + 1) * elems_per_group - 1) - sKernel constants "scan_stage2" $ do - init_constants - - local_arrs <- makeLocalArrays (spaceNumGroups space) (spaceNumGroups space) + sKernelThread "scan_stage2" 1 group_size' (segFlat space) $ \constants -> do + local_arrs <- makeLocalArrays group_size (unCount group_size) nes scan_op flat_idx <- dPrimV "flat_idx" $ @@ -195,15 +192,14 @@ (Var arr) [kernelLocalThreadId constants] scanStage3 :: Pattern ExplicitMemory - -> Imp.Exp -> CrossesSegment -> KernelSpace - -> Lambda InKernel -> [SubExp] + -> Imp.Exp -> CrossesSegment -> SegSpace + -> Lambda ExplicitMemory -> [SubExp] -> CallKernelGen () scanStage3 (Pattern _ pes) elems_per_group crossesSegment space scan_op nes = do - let (gtids, dims) = unzip $ spaceDimensions space + let (gtids, dims) = unzip $ unSegSpace space dims' <- mapM toExp dims - (constants, init_constants) <- simpleKernelConstants (product dims') "scan" - sKernel constants "scan_stage3" $ do - init_constants + sKernelSimple "scan_stage3" (product dims') $ \constants -> do + dPrimV_ (segFlat space) $ kernelGlobalThreadId constants -- Compute our logical index. zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ kernelGlobalThreadId constants -- Figure out which group this element was originally in. @@ -242,17 +238,18 @@ -- | Compile 'SegScan' instance to host-level code with calls to -- various kernels. compileSegScan :: Pattern ExplicitMemory - -> KernelSpace - -> Lambda InKernel -> [SubExp] - -> KernelBody InKernel + -> SegLevel -> SegSpace + -> Lambda ExplicitMemory -> [SubExp] + -> KernelBody ExplicitMemory -> CallKernelGen () -compileSegScan pat space scan_op nes kbody = do - (elems_per_group, crossesSegment) <- scanStage1 pat space scan_op nes kbody +compileSegScan pat lvl space scan_op nes kbody = do + (elems_per_group, crossesSegment) <- + scanStage1 pat (segNumGroups lvl) (segGroupSize lvl) space scan_op nes kbody emit $ Imp.DebugPrint "\n# SegScan" Nothing emit $ Imp.DebugPrint "elems_per_group" $ Just (int32, elems_per_group) scan_op' <- renameLambda scan_op scan_op'' <- renameLambda scan_op - scanStage2 pat elems_per_group crossesSegment space scan_op' nes + scanStage2 pat elems_per_group (segNumGroups lvl) crossesSegment space scan_op' nes scanStage3 pat elems_per_group crossesSegment space scan_op'' nes diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs b/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs @@ -28,6 +28,7 @@ import Futhark.CodeGen.ImpCode.OpenCL hiding (Program) import qualified Futhark.CodeGen.ImpCode.OpenCL as ImpOpenCL import Futhark.MonadFreshNames +import Futhark.Representation.ExplicitMemory (allScalarMemory) import Futhark.Util (zEncodeString) import Futhark.Util.Pretty (pretty) @@ -169,7 +170,7 @@ mem_aligned <- newVName $ baseString mem ++ "_aligned" return (Just $ SharedMemoryKArg size, Just [C.cparam|__local volatile typename int64_t* $id:mem_aligned|], - [C.citem|__local volatile char* restrict $id:mem = $id:mem_aligned;|]) + [C.citem|__local volatile char* restrict $id:mem = (__local volatile char*)$id:mem_aligned;|]) prepareLocalMemory TargetOpenCL (mem, Right size) = do let size' = compilePrimExp size return (Nothing, Nothing, @@ -262,7 +263,7 @@ cudaAtomicOps :: [C.Definition] -cudaAtomicOps = (return mkOp <*> opNames <*> types) ++ extraOps +cudaAtomicOps = (mkOp <$> opNames <*> types) ++ extraOps where mkOp (clName, cuName) t = [C.cedecl|static inline $ty:t $id:clName(volatile $ty:t *p, $ty:t val) { @@ -294,7 +295,7 @@ typedef char int8_t; typedef short int16_t; typedef int int32_t; -typedef long int64_t; +typedef long long int64_t; typedef unsigned char uint8_t; typedef unsigned short uint16_t; typedef unsigned int uint32_t; @@ -408,8 +409,8 @@ inKernelOperations = GenericC.Operations { GenericC.opsCompiler = kernelOps , GenericC.opsMemoryType = kernelMemoryType - , GenericC.opsWriteScalar = GenericC.writeScalarPointerWithQuals pointerQuals - , GenericC.opsReadScalar = GenericC.readScalarPointerWithQuals pointerQuals + , GenericC.opsWriteScalar = kernelWriteScalar + , GenericC.opsReadScalar = kernelReadScalar , GenericC.opsAllocate = cannotAllocate , GenericC.opsDeallocate = cannotDeallocate , GenericC.opsCopy = copyInKernel @@ -438,7 +439,7 @@ kernelOps MemFenceGlobal = GenericC.stm [C.cstm|mem_fence_global();|] kernelOps (PrivateAlloc name size) = do - size' <- GenericC.compileExp $ innerExp size + size' <- GenericC.compileExp $ unCount size name' <- newVName $ pretty name ++ "_backing" GenericC.item [C.citem|__private char $id:name'[$exp:size'];|] GenericC.stm [C.cstm|$id:name = $id:name';|] @@ -455,7 +456,7 @@ return [C.cty|$tyquals:(volatile++quals) $ty:t|] doAtomic s old arr ind val op ty = do - ind' <- GenericC.compileExp $ innerExp ind + ind' <- GenericC.compileExp $ unCount ind val' <- GenericC.compileExp val cast <- atomicCast s ty GenericC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], ($ty:ty) $exp:val');|] @@ -485,14 +486,14 @@ doAtomic s old arr ind val "atomic_xor" [C.cty|unsigned int|] atomicOps s (AtomicCmpXchg old arr ind cmp val) = do - ind' <- GenericC.compileExp $ innerExp ind + ind' <- GenericC.compileExp $ unCount ind cmp' <- GenericC.compileExp cmp val' <- GenericC.compileExp val cast <- atomicCast s [C.cty|int|] GenericC.stm [C.cstm|$id:old = atomic_cmpxchg(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:cmp', $exp:val');|] atomicOps s (AtomicXchg old arr ind val) = do - ind' <- GenericC.compileExp $ innerExp ind + ind' <- GenericC.compileExp $ unCount ind val' <- GenericC.compileExp val cast <- atomicCast s [C.cty|int|] GenericC.stm [C.cstm|$id:old = atomic_xchg(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:val');|] @@ -513,9 +514,29 @@ noStaticArrays _ _ _ _ = fail "Cannot create static array in kernel." + kernelMemoryType space + | Just t <- M.lookup space allScalarMemory = + return $ GenericC.primTypeToCType t + kernelMemoryType space = do quals <- pointerQuals space return [C.cty|$tyquals:quals $ty:defaultMemBlockType|] + + kernelWriteScalar dest _ _ space _ v + | space `M.member` allScalarMemory = + GenericC.stm [C.cstm|$exp:dest = $exp:v;|] + + kernelWriteScalar dest i elemtype space vol v = + GenericC.writeScalarPointerWithQuals pointerQuals + dest i elemtype space vol v + + kernelReadScalar dest _ _ space _ + | space `M.member` allScalarMemory = + return dest + + kernelReadScalar dest i elemtype space vol = + GenericC.readScalarPointerWithQuals pointerQuals + dest i elemtype space vol --- Checking requirements diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs b/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs @@ -6,13 +6,11 @@ ) where -import qualified Data.Set as S - import Prelude hiding (quot, rem) import Futhark.CodeGen.ImpCode.Kernels import Futhark.Representation.AST.Attributes.Types -import Futhark.Representation.AST.Attributes.Names (freeIn) +import Futhark.Representation.AST.Attributes.Names (freeIn, namesToList) import Futhark.Util.IntegralExp (IntegralExp, quot, rem, quotRoundingUp) -- | Which form of transposition to generate code for. @@ -256,7 +254,7 @@ , [actual_dim, actual_dim `quot` elemsPerThread, 1]) uses = map (`ScalarUse` int32) - (S.toList $ mconcat $ map freeIn + (namesToList $ mconcat $ map freeIn [basic_odata_offset, basic_idata_offset, num_arrays, width, height, input_size, output_size, mulx, muly]) ++ map MemoryUse [odata, idata] diff --git a/src/Futhark/CodeGen/SetDefaultSpace.hs b/src/Futhark/CodeGen/SetDefaultSpace.hs --- a/src/Futhark/CodeGen/SetDefaultSpace.hs +++ b/src/Futhark/CodeGen/SetDefaultSpace.hs @@ -40,7 +40,7 @@ setBodySpace :: Space -> Code op -> Code op setBodySpace space (Allocate v e old_space) = - Allocate v (setCountSpace space e) $ setSpace space old_space + Allocate v (fmap (setExpSpace space) e) $ setSpace space old_space setBodySpace space (Free v old_space) = Free v $ setSpace space old_space setBodySpace space (DeclareMem name old_space) = @@ -49,13 +49,13 @@ DeclareArray name space t vs setBodySpace space (Copy dest dest_offset dest_space src src_offset src_space n) = Copy - dest (setCountSpace space dest_offset) dest_space' - src (setCountSpace space src_offset) src_space' $ - setCountSpace space n + dest (fmap (setExpSpace space) dest_offset) dest_space' + src (fmap (setExpSpace space) src_offset) src_space' $ + fmap (setExpSpace space) n where dest_space' = setSpace space dest_space src_space' = setSpace space src_space setBodySpace space (Write dest dest_offset bt dest_space vol e) = - Write dest (setCountSpace space dest_offset) bt (setSpace space dest_space) + Write dest (fmap (setExpSpace space) dest_offset) bt (setSpace space dest_space) vol (setExpSpace space e) setBodySpace space (c1 :>>: c2) = setBodySpace space c1 :>>: setBodySpace space c2 @@ -85,10 +85,6 @@ DebugPrint s $ fmap (fmap (setExpSpace space)) v setBodySpace _ (Op op) = Op op - -setCountSpace :: Space -> Count a -> Count a -setCountSpace space (Count e) = - Count $ setExpSpace space e setExpSpace :: Space -> Exp -> Exp setExpSpace space = fmap setLeafSpace diff --git a/src/Futhark/Compiler.hs b/src/Futhark/Compiler.hs --- a/src/Futhark/Compiler.hs +++ b/src/Futhark/Compiler.hs @@ -119,7 +119,7 @@ , pipelineValidate = True } -typeCheckInternalProgram :: I.Prog -> FutharkM () +typeCheckInternalProgram :: I.Prog I.SOACS -> FutharkM () typeCheckInternalProgram prog = case I.checkProg prog of Left err -> internalErrorS ("After internalisation:\n" ++ show err) (Just prog) diff --git a/src/Futhark/Construct.hs b/src/Futhark/Construct.hs --- a/src/Futhark/Construct.hs +++ b/src/Futhark/Construct.hs @@ -1,6 +1,5 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeSynonymInstances #-} {-# LANGUAGE FlexibleInstances #-} module Futhark.Construct ( letSubExp @@ -28,7 +27,7 @@ , eDivRoundingUp , eRoundToMultipleOf , eSliceArray - , eSplitArray + , eBlank , eOutOfBounds , eWriteArray @@ -46,6 +45,7 @@ , fullSlice , fullSliceNum , isFullSlice + , sliceAt , ifCommon , module Futhark.Binder @@ -53,9 +53,6 @@ -- * Result types , instantiateShapes , instantiateShapes' - , instantiateShapesFromIdentList - , instantiateExtTypes - , instantiateIdents , removeExistentials -- * Convenience @@ -75,7 +72,6 @@ import Futhark.Representation.AST import Futhark.MonadFreshNames import Futhark.Binder -import Futhark.Util letSubExp :: MonadBinder m => String -> Exp (Lore m) -> m SubExp @@ -274,18 +270,6 @@ return $ BasicOp $ Index arr $ fullSlice arr_t $ skips ++ [slice i' n'] where slice j m = DimSlice j m (constant (1::Int32)) --- | Construct an 'Index' expressions that splits an array in different parts along the outer dimension. -eSplitArray :: MonadBinder m => - VName -> [m (Exp (Lore m))] -> m [Exp (Lore m)] -eSplitArray arr sizes = do - sizes' <- mapM (letSubExp "split_size") =<< sequence sizes - -- Compute the starting offset for each slice. - (_, offsets) <- mapAccumLM increase (intConst Int32 0) sizes' - zipWithM (eSliceArray 0 arr) (map eSubExp offsets) (map eSubExp sizes') - where increase offset size = do - offset' <- letSubExp "offset" $ BasicOp $ BinOp (Add Int32) offset size - return (offset', offset) - -- | Are these indexes out-of-bounds for the array? eOutOfBounds :: MonadBinder m => VName -> [m (Exp (Lore m))] -> m (Exp (Lore m)) @@ -325,6 +309,12 @@ If outside_bounds outside_bounds_branch in_bounds_branch $ ifCommon [arr_t] +-- | Construct an unspecified value of the given type. +eBlank :: MonadBinder m => Type -> m (Exp (Lore m)) +eBlank (Prim t) = return $ BasicOp $ SubExp $ Constant $ blankPrimValue t +eBlank (Array pt shape _) = return $ BasicOp $ Scratch pt $ shapeDims shape +eBlank Mem{} = fail "eBlank: cannot create blank memory" + -- | Sign-extend to the given integer type. asIntS :: MonadBinder m => IntType -> SubExp -> m SubExp asIntS = asInt SExt @@ -383,16 +373,24 @@ , lambdaBody = body } +sliceDim :: SubExp -> DimIndex SubExp +sliceDim d = DimSlice (constant (0::Int32)) d (constant (1::Int32)) + -- | @fullSlice t slice@ returns @slice@, but with 'DimSlice's of -- entire dimensions appended to the full dimensionality of @t@. This -- function is used to turn incomplete indexing complete, as required -- by 'Index'. fullSlice :: Type -> [DimIndex SubExp] -> Slice SubExp fullSlice t slice = - slice ++ - map (\d -> DimSlice (constant (0::Int32)) d (constant (1::Int32))) - (drop (length slice) $ arrayDims t) + slice ++ map sliceDim (drop (length slice) $ arrayDims t) +-- | @ sliceAt t n slice@ returns @slice@ but with 'DimSlice's of the +-- outer @n@ dimensions prepended, and as many appended as to make it +-- a full slice. This is a generalisation of 'fullSlice'. +sliceAt :: Type -> Int -> [DimIndex SubExp] -> Slice SubExp +sliceAt t n slice = + fullSlice t $ map sliceDim (take n $ arrayDims t) ++ slice + -- | Like 'fullSlice', but the dimensions are simply numeric. fullSliceNum :: Num d => [d] -> [DimIndex d] -> Slice d fullSliceNum dims slice = @@ -466,41 +464,6 @@ where instantiate _ = do v <- lift $ newIdent "size" $ Prim int32 tell [v] return $ Var $ identName v - -instantiateShapesFromIdentList :: [Ident] -> [ExtType] -> [Type] -instantiateShapesFromIdentList idents ts = - evalState (instantiateShapes instantiate ts) idents - where instantiate _ = do - idents' <- get - case idents' of - [] -> fail "instantiateShapesFromIdentList: insufficiently sized context" - ident:idents'' -> do put idents'' - return $ Var $ identName ident - -instantiateExtTypes :: [VName] -> [ExtType] -> [Ident] -instantiateExtTypes names rt = - let (shapenames,valnames) = splitAt (shapeContextSize rt) names - shapes = [ Ident name (Prim int32) | name <- shapenames ] - valts = instantiateShapesFromIdentList shapes rt - vals = [ Ident name t | (name,t) <- zip valnames valts ] - in shapes ++ vals - -instantiateIdents :: [VName] -> [ExtType] - -> Maybe ([Ident], [Ident]) -instantiateIdents names ts - | let n = shapeContextSize ts, - n + length ts == length names = do - let (context, vals) = splitAt n names - nextShape _ = do - (context', remaining) <- get - case remaining of [] -> lift Nothing - x:xs -> do let ident = Ident x (Prim int32) - put (context'++[ident], xs) - return $ Var x - (ts', (context', _)) <- - runStateT (instantiateShapes nextShape ts) ([],context) - return (context', zipWith Ident vals ts') - | otherwise = Nothing removeExistentials :: ExtType -> Type -> Type removeExistentials t1 t2 = diff --git a/src/Futhark/Doc/Generator.hs b/src/Futhark/Doc/Generator.hs --- a/src/Futhark/Doc/Generator.hs +++ b/src/Futhark/Doc/Generator.hs @@ -4,7 +4,7 @@ import Control.Arrow ((***)) import Control.Monad import Control.Monad.Reader -import Control.Monad.Writer +import Control.Monad.Writer hiding (Sum) import Data.List (sort, sortOn, intersperse, inits, tails, isPrefixOf, find, groupBy, partition) import Data.Char (isSpace, isAlpha, toUpper) import Data.Loc @@ -380,13 +380,14 @@ keyword "val " <> vnameHtml name <> mconcat (map (" "<>) tps') <> ": " <> t' -prettyEnum :: [Name] -> Html -prettyEnum cs = pipes $ map (("#"<>) . renderName) cs - typeHtml :: StructType -> DocM Html typeHtml t = case t of - Prim et -> return $ primTypeHtml et - Record fs + Array _ u et shape -> do + shape' <- prettyShapeDecl shape + et' <- typeHtml $ Scalar et + return $ prettyU u <> shape' <> et' + Scalar (Prim et) -> return $ primTypeHtml et + Scalar (Record fs) | Just ts <- areTupleFields fs -> parens . commas <$> mapM typeHtml ts | otherwise -> @@ -394,43 +395,21 @@ where ppField (name, tp) = do tp' <- typeHtml tp return $ toHtml (nameToString name) <> ": " <> tp' - TypeVar _ u et targs -> do + Scalar (TypeVar _ u et targs) -> do targs' <- mapM typeArgHtml targs et' <- typeNameHtml et return $ prettyU u <> et' <> joinBy " " targs' - Array _ u et shape -> do - shape' <- prettyShapeDecl shape - et' <- prettyElem et - return $ prettyU u <> shape' <> et' - Arrow _ pname t1 t2 -> do + Scalar (Arrow _ pname t1 t2) -> do t1' <- typeHtml t1 t2' <- typeHtml t2 return $ case pname of - Just v -> + Named v -> parens (vnameHtml v <> ": " <> t1') <> " -> " <> t2' - Nothing -> + Unnamed -> t1' <> " -> " <> t2' - Enum cs -> return $ prettyEnum cs - -prettyElem :: ArrayElemTypeBase (DimDecl VName) -> DocM Html -prettyElem (ArrayPrimElem et) = return $ primTypeHtml et -prettyElem (ArrayPolyElem et targs) = do - targs' <- mapM typeArgHtml targs - return $ prettyTypeName et <> joinBy " " targs' -prettyElem (ArrayRecordElem fs) - | Just ts <- areTupleFields fs = - parens . commas <$> mapM prettyRecordElem ts - | otherwise = - braces . commas <$> mapM ppField (M.toList fs) - where ppField (name, tp) = do - tp' <- prettyRecordElem tp - return $ toHtml (nameToString name) <> ": " <> tp' -prettyElem (ArrayEnumElem cs) = return $ braces $ prettyEnum cs - -prettyRecordElem :: RecordArrayElemTypeBase (DimDecl VName) -> DocM Html -prettyRecordElem (RecordArrayElem et) = prettyElem et -prettyRecordElem (RecordArrayArrayElem et shape) = - typeHtml $ Array () Nonunique et shape + Scalar (Sum cs) -> pipes <$> mapM ppClause (sortConstrs cs) + where ppClause (n, ts) = joinBy " " . (ppConstr n :) <$> mapM typeHtml ts + ppConstr name = "#" <> toHtml (nameToString name) prettyShapeDecl :: ShapeDecl (DimDecl VName) -> DocM Html prettyShapeDecl (ShapeDecl ds) = @@ -537,7 +516,9 @@ parens (vnameHtml v <> ": " <> t1') <> " -> " <> t2' Nothing -> t1' <> " -> " <> t2' - TEEnum cs _ -> return $ prettyEnum cs + TESum cs _ -> pipes <$> mapM ppClause cs + where ppClause (n, ts) = joinBy " " . (ppConstr n :) <$> mapM typeExpHtml ts + ppConstr name = "#" <> toHtml (nameToString name) qualNameHtml :: QualName VName -> DocM Html qualNameHtml (QualName names vname@(VName name tag)) = @@ -574,8 +555,8 @@ let (pat_param, t) = patternParam pat t' <- typeHtml t return $ case pat_param of - Just v -> parens (vnameHtml v <> ": " <> t') - Nothing -> t' + Named v -> parens (vnameHtml v <> ": " <> t') + Unnamed -> t' relativise :: FilePath -> FilePath -> FilePath relativise dest src = diff --git a/src/Futhark/Internalise.hs b/src/Futhark/Internalise.hs --- a/src/Futhark/Internalise.hs +++ b/src/Futhark/Internalise.hs @@ -16,9 +16,9 @@ import qualified Data.Map.Strict as M import qualified Data.Set as S import Data.List +import qualified Data.List.NonEmpty as NE import Data.Loc import Data.Char (chr) -import Data.Maybe import Language.Futhark as E hiding (TypeArg) import Language.Futhark.Semantic (Imports) @@ -42,7 +42,7 @@ -- | Convert a program in source Futhark to a program in the Futhark -- core language. internaliseProg :: MonadFreshNames m => - Bool -> Imports -> m (Either String I.Prog) + Bool -> Imports -> m (Either String (I.Prog SOACS)) internaliseProg always_safe prog = do prog_decs <- Defunctorise.transformProg prog prog_decs' <- Monomorphise.transformProg prog_decs @@ -80,7 +80,7 @@ shapenames = map I.paramName shapeparams normal_params = map I.paramName constparams ++ shapenames ++ map I.paramName (concat params') - normal_param_names = S.fromList normal_params + normal_param_names = namesFromList normal_params fname' <- internaliseFunName fname params @@ -93,9 +93,9 @@ internaliseBody body >>= ensureResultExtShape asserting msg loc (map I.fromDecl rettype') - let free_in_fun = freeIn body' `S.difference` normal_param_names + let free_in_fun = freeIn body' `namesSubtract` normal_param_names - used_free_params <- forM (S.toList free_in_fun) $ \v -> do + used_free_params <- forM (namesToList free_in_fun) $ \v -> do v_t <- lookupType v return $ Param v $ toDecl v_t Nonunique @@ -149,6 +149,9 @@ 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 generateEntryPoint :: E.StructType -> E.ValBind -> InternaliseM () generateEntryPoint ftype (E.ValBind _ ofname retdecl (Info rettype) _ params _ _ loc) = do @@ -196,13 +199,13 @@ E.StructType, [I.TypeBase ExtShape Uniqueness]) -> [EntryPointType] - entryPointType (_, E.Prim E.Unsigned{}, _) = + entryPointType (_, E.Scalar (E.Prim E.Unsigned{}), _) = [I.TypeUnsigned] - entryPointType (_, E.Array _ _ (ArrayPrimElem Unsigned{}) _, _) = + entryPointType (_, E.Array _ _ (E.Prim E.Unsigned{}) _, _) = [I.TypeUnsigned] - entryPointType (_, E.Prim{}, _) = + entryPointType (_, E.Scalar E.Prim{}, _) = [I.TypeDirect] - entryPointType (_, E.Array _ _ ArrayPrimElem{} _, _) = + entryPointType (_, E.Array _ _ E.Prim{} _, _) = [I.TypeDirect] entryPointType (te, t, ts) = [I.TypeOpaque desc $ length ts] @@ -230,10 +233,10 @@ internaliseIdent :: E.Ident -> InternaliseM I.VName internaliseIdent (E.Ident name (Info tp) loc) = case tp of - E.Prim{} -> return name - _ -> fail $ "Futhark.Internalise.internaliseIdent: asked to internalise non-prim-typed ident '" - ++ pretty name ++ " of type " ++ pretty tp ++ - " at " ++ locStr loc ++ "." + E.Scalar E.Prim{} -> return name + _ -> fail $ "Futhark.Internalise.internaliseIdent: asked to internalise non-prim-typed ident '" + ++ pretty name ++ " of type " ++ pretty tp ++ + " at " ++ locStr loc ++ "." internaliseBody :: E.Exp -> InternaliseM Body internaliseBody e = insertStmsM $ resultBody <$> internaliseExp "res" e @@ -274,11 +277,10 @@ return $ I.BasicOp $ I.Index v $ fullSlice v_t idxs' certifying cs $ letSubExps desc =<< mapM index vs -internaliseExp desc (E.TupLit es _) = - concat <$> mapM (internaliseExp desc) es +internaliseExp desc (E.TupLit es _) = concat <$> mapM (internaliseExp desc) es internaliseExp desc (E.RecordLit orig_fields _) = - concatMap snd . sortFields . M.unions . reverse <$> mapM internaliseField orig_fields + concatMap snd . sortFields . M.unions <$> mapM internaliseField orig_fields where internaliseField (E.RecordFieldExplicit name e _) = M.singleton name <$> internaliseExp desc e internaliseField (E.RecordFieldImplicit name t loc) = @@ -340,8 +342,8 @@ (it, le_op, lt_op) <- case E.typeOf start of - E.Prim (E.Signed it) -> return (it, CmpSle it, CmpSlt it) - E.Prim (E.Unsigned it) -> return (it, CmpUle it, CmpUlt it) + E.Scalar (E.Prim (E.Signed it)) -> return (it, CmpSle it, CmpSlt it) + E.Scalar (E.Prim (E.Unsigned it)) -> return (it, CmpUle it, CmpUlt it) start_t -> fail $ "Start value in range has type " ++ pretty start_t let one = intConst it 1 @@ -639,7 +641,7 @@ src' <- internaliseExp desc src ve' <- internaliseExp desc ve replace (E.typeOf src `setAliases` ()) fields ve' src' - where replace (E.Record m) (f:fs) ve' src' + where replace (E.Scalar (E.Record m)) (f:fs) ve' src' | Just t <- M.lookup f m = do i <- fmap sum $ mapM (internalisedTypeSize . snd) $ takeWhile ((/=f) . fst) $ sortFields m @@ -664,28 +666,47 @@ letBindNames_ [v'] $ I.BasicOp $ I.SubExp v return $ I.Var v' -internaliseExp _ (E.VConstr0 c (Info t) loc) = - case t of - Enum cs -> - case elemIndex c $ sort cs of - Just i -> return [I.Constant $ I.IntValue $ intValue I.Int8 i] - _ -> fail $ "internaliseExp: invalid constructor: #" ++ nameToString c ++ - "\nfor enum at " ++ locStr loc ++ ": " ++ pretty t - _ -> fail $ "internaliseExp: nonsensical type for enum at " - ++ locStr loc ++ ": " ++ pretty t +internaliseExp _ (E.Constr c es (Info (E.Scalar (E.Sum fs))) loc) = do + ((ts, constr_map), cm) <- internaliseSumType $ M.map (map E.toStruct) fs + mapM_ (uncurry (internaliseDimConstant loc)) cm + es' <- concat <$> mapM (internaliseExp "payload") es -internaliseExp desc (E.Match e cs _ loc) = - case cs of - [CasePat _ eCase _] -> internaliseExp desc eCase - (c:cs') -> do - bFalse <- bFalseM - letTupExp' desc =<< generateCaseIf desc e c bFalse - where bFalseM = do - eLast' <- internalisePat desc pLast e eLast locLast internaliseBody - foldM (\bf c' -> eBody $ return $ generateCaseIf desc e c' bf) eLast' (reverse $ init cs') - CasePat pLast eLast locLast = last cs' - [] -> fail $ "internaliseExp: match with no cases at: " ++ locStr loc + let noExt _ = return $ intConst Int32 0 + ts' <- instantiateShapes noExt $ map fromDecl ts + case M.lookup c constr_map of + Just (i, js) -> + (intConst Int8 (toInteger i):) <$> clauses 0 ts' (zip js es') + Nothing -> + fail "internaliseExp Constr: missing constructor" + + where clauses j (t:ts) js_to_es + | Just e <- j `lookup` js_to_es = + (e:) <$> clauses (j+1) ts js_to_es + | otherwise = do + blank <- letSubExp "zero" =<< eBlank t + (blank:) <$> clauses (j+1) ts js_to_es + clauses _ [] _ = + return [] + +internaliseExp _ (E.Constr _ _ (Info t) loc) = + fail $ "internaliseExp: constructor with type " ++ pretty t ++ " at " ++ locStr loc + +internaliseExp desc (E.Match e cs _ _) = do + ses <- internaliseExp (desc ++ "_scrutinee") e + case NE.uncons cs of + (CasePat pCase eCase locCase, Nothing) -> do + (_, pertinent) <- generateCond pCase ses + internalisePat' pCase pertinent eCase locCase (internaliseExp desc) + (c, Just cs') -> do + let CasePat pLast eLast locLast = NE.last cs' + bFalse <- do + (_, pertinent) <- generateCond pLast ses + eLast' <- internalisePat' pLast pertinent eLast locLast internaliseBody + foldM (\bf c' -> eBody $ return $ generateCaseIf ses c' bf) eLast' $ + reverse $ NE.init cs' + letTupExp' desc =<< generateCaseIf ses c bFalse + -- The "interesting" cases are over, now it's mostly boilerplate. internaliseExp _ (E.Literal v _) = @@ -693,17 +714,17 @@ internaliseExp _ (E.IntLit v (Info t) _) = case t of - E.Prim (E.Signed it) -> + E.Scalar (E.Prim (E.Signed it)) -> return [I.Constant $ I.IntValue $ intValue it v] - E.Prim (E.Unsigned it) -> + E.Scalar (E.Prim (E.Unsigned it)) -> return [I.Constant $ I.IntValue $ intValue it v] - E.Prim (E.FloatType ft) -> + E.Scalar (E.Prim (E.FloatType ft)) -> return [I.Constant $ I.FloatValue $ floatValue ft v] _ -> fail $ "internaliseExp: nonsensical type for integer literal: " ++ pretty t internaliseExp _ (E.FloatLit v (Info t) _) = case t of - E.Prim (E.FloatType ft) -> + E.Scalar (E.Prim (E.FloatType ft)) -> return [I.Constant $ I.FloatValue $ floatValue ft v] _ -> fail $ "internaliseExp: nonsensical type for float literal: " ++ pretty t @@ -728,8 +749,9 @@ n <- internalisedTypeSize $ rt `setAliases` () i' <- fmap sum $ mapM internalisedTypeSize $ case E.typeOf e `setAliases` () of - Record fs -> map snd $ takeWhile ((/=k) . fst) $ sortFields fs - t -> [t] + E.Scalar (Record fs) -> + map snd $ takeWhile ((/=k) . fst) $ sortFields fs + t -> [t] take n . drop i' <$> internaliseExp desc e internaliseExp _ e@E.Lambda{} = @@ -750,53 +772,85 @@ internaliseExp _ e@E.IndexSection{} = fail $ "internaliseExp: Unexpected index section at " ++ locStr (srclocOf e) -andExp :: E.Exp -> E.Exp -> E.Exp -andExp l r = E.If l r (E.Literal (E.BoolValue False) noLoc) (Info (E.Prim E.Bool)) noLoc +generateCond :: E.Pattern -> [I.SubExp] -> InternaliseM (I.SubExp, [I.SubExp]) +generateCond orig_p orig_ses = do + (cmps, pertinent, _) <- compares orig_p orig_ses + cmp <- letSubExp "matches" =<< foldBinOp I.LogAnd (constant True) cmps + return (cmp, pertinent) + where + -- Literals are always primitive values. + compares (E.PatternLit e _ _) (se:ses) = do + e' <- internaliseExp1 "constant" e + I.Prim t' <- subExpType se + cmp <- letSubExp "match_lit" $ I.BasicOp $ I.CmpOp (I.CmpEq t') e' se + return ([cmp], [se], ses) -eqExp :: E.Exp -> E.Exp -> E.Exp -eqExp l r = E.BinOp eq (Info ft) - (l, sType l) (r, sType r) (Info (E.Prim E.Bool)) noLoc - where sType e = Info $ toStruct $ E.typeOf e - arrow = Arrow S.empty Nothing - ft = E.typeOf l `arrow` E.typeOf r `arrow` E.Prim E.Bool - eq = qualName $ VName "==" (-1) + compares (E.PatternConstr c (Info (E.Scalar (E.Sum fs))) pats _) (se:ses) = do + ((payload_ts, m), _) <- internaliseSumType $ M.map (map toStruct) fs + case M.lookup c m of + Just (i, payload_is) -> do + let i' = intConst Int8 $ toInteger i + let (payload_ses, ses') = splitAt (length payload_ts) ses + cmp <- letSubExp "match_constr" $ I.BasicOp $ I.CmpOp (I.CmpEq int8) i' se + (cmps, pertinent, _) <- comparesMany pats $ map (payload_ses!!) payload_is + return (cmp : cmps, pertinent, ses') + Nothing -> + fail "generateCond: missing constructor" -generateCond :: E.Pattern -> E.Exp -> E.Exp -generateCond p e = foldr andExp (E.Literal (E.BoolValue True) noLoc) conds - where conds = mapMaybe ((<*> pure e) . fst) $ generateCond' p + compares (E.PatternConstr _ (Info t) _ _) _ = + fail $ "generateCond: PatternConstr has nonsensical type: " ++ pretty t - generateCond' :: E.Pattern -> [(Maybe (E.Exp -> E.Exp), PatternType)] - generateCond' (E.TuplePattern ps loc) = generateCond' (E.RecordPattern fs loc) - where fs = zipWith (\i p' -> (nameFromString (show i), p')) ([1..] :: [Integer]) ps - generateCond' (E.RecordPattern fs _) = concatMap instCond holes - where holes = map (\(n, p') -> (generateCond' p', n)) fs - field ([],_) = Nothing - field ((_, t):_, f) = Just (f, t) - t' = Record $ M.fromList $ mapMaybe field holes - projectHole _ (Nothing, _) = (Nothing, t') - projectHole f (Just condHole, t) = - (Just (\e' -> condHole $ Project f e' (Info t) noLoc), t') - instCond (condHoles, f) = map (projectHole f) condHoles - generateCond' (E.PatternParens p' _) = generateCond' p' - generateCond' (E.Id _ (Info t) _) = - [(Nothing, t)] - generateCond' (E.Wildcard (Info t) _)= - [(Nothing, t)] - generateCond' (E.PatternAscription p' _ _) = generateCond' p' - generateCond' (E.PatternLit ePat (Info t) _) = - [(Just (eqExp ePat), t)] + compares (E.Id _ t loc) ses = + compares (E.Wildcard t loc) ses + compares (E.Wildcard (Info t) _) ses = do + n <- internalisedTypeSize $ E.toStruct t + let (id_ses, rest_ses) = splitAt n ses + return ([], id_ses, rest_ses) -generateCaseIf :: String -> E.Exp -> Case -> I.Body -> InternaliseM I.Exp -generateCaseIf desc e (CasePat p eCase loc) bFail = do - eCase' <- internalisePat desc p e eCase loc internaliseBody - eIf cond (return eCase') (return bFail) - where cond = BasicOp . SubExp <$> internaliseExp1 "cond" (generateCond p e) + compares (E.PatternParens pat _) ses = + compares pat ses + compares (E.TuplePattern pats _) ses = + comparesMany pats ses + + compares (E.RecordPattern fs _) ses = + comparesMany (map snd $ E.sortFields $ M.fromList fs) ses + + compares (E.PatternAscription pat _ _) ses = + compares pat ses + + compares pat [] = + error $ "generateCond: No values left for pattern " ++ pretty pat + + comparesMany [] ses = return ([], [], ses) + comparesMany (pat:pats) ses = do + (cmps1, pertinent1, ses') <- compares pat ses + (cmps2, pertinent2, ses'') <- comparesMany pats ses' + return (cmps1 <> cmps2, + pertinent1 <> pertinent2, + ses'') + +generateCaseIf :: [I.SubExp] -> Case -> I.Body -> InternaliseM I.Exp +generateCaseIf ses (CasePat p eCase loc) bFail = do + (cond, pertinent) <- generateCond p ses + eCase' <- internalisePat' p pertinent eCase loc internaliseBody + eIf (eSubExp cond) (return eCase') (return bFail) + internalisePat :: String -> E.Pattern -> E.Exp - -> E.Exp -> SrcLoc -> (E.Exp -> InternaliseM a) -> InternaliseM a + -> E.Exp -> SrcLoc -> (E.Exp -> InternaliseM a) + -> InternaliseM a internalisePat desc p e body loc m = do ses <- internaliseExp desc' e + internalisePat' p ses body loc m + where desc' = case S.toList $ E.patternIdents p of + [v] -> baseString $ E.identName v + _ -> desc + +internalisePat' :: E.Pattern -> [I.SubExp] + -> E.Exp -> SrcLoc -> (E.Exp -> InternaliseM a) + -> InternaliseM a +internalisePat' p ses body loc m = do t <- I.staticShapes <$> mapM I.subExpType ses stmPattern p t $ \cm pat_names match -> do mapM_ (uncurry (internaliseDimConstant loc)) cm @@ -804,9 +858,6 @@ forM_ (zip pat_names ses') $ \(v,se) -> letBindNames_ [v] $ I.BasicOp $ I.SubExp se m body - where desc' = case S.toList $ E.patternIdents p of - [v] -> baseString $ E.identName v - _ -> desc internaliseSlice :: SrcLoc -> [SubExp] @@ -1017,7 +1068,7 @@ let consumed = consumedByLambda $ Alias.analyseLambda lam0' copyIfConsumed p (I.Var v) - | I.paramName p `S.member` consumed = + | I.paramName p `nameIn` consumed = letSubExp "acc_copy" $ I.BasicOp $ I.Copy v copyIfConsumed _ x = return x @@ -1048,9 +1099,9 @@ internaliseDimExp s e = do e' <- internaliseExp1 s e case E.typeOf e of - E.Prim (Signed it) -> (,it) <$> asIntS Int32 e' - E.Prim (Unsigned it) -> (,it) <$> asIntZ Int32 e' - _ -> fail "internaliseDimExp: bad type" + E.Scalar (E.Prim (Signed it)) -> (,it) <$> asIntS Int32 e' + E.Scalar (E.Prim (Unsigned it)) -> (,it) <$> asIntZ Int32 e' + _ -> fail "internaliseDimExp: bad type" internaliseExpToVars :: String -> E.Exp -> InternaliseM [I.VName] internaliseExpToVars desc e = @@ -1241,10 +1292,7 @@ handle [x] "unsign_i32" = Just $ toUnsigned I.Int32 x handle [x] "unsign_i64" = Just $ toUnsigned I.Int64 x - handle [x] "sgn" = Just $ signumF x - handle [x] "abs" = Just $ absF x - handle [x] "!" = Just $ notF x - handle [x] "~" = Just $ complementF x + handle [x] "!" = Just $ complementF x handle [x] "opaque" = Just $ \desc -> mapM (letSubExp desc . BasicOp . Opaque) =<< internaliseExp "opaque_arg" x @@ -1271,10 +1319,10 @@ -- Short-circuiting operators are magical. handle [x,y] "&&" = Just $ \desc -> internaliseExp desc $ - E.If x y (E.Literal (E.BoolValue False) noLoc) (Info (E.Prim E.Bool)) noLoc + E.If x y (E.Literal (E.BoolValue False) noLoc) (Info $ E.Scalar $ E.Prim E.Bool) noLoc handle [x,y] "||" = Just $ \desc -> internaliseExp desc $ - E.If x (E.Literal (E.BoolValue True) noLoc) y (Info (E.Prim E.Bool)) noLoc + E.If x (E.Literal (E.BoolValue True) noLoc) y (Info $ E.Scalar $ E.Prim E.Bool) noLoc -- Handle equality and inequality specially, to treat the case of -- arrays. @@ -1332,7 +1380,7 @@ x' <- internaliseExp1 "x" x y' <- internaliseExp1 "y" y case (E.typeOf x, E.typeOf y) of - (E.Prim t1, E.Prim t2) -> + (E.Scalar (E.Prim t1), E.Scalar (E.Prim t2)) -> internaliseBinOp desc bop x' y' t1 t2 _ -> fail "Futhark.Internalise.internaliseExp: non-primitive type in BinOp." @@ -1422,7 +1470,7 @@ lam' <- internalisePartitionLambda internaliseLambda k' lam $ map I.Var arrs uncurry (++) <$> partitionWithSOACS k' lam' arrs where isInt32 (Literal (SignedValue (Int32Value k')) _) = Just k' - isInt32 (IntLit k' (Info (E.Prim (Signed Int32))) _) = Just $ fromInteger k' + isInt32 (IntLit k' (Info (E.Scalar (E.Prim (Signed Int32)))) _) = Just $ fromInteger k' isInt32 _ = Nothing handle [TupLit [lam, ne, arr] _] "reduce" = Just $ \desc -> @@ -1442,16 +1490,16 @@ where reduce w scan_lam nes arrs = I.Screma w <$> I.scanSOAC scan_lam nes <*> pure arrs - handle [TupLit [op, f, arr] _] "stream_red" = Just $ \desc -> + handle [TupLit [op, f, arr] _] "reduce_stream" = Just $ \desc -> internaliseStreamRed desc InOrder Noncommutative op f arr - handle [TupLit [op, f, arr] _] "stream_red_per" = Just $ \desc -> + handle [TupLit [op, f, arr] _] "reduce_stream_per" = Just $ \desc -> internaliseStreamRed desc Disorder Commutative op f arr - handle [TupLit [f, arr] _] "stream_map" = Just $ \desc -> + handle [TupLit [f, arr] _] "map_stream" = Just $ \desc -> internaliseStreamMap desc InOrder f arr - handle [TupLit [f, arr] _] "stream_map_per" = Just $ \desc -> + handle [TupLit [f, arr] _] "map_stream_per" = Just $ \desc -> internaliseStreamMap desc Disorder f arr handle [TupLit [dest, op, ne, buckets, img] _] "gen_reduce" = Just $ \desc -> @@ -1466,64 +1514,42 @@ toSigned int_to e desc = do e' <- internaliseExp1 "trunc_arg" e case E.typeOf e of - E.Prim E.Bool -> + E.Scalar (E.Prim E.Bool) -> letTupExp' desc $ I.If e' (resultBody [intConst int_to 1]) (resultBody [intConst int_to 0]) $ ifCommon [I.Prim $ I.IntType int_to] - E.Prim (E.Signed int_from) -> + E.Scalar (E.Prim (E.Signed int_from)) -> letTupExp' desc $ I.BasicOp $ I.ConvOp (I.SExt int_from int_to) e' - E.Prim (E.Unsigned int_from) -> + E.Scalar (E.Prim (E.Unsigned int_from)) -> letTupExp' desc $ I.BasicOp $ I.ConvOp (I.ZExt int_from int_to) e' - E.Prim (E.FloatType float_from) -> + E.Scalar (E.Prim (E.FloatType float_from)) -> letTupExp' desc $ I.BasicOp $ I.ConvOp (I.FPToSI float_from int_to) e' _ -> fail "Futhark.Internalise.handle: non-numeric type in ToSigned" toUnsigned int_to e desc = do e' <- internaliseExp1 "trunc_arg" e case E.typeOf e of - E.Prim E.Bool -> + E.Scalar (E.Prim E.Bool) -> letTupExp' desc $ I.If e' (resultBody [intConst int_to 1]) (resultBody [intConst int_to 0]) $ ifCommon [I.Prim $ I.IntType int_to] - E.Prim (E.Signed int_from) -> + E.Scalar (E.Prim (E.Signed int_from)) -> letTupExp' desc $ I.BasicOp $ I.ConvOp (I.ZExt int_from int_to) e' - E.Prim (E.Unsigned int_from) -> + E.Scalar (E.Prim (E.Unsigned int_from)) -> letTupExp' desc $ I.BasicOp $ I.ConvOp (I.ZExt int_from int_to) e' - E.Prim (E.FloatType float_from) -> + E.Scalar (E.Prim (E.FloatType float_from)) -> letTupExp' desc $ I.BasicOp $ I.ConvOp (I.FPToUI float_from int_to) e' _ -> fail "Futhark.Internalise.internaliseExp: non-numeric type in ToUnsigned" - signumF e desc = do - e' <- internaliseExp1 "signum_arg" e - case E.typeOf e of - E.Prim (E.Signed t) -> - letTupExp' desc $ I.BasicOp $ I.UnOp (I.SSignum t) e' - E.Prim (E.Unsigned t) -> - letTupExp' desc $ I.BasicOp $ I.UnOp (I.USignum t) e' - _ -> fail "Futhark.Internalise.internaliseExp: non-integer type in Signum" - - absF e desc = do - e' <- internaliseExp1 "abs_arg" e - case E.typeOf e of - E.Prim (E.Signed t) -> - letTupExp' desc $ I.BasicOp $ I.UnOp (I.Abs t) e' - E.Prim (E.Unsigned _) -> - return [e'] - E.Prim (E.FloatType t) -> - letTupExp' desc $ I.BasicOp $ I.UnOp (I.FAbs t) e' - _ -> fail "Futhark.Internalise.internaliseExp: non-integer type in Abs" - - notF e desc = do - e' <- internaliseExp1 "not_arg" e - letTupExp' desc $ I.BasicOp $ I.UnOp I.Not e' - complementF e desc = do e' <- internaliseExp1 "complement_arg" e et <- subExpType e' case et of I.Prim (I.IntType t) -> letTupExp' desc $ I.BasicOp $ I.UnOp (I.Complement t) e' + I.Prim I.Bool -> + letTupExp' desc $ I.BasicOp $ I.UnOp I.Not e' _ -> - fail "Futhark.Internalise.internaliseExp: non-integer type in Complement" + fail "Futhark.Internalise.internaliseExp: non-int/bool type in Complement" scatterF a si v desc = do si' <- letExp "write_si" . BasicOp . SubExp =<< internaliseExp1 "write_arg_i" si @@ -1748,8 +1774,12 @@ arg' <- case arg of TypeArgExpType argt -> typeExpForError cm argt TypeArgExpDim d _ -> pure <$> dimDeclForError cm d return $ t' ++ [" "] ++ arg' -typeExpForError _ e@E.TEEnum{} = - return [ErrorString $ pretty e] +typeExpForError cm (E.TESum cs _) = do + cs' <- mapM (onClause . snd) cs + return $ intercalate [" | "] cs' + where onClause c = do + c' <- mapM (typeExpForError cm) c + return $ intercalate [" "] c' dimDeclForError :: ConstParams -> E.DimDecl VName -> InternaliseM (ErrorMsgPart SubExp) dimDeclForError cm (NamedDim d) = do diff --git a/src/Futhark/Internalise/Bindings.hs b/src/Futhark/Internalise/Bindings.hs --- a/src/Futhark/Internalise/Bindings.hs +++ b/src/Futhark/Internalise/Bindings.hs @@ -139,6 +139,8 @@ flattenPattern' p flattenPattern' (E.PatternLit _ t loc) = flattenPattern' $ E.Wildcard t loc + flattenPattern' (E.PatternConstr _ _ ps _) = + concat <$> mapM flattenPattern' ps type MatchPattern = SrcLoc -> [I.SubExp] -> InternaliseM [I.SubExp] diff --git a/src/Futhark/Internalise/Defunctionalise.hs b/src/Futhark/Internalise/Defunctionalise.hs --- a/src/Futhark/Internalise/Defunctionalise.hs +++ b/src/Futhark/Internalise/Defunctionalise.hs @@ -4,10 +4,11 @@ ( transformProg ) where import Control.Arrow (first, second) -import Control.Monad.RWS +import Control.Monad.RWS hiding (Sum) import Data.Bifunctor hiding (first, second) import Data.Foldable import Data.List +import qualified Data.List.NonEmpty as NE import Data.Loc import qualified Data.Map.Strict as M import qualified Data.Set as S @@ -30,6 +31,9 @@ -- ^ The 'VName's are shape parameters that are bound -- by the 'Pattern'. | RecordSV [(Name, StaticVal)] + | SumSV Name [StaticVal] [(Name, [PatternType])] + -- ^ The constructor that is actually present, plus + -- the others that are not. | DynamicFun (Exp, StaticVal) StaticVal | IntrinsicSV deriving (Show) @@ -71,6 +75,8 @@ LambdaSV dims pat t e $ M.map (restrict' u) env restrict' u (RecordSV fields) = RecordSV $ map (fmap $ restrict' u) fields + restrict' u (SumSV c svs fields) = + SumSV c (map (restrict' u) svs) fields restrict' u (DynamicFun (e, sv1) sv2) = DynamicFun (e, restrict' u sv1) $ restrict' u sv2 restrict' _ IntrinsicSV = IntrinsicSV @@ -125,7 +131,7 @@ let bound_by_pat = (`S.member` patternDimNames pat') . typeParamName (pat_dims, rest_dims) = partition bound_by_pat tparams in (map typeParamName pat_dims, pat', - foldFunType (map (toStruct . patternPatternType) pats') ret, + foldFunType (map (toStruct . patternType) pats') ret, ExtLambda rest_dims pats' e0 (closure, ret) loc) -- Construct a record literal that closes over the environment of @@ -288,7 +294,7 @@ -- We handle BinOps by turning them into ordinary function applications. defuncExp (BinOp qn (Info t) (e1, Info pt1) (e2, Info pt2) (Info ret) loc) = defuncExp $ Apply (Apply (Var qn (Info t) loc) - e1 (Info (diet pt1)) (Info (Arrow mempty Nothing (fromStruct pt2) ret)) loc) + e1 (Info (diet pt1)) (Info (Scalar $ Arrow mempty Unnamed (fromStruct pt2) ret)) loc) e2 (Info (diet pt2)) (Info ret) loc defuncExp (Project vn e0 tp@(Info tp') loc) = do @@ -332,7 +338,7 @@ Just sv -> RecordSV $ (f, staticField sv sv2 fs') : filter ((/=f) . fst) svs Nothing -> error "Invalid record projection." - staticField (Dynamic t@Record{}) sv2 fs'@(_:_) = + staticField (Dynamic t@(Scalar Record{})) sv2 fs'@(_:_) = staticField (svFromType t) sv2 fs' staticField _ sv2 _ = sv2 @@ -345,15 +351,35 @@ (e2', sv) <- defuncExp e2 return (Assert e1' e2' desc loc, sv) -defuncExp e@VConstr0{} = return (e, Dynamic $ typeOf e) +defuncExp (Constr name es (Info (Scalar (Sum all_fs))) loc) = do + (es', svs) <- unzip <$> mapM defuncExp es + let sv = SumSV name svs $ M.toList $ + name `M.delete` M.map (map defuncType) all_fs + return (Constr name es' (Info (typeFromSV sv)) loc, sv) + where defuncType :: Monoid als => + TypeBase (DimDecl VName) als + -> TypeBase (DimDecl VName) als + defuncType (Array as u t shape) = Array as u (defuncScalar t) shape + defuncType (Scalar t) = Scalar $ defuncScalar t + defuncScalar :: Monoid als => + ScalarTypeBase (DimDecl VName) als + -> ScalarTypeBase (DimDecl VName) als + defuncScalar (Record fs) = Record $ M.map defuncType fs + defuncScalar Arrow{} = Record mempty + defuncScalar (Sum fs) = Sum $ M.map (map defuncType) fs + defuncScalar (Prim t) = Prim t + defuncScalar (TypeVar as u tn targs) = TypeVar as u tn targs + +defuncExp (Constr name _ (Info t) loc) = + error $ "Constructor " ++ pretty name ++ " given type " ++ + pretty t ++ " at " ++ locStr loc + defuncExp (Match e cs t loc) = do (e', sv) <- defuncExp e csPairs <- mapM (defuncCase sv) cs - let cs' = map fst csPairs - sv' = case csPairs of - [] -> error "Matches must always have at least one case." - c':_ -> snd c' + let cs' = fmap fst csPairs + sv' = snd $ NE.head csPairs return (Match e' cs' t loc, sv') -- | Same as 'defuncExp', except it ignores the static value. @@ -389,7 +415,7 @@ return $ Lambda params e0' decl tp loc defuncSoacExp e - | Arrow{} <- typeOf e = do + | Scalar Arrow{} <- typeOf e = do (pats, body, tp) <- etaExpand e let env = foldMap envFromPattern pats body' <- localEnv env $ defuncExp' body @@ -409,7 +435,7 @@ e $ zip3 vars ps (drop 1 $ tails ps) return (pats, e', toStruct ret) - where getType (Arrow _ _ t1 t2) = + where getType (Scalar (Arrow _ _ t1 t2)) = let (ps, r) = getType t2 in (t1 : ps, r) getType t = ([], t) @@ -440,7 +466,7 @@ return ([], body', imposeType sv rettype ) where imposeType Dynamic{} t = Dynamic $ fromStruct t - imposeType (RecordSV fs1) (Record fs2) = + imposeType (RecordSV fs1) (Scalar (Record fs2)) = RecordSV $ M.toList $ M.intersectionWith imposeType (M.fromList fs1) fs2 imposeType sv _ = sv @@ -488,9 +514,9 @@ let t1 = toStruct $ typeOf e1' t2 = toStruct $ typeOf e2' fname' = qualName fname - return (Parens (Apply (Apply (Var fname' (Info (Arrow mempty Nothing (fromStruct t1) $ - Arrow mempty Nothing (fromStruct t2) rettype)) loc) - e1' (Info Observe) (Info $ Arrow mempty Nothing (fromStruct t2) rettype) loc) + return (Parens (Apply (Apply (Var fname' (Info (Scalar $ Arrow mempty Unnamed (fromStruct t1) $ + Scalar $ Arrow mempty Unnamed (fromStruct t2) rettype)) loc) + e1' (Info Observe) (Info $ Scalar $ Arrow mempty Unnamed (fromStruct t2) rettype) loc) e2' d (Info rettype) loc) noLoc, sv) -- If e1 is a dynamic function, we just leave the application in place, @@ -534,6 +560,8 @@ _ -> return (Var qn (Info (typeFromSV sv)) loc, sv) +defuncApply depth (Parens e _) = defuncApply depth e + defuncApply _ expr = defuncExp expr -- | Check if a 'StaticVal' and a given application depth corresponds @@ -567,6 +595,7 @@ 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 @@ -578,7 +607,7 @@ " at " ++ locStr (srclocOf tparam) ++ "." envFromDimNames :: [VName] -> Env -envFromDimNames = M.fromList . flip zip (repeat $ Dynamic $ Prim $ Signed Int32) +envFromDimNames = M.fromList . flip zip (repeat $ Dynamic $ Scalar $ Prim $ Signed Int32) -- | Create a new top-level value declaration with the given function name, -- return type, list of parameters, and body expression. @@ -619,27 +648,29 @@ maybe False (unique . unInfo . identType) $ find ((==v) . identName) $ S.toList $ foldMap patternIdents pats problematic v = (v `member` bound) && not (boundAsUnique v) - comb (Record fs_annot) (Record fs_got) = - Record $ M.intersectionWith comb fs_annot fs_got - comb Arrow{} t = descend t + comb (Scalar (Record fs_annot)) (Scalar (Record fs_got)) = + Scalar $ Record $ M.intersectionWith comb fs_annot fs_got + comb (Scalar Arrow{}) t = descend t comb got et = descend $ fromStruct got `setUniqueness` uniqueness et `setAliases` aliases et descend t@Array{} | any (problematic . aliasVar) (aliases t) = t `setUniqueness` Nonunique - descend (Record t) = Record $ fmap descend t + descend (Scalar (Record t)) = Scalar $ Record $ fmap descend t descend t = t -- | Compute the corresponding type for a given static value. typeFromSV :: StaticVal -> PatternType typeFromSV (Dynamic tp) = anyDimShapeAnnotations tp typeFromSV (LambdaSV _ _ _ _ env) = typeFromEnv env -typeFromSV (RecordSV ls) = Record $ M.fromList $ map (fmap typeFromSV) ls +typeFromSV (RecordSV ls) = Scalar $ Record $ M.fromList $ map (fmap typeFromSV) ls typeFromSV (DynamicFun (_, sv) _) = typeFromSV sv +typeFromSV (SumSV name svs fields) = + Scalar $ Sum $ M.insert name (map typeFromSV svs) $ M.fromList fields typeFromSV IntrinsicSV = error $ "Tried to get the type from the " ++ "static value of an intrinsic." typeFromEnv :: Env -> PatternType -typeFromEnv = Record . M.fromList . +typeFromEnv = Scalar . Record . M.fromList . map (bimap (nameFromString . pretty) typeFromSV) . M.toList -- | Construct the type for a fully-applied dynamic function from its @@ -670,6 +701,18 @@ matchPatternSV (Wildcard _ _) _ = mempty matchPatternSV (PatternAscription pat _ _) sv = matchPatternSV pat sv matchPatternSV PatternLit{} _ = mempty +matchPatternSV (PatternConstr c1 _ ps _) (SumSV c2 ls fs) + | c1 == c2 = + mconcat $ zipWith matchPatternSV ps ls + | Just ts <- lookup c1 fs = + mconcat $ zipWith matchPatternSV ps $ map svFromType ts + | otherwise = + error $ "matchPatternSV: missing constructor in type: " ++ pretty c1 +matchPatternSV (PatternConstr c1 _ ps _) (Dynamic (Scalar (Sum fs))) + | Just ts <- M.lookup c1 fs = + mconcat $ zipWith matchPatternSV ps $ map svFromType ts + | otherwise = + error $ "matchPatternSV: missing constructor in type: " ++ pretty c1 matchPatternSV pat (Dynamic t) = matchPatternSV pat $ svFromType t matchPatternSV pat sv = error $ "Tried to match pattern " ++ pretty pat ++ " with static value " ++ show sv ++ "." @@ -701,6 +744,13 @@ PatternAscription (updatePattern pat sv) tydecl loc | otherwise = updatePattern pat sv updatePattern p@PatternLit{} _ = p +updatePattern pat@(PatternConstr c1 (Info t) ps loc) sv@(SumSV _ svs _) + | orderZero t = pat + | otherwise = PatternConstr c1 (Info t') ps' loc + where t' = typeFromSV sv `setUniqueness` Nonunique + ps' = zipWith updatePattern ps svs +updatePattern (PatternConstr c1 _ ps loc) (Dynamic t) = + PatternConstr c1 (Info t) ps loc updatePattern pat (Dynamic t) = updatePattern pat (svFromType t) updatePattern pat sv = error $ "Tried to update pattern " ++ pretty pat @@ -709,8 +759,8 @@ -- | 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 (Record fs) = RecordSV . M.toList $ M.map svFromType fs -svFromType t = Dynamic t +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) @@ -795,7 +845,7 @@ Unsafe e _ -> freeVars e Assert e1 e2 _ _ -> freeVars e1 <> freeVars e2 - VConstr0{} -> mempty + Constr _ es _ _ -> foldMap freeVars es Match e cs _ _ -> freeVars e <> foldMap caseFV cs where caseFV (CasePat p eCase _) = (names (patternDimNames p) <> freeVars eCase) `without` patternVars p diff --git a/src/Futhark/Internalise/Lambdas.hs b/src/Futhark/Internalise/Lambdas.hs --- a/src/Futhark/Internalise/Lambdas.hs +++ b/src/Futhark/Internalise/Lambdas.hs @@ -11,7 +11,6 @@ import Control.Monad import Data.Loc -import qualified Data.Set as S import Language.Futhark as E import Futhark.Representation.SOACS as I @@ -107,7 +106,7 @@ let sizefun_safe = all (I.safeExp . I.stmExp) $ I.bodyStms $ I.lambdaBody sizefun' sizefun_arg_invariant = - not $ any (`S.member` freeIn (I.lambdaBody sizefun')) $ + not $ any (`nameIn` freeIn (I.lambdaBody sizefun')) $ map I.paramName $ lambdaParams sizefun' if sizefun_safe && sizefun_arg_invariant then do ses <- bodyBind $ lambdaBody sizefun' diff --git a/src/Futhark/Internalise/Monad.hs b/src/Futhark/Internalise/Monad.hs --- a/src/Futhark/Internalise/Monad.hs +++ b/src/Futhark/Internalise/Monad.hs @@ -73,7 +73,7 @@ , stateFunTable :: FunTable } -newtype InternaliseResult = InternaliseResult [FunDef] +newtype InternaliseResult = InternaliseResult [FunDef SOACS] deriving (Semigroup, Monoid) newtype InternaliseM a = InternaliseM (BinderT SOACS @@ -110,7 +110,7 @@ runInternaliseM :: MonadFreshNames m => Bool -> InternaliseM () - -> m (Either String [FunDef]) + -> m (Either String [FunDef SOACS]) runInternaliseM safe (InternaliseM m) = modifyNameSource $ \src -> do let onError e = (Left e, src) @@ -132,7 +132,7 @@ substitutingVars substs = local $ \env -> env { envSubsts = substs <> envSubsts env } -- | Add a function definition to the program being constructed. -addFunction :: FunDef -> InternaliseM () +addFunction :: FunDef SOACS -> InternaliseM () addFunction = InternaliseM . lift . tell . InternaliseResult . pure lookupFunction' :: VName -> InternaliseM (Maybe FunInfo) diff --git a/src/Futhark/Internalise/Monomorphise.hs b/src/Futhark/Internalise/Monomorphise.hs --- a/src/Futhark/Internalise/Monomorphise.hs +++ b/src/Futhark/Internalise/Monomorphise.hs @@ -25,12 +25,11 @@ module Futhark.Internalise.Monomorphise ( transformProg , transformDecs - , runMonoM ) where -import Control.Monad.RWS +import Control.Monad.RWS hiding (Sum) import Control.Monad.State -import Control.Monad.Writer +import Control.Monad.Writer hiding (Sum) import Data.Bitraversable import Data.Bifunctor import Data.Loc @@ -210,7 +209,7 @@ Ascript <$> transformExp e <*> pure tp <*> pure t <*> pure loc transformExp (LetPat pat e1 e2 (Info t) loc) = do - (pat', rr) <- expandRecordPattern pat + (pat', rr) <- transformPattern pat t' <- transformType t LetPat pat' <$> transformExp e1 <*> withRecordReplacements rr (transformExp e2) <*> @@ -324,13 +323,15 @@ transformExp (Assert e1 e2 desc loc) = Assert <$> transformExp e1 <*> transformExp e2 <*> pure desc <*> pure loc -transformExp e@VConstr0{} = return e +transformExp (Constr name all_es t loc) = + Constr name <$> mapM transformExp all_es <*> pure t <*> pure loc + transformExp (Match e cs t loc) = Match <$> transformExp e <*> mapM transformCase cs <*> traverse transformType t <*> pure loc transformCase :: Case -> MonoM Case transformCase (CasePat p e loc) = do - (p', rr) <- expandRecordPattern p + (p', rr) <- transformPattern p CasePat <$> pure p' <*> withRecordReplacements rr (transformExp e) <*> pure loc transformDimIndex :: DimIndexBase Info VName -> MonoM (DimIndexBase Info VName) @@ -356,20 +357,21 @@ [Id x (Info $ fromStruct argtype) noLoc]) desugarProjectSection :: [Name] -> PatternType -> SrcLoc -> MonoM Exp -desugarProjectSection fields (Arrow _ _ t1 t2) loc = do +desugarProjectSection fields (Scalar (Arrow _ _ t1 t2)) loc = do p <- newVName "project_p" let body = foldl project (Var (qualName p) (Info t1) noLoc) fields return $ Lambda [Id p (Info t1) noLoc] body Nothing (Info (mempty, toStruct t2)) loc where project e field = case typeOf e of - Record fs | Just t <- M.lookup field fs -> - Project field e (Info t) noLoc + Scalar (Record fs) + | Just t <- M.lookup field fs -> + Project field e (Info t) noLoc t -> error $ "desugarOpSection: type " ++ pretty t ++ " does not have field " ++ pretty field desugarProjectSection _ t _ = error $ "desugarOpSection: not a function type: " ++ pretty t desugarIndexSection :: [DimIndex] -> PatternType -> SrcLoc -> MonoM Exp -desugarIndexSection idxs (Arrow _ _ t1 t2) loc = do +desugarIndexSection idxs (Scalar (Arrow _ _ t1 t2)) loc = do p <- newVName "index_i" let body = Index (Var (qualName p) (Info t1) loc) idxs (Info t2) loc return $ Lambda [Id p (Info t1) noLoc] body Nothing (Info (mempty, toStruct t2)) loc @@ -377,7 +379,7 @@ noticeDims :: TypeBase (DimDecl VName) as -> MonoM () noticeDims = mapM_ notice . nestedDims - where notice (NamedDim v) = void $ transformFName (qualLeaf v) $ Prim $ Signed Int32 + where notice (NamedDim v) = void $ transformFName (qualLeaf v) $ Scalar $ Prim $ Signed Int32 notice _ = return () -- | Convert a collection of 'ValBind's to a nested sequence of let-bound, @@ -388,8 +390,8 @@ LetFun fname (dim_params, params, Nothing, rettype, body) e' loc where e' = unfoldLetFuns rest e -expandRecordPattern :: Pattern -> MonoM (Pattern, RecordReplacements) -expandRecordPattern (Id v (Info (Record fs)) loc) = do +transformPattern :: Pattern -> MonoM (Pattern, RecordReplacements) +transformPattern (Id v (Info (Scalar (Record fs))) loc) = do let fs' = M.toList fs (fs_ks, fs_ts) <- fmap unzip $ forM fs' $ \(f, ft) -> (,) <$> newVName (nameToString f) <*> transformType ft @@ -397,25 +399,34 @@ (zipWith3 Id fs_ks (map Info fs_ts) $ repeat loc)) loc, M.singleton v $ M.fromList $ zip (map fst fs') $ zip fs_ks fs_ts) -expandRecordPattern (Id v t loc) = return (Id v t loc, mempty) -expandRecordPattern (TuplePattern pats loc) = do - (pats', rrs) <- unzip <$> mapM expandRecordPattern pats +transformPattern (Id v t loc) = return (Id v t loc, mempty) +transformPattern (TuplePattern pats loc) = do + (pats', rrs) <- unzip <$> mapM transformPattern pats return (TuplePattern pats' loc, mconcat rrs) -expandRecordPattern (RecordPattern fields loc) = do +transformPattern (RecordPattern fields loc) = do let (field_names, field_pats) = unzip fields - (field_pats', rrs) <- unzip <$> mapM expandRecordPattern field_pats + (field_pats', rrs) <- unzip <$> mapM transformPattern field_pats return (RecordPattern (zip field_names field_pats') loc, mconcat rrs) -expandRecordPattern (PatternParens pat loc) = do - (pat', rr) <- expandRecordPattern pat +transformPattern (PatternParens pat loc) = do + (pat', rr) <- transformPattern pat return (PatternParens pat' loc, rr) -expandRecordPattern (Wildcard t loc) = do - t' <- traverse transformType t - return (Wildcard t' loc, mempty) -expandRecordPattern (PatternAscription pat td loc) = do - (pat', rr) <- expandRecordPattern pat +transformPattern (Wildcard (Info t) loc) = do + t' <- transformType t + return (wildcard t' loc, mempty) +transformPattern (PatternAscription pat td loc) = do + (pat', rr) <- transformPattern pat return (PatternAscription pat' td loc, rr) -expandRecordPattern (PatternLit e t loc) = return (PatternLit e t loc, mempty) +transformPattern (PatternLit e t loc) = return (PatternLit e t loc, mempty) +transformPattern (PatternConstr name t all_ps loc) = do + (all_ps', rrs) <- unzip <$> mapM transformPattern all_ps + return (PatternConstr name t all_ps' loc, mconcat rrs) +wildcard :: PatternType -> SrcLoc -> Pattern +wildcard (Scalar (Record fs)) loc = + RecordPattern (zip (M.keys fs) $ map ((`Wildcard` loc) . Info) $ M.elems fs) loc +wildcard t loc = + Wildcard (Info t) loc + -- | Monomorphize a polymorphic function at the types given in the instance -- list. Monomorphizes the body of the function as well. Returns the fresh name -- of the generated monomorphic function and its 'ValBind' representation. @@ -432,17 +443,26 @@ substTypesAny (fmap (fmap fromStruct) . (`M.lookup` substs')) params' = map (substPattern entry substPatternType) params - (params'', rrs) <- unzip <$> mapM expandRecordPattern params' + (params'', rrs) <- unzip <$> mapM transformPattern params' mapM_ noticeDims $ rettype : map patternStructType params'' body' <- updateExpTypes (`M.lookup` M.map (fmap toStructural) substs') body body'' <- withRecordReplacements (mconcat rrs) $ transformExp body' + body''' <- astMap noMoreSumTypes body'' + params''' <- astMap noMoreSumTypes params'' name' <- if null tparams then return name else newName name - return (name', toValBinding t_shape_params name' params'' rettype' body'') + return (name', toValBinding t_shape_params name' params''' rettype' body''') where shape_params = filter (not . isTypeParam) tparams + noMoreSumTypes = ASTMapper { mapOnExp = pure + , mapOnName = pure + , mapOnQualName = pure + , mapOnStructType = pure + , mapOnPatternType = pure + } + updateExpTypes substs = astMap $ mapper substs mapper substs = ASTMapper { mapOnExp = astMap $ mapper substs , mapOnName = pure @@ -475,19 +495,23 @@ sub False t1_rt t2_rt in runWriterT $ execStateT m mempty - where sub pos (TypeVar _ _ v _) t = addSubst pos v t - sub pos (Record fields1) (Record fields2) = - zipWithM_ (sub pos) - (map snd $ sortFields fields1) (map snd $ sortFields fields2) - sub _ Prim{} Prim{} = return () - sub _ Enum{} Enum{} = return () - sub pos t1@Array{} t2@Array{} + where sub pos t1@Array{} t2@Array{} | Just t1' <- peelArray (arrayRank t1) t1, Just t2' <- peelArray (arrayRank t1) t2 = sub pos t1' t2' - sub _ (Arrow _ _ t1a t1b) (Arrow _ _ t2a t2b) = do + sub pos (Scalar (TypeVar _ _ v _)) t = addSubst pos v t + sub pos (Scalar (Record fields1)) (Scalar (Record fields2)) = + zipWithM_ (sub pos) + (map snd $ sortFields fields1) (map snd $ sortFields fields2) + sub _ (Scalar Prim{}) (Scalar Prim{}) = return () + sub _ (Scalar (Arrow _ _ t1a t1b)) (Scalar (Arrow _ _ t2a t2b)) = do sub False t1a t2a sub False t1b t2b + sub pos (Scalar (Sum cs1)) (Scalar (Sum cs2)) = + zipWithM_ typeSubstClause (sortConstrs cs1) (sortConstrs cs2) + where typeSubstClause (_, ts1) (_, ts2) = zipWithM (sub pos) ts1 ts2 + sub pos t1@(Scalar Sum{}) t2 = sub pos t1 t2 + sub pos t1 t2@(Scalar Sum{}) = sub pos t1 t2 sub _ t1 t2 = error $ unlines ["typeSubstsM: mismatched types:", pretty t1, pretty t2] @@ -515,6 +539,7 @@ PatternAscription p td loc | entry -> PatternAscription (substPattern False f p) td loc | otherwise -> substPattern False f p PatternLit e (Info tp) loc -> PatternLit e (Info $ f tp) loc + PatternConstr n (Info tp) ps loc -> PatternConstr n (Info $ f tp) ps loc toPolyBinding :: ValBind -> PolyBinding toPolyBinding (ValBind _ name retdecl (Info rettype) tparams params body _ loc) = diff --git a/src/Futhark/Internalise/TypesValues.hs b/src/Futhark/Internalise/TypesValues.hs --- a/src/Futhark/Internalise/TypesValues.hs +++ b/src/Futhark/Internalise/TypesValues.hs @@ -11,6 +11,7 @@ , internaliseType , internalisePrimType , internalisedTypeSize + , internaliseSumType -- * Internalising values , internalisePrimValue @@ -108,33 +109,52 @@ -> InternaliseTypeM [I.TypeBase ExtShape Uniqueness] internaliseTypeM orig_t = case orig_t of - E.Prim bt -> return [I.Prim $ internalisePrimType bt] - E.TypeVar{} -> - fail "internaliseTypeM: cannot handle type variable." - E.Record ets -> - concat <$> mapM (internaliseTypeM . snd) (E.sortFields ets) E.Array _ u et shape -> do dims <- internaliseShape shape - ets <- internaliseElemType et + ets <- internaliseTypeM $ E.Scalar et return [I.arrayOf et' (Shape dims) $ internaliseUniqueness u | et' <- ets ] - E.Arrow{} -> fail $ "internaliseTypeM: cannot handle function type: " ++ pretty orig_t - E.Enum{} -> return [I.Prim $ I.IntType I.Int8] + E.Scalar (E.Prim bt) -> + return [I.Prim $ internalisePrimType bt] + E.Scalar (E.Record ets) -> + concat <$> mapM (internaliseTypeM . snd) (E.sortFields ets) + E.Scalar E.TypeVar{} -> + fail "internaliseTypeM: cannot handle type variable." + E.Scalar E.Arrow{} -> + fail $ "internaliseTypeM: cannot handle function type: " ++ pretty orig_t + E.Scalar (E.Sum cs) -> do + (ts, _) <- internaliseConstructors <$> + traverse (fmap concat . mapM internaliseTypeM) cs + return $ I.Prim (I.IntType I.Int8) : ts - where internaliseElemType E.ArrayPolyElem{} = - fail "internaliseElemType: cannot handle type variable." - internaliseElemType (E.ArrayPrimElem bt) = - return [I.Prim $ internalisePrimType bt] - internaliseElemType (E.ArrayRecordElem elemts) = - concat <$> mapM (internaliseRecordElem . snd) (E.sortFields elemts) - internaliseElemType E.ArrayEnumElem{} = - return [I.Prim $ I.IntType I.Int8] + where internaliseShape = mapM internaliseDim . E.shapeDims - internaliseRecordElem (E.RecordArrayElem et) = - internaliseElemType et - internaliseRecordElem (E.RecordArrayArrayElem et shape) = - internaliseTypeM $ E.Array mempty Nonunique et shape +internaliseConstructors :: M.Map Name [I.TypeBase ExtShape Uniqueness] + -> ([I.TypeBase ExtShape Uniqueness], + M.Map Name (Int, [Int])) +internaliseConstructors cs = + foldl' onConstructor mempty $ zip (E.sortConstrs cs) [0..] + where onConstructor (ts, mapping) ((c, c_ts), i) = + let (_, js, new_ts) = + foldl' f (zip ts [0..], mempty, mempty) c_ts + in (ts ++ new_ts, M.insert c (i, js) mapping) + where f (ts', js, new_ts) t + | primType t, + Just (_, j) <- find ((==t) . fst) ts' = + (delete (t, j) ts', + js ++ [j], + new_ts) + | otherwise = + (ts', + js ++ [length ts + length new_ts], + new_ts ++ [t]) - internaliseShape = mapM internaliseDim . E.shapeDims +internaliseSumType :: M.Map Name [E.StructType] + -> InternaliseM (([I.TypeBase ExtShape Uniqueness], + M.Map Name (Int, [Int])), + ConstParams) +internaliseSumType cs = + runInternaliseTypeM $ internaliseConstructors <$> + traverse (fmap concat . mapM internaliseTypeM) cs -- | How many core language values are needed to represent one source -- language value of the given type? diff --git a/src/Futhark/MonadFreshNames.hs b/src/Futhark/MonadFreshNames.hs --- a/src/Futhark/MonadFreshNames.hs +++ b/src/Futhark/MonadFreshNames.hs @@ -10,8 +10,6 @@ , modifyNameSource , newName , newNameFromString - , newID - , newIDFromString , newVName , newVName' , newIdent @@ -86,10 +84,6 @@ -- | Produce a fresh 'ID', using the given base name as a template. newID :: MonadFreshNames m => Name -> m VName newID s = newName $ VName s 0 - --- | As 'newID', but takes a 'String' for the name template. -newIDFromString :: MonadFreshNames m => String -> m VName -newIDFromString = newID . nameFromString -- | Produce a fresh 'VName', using the given base name as a template. newVName :: MonadFreshNames m => String -> m VName diff --git a/src/Futhark/Optimise/CSE.hs b/src/Futhark/Optimise/CSE.hs --- a/src/Futhark/Optimise/CSE.hs +++ b/src/Futhark/Optimise/CSE.hs @@ -33,7 +33,6 @@ where import Control.Monad.Reader -import qualified Data.Set as S import qualified Data.Map.Strict as M import Futhark.Analysis.Alias @@ -42,7 +41,6 @@ import Futhark.Representation.Aliases (removeFunDefAliases, Aliases, consumedInStms) import qualified Futhark.Representation.Kernels.Kernel as Kernel -import qualified Futhark.Representation.Kernels.KernelExp as KernelExp import qualified Futhark.Representation.SOACS.SOAC as SOAC import qualified Futhark.Representation.ExplicitMemory as ExplicitMemory import Futhark.Transform.Substitute @@ -104,8 +102,8 @@ , mapOnOp = cseInOp } - patElemDiet pe | patElemName pe `S.member` consumed = Consume - | otherwise = Observe + patElemDiet pe | patElemName pe `nameIn` consumed = Consume + | otherwise = Observe cseInStm :: Attributes lore => Names -> Stm lore @@ -134,7 +132,7 @@ where bad cse_arrays pe | Mem{} <- patElemType pe = True | Array{} <- patElemType pe, not cse_arrays = True - | patElemName pe `S.member` consumed = True + | patElemName pe `nameIn` consumed = True | otherwise = False type ExpressionSubstitutions lore = M.Map @@ -177,43 +175,22 @@ CSEState _ cse_arrays <- ask return $ runReader m $ newCSEState cse_arrays -instance CSEInOp op => CSEInOp (Kernel.HostOp lore op) where - cseInOp (Kernel.HostOp op) = Kernel.HostOp <$> cseInOp op +instance (Attributes lore, Aliased lore, + CSEInOp (Op lore), CSEInOp op) => CSEInOp (Kernel.HostOp lore op) where + cseInOp (Kernel.SegOp op) = Kernel.SegOp <$> cseInOp op + cseInOp (Kernel.OtherOp op) = Kernel.OtherOp <$> cseInOp op cseInOp x = return x -instance (Attributes lore, Aliased lore, CSEInOp (Op lore)) => CSEInOp (Kernel.Kernel lore) where +instance (Attributes lore, Aliased lore, CSEInOp (Op lore)) => CSEInOp (Kernel.SegOp lore) where cseInOp = subCSE . - Kernel.mapKernelM - (Kernel.KernelMapper return cseInLambda - (\b -> cseInBody (map (const Observe) $ bodyResult b) b) - return return cseInKernelBody) + Kernel.mapSegOpM + (Kernel.SegOpMapper return cseInLambda cseInKernelBody return) cseInKernelBody :: (Attributes lore, Aliased lore, CSEInOp (Op lore)) => Kernel.KernelBody lore -> CSEM lore (Kernel.KernelBody lore) cseInKernelBody (Kernel.KernelBody bodyattr bnds res) = do Body _ bnds' _ <- cseInBody (map (const Observe) res) $ Body bodyattr bnds [] return $ Kernel.KernelBody bodyattr bnds' res - -instance (Attributes lore, Aliased lore, CSEInOp (Op lore)) => CSEInOp (KernelExp.KernelExp lore) where - cseInOp (KernelExp.Combine cspace ts active body) = - subCSE $ KernelExp.Combine cspace ts active <$> - cseInBody (map (const Observe) ts) body - cseInOp (KernelExp.GroupReduce w lam input) = - subCSE $ KernelExp.GroupReduce w <$> - cseInLambda lam <*> pure input - cseInOp (KernelExp.GroupStream w max_chunk lam nes arrs) = - subCSE $ KernelExp.GroupStream w max_chunk <$> - cseInGroupStreamLambda lam <*> pure nes <*> pure arrs - cseInOp op = return op - -cseInGroupStreamLambda :: (Attributes lore, Aliased lore, CSEInOp (Op lore)) => - KernelExp.GroupStreamLambda lore - -> CSEM lore (KernelExp.GroupStreamLambda lore) -cseInGroupStreamLambda lam = do - body' <- cseInBody (map (const Observe) $ KernelExp.groupStreamAccParams lam) $ - KernelExp.groupStreamLambdaBody lam - return lam { KernelExp.groupStreamLambdaBody = body' } - instance CSEInOp op => CSEInOp (ExplicitMemory.MemOp op) where cseInOp o@ExplicitMemory.Alloc{} = return o diff --git a/src/Futhark/Optimise/DoubleBuffer.hs b/src/Futhark/Optimise/DoubleBuffer.hs --- a/src/Futhark/Optimise/DoubleBuffer.hs +++ b/src/Futhark/Optimise/DoubleBuffer.hs @@ -29,7 +29,6 @@ import Control.Monad.Writer import Control.Monad.Reader import qualified Data.Map.Strict as M -import qualified Data.Set as S import Data.Maybe import Data.List @@ -50,77 +49,41 @@ , passFunction = intraproceduralTransformation optimiseFunDef } --- This pass is written in a slightly weird way because we want to --- apply essentially the same transformation both outside and inside --- kernel bodies, which are different (but similar) representations. --- Thus, the environment is parametrised by the lore and contains the --- function used to transform 'Op's for the lore. - optimiseFunDef :: FunDef ExplicitMemory -> PassM (FunDef ExplicitMemory) optimiseFunDef fundec = modifyNameSource $ \src -> let m = runDoubleBufferM $ inScopeOf fundec $ optimiseBody $ funDefBody fundec (body', src') = runState (runReaderT m env) src in (fundec { funDefBody = body' }, src') - where env = Env mempty optimiseKernelOp doNotTouchLoop - - optimiseKernelOp (Inner (HostOp k)) = do - scope <- castScope <$> askScope - modifyNameSource $ - runState (runReaderT (runDoubleBufferM $ Inner . HostOp <$> optimiseKernel k) $ - Env scope optimiseInKernelOp optimiseLoop) - where optimiseKernel = - mapKernelM identityKernelMapper - { mapOnKernelBody = optimiseBody - , mapOnKernelKernelBody = optimiseKernelBody - , mapOnKernelLambda = optimiseLambda - } - optimiseKernelOp op = return op - - optimiseInKernelOp (Inner (GroupStream w maxchunk lam accs arrs)) = do - lam' <- optimiseGroupStreamLambda lam - return $ Inner $ GroupStream w maxchunk lam' accs arrs - optimiseInKernelOp op = return op - + where env = Env mempty doNotTouchLoop doNotTouchLoop ctx val body = return (mempty, ctx, val, body) -data Env lore = Env { envScope :: Scope lore - , envOptimiseOp :: Op lore -> DoubleBufferM lore (Op lore) - , envOptimiseLoop :: OptimiseLoop lore - } +data Env = Env { envScope :: Scope ExplicitMemory + , envOptimiseLoop :: OptimiseLoop + } -newtype DoubleBufferM lore a = - DoubleBufferM { runDoubleBufferM :: ReaderT (Env lore) (State VNameSource) a } - deriving (Functor, Applicative, Monad, MonadReader (Env lore), MonadFreshNames) +newtype DoubleBufferM a = + DoubleBufferM { runDoubleBufferM :: ReaderT Env (State VNameSource) a } + deriving (Functor, Applicative, Monad, MonadReader Env, MonadFreshNames) -instance Annotations lore => HasScope lore (DoubleBufferM lore) where +instance HasScope ExplicitMemory DoubleBufferM where askScope = asks envScope -instance Annotations lore => LocalScope lore (DoubleBufferM lore) where +instance LocalScope ExplicitMemory DoubleBufferM where localScope scope = local $ \env -> env { envScope = envScope env <> scope } --- | Bunch up all the constraints for less typing. -type LoreConstraints lore inner = - (ExpAttr lore ~ (), BodyAttr lore ~ (), - ExplicitMemorish lore, Op lore ~ MemOp inner, - BinderOps lore) - -optimiseBody :: LoreConstraints lore inner => - Body lore -> DoubleBufferM lore (Body lore) +optimiseBody :: Body ExplicitMemory -> DoubleBufferM (Body ExplicitMemory) optimiseBody body = do bnds' <- optimiseStms $ stmsToList $ bodyStms body return $ body { bodyStms = stmsFromList bnds' } -optimiseStms :: LoreConstraints lore inner => - [Stm lore] -> DoubleBufferM lore [Stm lore] +optimiseStms :: [Stm ExplicitMemory] -> DoubleBufferM [Stm ExplicitMemory] optimiseStms [] = return [] optimiseStms (e:es) = do e_es <- optimiseStm e es' <- localScope (castScope $ scopeOf e_es) $ optimiseStms es return $ e_es ++ es' -optimiseStm :: forall lore inner. - LoreConstraints lore inner => - Stm lore -> DoubleBufferM lore [Stm lore] +optimiseStm :: Stm ExplicitMemory -> DoubleBufferM [Stm ExplicitMemory] optimiseStm (Let pat aux (DoLoop ctx val form body)) = do body' <- localScope (scopeOf form <> scopeOfFParams (map fst $ ctx++val)) $ optimiseBody body @@ -130,44 +93,40 @@ optimiseStm (Let pat aux e) = pure . Let pat aux <$> mapExpM optimise e where optimise = identityMapper { mapOnBody = \_ x -> - -- This type annotation is - -- necessary to prevent the GHC - -- 8.4 type checker from going - -- nuts. - optimiseBody x :: DoubleBufferM lore (Body lore) + optimiseBody x :: DoubleBufferM (Body ExplicitMemory) , mapOnOp = optimiseOp } -optimiseOp :: Op lore -> DoubleBufferM lore (Op lore) -optimiseOp op = do f <- asks envOptimiseOp - f op +optimiseOp :: Op ExplicitMemory + -> DoubleBufferM (Op ExplicitMemory) +optimiseOp (Inner (SegOp op)) = + local inSegOp $ Inner . SegOp <$> mapSegOpM mapper op + where mapper = identitySegOpMapper + { mapOnSegOpLambda = optimiseLambda + , mapOnSegOpBody = optimiseKernelBody + } + inSegOp env = env { envOptimiseLoop = optimiseLoop } +optimiseOp op = return op -optimiseKernelBody :: KernelBody InKernel - -> DoubleBufferM InKernel (KernelBody InKernel) +optimiseKernelBody :: KernelBody ExplicitMemory + -> DoubleBufferM (KernelBody ExplicitMemory) optimiseKernelBody kbody = do stms' <- optimiseStms $ stmsToList $ kernelBodyStms kbody return $ kbody { kernelBodyStms = stmsFromList stms' } -optimiseLambda :: Lambda InKernel -> DoubleBufferM InKernel (Lambda InKernel) +optimiseLambda :: Lambda ExplicitMemory -> DoubleBufferM (Lambda ExplicitMemory) optimiseLambda lam = do body <- localScope (castScope $ scopeOf lam) $ optimiseBody $ lambdaBody lam return lam { lambdaBody = body } -optimiseGroupStreamLambda :: GroupStreamLambda InKernel - -> DoubleBufferM InKernel (GroupStreamLambda InKernel) -optimiseGroupStreamLambda lam = do - body <- localScope (scopeOf lam) $ - optimiseBody $ groupStreamLambdaBody lam - return lam { groupStreamLambdaBody = body } - -type OptimiseLoop lore = - [(FParam lore, SubExp)] -> [(FParam lore, SubExp)] -> Body lore - -> DoubleBufferM lore ([Stm lore], - [(FParam lore, SubExp)], - [(FParam lore, SubExp)], - Body lore) +type OptimiseLoop = + [(FParam ExplicitMemory, SubExp)] -> [(FParam ExplicitMemory, SubExp)] -> Body ExplicitMemory + -> DoubleBufferM ([Stm ExplicitMemory], + [(FParam ExplicitMemory, SubExp)], + [(FParam ExplicitMemory, SubExp)], + Body ExplicitMemory) -optimiseLoop :: LoreConstraints lore inner => OptimiseLoop lore +optimiseLoop :: OptimiseLoop optimiseLoop ctx val body = do -- We start out by figuring out which of the merge variables should -- be double-buffered. @@ -186,20 +145,20 @@ -- | The booleans indicate whether we should also play with the -- initial merge values. -data DoubleBuffer lore = BufferAlloc VName (PrimExp VName) Space Bool - | BufferCopy VName IxFun VName Bool - -- ^ First name is the memory block to copy to, - -- second is the name of the array copy. - | NoBuffer +data DoubleBuffer = BufferAlloc VName (PrimExp VName) Space Bool + | BufferCopy VName IxFun VName Bool + -- ^ First name is the memory block to copy to, + -- second is the name of the array copy. + | NoBuffer deriving (Show) -doubleBufferMergeParams :: (ExplicitMemorish lore, MonadFreshNames m) => - [(FParam lore,SubExp)] - -> [FParam lore] -> Names - -> m [DoubleBuffer lore] +doubleBufferMergeParams :: MonadFreshNames m => + [(FParam ExplicitMemory, SubExp)] + -> [FParam ExplicitMemory] -> Names + -> m [DoubleBuffer] doubleBufferMergeParams ctx_and_res val_params bound_in_loop = evalStateT (mapM buffer val_params) M.empty - where loopVariant v = v `S.member` bound_in_loop || + where loopVariant v = v `nameIn` bound_in_loop || v `elem` map (paramName . fst) ctx_and_res loopInvariantSize (Constant v) = @@ -244,9 +203,8 @@ return NoBuffer _ -> return NoBuffer -allocStms :: LoreConstraints lore inner => - [(FParam lore,SubExp)] -> [DoubleBuffer lore] - -> DoubleBufferM lore ([(FParam lore, SubExp)], [Stm lore]) +allocStms :: [(FParam ExplicitMemory, SubExp)] -> [DoubleBuffer] + -> DoubleBufferM ([(FParam ExplicitMemory, SubExp)], [Stm ExplicitMemory]) allocStms merge = runWriterT . zipWithM allocation merge where allocation m@(Param pname _, _) (BufferAlloc name size space b) = do stms <- lift $ runBinder_ $ do @@ -267,10 +225,8 @@ allocation (f, se) _ = return (f, se) -doubleBufferResult :: (ExplicitMemorish lore, - ExpAttr lore ~ (), BodyAttr lore ~ ()) => - [FParam lore] -> [DoubleBuffer lore] - -> Body lore -> Body lore +doubleBufferResult :: [FParam ExplicitMemory] -> [DoubleBuffer] + -> Body ExplicitMemory -> Body ExplicitMemory doubleBufferResult valparams buffered (Body () bnds res) = let (ctx_res, val_res) = splitAt (length res - length valparams) res (copybnds,val_res') = diff --git a/src/Futhark/Optimise/Fusion.hs b/src/Futhark/Optimise/Fusion.hs --- a/src/Futhark/Optimise/Fusion.hs +++ b/src/Futhark/Optimise/Fusion.hs @@ -84,7 +84,7 @@ Array {} -> IsArray name (LetInfo t) aliases' $ SOAC.identInput $ Ident name t _ -> IsNotArray name $ LetInfo t expand = maybe mempty varEntryAliases . flip M.lookup (varsInScope env) - aliases' = aliases <> mconcat (map expand $ S.toList aliases) + aliases' = aliases <> mconcat (map expand $ namesToList aliases) bindVars :: FusionGEnv -> [(Ident, Names)] -> FusionGEnv bindVars = foldl bindVar @@ -114,16 +114,16 @@ } varAliases :: VName -> FusionGM Names -varAliases v = asks $ S.insert v . maybe mempty varEntryAliases . +varAliases v = asks $ (oneName v<>) . maybe mempty varEntryAliases . M.lookup v . varsInScope varsAliases :: Names -> FusionGM Names -varsAliases = fmap mconcat . mapM varAliases . S.toList +varsAliases = fmap mconcat . mapM varAliases . namesToList checkForUpdates :: FusedRes -> Exp -> FusionGM FusedRes checkForUpdates res (BasicOp (Update src is _)) = do res' <- foldM addVarToInfusible res $ - src : S.toList (mconcat $ map freeIn is) + src : namesToList (mconcat $ map freeIn is) aliases <- varAliases src let inspectKer k = k { inplace = aliases <> inplace k } return res' { kernels = M.map inspectKer $ kernels res' } @@ -146,11 +146,11 @@ Just (IsArray src' _ aliases input) -> env { varsInScope = M.insert vname - (IsArray src' (LetInfo attr) (srcname `S.insert` aliases) $ + (IsArray src' (LetInfo attr) (oneName srcname <> aliases) $ trns `SOAC.addTransform` input) $ varsInScope env } - _ -> bindVar env (patElemIdent pe, S.singleton vname) + _ -> bindVar env (patElemIdent pe, oneName vname) where vname = patElemName pe attr = patElemAttr pe @@ -179,7 +179,7 @@ , passFunction = simplifySOACS <=< renameProg <=< intraproceduralTransformation fuseFun } -fuseFun :: FunDef -> PassM FunDef +fuseFun :: FunDef SOACS -> PassM (FunDef SOACS) fuseFun fun = do let env = FusionGEnv { soacs = M.empty , varsInScope = M.empty @@ -191,12 +191,12 @@ then return fun else liftEitherM $ runFusionGatherM (fuseInFun k fun) env -fusionGatherFun :: FunDef -> FusionGM FusedRes +fusionGatherFun :: FunDef SOACS -> FusionGM FusedRes fusionGatherFun fundec = bindingParams (funDefParams fundec) $ fusionGatherBody mempty $ funDefBody fundec -fuseInFun :: FusedRes -> FunDef -> FusionGM FunDef +fuseInFun :: FusedRes -> FunDef SOACS -> FusionGM (FunDef SOACS) fuseInFun res fundec = do body' <- bindingParams (funDefParams fundec) $ bindRes res $ @@ -247,12 +247,12 @@ FusedRes (rsucc res1 || rsucc res2) (outArr res1 `M.union` outArr res2) (M.unionWith S.union (inpArr res1) (inpArr res2) ) - (infusible res1 `S.union` infusible res2) + (infusible res1 <> infusible res2) (kernels res1 `M.union` kernels res2) instance Monoid FusedRes where mempty = FusedRes { rsucc = False, outArr = M.empty, inpArr = M.empty, - infusible = S.empty, kernels = M.empty } + infusible = mempty, kernels = M.empty } isInpArrInResModKers :: FusedRes -> S.Set KernName -> VName -> Bool isInpArrInResModKers ress kers nm = @@ -332,7 +332,7 @@ (inp_nms, other_nms) <- soacInputs soac -- Assumption: the free vars in lambda are already in @infusible res@. let out_nms = patternNames out_idds - isInfusible = (`S.member` infusible res) + isInfusible = (`nameIn` infusible res) is_screma = case orig_soac of SOAC.Screma _ form _ -> (isJust (isRedomapSOAC form) || isJust (isScanomapSOAC form)) && @@ -353,8 +353,8 @@ -- -- (ii) check whether fusing @soac@ will violate any in-place update -- restriction, e.g., would move an input array past its in-place update. - let all_used_names = S.toList $ S.unions [lam_used_nms, S.fromList inp_nms, S.fromList other_nms] - has_inplace ker = any (`S.member` inplace ker) all_used_names + let all_used_names = namesToList $ mconcat [lam_used_nms, namesFromList inp_nms, namesFromList other_nms] + has_inplace ker = any (`nameIn` inplace ker) all_used_names ok_inplace = not $ any has_inplace old_kers -- -- (iii) there are some kernels that use some of `out_idds' as inputs @@ -366,11 +366,11 @@ -- (ii) will also become part of the infusible set the inparr vars -- that also appear as inparr of another kernel, -- BUT which said kernel is not the one we are fusing with (now)! - let mod_kerS = if fusible_ker then S.fromList oldker_nms else S.empty + let mod_kerS = if fusible_ker then S.fromList oldker_nms else mempty let used_inps = filter (isInpArrInResModKers res mod_kerS) inp_nms - let ufs = S.unions [infusible res, S.fromList used_inps, - S.fromList other_nms `S.difference` - S.fromList (map SOAC.inputArray $ SOAC.inputs soac)] + let ufs = mconcat [infusible res, namesFromList used_inps, + namesFromList other_nms `namesSubtract` + namesFromList (map SOAC.inputArray $ SOAC.inputs soac)] let comb = M.unionWith S.union if not fusible_ker then @@ -422,7 +422,7 @@ -- try producer-consumer fusion (ok_kers_compat, fused_kers) <- do kers <- forM to_fuse_kers $ - attemptFusion S.empty (patternNames out_idds) soac consumed + attemptFusion mempty (patternNames out_idds) soac consumed case sequence kers of Nothing -> return (False, []) Just kers' -> return (True, map certifyKer kers') @@ -434,7 +434,7 @@ horizontGreedyFuse rem_bnds res (out_idds, cs, soac, consumed) = do (inp_nms, _) <- soacInputs soac let out_nms = patternNames out_idds - infusible_nms = S.fromList $ filter (`S.member` infusible res) out_nms + infusible_nms = namesFromList $ filter (`nameIn` infusible res) out_nms out_arr_nms = case soac of -- the accumulator result cannot be fused! SOAC.Screma _ (ScremaForm (_, scan_nes) reds _) _ -> @@ -472,15 +472,15 @@ -- check that we still try fusion and that the intermediate -- bindings do not use the results of cur_ker let curker_outnms = outNames cur_ker - curker_outset = S.fromList curker_outnms - new_ufus_nms = S.fromList $ outNames ker ++ S.toList ufus_nms + curker_outset = namesFromList curker_outnms + new_ufus_nms = namesFromList $ outNames ker ++ namesToList ufus_nms -- disable horizontal fusion in the case when an output array of -- producer SOAC is a non-trivially transformed input of the consumer out_transf_ok = let ker_inp = SOAC.inputs $ fsoac ker - unfuse1 = S.fromList (map SOAC.inputArray ker_inp) `S.difference` - S.fromList (mapMaybe SOAC.isVarInput ker_inp) - unfuse2 = S.intersection curker_outset ufus_nms - in S.null $ S.intersection unfuse1 unfuse2 + unfuse1 = namesFromList (map SOAC.inputArray ker_inp) `namesSubtract` + namesFromList (mapMaybe SOAC.isVarInput ker_inp) + unfuse2 = namesIntersection curker_outset ufus_nms + in not $ unfuse1 `namesIntersect` unfuse2 -- Disable horizontal fusion if consumer has any -- output transforms. cons_no_out_transf = SOAC.nullTransforms $ outputTransform ker @@ -490,16 +490,17 @@ case maybesoac of -- check that consumer's lambda body does not use -- directly the produced arrays (e.g., see noFusion3.fut). - Right conssoac -> return $ S.null $ S.intersection curker_outset $ - freeIn $ lambdaBody $ SOAC.lambda conssoac + Right conssoac -> return $ not $ + curker_outset + `namesIntersect` + freeIn (lambdaBody $ SOAC.lambda conssoac) Left _ -> return True let interm_bnds_ok = cur_ok && consumer_ok && out_transf_ok && cons_no_out_transf && foldl (\ok bnd-> ok && -- hardwired to False after first fail -- (i) check that the in-between bindings do -- not use the result of current kernel OR - S.null ( S.intersection curker_outset $ - freeIn (stmExp bnd) ) || + not (curker_outset `namesIntersect` freeIn (stmExp bnd)) || --(ii) that the pattern-binding corresponds to -- the result of the consumer kernel; in the -- latter case it means it corresponds to a @@ -508,11 +509,11 @@ not ( null $ curker_outnms `L.intersect` patternNames (stmPattern bnd)) ) True (drop (prev_ind+1) $ take bnd_ind rem_bnds) - if not interm_bnds_ok then return (False,n,bnd_ind,cur_ker,S.empty) + if not interm_bnds_ok then return (False,n,bnd_ind,cur_ker,mempty) else do new_ker <- attemptFusion ufus_nms (outNames cur_ker) (fsoac cur_ker) (fusedConsumed cur_ker) ker case new_ker of - Nothing -> return (False, n,bnd_ind,cur_ker,S.empty) + Nothing -> return (False, n,bnd_ind,cur_ker,mempty) Just krn-> return (True,n+1,bnd_ind,krn,new_ufus_nms) ) (True,0,0,soac_kernel,infusible_nms) kernminds' @@ -613,14 +614,14 @@ -- We put the variables produced by Scatter into the infusible -- set to force horizontal fusion. It is not possible to -- producer/consumer-fuse Scatter anyway. - fres' <- addNamesToInfusible fres $ S.fromList $ patternNames pat + fres' <- addNamesToInfusible fres $ namesFromList $ patternNames pat mapLike fres' soac lam Right soac@(SOAC.GenReduce _ _ lam _) -> do -- We put the variables produced by GenReduce into the infusible -- set to force horizontal fusion. It is not possible to -- producer/consumer-fuse GenReduce anyway. - fres' <- addNamesToInfusible fres $ S.fromList $ patternNames pat + fres' <- addNamesToInfusible fres $ namesFromList $ patternNames pat mapLike fres' soac lam Right soac@(SOAC.Screma _ (ScremaForm (scan_lam, scan_nes) reds map_lam) _) -> @@ -652,7 +653,7 @@ consumed = consumedInExp $ Alias.analyseExp e reduceLike soac lambdas nes = do - (used_lam, lres) <- foldM fusionGatherLam (S.empty, fres) lambdas + (used_lam, lres) <- foldM fusionGatherLam (mempty, fres) lambdas bres <- bindingFamily pat $ fusionGatherStms lres bnds res bres' <- foldM fusionGatherSubExp bres nes consumed' <- varsAliases consumed @@ -660,7 +661,7 @@ mapLike fres' soac lambda = do bres <- bindingFamily pat $ fusionGatherStms fres' bnds res - (used_lam, blres) <- fusionGatherLam (S.empty, bres) lambda + (used_lam, blres) <- fusionGatherLam (mempty, bres) lambda consumed' <- varsAliases consumed greedyFuse rem_bnds used_lam blres (pat, cs, soac, consumed') @@ -687,7 +688,7 @@ -- make the inpArr infusible, so that they -- cannot be fused from outside the loop: let (inp_arrs, _) = unzip $ M.toList $ inpArr new_res - let new_res' = new_res { infusible = foldl (flip S.insert) (infusible new_res) inp_arrs } + let new_res' = new_res { infusible = infusible new_res <> mconcat (map oneName inp_arrs) } -- merge new_res with fres' return $ new_res' <> fres' @@ -717,7 +718,7 @@ fusionGatherSubExp fres _ = return fres addNamesToInfusible :: FusedRes -> Names -> FusionGM FusedRes -addNamesToInfusible fres = foldM addVarToInfusible fres . S.toList +addNamesToInfusible fres = foldM addVarToInfusible fres . namesToList addVarToInfusible :: FusedRes -> VName -> FusionGM FusedRes addVarToInfusible fres name = do @@ -725,23 +726,23 @@ let name' = case trns of Nothing -> name Just (SOAC.Input _ orig _) -> orig - return fres { infusible = S.insert name' $ infusible fres } + return fres { infusible = oneName name' <> infusible fres } -- Lambdas create a new scope. Disallow fusing from outside lambda by -- adding inp_arrs to the infusible set. -fusionGatherLam :: (Names, FusedRes) -> Lambda -> FusionGM (S.Set VName, FusedRes) +fusionGatherLam :: (Names, FusedRes) -> Lambda -> FusionGM (Names, FusedRes) fusionGatherLam (u_set,fres) (Lambda idds body _) = do new_res <- bindingParams idds $ fusionGatherBody mempty body -- make the inpArr infusible, so that they -- cannot be fused from outside the lambda: - let inp_arrs = S.fromList $ M.keys $ inpArr new_res - let unfus = infusible new_res `S.union` inp_arrs + let inp_arrs = namesFromList $ M.keys $ inpArr new_res + let unfus = infusible new_res <> inp_arrs bnds <- M.keys <$> asks varsInScope - let unfus' = unfus `S.intersection` S.fromList bnds + let unfus' = unfus `namesIntersection` namesFromList bnds -- merge fres with new_res' let new_res' = new_res { infusible = unfus' } -- merge new_res with fres' - return (u_set `S.union` unfus', new_res' <> fres) + return (u_set <> unfus', new_res' <> fres) ------------------------------------------------------------- ------------------------------------------------------------- @@ -843,7 +844,7 @@ simplifyAndFuseInLambda lam = do let args = replicate (length $ lambdaParams lam) Nothing lam' <- simplifyLambda lam args - (_, nfres) <- fusionGatherLam (S.empty, mkFreshFusionRes) lam' + (_, nfres) <- fusionGatherLam (mempty, mkFreshFusionRes) lam' let nfres' = cleanFusionResult nfres bindRes nfres' $ fuseInLambda lam' @@ -869,18 +870,18 @@ _ -> return $ removeOpAliases soac where consumed = consumedInOp soac - newly_consumed = consumed `S.difference` was_consumed + newly_consumed = consumed `namesSubtract` was_consumed copyConsumedArr a - | a `S.member` newly_consumed = + | a `nameIn` newly_consumed = letExp (baseString a <> "_copy") $ BasicOp $ Copy a | otherwise = return a copyFreeInLambda lam = do - let free_consumed = consumedByLambda lam `S.difference` - S.fromList (map paramName $ lambdaParams lam) + let free_consumed = consumedByLambda lam `namesSubtract` + namesFromList (map paramName $ lambdaParams lam) (bnds, subst) <- - foldM copyFree (mempty, mempty) $ S.toList free_consumed + foldM copyFree (mempty, mempty) $ namesToList free_consumed let lam' = Aliases.removeLambdaAliases lam return $ if null bnds then lam' @@ -905,13 +906,13 @@ mkFreshFusionRes :: FusedRes mkFreshFusionRes = FusedRes { rsucc = False, outArr = M.empty, inpArr = M.empty, - infusible = S.empty, kernels = M.empty } + infusible = mempty, kernels = M.empty } mergeFusionRes :: FusedRes -> FusedRes -> FusionGM FusedRes mergeFusionRes res1 res2 = do - let ufus_mres = infusible res1 `S.union` infusible res2 - inp_both <- expandSoacInpArr $ M.keys $ inpArr res1 `M.intersection` inpArr res2 - let m_unfus = foldl (flip S.insert) ufus_mres inp_both + let ufus_mres = infusible res1 <> infusible res2 + inp_both <- expandSoacInpArr $ M.keys $ inpArr res1 `M.intersection` inpArr res2 + let m_unfus = ufus_mres <> mconcat (map oneName inp_both) return $ FusedRes (rsucc res1 || rsucc res2) (outArr res1 `M.union` outArr res2) (M.unionWith S.union (inpArr res1) (inpArr res2) ) diff --git a/src/Futhark/Optimise/Fusion/Composing.hs b/src/Futhark/Optimise/Fusion/Composing.hs --- a/src/Futhark/Optimise/Fusion/Composing.hs +++ b/src/Futhark/Optimise/Fusion/Composing.hs @@ -19,7 +19,6 @@ import Data.List import qualified Data.Map.Strict as M -import qualified Data.Set as S import Data.Maybe import qualified Futhark.Analysis.HORepresentation.SOAC as SOAC @@ -106,7 +105,7 @@ Nothing -> Nothing --should not be reached! outinsrev = M.fromList $ mapMaybe getVarParPair $ M.toList outins unfusible outname - | outname `S.member` unfus_nms = + | outname `nameIn` unfus_nms = outname `M.lookup` M.union outinsrev (M.fromList out1) unfusible _ = Nothing unfus_vars= mapMaybe (unfusible . fst) out1 @@ -162,7 +161,7 @@ -- (i) we remove the accumulator formal paramter and corresponding -- (body) result from from redomap's fold-lambda body let p_num_nes = length p_scan_nes + length p_red_nes - unfus_arrs = filter (`S.member` unfus_nms) outVars + unfus_arrs = filter (`nameIn` unfus_nms) outVars p_lam_body = lambdaBody p_lam (p_lam_scan_ts, p_lam_red_ts, p_lam_map_ts) = splitAt3 (length p_scan_nes) (length p_red_nes) $ lambdaReturnType p_lam @@ -176,11 +175,11 @@ -- @outPairs@, then ``map o redomap'' fuse the two lambdas -- (in the usual way), and construct the extra return types -- for the arrays that fall through. - (res_lam, new_inp) = fuseMaps (S.fromList unfus_arrs) p_lam_hacked p_inparr + (res_lam, new_inp) = fuseMaps (namesFromList unfus_arrs) p_lam_hacked p_inparr (drop p_num_nes outPairs) c_lam c_inparr (res_lam_scan_ts, res_lam_red_ts, res_lam_map_ts) = splitAt3 (length c_scan_nes) (length c_red_nes) $ lambdaReturnType res_lam - (_,extra_map_ts) = unzip $ filter (\(nm,_)->elem nm unfus_arrs) $ + (_,extra_map_ts) = unzip $ filter (\(nm,_)-> nm `elem` unfus_arrs) $ zip (drop p_num_nes outVars) $ drop p_num_nes $ lambdaReturnType p_lam diff --git a/src/Futhark/Optimise/Fusion/LoopKernel.hs b/src/Futhark/Optimise/Fusion/LoopKernel.hs --- a/src/Futhark/Optimise/Fusion/LoopKernel.hs +++ b/src/Futhark/Optimise/Fusion/LoopKernel.hs @@ -208,7 +208,7 @@ (ps', inps') = case (unzip $ filter (used . fst) pInps, pInps) of (([], []), (p,inp):_) -> ([p], [inp]) ((ps_, inps_), _) -> (ps_, inps_) - used p = paramName p `S.member` freeVars + used p = paramName p `nameIn` freeVars freeVars = freeIn $ lambdaBody l -- | Check that the consumer uses at least one output of the producer @@ -230,13 +230,13 @@ -- into soac_c. let soac_c = fsoac ker inp_p_arr = SOAC.inputs soac_p - horizFuse= not (S.null unfus_set) && + horizFuse= unfus_set /= mempty && SOAC.width soac_p == SOAC.width soac_c inp_c_arr = SOAC.inputs soac_c lam_p = SOAC.lambda soac_p lam_c = SOAC.lambda soac_c w = SOAC.width soac_p - returned_outvars = filter (`S.member` unfus_set) outVars + returned_outvars = filter (`nameIn` unfus_set) outVars success res_outnms res_soac = do let fusedVars_new = fusedVars ker++outVars -- Avoid name duplication, because the producer lambda is not @@ -255,7 +255,7 @@ let mapLikeFusionCheck = let (res_lam, new_inp) = fuseMaps unfus_set lam_p inp_p_arr outPairs lam_c inp_c_arr - (extra_nms,extra_rtps) = unzip $ filter ((`S.member` unfus_set) . fst) $ + (extra_nms,extra_rtps) = unzip $ filter ((`nameIn` unfus_set) . fst) $ zip outVars $ map (stripArray 1) $ SOAC.typeOf soac_p res_lam' = res_lam { lambdaReturnType = lambdaReturnType res_lam ++ extra_rtps } in (extra_nms, res_lam', new_inp) @@ -302,7 +302,7 @@ | isJust $ isMapSOAC form, -- 1. all arrays produced by the map are ONLY used (consumed) -- by the scatter, i.e., not used elsewhere. - not (any (`S.member` unfus_set) outVars), + not (any (`nameIn` unfus_set) outVars), -- 2. all arrays produced by the map are input to the scatter. mapWriteFusionOK outVars ker -> do let (extra_nms, res_lam', new_inp) = mapLikeFusionCheck @@ -319,7 +319,7 @@ | isJust $ isMapSOAC form, -- 1. all arrays produced by the map are ONLY used (consumed) -- by the genreduce, i.e., not used elsewhere. - not (any (`S.member` unfus_set) outVars), + not (any (`nameIn` unfus_set) outVars), -- 2. all arrays produced by the map are input to the scatter. mapWriteFusionOK outVars ker -> do let (extra_nms, res_lam', new_inp) = mapLikeFusionCheck @@ -481,7 +481,7 @@ inp2_arr res_lam'' = res_lam' { lambdaParams = chunk1 : lambdaParams res_lam' } unfus_accs = take (length nes1) outVars - unfus_arrs = filter (`S.member` unfus_set) outVars + unfus_arrs = filter (`nameIn` unfus_set) outVars res_form <- mergeForms form2 form1 return (unfus_accs ++ out_kernms ++ unfus_arrs, SOAC.Stream w2 res_form res_lam'' new_inp ) diff --git a/src/Futhark/Optimise/InPlaceLowering.hs b/src/Futhark/Optimise/InPlaceLowering.hs --- a/src/Futhark/Optimise/InPlaceLowering.hs +++ b/src/Futhark/Optimise/InPlaceLowering.hs @@ -69,7 +69,6 @@ import Control.Monad.RWS import qualified Data.Map.Strict as M -import qualified Data.Set as S import Futhark.Analysis.Alias import Futhark.Representation.Aliases @@ -161,26 +160,16 @@ where optimise = identityMapper { mapOnBody = const optimiseBody } onKernelOp :: OnOp Kernels -onKernelOp (HostOp op) = do - old_scope <- askScope - modifyNameSource $ runForwardingM lowerUpdateInKernel onKernelExp $ - bindingScope (castScope old_scope <> scopeOfKernelSpace (kernelSpace op)) $ do - let mapper = identityKernelMapper { mapOnKernelKernelBody = onKernelBody } - onKernelBody kbody = do - stms <- deepen $ optimiseStms (stmsToList (kernelBodyStms kbody)) $ - mapM_ seenVar $ freeIn $ kernelBodyResult kbody - return kbody { kernelBodyStms = stmsFromList stms } - HostOp <$> mapKernelM mapper op +onKernelOp (SegOp op) = + bindingScope (scopeOfSegSpace (segSpace op)) $ do + let mapper = identitySegOpMapper { mapOnSegOpBody = onKernelBody } + onKernelBody kbody = do + stms <- deepen $ optimiseStms (stmsToList (kernelBodyStms kbody)) $ + mapM_ seenVar $ namesToList $ freeIn $ kernelBodyResult kbody + return kbody { kernelBodyStms = stmsFromList stms } + SegOp <$> mapSegOpM mapper op onKernelOp op = return op -onKernelExp :: OnOp InKernel -onKernelExp (GroupStream w maxchunk lam accs arrs) = do - lam_body <- bindingScope (scopeOf lam) $ - optimiseBody $ groupStreamLambdaBody lam - let lam' = lam { groupStreamLambdaBody = lam_body } - return $ GroupStream w maxchunk lam' accs arrs -onKernelExp op = return op - data Entry lore = Entry { entryNumber :: Int , entryAliases :: Names , entryDepth :: Int @@ -314,7 +303,7 @@ aliases <- asks $ maybe mempty entryAliases . M.lookup name . topDownTable - tell $ mempty { bottomUpSeen = S.insert name aliases } + tell $ mempty { bottomUpSeen = oneName name <> aliases } tapBottomUp :: ForwardingM lore a -> ForwardingM lore (a, BottomUp lore) tapBottomUp m = do (x,bup) <- listen m @@ -328,7 +317,7 @@ -- Checks condition (2) available <- [i,Var src] `areAvailableBefore` v -- ...subcondition, the certificates must also. - certs_available <- map Var (S.toList $ freeIn cs) `areAvailableBefore` v + certs_available <- map Var (namesToList $ freeIn cs) `areAvailableBefore` v -- Check condition (3) samebody <- isInCurrentBody v -- Check condition (6) diff --git a/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs b/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs --- a/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs +++ b/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs @@ -1,12 +1,10 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} module Futhark.Optimise.InPlaceLowering.LowerIntoStm - ( - lowerUpdateInKernel - , lowerUpdateKernels - , LowerUpdate - , DesiredUpdate (..) - ) where + ( lowerUpdateKernels + , LowerUpdate + , DesiredUpdate (..) + ) where import Control.Monad import Control.Monad.Writer @@ -14,7 +12,6 @@ import Data.Maybe (mapMaybe) import Data.Either import qualified Data.Map as M -import qualified Data.Set as S import Futhark.Representation.AST.Attributes.Aliases import Futhark.Representation.Aliases @@ -67,30 +64,27 @@ lowerUpdateKernels :: MonadFreshNames m => LowerUpdate Kernels m lowerUpdateKernels _ - (Let (Pattern [] [PatElem v v_attr]) aux (Op (HostOp (Kernel debug kspace ts kbody)))) + (Let (Pattern [] [PatElem v v_attr]) aux (Op (SegOp (SegMap lvl space ts kbody)))) [update@(DesiredUpdate bindee_nm bindee_attr cs _src is val)] | v == val = do - kbody' <- lowerUpdateIntoKernel update kspace kbody + kbody' <- lowerUpdateIntoKernel update space kbody let is' = fullSlice (typeOf bindee_attr) is Just $ return [certify (stmAuxCerts aux <> cs) $ mkLet [] [Ident bindee_nm $ typeOf bindee_attr] $ - Op $ HostOp $ Kernel debug kspace ts kbody', + Op $ SegOp $ SegMap lvl space ts kbody', mkLet [] [Ident v $ typeOf v_attr] $ BasicOp $ Index bindee_nm is'] lowerUpdateKernels scope stm updates = lowerUpdate scope stm updates -lowerUpdateInKernel :: MonadFreshNames m => LowerUpdate InKernel m -lowerUpdateInKernel = lowerUpdate - lowerUpdateIntoKernel :: DesiredUpdate (LetAttr (Aliases Kernels)) - -> KernelSpace -> KernelBody (Aliases InKernel) - -> Maybe (KernelBody (Aliases InKernel)) + -> SegSpace -> KernelBody (Aliases Kernels) + -> Maybe (KernelBody (Aliases Kernels)) lowerUpdateIntoKernel update kspace kbody = do - [ThreadsReturn se] <- Just $ kernelBodyResult kbody + [Returns se] <- Just $ kernelBodyResult kbody is' <- mapM dimFix is - let ret = WriteReturn (arrayDims $ snd bindee_attr) src [(is'++map Var gtids, se)] + let ret = WriteReturns (arrayDims $ snd bindee_attr) src [(is'++map Var gtids, se)] return kbody { kernelBodyResult = [ret] } where DesiredUpdate _bindee_nm bindee_attr _cs src is _val = update - gtids = map fst $ spaceDimensions kspace + gtids = map fst $ unSegSpace kspace lowerUpdateIntoLoop :: (Bindable lore, BinderOps lore, Aliased lore, LetAttr lore ~ (als, Type), @@ -138,7 +132,7 @@ -- Safety condition (8). forM_ (zip val $ bodyAliases body) $ \((p, _), als) -> - guard $ not $ paramName p `S.member` als + guard $ not $ paramName p `nameIn` als mk_in_place_map <- summariseLoop updates usedInBody resmap val @@ -151,10 +145,10 @@ (body_res, res_bnds) <- manipulateResult in_place_map idxsubsts' let body' = mkBody (newbnds<>res_bnds) body_res return (prebnds, postbnds, ctxpat, valpat, ctx, val', body') - where usedInBody = S.unions $ map expandAliases $ S.toList $ freeIn body <> freeIn form + where usedInBody = mconcat $ map expandAliases $ namesToList $ freeIn body <> freeIn form expandAliases v = case M.lookup v scope of - Just (LetInfo attr) -> S.insert v $ aliasesOf attr - _ -> S.singleton v + Just (LetInfo attr) -> oneName v <> aliasesOf attr + _ -> oneName v resmap = zip (bodyResult body) $ patternValueIdents pat mkMerges :: (MonadFreshNames m, Bindable lore) => @@ -203,7 +197,7 @@ sequence <$> zipWithM summariseLoopResult resmap merge where summariseLoopResult (se, v) (fparam, mergeinit) | Just update <- find (updateHasValue $ identName v) updates = - if updateSource update `S.member` usedInBody + if updateSource update `nameIn` usedInBody then Nothing else if hasLoopInvariantShape fparam then Just $ do lowered_array <- newVName "lowered_array" diff --git a/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs b/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs --- a/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs +++ b/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs @@ -13,7 +13,6 @@ import Control.Monad import qualified Data.Map.Strict as M -import qualified Data.Set as S import Futhark.Representation.AST.Attributes.Aliases import Futhark.Representation.AST @@ -92,7 +91,7 @@ []) substs' consumingSubst substs' _ = return substs' - in foldM consumingSubst substs . S.toList . consumedInExp + in foldM consumingSubst substs . namesToList . consumedInExp substituteIndicesInSubExp :: MonadBinder m => IndexSubstitutions (LetAttr (Lore m)) diff --git a/src/Futhark/Optimise/InliningDeadFun.hs b/src/Futhark/Optimise/InliningDeadFun.hs --- a/src/Futhark/Optimise/InliningDeadFun.hs +++ b/src/Futhark/Optimise/InliningDeadFun.hs @@ -15,29 +15,37 @@ import qualified Data.Set as S import Futhark.Representation.SOACS +import Futhark.Representation.SOACS.Simplify (simplifyFun) import Futhark.Transform.Rename import Futhark.Analysis.CallGraph import Futhark.Binder import Futhark.Pass -aggInlining :: CallGraph -> [FunDef] -> [FunDef] -aggInlining cg = filter keep . recurse - where noInterestingCalls :: S.Set Name -> FunDef -> Bool +aggInlining :: MonadFreshNames m => CallGraph -> [FunDef SOACS] -> m [FunDef SOACS] +aggInlining cg = fmap (filter keep) . recurse + where noInterestingCalls :: S.Set Name -> FunDef SOACS -> Bool noInterestingCalls interesting fundec = case M.lookup (funDefName fundec) cg of Just calls | not $ any (`elem` interesting') calls -> True _ -> False where interesting' = funDefName fundec `S.insert` interesting - recurse funs = + -- We apply simplification after every round of inlining, + -- because it is more efficient to shrink the program as soon + -- as possible, rather than wait until it has balooned after + -- full inlining. + recurse funs = do let interesting = S.fromList $ map funDefName funs (to_be_inlined, to_inline_in) = partition (noInterestingCalls interesting) funs inlined_but_entry_points = filter (isJust . funDefEntryPoint) to_be_inlined - in if null to_be_inlined then funs - else inlined_but_entry_points ++ - recurse (map (`doInlineInCaller` to_be_inlined) to_inline_in) + if null to_be_inlined + then return funs + else do let onFun = simplifyFun <=< renameFun . + (`doInlineInCaller` to_be_inlined) + to_inline_in' <- recurse =<< mapM onFun to_inline_in + return $ inlined_but_entry_points ++ to_inline_in' keep fundec = isJust (funDefEntryPoint fundec) || callsRecursive fundec @@ -54,12 +62,12 @@ -- not call any other functions. Further extensions that transform a -- tail-recursive function to a do or while loop, should do the transformation -- first and then do the inlining. -doInlineInCaller :: FunDef -> [FunDef] -> FunDef +doInlineInCaller :: FunDef SOACS -> [FunDef SOACS] -> FunDef SOACS doInlineInCaller (FunDef entry name rtp args body) inlcallees = let body' = inlineInBody inlcallees body in FunDef entry name rtp args body' -inlineInBody :: [FunDef] -> Body -> Body +inlineInBody :: [FunDef SOACS] -> Body -> Body inlineInBody inlcallees (Body attr stms res) = Body attr stms' res where stms' = stmsFromList (concatMap inline $ stmsToList stms) @@ -84,21 +92,21 @@ notNoLoc :: SrcLoc -> Bool notNoLoc = (/=NoLoc) . locOf -inliner :: Monad m => [FunDef] -> Mapper SOACS SOACS m +inliner :: Monad m => [FunDef SOACS] -> Mapper SOACS SOACS m inliner funs = identityMapper { mapOnBody = const $ return . inlineInBody funs , mapOnOp = return . inlineInSOAC funs } -inlineInSOAC :: [FunDef] -> SOAC SOACS -> SOAC SOACS +inlineInSOAC :: [FunDef SOACS] -> SOAC SOACS -> SOAC SOACS inlineInSOAC inlcallees = runIdentity . mapSOACM identitySOACMapper { mapOnSOACLambda = return . inlineInLambda inlcallees } -inlineInStm :: [FunDef] -> Stm -> Stm +inlineInStm :: [FunDef SOACS] -> Stm -> Stm inlineInStm inlcallees (Let pat aux e) = Let pat aux $ mapExp (inliner inlcallees) e -inlineInLambda :: [FunDef] -> Lambda -> Lambda +inlineInLambda :: [FunDef SOACS] -> Lambda -> Lambda inlineInLambda inlcallees (Lambda params body ret) = Lambda params (inlineInBody inlcallees body) ret @@ -131,7 +139,7 @@ } where pass prog = do let cg = buildCallGraph prog - renameProg . Prog . aggInlining cg . progFunctions =<< renameProg prog + Prog <$> aggInlining cg (progFunctions prog) -- | @removeDeadFunctions prog@ removes the functions that are unreachable from -- the main function from the program. diff --git a/src/Futhark/Optimise/Simplify/ClosedForm.hs b/src/Futhark/Optimise/Simplify/ClosedForm.hs --- a/src/Futhark/Optimise/Simplify/ClosedForm.hs +++ b/src/Futhark/Optimise/Simplify/ClosedForm.hs @@ -9,13 +9,13 @@ module Futhark.Optimise.Simplify.ClosedForm ( foldClosedForm , loopClosedForm + , VarLookup ) where import Control.Monad import Data.Maybe import qualified Data.Map.Strict as M -import qualified Data.Set as S import Futhark.Construct import Futhark.Representation.AST @@ -110,9 +110,7 @@ (accparams, _) = splitAt (length accs) params res = bodyResult body - nonFree = boundInBody body <> - S.fromList params <> - untouchable + nonFree = boundInBody body <> namesFromList params <> untouchable checkResult (p, Var v) (accparam, acc) | Just (BasicOp (BinOp bop x y)) <- M.lookup v bndMap = do @@ -145,7 +143,7 @@ asFreeSubExp :: SubExp -> Maybe SubExp asFreeSubExp (Var v) - | S.member v nonFree = M.lookup v knownBnds + | v `nameIn` nonFree = M.lookup v knownBnds asFreeSubExp se = Just se properIntSize Int32 = Just $ return size diff --git a/src/Futhark/Optimise/Simplify/Engine.hs b/src/Futhark/Optimise/Simplify/Engine.hs --- a/src/Futhark/Optimise/Simplify/Engine.hs +++ b/src/Futhark/Optimise/Simplify/Engine.hs @@ -53,25 +53,20 @@ , simplifyLambdaNoHoisting , simplifyParam , bindLParams - , bindChunkLParams - , bindLoopVar - , enterLoop , simplifyBody , SimplifiedBody , blockIf - , constructBody - , protectIf , module Futhark.Optimise.Simplify.Lore ) where import Control.Monad.Writer -import Control.Monad.RWS.Strict +import Control.Monad.Reader +import Control.Monad.State.Strict import Data.Either import Data.List import Data.Maybe -import qualified Data.Set as S import Futhark.Representation.AST import Futhark.Representation.AST.Attributes.Aliases @@ -96,7 +91,7 @@ } noExtraHoistBlockers :: HoistBlockers lore -noExtraHoistBlockers = HoistBlockers neverBlocks neverBlocks neverBlocks (const S.empty) (const False) +noExtraHoistBlockers = HoistBlockers neverBlocks neverBlocks neverBlocks (const mempty) (const False) data Env lore = Env { envRules :: RuleBook (Wise lore) , envHoistBlockers :: HoistBlockers lore @@ -120,28 +115,27 @@ , mkLetNamesS :: ST.SymbolTable (Wise lore) -> [VName] -> Exp (Wise lore) -> SimpleM lore (Stm (Wise lore), Stms (Wise lore)) - , simplifyOpS :: SimplifyOp lore + , simplifyOpS :: SimplifyOp lore (Op lore) } -type SimplifyOp lore = Op lore -> SimpleM lore (OpWithWisdom (Op lore), Stms (Wise lore)) +type SimplifyOp lore op = op -> SimpleM lore (OpWithWisdom op, Stms (Wise lore)) bindableSimpleOps :: (SimplifiableLore lore, Bindable lore) => - SimplifyOp lore -> SimpleOps lore + SimplifyOp lore (Op lore) -> SimpleOps lore bindableSimpleOps = SimpleOps mkExpAttrS' mkBodyS' mkLetNamesS' where mkExpAttrS' _ pat e = return $ mkExpAttr pat e mkBodyS' _ bnds res = return $ mkBody bnds res mkLetNamesS' _ name e = (,) <$> mkLetNames name e <*> pure mempty newtype SimpleM lore a = - SimpleM (RWS (SimpleOps lore, Env lore) Certificates (VNameSource, Bool) a) + SimpleM (ReaderT (SimpleOps lore, Env lore) (State (VNameSource, Bool, Certificates)) a) deriving (Applicative, Functor, Monad, MonadReader (SimpleOps lore, Env lore), - MonadState (VNameSource, Bool), - MonadWriter Certificates) + MonadState (VNameSource, Bool, Certificates)) instance MonadFreshNames (SimpleM lore) where - putNameSource src = modify $ \(_, b) -> (src, b) - getNameSource = gets fst + putNameSource src = modify $ \(_, b, c) -> (src, b, c) + getNameSource = gets $ \(a, _, _) -> a instance SimplifiableLore lore => HasScope (Wise lore) (SimpleM lore) where askScope = ST.toScope <$> askVtable @@ -163,30 +157,21 @@ -> VNameSource -> ((a, Bool), VNameSource) runSimpleM (SimpleM m) simpl env src = - let (x, (src', b), _) = runRWS m (simpl, env) (src, False) + let (x, (src', b, _)) = runState (runReaderT m (simpl, env)) (src, False, mempty) in ((x, b), src') -subSimpleM :: (SameScope outerlore lore, - ExpAttr outerlore ~ ExpAttr lore, - BodyAttr outerlore ~ BodyAttr lore, - RetType outerlore ~ RetType lore, - BranchType outerlore ~ BranchType lore) => - SimpleOps lore - -> Env lore - -> ST.SymbolTable (Wise outerlore) +subSimpleM :: RuleBook (Wise lore) + -> HoistBlockers lore -> SimpleM lore a - -> SimpleM outerlore a -subSimpleM simpl env outer_vtable m = do - let inner_vtable = ST.castSymbolTable outer_vtable - src <- getNameSource - let SimpleM m' = localVtable (<>inner_vtable) m - (x, (src', b), _) = runRWS m' (simpl, env) (src, False) - putNameSource src' - when b changed - return x + -> SimpleM lore a +subSimpleM rules blockers = + local $ \(ops, env) -> (ops, + env { envRules = rules + , envHoistBlockers = blockers + }) askEngineEnv :: SimpleM lore (Env lore) -askEngineEnv = snd <$> ask +askEngineEnv = asks snd asksEngineEnv :: (Env lore -> a) -> SimpleM lore a asksEngineEnv f = f <$> askEngineEnv @@ -199,16 +184,18 @@ localVtable f = local $ \(ops, env) -> (ops, env { envVtable = f $ envVtable env }) collectCerts :: SimpleM lore a -> SimpleM lore (a, Certificates) -collectCerts m = pass $ do (x, cs) <- listen m - return ((x, cs), const mempty) +collectCerts m = do x <- m + (a, b, cs) <- get + put (a, b, mempty) + return (x, cs) -- | Mark that we have changed something and it would be a good idea -- to re-run the simplifier. changed :: SimpleM lore () -changed = modify $ \(src, _) -> (src, True) +changed = modify $ \(src, _, cs) -> (src, True, cs) usedCerts :: Certificates -> SimpleM lore () -usedCerts = tell +usedCerts cs = modify $ \(a, b, c) -> (a, b, cs <> c) enterLoop :: SimpleM lore a -> SimpleM lore a enterLoop = localVtable ST.deepen @@ -230,12 +217,6 @@ localVtable $ \vtable -> foldr (uncurry ST.insertArrayLParam) vtable params -bindChunkLParams :: SimplifiableLore lore => - VName -> [(LParam (Wise lore),VName)] -> SimpleM lore a -> SimpleM lore a -bindChunkLParams offset params = - localVtable $ \vtable -> - foldr (uncurry $ ST.insertChunkLParam offset) vtable params - bindLoopVar :: SimplifiableLore lore => VName -> IntType -> SubExp -> SimpleM lore a -> SimpleM lore a bindLoopVar var it bound = @@ -328,7 +309,7 @@ -- are unsafe, and added safety (by 'protectLoopHoisted') may inhibit -- further optimisation.. notWorthHoisting :: Attributes lore => BlockPred lore -notWorthHoisting _ (Let pat _ e) = +notWorthHoisting _ _ (Let pat _ e) = not (safeExp e) && any (>0) (map arrayRank $ patternTypes pat) hoistStms :: SimplifiableLore lore => @@ -360,7 +341,7 @@ bottomUpSimplifyStm rules (vtable', uses') stm case res of Nothing -- Nothing to optimise - see if hoistable. - | block uses' stm -> + | block vtable' uses' stm -> return (expandUsage vtable' uses' stm `UT.without` provides stm, Left stm : stms) | otherwise -> @@ -373,12 +354,12 @@ blockUnhoistedDeps :: Attributes lore => [Either (Stm lore) (Stm lore)] -> [Either (Stm lore) (Stm lore)] -blockUnhoistedDeps = snd . mapAccumL block S.empty +blockUnhoistedDeps = snd . mapAccumL block mempty where block blocked (Left need) = - (blocked <> S.fromList (provides need), Left need) + (blocked <> namesFromList (provides need), Left need) block blocked (Right need) - | blocked `intersects` freeIn need = - (blocked <> S.fromList (provides need), Left need) + | blocked `namesIntersect` freeIn need = + (blocked <> namesFromList (provides need), Left need) | otherwise = (blocked, Right need) @@ -396,31 +377,28 @@ zip (patternNames pat) (patternAliases pat) usageThroughBindeeAliases (name, aliases) = do uses <- UT.lookup name utable - return $ mconcat $ map (`UT.usage` uses) $ S.toList aliases - -intersects :: Ord a => S.Set a -> S.Set a -> Bool -intersects a b = not $ S.null $ a `S.intersection` b + return $ mconcat $ map (`UT.usage` uses) $ namesToList aliases -type BlockPred lore = UT.UsageTable -> Stm lore -> Bool +type BlockPred lore = ST.SymbolTable lore -> UT.UsageTable -> Stm lore -> Bool neverBlocks :: BlockPred lore -neverBlocks _ _ = False +neverBlocks _ _ _ = False isFalse :: Bool -> BlockPred lore -isFalse b _ _ = not b +isFalse b _ _ _ = not b orIf :: BlockPred lore -> BlockPred lore -> BlockPred lore -orIf p1 p2 body need = p1 body need || p2 body need +orIf p1 p2 body vtable need = p1 body vtable need || p2 body vtable need andAlso :: BlockPred lore -> BlockPred lore -> BlockPred lore -andAlso p1 p2 body need = p1 body need && p2 body need +andAlso p1 p2 body vtable need = p1 body vtable need && p2 body vtable need isConsumed :: BlockPred lore -isConsumed utable = any (`UT.isConsumed` utable) . patternNames . stmPattern +isConsumed _ utable = any (`UT.isConsumed` utable) . patternNames . stmPattern isOp :: BlockPred lore -isOp _ (Let _ _ Op{}) = True -isOp _ _ = False +isOp _ _ (Let _ _ Op{}) = True +isOp _ _ _ = False constructBody :: SimplifiableLore lore => Stms (Wise lore) -> Result -> SimpleM lore (Body (Wise lore)) @@ -447,18 +425,18 @@ insertAllStms = uncurry constructBody . fst <=< blockIf (isFalse False) hasFree :: Attributes lore => Names -> BlockPred lore -hasFree ks _ need = ks `intersects` freeIn need +hasFree ks _ _ need = ks `namesIntersect` freeIn need isNotSafe :: Attributes lore => BlockPred lore -isNotSafe _ = not . safeExp . stmExp +isNotSafe _ _ = not . safeExp . stmExp isInPlaceBound :: BlockPred m -isInPlaceBound _ = isUpdate . stmExp +isInPlaceBound _ _ = isUpdate . stmExp where isUpdate (BasicOp Update{}) = True isUpdate _ = False isNotCheap :: Attributes lore => BlockPred lore -isNotCheap _ = not . cheapStm +isNotCheap _ _ = not . cheapStm cheapStm :: Attributes lore => Stm lore -> Bool cheapStm = cheapExp . stmExp @@ -478,11 +456,11 @@ -- let's try it out. stmIs :: (Stm lore -> Bool) -> BlockPred lore -stmIs f _ = f +stmIs f _ _ = f loopInvariantStm :: Attributes lore => ST.SymbolTable lore -> Stm lore -> Bool loopInvariantStm vtable = - all (`S.member` ST.availableAtClosestLoop vtable) . freeIn + all (`nameIn` ST.availableAtClosestLoop vtable) . namesToList . freeIn hoistCommon :: SimplifiableLore lore => SubExp -> IfSort @@ -504,18 +482,25 @@ -- shape computations, and expensive loop-invariant operations -- are if-hoistable. cond_loop_invariant = - all (`S.member` ST.availableAtClosestLoop vtable) $ freeIn cond + all (`nameIn` ST.availableAtClosestLoop vtable) $ namesToList $ freeIn cond + desirableToHoist stm = is_alloc_fun stm || (ST.loopDepth vtable > 0 && cond_loop_invariant && ifsort /= IfFallback && loopInvariantStm vtable stm) + hoistbl_nms = filterBnds desirableToHoist getArrSz_fun $ stmsToList $ stms1<>stms2 + + isNotHoistableBnd _ _ _ (Let _ _ (BasicOp ArrayLit{})) = False + isNotHoistableBnd nms _ _ stm = not (hasPatName nms stm) + block = branch_blocker `orIf` ((isNotSafe `orIf` isNotCheap) `andAlso` stmIs (not . desirableToHoist)) `orIf` isInPlaceBound `orIf` isNotHoistableBnd hoistbl_nms + rules <- asksEngineEnv envRules (body1_bnds', safe1) <- protectIfHoisted cond True $ hoistStms rules block vtable usages1 stms1 @@ -529,7 +514,7 @@ let sz_nms = mconcat $ map getArrSz_fn all_bnds sz_needs = transClosSizes all_bnds sz_nms [] alloc_bnds = filter interesting all_bnds - sel_nms = S.fromList $ + sel_nms = namesFromList $ concatMap (patternNames . stmPattern) (sz_needs ++ alloc_bnds) in sel_nms @@ -539,10 +524,7 @@ in if null new_bnds then hoist_bnds else transClosSizes all_bnds new_nms (new_bnds ++ hoist_bnds) - hasPatName nms bnd = intersects nms $ S.fromList $ - patternNames $ stmPattern bnd - isNotHoistableBnd _ _ (Let _ _ (BasicOp ArrayLit{})) = False - isNotHoistableBnd nms _ stm = not (hasPatName nms stm) + hasPatName nms bnd = any (`nameIn` nms) $ patternNames $ stmPattern bnd -- | Simplify a single 'Body'. The @[Diet]@ only covers the value -- elements, because the context cannot be consumed. @@ -665,7 +647,7 @@ loop_arrs' <- mapM simplify loop_arrs let form' = ForLoop loopvar it boundexp' (zip loop_params' loop_arrs') return (form', - S.fromList (loopvar : map paramName loop_params') <> fparamnames, + namesFromList (loopvar : map paramName loop_params') <> fparamnames, bindLoopVar loopvar it boundexp' . protectLoopHoisted ctx' val' form' . bindArrayLParams (zip loop_params' (map Just loop_arrs'))) @@ -686,9 +668,9 @@ loopbody' <- constructBody loopstms loopres return (DoLoop ctx' val' form' loopbody', hoisted) where fparamnames = - S.fromList (map (paramName . fst) $ ctx++val) + namesFromList (map (paramName . fst) $ ctx++val) consumeMerge = - localVtable $ flip (foldl' (flip ST.consume)) consumed_by_merge + localVtable $ flip (foldl' (flip ST.consume)) $ namesToList consumed_by_merge consumed_by_merge = freeIn $ map snd $ filter (unique . paramDeclType . fst) val @@ -782,7 +764,7 @@ mapM inspect (patternValueElements pat) where inspect (PatElem name lore) = PatElem name <$> simplify lore -simplifyParam :: (attr -> SimpleM lore attr) -> ParamT attr -> SimpleM lore (ParamT attr) +simplifyParam :: (attr -> SimpleM lore attr) -> Param attr -> SimpleM lore (Param attr) simplifyParam simplifyAttribute (Param name attr) = Param name <$> simplifyAttribute attr @@ -844,7 +826,7 @@ params' <- mapM (simplifyParam simplify) params let (nonarrayparams, arrayparams) = splitAt (length params' - length arrs) params' - paramnames = S.fromList $ boundByLambda lam + paramnames = namesFromList $ boundByLambda lam ((lamstms, lamres), hoisted) <- enterLoop $ bindLParams nonarrayparams $ @@ -858,7 +840,7 @@ consumeResult :: [(Diet, SubExp)] -> UT.UsageTable consumeResult = mconcat . map inspect where inspect (Consume, se) = - mconcat $ map UT.consumedUsage $ S.toList $ subExpAliases se + mconcat $ map UT.consumedUsage $ namesToList $ subExpAliases se inspect _ = mempty instance Simplifiable Certificates where diff --git a/src/Futhark/Optimise/Simplify/Lore.hs b/src/Futhark/Optimise/Simplify/Lore.hs --- a/src/Futhark/Optimise/Simplify/Lore.hs +++ b/src/Futhark/Optimise/Simplify/Lore.hs @@ -59,7 +59,7 @@ VarWisdom (substituteNames substs als) (substituteNames substs range) instance FreeIn VarWisdom where - freeIn (VarWisdom als range) = freeIn als <> freeIn range + freeIn' (VarWisdom als range) = freeIn' als <> freeIn' range -- | Wisdom about an expression. data ExpWisdom = ExpWisdom { _expWisdomConsumed :: ConsumedInExp @@ -68,10 +68,10 @@ deriving (Eq, Ord, Show) instance FreeIn ExpWisdom where - freeIn = mempty + freeIn' = mempty instance FreeAttr ExpWisdom where - precomputed = const . unNames . expWisdomFree + precomputed = const . fvNames . unNames . expWisdomFree instance Substitute ExpWisdom where substituteNames substs (ExpWisdom cons free) = @@ -102,11 +102,11 @@ (substituteNames substs free) instance FreeIn BodyWisdom where - freeIn (BodyWisdom als cons rs free) = - freeIn als <> freeIn cons <> freeIn rs <> freeIn free + freeIn' (BodyWisdom als cons rs free) = + freeIn' als <> freeIn' cons <> freeIn' rs <> freeIn' free instance FreeAttr BodyWisdom where - precomputed = const . unNames . bodyWisdomFree + precomputed = const . fvNames . unNames . bodyWisdomFree instance (Annotations lore, CanBeWise (Op lore)) => Annotations (Wise lore) where @@ -215,7 +215,7 @@ mkWiseBody :: (Attributes lore, CanBeWise (Op lore)) => BodyAttr lore -> Stms (Wise lore) -> Result -> Body (Wise lore) mkWiseBody innerlore bnds res = - Body (BodyWisdom aliases consumed ranges (Names' $ freeInStmsAndRes bnds res), + Body (BodyWisdom aliases consumed ranges (Names' $ freeIn $ freeInStmsAndRes bnds res), innerlore) bnds res where (aliases, consumed) = Aliases.mkBodyAliases bnds res ranges = Ranges.mkBodyRanges bnds res diff --git a/src/Futhark/Optimise/Simplify/Rule.hs b/src/Futhark/Optimise/Simplify/Rule.hs --- a/src/Futhark/Optimise/Simplify/Rule.hs +++ b/src/Futhark/Optimise/Simplify/Rule.hs @@ -18,6 +18,7 @@ , liftMaybe -- * Rule definition + , Rule(..) , SimplificationRule(..) , RuleGeneric , RuleBasicOp @@ -86,16 +87,14 @@ -- list of new bindings. Even if the action fail, there may still be -- a monadic effect - particularly, the name source may have been -- modified. -simplify :: (MonadFreshNames m, HasScope lore m) => - RuleM lore a - -> m (Maybe (a, Stms lore)) -simplify (RuleM m) = do - scope <- askScope - modifyNameSource $ \src -> - case runExcept $ runStateT (runBinderT m scope) src of - Left CannotSimplify -> (Nothing, src) - Left (OtherError err) -> error $ "simplify: " ++ err - Right (x, src') -> (Just x, src') +simplify :: Scope lore -> VNameSource -> Rule lore + -> Maybe (Stms lore, VNameSource) +simplify _ _ Skip = Nothing +simplify scope src (Simplify (RuleM m)) = + case runExcept $ runStateT (runBinderT m scope) src of + Left CannotSimplify -> Nothing + Left (OtherError err) -> error $ "simplify: " ++ err + Right (((), x), src') -> Just (x, src') cannotSimplify :: RuleM lore a cannotSimplify = throwError CannotSimplify @@ -104,19 +103,23 @@ liftMaybe Nothing = cannotSimplify liftMaybe (Just x) = return x -type RuleGeneric lore a = a -> Stm lore -> RuleM lore () +-- | An efficient way of encoding whether a simplification rule should even be attempted. +data Rule lore = Simplify (RuleM lore ()) -- ^ Give it a shot. + | Skip -- ^ Don't bother. + +type RuleGeneric lore a = a -> Stm lore -> Rule lore type RuleBasicOp lore a = (a -> Pattern lore -> StmAux (ExpAttr lore) -> - BasicOp lore -> RuleM lore ()) + BasicOp lore -> Rule lore) type RuleIf lore a = a -> Pattern lore -> StmAux (ExpAttr lore) -> (SubExp, BodyT lore, BodyT lore, IfAttr (BranchType lore)) -> - RuleM lore () + Rule lore type RuleDoLoop lore a = a -> Pattern lore -> StmAux (ExpAttr lore) -> ([(FParam lore, SubExp)], [(FParam lore, SubExp)], LoopForm lore, BodyT lore) -> - RuleM lore () + Rule lore type RuleOp lore a = a -> Pattern lore -> StmAux (ExpAttr lore) -> - Op lore -> RuleM lore () + Op lore -> Rule lore -- | A simplification rule takes some argument and a statement, and -- tries to simplify the statement. @@ -214,7 +217,7 @@ -- binding @bnd@. If simplification is possible, a replacement list -- of bindings is returned, that bind at least the same names as the -- original binding (and possibly more, for intermediate results). -topDownSimplifyStm :: (MonadFreshNames m, HasScope lore m, BinderOps lore) => +topDownSimplifyStm :: (MonadFreshNames m, HasScope lore m) => RuleBook lore -> ST.SymbolTable lore -> Stm lore @@ -226,7 +229,7 @@ -- bindings is returned, that bind at least the same names as the -- original binding (and possibly more, for intermediate results). -- The first argument is the set of names used after this binding. -bottomUpSimplifyStm :: (MonadFreshNames m, HasScope lore m, BinderOps lore) => +bottomUpSimplifyStm :: (MonadFreshNames m, HasScope lore m) => RuleBook lore -> (ST.SymbolTable lore, UT.UsageTable) -> Stm lore @@ -240,7 +243,7 @@ If{} -> rulesIf _ -> rulesAny -applyRule :: SimplificationRule lore a -> a -> Stm lore -> RuleM lore () +applyRule :: SimplificationRule lore a -> a -> Stm lore -> Rule lore applyRule (RuleGeneric f) a stm = f a stm applyRule (RuleBasicOp f) a (Let pat aux (BasicOp e)) = f a pat aux e applyRule (RuleDoLoop f) a (Let pat aux (DoLoop ctx val form body)) = @@ -250,18 +253,21 @@ applyRule (RuleOp f) a (Let pat aux (Op op)) = f a pat aux op applyRule _ _ _ = - cannotSimplify + Skip -applyRules :: (MonadFreshNames m, HasScope lore m, BinderOps lore) => +applyRules :: (MonadFreshNames m, HasScope lore m) => Rules lore a -> a -> Stm lore -> m (Maybe (Stms lore)) -applyRules rules context stm = applyRules' (rulesForStm stm rules) context stm +applyRules all_rules context stm = do + scope <- askScope -applyRules' :: (MonadFreshNames m, HasScope lore m, BinderOps lore) => - [SimplificationRule lore a] -> a -> Stm lore - -> m (Maybe (Stms lore)) -applyRules' [] _ _ = return Nothing -applyRules' (rule:rules) context bnd = do - res <- simplify $ applyRule rule context bnd - case res of Just ((), bnds) -> return $ Just bnds - Nothing -> applyRules' rules context bnd + modifyNameSource $ \src -> + let applyRules' [] = Nothing + applyRules' (rule:rules) = + case simplify scope src (applyRule rule context stm) of + Just x -> Just x + Nothing -> applyRules' rules + + in case applyRules' $ rulesForStm stm all_rules of + Just (stms, src') -> (Just stms, src') + Nothing -> (Nothing, src) diff --git a/src/Futhark/Optimise/Simplify/Rules.hs b/src/Futhark/Optimise/Simplify/Rules.hs --- a/src/Futhark/Optimise/Simplify/Rules.hs +++ b/src/Futhark/Optimise/Simplify/Rules.hs @@ -24,9 +24,7 @@ import Data.Foldable (all) import Data.List hiding (all) import Data.Maybe - import qualified Data.Map.Strict as M -import qualified Data.Set as S import qualified Futhark.Analysis.SymbolTable as ST import qualified Futhark.Analysis.UsageTable as UT @@ -90,7 +88,7 @@ resIsNecessary ((v,_), _) = usedAfterLoop v || - paramName v `S.member` necessaryForReturned || + paramName v `nameIn` necessaryForReturned || referencedInPat v || referencedInForm v @@ -109,7 +107,7 @@ free_in_keeps = freeIn keep_valpatelems stillUsedContext pat_elem = - patElemName pat_elem `S.member` + patElemName pat_elem `nameIn` (free_in_keeps <> freeIn (filter (/=pat_elem) $ patternContextElements pat)) @@ -117,8 +115,8 @@ , patternContextElements = filter stillUsedContext $ patternContextElements pat } in if ctx' ++ val' == ctx ++ val - then cannotSimplify - else do + then Skip + else Simplify $ do -- We can't just remove the bindings in 'discard', since the loop -- body may still use their names in (now-dead) expressions. -- Hence, we add them inside the loop, fully aware that dead-code @@ -133,10 +131,10 @@ used_vals = map fst $ filter snd $ zip (map (paramName . fst) val) pat_used usedAfterLoop = flip elem used_vals . paramName usedAfterLoopOrInForm p = - usedAfterLoop p || paramName p `S.member` freeIn form + usedAfterLoop p || paramName p `nameIn` freeIn form patAnnotNames = freeIn $ map fst $ ctx++val - referencedInPat = (`S.member` patAnnotNames) . paramName - referencedInForm = (`S.member` freeIn form) . paramName + referencedInPat = (`nameIn` patAnnotNames) . paramName + referencedInForm = (`nameIn` freeIn form) . paramName dummyStms = map dummyStm dummyStm ((p,e), _) @@ -144,7 +142,7 @@ Var v <- e = ([paramName p], BasicOp $ Copy v) | otherwise = ([paramName p], BasicOp $ SubExp e) removeRedundantMergeVariables _ _ _ _ = - cannotSimplify + Skip -- We may change the type of the loop if we hoist out a shape -- annotation, in which case we also need to tweak the bound pattern. @@ -156,8 +154,8 @@ zip merge res of ([], _, _, _) -> -- Nothing is invariant. - cannotSimplify - (invariant, explpat', merge', res') -> do + Skip + (invariant, explpat', merge', res') -> Simplify $ do -- We have moved something invariant out of the loop. let loopbody' = loopbody { bodyResult = res' } invariantShape :: (a, VName) -> Bool @@ -180,7 +178,7 @@ explpat = zip (patternValueElements pat) $ map (paramName . fst) val - namesOfMergeParams = S.fromList $ map (paramName . fst) $ ctx++val + namesOfMergeParams = namesFromList $ map (paramName . fst) $ ctx++val removeFromResult (mergeParam,mergeInit) explpat' = case partition ((==paramName mergeParam) . snd) explpat' of @@ -196,7 +194,7 @@ arrayRank (paramDeclType mergeParam) == 1, isInvariant resExp, -- Also do not remove the condition in a while-loop. - not $ paramName mergeParam `S.member` freeIn form = + not $ paramName mergeParam `nameIn` freeIn form = let (bnd, explpat'') = removeFromResult (mergeParam,mergeInit) explpat' in (maybe id (:) bnd $ (paramIdent mergeParam, mergeInit) : invariant, @@ -211,7 +209,7 @@ isInvariant (Var v2) | paramName mergeParam == v2 = allExistentialInvariant - (S.fromList $ map (identName . fst) invariant) mergeParam + (namesFromList $ map (identName . fst) invariant) mergeParam -- (1) or identical to the initial value of the parameter. isInvariant _ = mergeInit == resExp @@ -219,14 +217,11 @@ (invariant, explpat', (mergeParam,mergeInit):merge', resExp:resExps) allExistentialInvariant namesOfInvariant mergeParam = - all (invariantOrNotMergeParam namesOfInvariant) - (paramName mergeParam `S.delete` freeIn mergeParam) + all (invariantOrNotMergeParam namesOfInvariant) $ namesToList $ + freeIn mergeParam `namesSubtract` oneName (paramName mergeParam) invariantOrNotMergeParam namesOfInvariant name = - not (name `S.member` namesOfMergeParams) || - name `S.member` namesOfInvariant - --- | A function that, given a variable name, returns its definition. -type VarLookup lore = VName -> Maybe (Exp lore, Certificates) + not (name `nameIn` namesOfMergeParams) || + name `nameIn` namesOfInvariant -- | A function that, given a subexpression, returns its type. type TypeLookup = SubExp -> Maybe Type @@ -251,13 +246,13 @@ simplifyClosedFormLoop :: BinderOps lore => TopDownRuleDoLoop lore simplifyClosedFormLoop _ pat _ ([], val, ForLoop i _ bound [], body) = - loopClosedForm pat val (S.singleton i) bound body -simplifyClosedFormLoop _ _ _ _ = cannotSimplify + Simplify $ loopClosedForm pat val (oneName i) bound body +simplifyClosedFormLoop _ _ _ _ = Skip simplifyLoopVariables :: (BinderOps lore, Aliased lore) => TopDownRuleDoLoop lore simplifyLoopVariables vtable pat _ (ctx, val, form@(ForLoop i it num_iters loop_vars), body) | simplifiable <- map checkIfSimplifiable loop_vars, - not $ all isNothing simplifiable = do + not $ all isNothing simplifiable = Simplify $ do -- Check if the simplifications throw away more information than -- we are comfortable with at this stage. (maybe_loop_vars, body_prefix_stms) <- @@ -282,7 +277,7 @@ checkIfSimplifiable (p,arr) = simplifyIndexing vtable' seType arr (DimFix (Var i) : fullSlice (paramType p) []) $ - paramName p `S.member` consumed_in_body + paramName p `nameIn` consumed_in_body -- We only want this simplification if the result does not refer -- to 'i' at all, or does not contain accesses. @@ -292,9 +287,9 @@ (x,x_stms) <- collectStms m case x of IndexResult cs arr' slice - | all (not . (i `S.member`) . freeIn) x_stms, + | all (not . (i `nameIn`) . freeIn) x_stms, DimFix (Var j) : slice' <- slice, - j == i, not $ i `S.member` freeIn slice -> do + j == i, not $ i `nameIn` freeIn slice -> do addStms x_stms w <- arraySize 0 <$> lookupType arr' for_in_partial <- @@ -313,16 +308,16 @@ notIndex (BasicOp Index{}) = False notIndex _ = True -simplifyLoopVariables _ _ _ _ = cannotSimplify +simplifyLoopVariables _ _ _ _ = Skip simplifKnownIterationLoop :: BinderOps lore => TopDownRuleDoLoop lore simplifKnownIterationLoop _ pat _ (ctx, val, ForLoop i it (Constant iters) loop_vars, body) - | zeroIsh iters = do + | zeroIsh iters = Simplify $ do let bindResult p r = letBindNames [patElemName p] $ BasicOp $ SubExp r zipWithM_ bindResult (patternContextElements pat) (map snd ctx) zipWithM_ bindResult (patternValueElements pat) (map snd val) - | oneIsh iters = do + | oneIsh iters = Simplify $ do forM_ (ctx++val) $ \(mergevar, mergeinit) -> letBindNames [paramName mergevar] $ BasicOp $ SubExp mergeinit @@ -339,7 +334,7 @@ forM_ (zip (patternNames pat) res) $ \(v, se) -> letBindNames_ [v] $ BasicOp $ SubExp se simplifKnownIterationLoop _ _ _ _ = - cannotSimplify + Skip -- | Turn @copy(x)@ into @x@ iff @x@ is not used after this copy -- statement and it can be consumed. @@ -349,7 +344,7 @@ removeUnnecessaryCopy (vtable,used) (Pattern [] [d]) _ (Copy v) | not (v `UT.isConsumed` used), (not (v `UT.used` used) && consumable) || not (patElemName d `UT.isConsumed` used) = - letBind_ (Pattern [] [d]) $ BasicOp $ SubExp $ Var v + Simplify $ letBind_ (Pattern [] [d]) $ BasicOp $ SubExp $ Var v where -- We need to make sure we can even consume the original. -- This is currently a hacky check, much too conservative, -- because we don't have the information conveniently @@ -357,7 +352,7 @@ consumable = case M.lookup v $ ST.toScope vtable of Just (FParamInfo info) -> unique $ declTypeOf info _ -> False -removeUnnecessaryCopy _ _ _ _ = cannotSimplify +removeUnnecessaryCopy _ _ _ _ = Skip simplifyCmpOp :: SimpleRule lore simplifyCmpOp _ _ (CmpOp cmp e1 e2) @@ -574,14 +569,14 @@ | Just args' <- mapM (isConst . fst) args, Just (_, _, fun) <- M.lookup (nameToString fname) primFuns, Just result <- fun args' = - certifying cs $ letBind_ pat $ BasicOp $ SubExp $ Constant result + Simplify $ certifying cs $ letBind_ pat $ BasicOp $ SubExp $ Constant result where isConst (Constant v) = Just v isConst _ = Nothing -constantFoldPrimFun _ _ = cannotSimplify +constantFoldPrimFun _ _ = Skip simplifyIndex :: BinderOps lore => BottomUpRuleBasicOp lore simplifyIndex (vtable, used) pat@(Pattern [] [pe]) (StmAux cs _) (Index idd inds) - | Just m <- simplifyIndexing vtable seType idd inds consumed = do + | Just m <- simplifyIndexing vtable seType idd inds consumed = Simplify $ do res <- m case res of SubExpResult cs' se -> @@ -594,7 +589,7 @@ seType (Var v) = ST.lookupType v vtable seType (Constant v) = Just $ Prim $ primValueType v -simplifyIndex _ _ _ _ = cannotSimplify +simplifyIndex _ _ _ _ = Skip data IndexResult = IndexResult Certificates VName (Slice SubExp) | SubExpResult Certificates SubExp @@ -780,7 +775,7 @@ | Just r <- arrayRank <$> ST.lookupType x vtable, let perm = [i] ++ [0..i-1] ++ [i+1..r-1], Just (x',x_cs) <- transposedBy perm x, - Just (xs',xs_cs) <- unzip <$> mapM (transposedBy perm) xs = do + Just (xs',xs_cs) <- unzip <$> mapM (transposedBy perm) xs = Simplify $ do concat_rearrange <- certifying (x_cs<>mconcat xs_cs) $ letExp "concat_rearrange" $ BasicOp $ Concat 0 x' xs' new_d @@ -793,7 +788,7 @@ -- concat xs (concat ys zs) == concat xs ys zs simplifyConcat (vtable, _) pat (StmAux cs _) (Concat i x xs new_d) - | x' /= x || concat xs' /= xs = + | x' /= x || concat xs' /= xs = Simplify $ certifying (cs<>x_cs<>mconcat xs_cs) $ letBind_ pat $ BasicOp $ Concat i x' (zs++concat xs') new_d where (x':zs, x_cs) = isConcat x @@ -805,7 +800,7 @@ -- If concatenating a bunch of array literals (or equivalent -- replicate), just construct the array literal instead. simplifyConcat (vtable, _) pat (StmAux cs _) (Concat 0 x xs _) - | Just (vs, vcs) <- unzip <$> mapM isArrayLit (x:xs) = do + | Just (vs, vcs) <- unzip <$> mapM isArrayLit (x:xs) = Simplify $ do rt <- rowType <$> lookupType x certifying (cs <> mconcat vcs) $ letBind_ pat $ BasicOp $ ArrayLit vs rt @@ -819,19 +814,17 @@ unitShape = (==Shape [Constant $ IntValue $ Int32Value 1]) -simplifyConcat _ _ _ _ = cannotSimplify +simplifyConcat _ _ _ _ = Skip ruleIf :: BinderOps lore => TopDownRuleIf lore -ruleIf _ pat _ (e1, tb, fb, IfAttr t ifsort) +ruleIf _ pat _ (e1, tb, fb, IfAttr _ ifsort) | Just branch <- checkBranch, - ifsort /= IfFallback || isCt1 e1 = do + ifsort /= IfFallback || isCt1 e1 = Simplify $ do let ses = bodyResult branch addStms $ bodyStms branch - ctx <- subExpShapeContext (bodyTypeValues t) ses - let ses' = ctx ++ ses sequence_ [ letBind (Pattern [] [p]) $ BasicOp $ SubExp se - | (p,se) <- zip (patternElements pat) ses'] + | (p,se) <- zip (patternElements pat) ses] where checkBranch | isCt1 e1 = Just tb @@ -847,14 +840,14 @@ (cond, Body _ tstms [Constant (BoolValue True)], Body _ fstms [se], IfAttr ts _) | null tstms, null fstms, [Prim Bool] <- bodyTypeValues ts = - letBind_ pat $ BasicOp $ BinOp LogOr cond se + Simplify $ letBind_ pat $ BasicOp $ BinOp LogOr cond se -- When type(x)==bool, if c then x else y == (c && x) || (!c && y) ruleIf _ pat _ (cond, tb, fb, IfAttr ts _) | Body _ tstms [tres] <- tb, Body _ fstms [fres] <- fb, all (safeExp . stmExp) $ tstms <> fstms, - all (==Prim Bool) $ bodyTypeValues ts = do + all (==Prim Bool) $ bodyTypeValues ts = Simplify $ do addStms tstms addStms fstms e <- eBinOp LogOr (pure $ BasicOp $ BinOp LogAnd cond tres) @@ -864,19 +857,19 @@ ruleIf _ pat _ (_, tbranch, _, IfAttr _ IfFallback) | null $ patternContextNames pat, - all (safeExp . stmExp) $ bodyStms tbranch = do + all (safeExp . stmExp) $ bodyStms tbranch = Simplify $ do let ses = bodyResult tbranch addStms $ bodyStms tbranch sequence_ [ letBind (Pattern [] [p]) $ BasicOp $ SubExp se | (p,se) <- zip (patternElements pat) ses] -ruleIf _ _ _ _ = cannotSimplify +ruleIf _ _ _ _ = Skip -- | Move out results of a conditional expression whose computation is -- either invariant to the branches (only done for results in the -- context), or the same in both branches. hoistBranchInvariant :: BinderOps lore => TopDownRuleIf lore -hoistBranchInvariant _ pat _ (cond, tb, fb, IfAttr ret ifsort) = do +hoistBranchInvariant _ pat _ (cond, tb, fb, IfAttr ret ifsort) = Simplify $ do let tses = bodyResult tb fses = bodyResult fb (hoistings, (pes, ts, res)) <- @@ -899,15 +892,15 @@ If cond tb'' fb'' (IfAttr ret' ifsort) else cannotSimplify where num_ctx = length $ patternContextElements pat - bound_in_branches = S.fromList $ concatMap (patternNames . stmPattern) $ + bound_in_branches = namesFromList $ concatMap (patternNames . stmPattern) $ bodyStms tb <> bodyStms fb mem_sizes = freeIn $ filter (isMem . patElemType) $ patternElements pat invariant Constant{} = True - invariant (Var v) = not $ v `S.member` bound_in_branches + invariant (Var v) = not $ v `nameIn` bound_in_branches isMem Mem{} = True isMem _ = False - sizeOfMem v = v `S.member` mem_sizes + sizeOfMem v = v `nameIn` mem_sizes branchInvariant (pe, t, (tse, fse)) -- Do both branches return the same value? @@ -1012,19 +1005,19 @@ -- Check all the simpleRules. ruleBasicOp vtable pat aux op | Just (op', cs) <- msum [ rule defOf seType op | rule <- simpleRules ] = - certifying (cs <> stmAuxCerts aux) $ letBind_ pat $ BasicOp op' + Simplify $ certifying (cs <> stmAuxCerts aux) $ letBind_ pat $ BasicOp op' where defOf = (`ST.lookupExp` vtable) seType (Var v) = ST.lookupType v vtable seType (Constant v) = Just $ Prim $ primValueType v ruleBasicOp vtable pat _ (Update src _ (Var v)) | Just (BasicOp Scratch{}, _) <- ST.lookupExp v vtable = - letBind_ pat $ BasicOp $ SubExp $ Var src + Simplify $ letBind_ pat $ BasicOp $ SubExp $ Var src ruleBasicOp vtable pat _ (Update dest destis (Var v)) | Just (e, _) <- ST.lookupExp v vtable, arrayFrom e = - letBind_ pat $ BasicOp $ SubExp $ Var dest + Simplify $ letBind_ pat $ BasicOp $ SubExp $ Var dest where arrayFrom (BasicOp (Copy copy_v)) | Just (e',_) <- ST.lookupExp copy_v vtable = arrayFrom e' @@ -1043,7 +1036,7 @@ ruleBasicOp vtable pat _ (Update dest is se) | Just dest_t <- ST.lookupType dest vtable, isFullSlice (arrayShape dest_t) is = - letBind_ pat $ BasicOp $ + Simplify $ letBind_ pat $ BasicOp $ case se of Var v | not $ null $ sliceDims is -> Reshape (map DimNew $ arrayDims dest_t) v @@ -1055,7 +1048,8 @@ | Just (Update dest2 is2 se2, cs2) <- ST.lookupBasicOp v1 vtable, Just (Copy v3, cs3) <- ST.lookupBasicOp dest2 vtable, Just (Index v4 is4, cs4) <- ST.lookupBasicOp v3 vtable, - is4 == is1, v4 == dest1 = certifying (cs1 <> cs2 <> cs3 <> cs4) $ do + is4 == is1, v4 == dest1 = + Simplify $ certifying (cs1 <> cs2 <> cs3 <> cs4) $ do is5 <- sliceSlice is1 is2 letBind_ pat $ BasicOp $ Update dest1 is5 se2 @@ -1067,15 +1061,15 @@ -- is actually the case, such as if we will obtain at least one -- constant-to-constant comparison? ruleBasicOp vtable pat _ (CmpOp (CmpEq t) se1 se2) - | Just m <- simplifyWith se1 se2 = m - | Just m <- simplifyWith se2 se1 = m + | Just m <- simplifyWith se1 se2 = Simplify m + | Just m <- simplifyWith se2 se1 = Simplify m where simplifyWith (Var v) x | Just bnd <- ST.entryStm =<< ST.lookup v vtable, If p tbranch fbranch _ <- stmExp bnd, Just (y, z) <- returns v (stmPattern bnd) tbranch fbranch, - S.null $ freeIn y `S.intersection` boundInBody tbranch, - S.null $ freeIn z `S.intersection` boundInBody fbranch = Just $ do + not $ boundInBody tbranch `namesIntersect` freeIn y, + not $ boundInBody fbranch `namesIntersect` freeIn z = Just $ do eq_x_y <- letSubExp "eq_x_y" $ BasicOp $ CmpOp (CmpEq t) x y eq_x_z <- @@ -1098,26 +1092,27 @@ zip (bodyResult tbranch) (bodyResult fbranch) ruleBasicOp _ pat _ (Replicate (Shape []) se@Constant{}) = - letBind_ pat $ BasicOp $ SubExp se -ruleBasicOp _ pat _ (Replicate (Shape []) (Var v)) = do + Simplify $ letBind_ pat $ BasicOp $ SubExp se +ruleBasicOp _ pat _ (Replicate (Shape []) (Var v)) = Simplify $ do v_t <- lookupType v letBind_ pat $ BasicOp $ if primType v_t - then SubExp $ Var v - else Copy v + then SubExp $ Var v + else Copy v ruleBasicOp vtable pat _ (Replicate shape (Var v)) | Just (BasicOp (Replicate shape2 se), cs) <- ST.lookupExp v vtable = - certifying cs $ letBind_ pat $ BasicOp $ Replicate (shape<>shape2) se + Simplify $ certifying cs $ letBind_ pat $ BasicOp $ Replicate (shape<>shape2) se -- | Turn array literals with identical elements into replicates. ruleBasicOp _ pat _ (ArrayLit (se:ses) _) | all (==se) ses = - let n = constant (fromIntegral (length ses) + 1 :: Int32) - in letBind_ pat $ BasicOp $ Replicate (Shape [n]) se + Simplify $ let n = constant (fromIntegral (length ses) + 1 :: Int32) + in letBind_ pat $ BasicOp $ Replicate (Shape [n]) se ruleBasicOp vtable pat (StmAux cs _) (Index idd slice) | Just inds <- sliceIndices slice, Just (BasicOp (Reshape newshape idd2), idd_cs) <- ST.lookupExp idd vtable, length newshape == length inds = + Simplify $ case shapeCoercion newshape of Just _ -> certifying (cs<>idd_cs) $ @@ -1136,22 +1131,22 @@ ruleBasicOp _ pat _ (BinOp (Pow t) e1 e2) | e1 == intConst t 2 = - letBind_ pat $ BasicOp $ BinOp (Shl t) (intConst t 1) e2 + Simplify $ letBind_ pat $ BasicOp $ BinOp (Shl t) (intConst t 1) e2 -- Handle identity permutation. ruleBasicOp _ pat _ (Rearrange perm v) | sort perm == perm = - letBind_ pat $ BasicOp $ SubExp $ Var v + Simplify $ letBind_ pat $ BasicOp $ SubExp $ Var v ruleBasicOp vtable pat (StmAux cs _) (Rearrange perm v) | Just (BasicOp (Rearrange perm2 e), v_cs) <- ST.lookupExp v vtable = -- Rearranging a rearranging: compose the permutations. - certifying (cs<>v_cs) $ + Simplify $ certifying (cs<>v_cs) $ letBind_ pat $ BasicOp $ Rearrange (perm `rearrangeCompose` perm2) e ruleBasicOp vtable pat (StmAux cs _) (Rearrange perm v) | Just (BasicOp (Rotate offsets v2), v_cs) <- ST.lookupExp v vtable, - Just (BasicOp (Rearrange perm3 v3), v2_cs) <- ST.lookupExp v2 vtable = do + Just (BasicOp (Rearrange perm3 v3), v2_cs) <- ST.lookupExp v2 vtable = Simplify $ do let offsets' = rearrangeShape (rearrangeInverse perm3) offsets rearrange_rotate <- letExp "rearrange_rotate" $ BasicOp $ Rotate offsets' v3 certifying (cs<>v_cs<>v2_cs) $ @@ -1163,18 +1158,18 @@ num_dims <- shapeRank dims, (rep_perm, rest_perm) <- splitAt num_dims perm, not $ null rest_perm, - rep_perm == [0..length rep_perm-1] = certifying (cs<>v1_cs) $ do + rep_perm == [0..length rep_perm-1] = Simplify $ certifying (cs<>v1_cs) $ do v <- letSubExp "rearrange_replicate" $ BasicOp $ Rearrange (map (subtract num_dims) rest_perm) v2 letBind_ pat $ BasicOp $ Replicate dims v -- A zero-rotation is identity. ruleBasicOp _ pat _ (Rotate offsets v) - | all isCt0 offsets = letBind_ pat $ BasicOp $ SubExp $ Var v + | all isCt0 offsets = Simplify $ letBind_ pat $ BasicOp $ SubExp $ Var v ruleBasicOp vtable pat (StmAux cs _) (Rotate offsets v) | Just (BasicOp (Rearrange perm v2), v_cs) <- ST.lookupExp v vtable, - Just (BasicOp (Rotate offsets2 v3), v2_cs) <- ST.lookupExp v2 vtable = do + Just (BasicOp (Rotate offsets2 v3), v2_cs) <- ST.lookupExp v2 vtable = Simplify $ do let offsets2' = rearrangeShape (rearrangeInverse perm) offsets2 addOffsets x y = letSubExp "summed_offset" $ BasicOp $ BinOp (Add Int32) x y offsets' <- zipWithM addOffsets offsets offsets2' @@ -1185,7 +1180,7 @@ -- Combining Rotates. ruleBasicOp vtable pat (StmAux cs _) (Rotate offsets1 v) - | Just (BasicOp (Rotate offsets2 v2), v_cs) <- ST.lookupExp v vtable = do + | Just (BasicOp (Rotate offsets2 v2), v_cs) <- ST.lookupExp v vtable = Simplify $ do offsets <- zipWithM add offsets1 offsets2 certifying (cs<>v_cs) $ letBind_ pat $ BasicOp $ Rotate offsets v2 @@ -1203,7 +1198,7 @@ -- XXX: we should check for proper aliasing here instead. arr_y /= arr_x, Just (slice_x_bef, DimFix i, []) <- focusNth (length slice_x - 1) slice_x, - Just (slice_y_bef, DimFix j, []) <- focusNth (length slice_y - 1) slice_y = do + Just (slice_y_bef, DimFix j, []) <- focusNth (length slice_y - 1) slice_y = Simplify $ do let slice_x' = slice_x_bef ++ [DimSlice i (intConst Int32 1) (intConst Int32 1)] slice_y' = slice_y_bef ++ [DimSlice j (intConst Int32 1) (intConst Int32 1)] v' <- letExp (baseString v ++ "_slice") $ BasicOp $ Index arr_y slice_y' @@ -1211,7 +1206,7 @@ letBind_ pat $ BasicOp $ Update arr_x slice_x' $ Var v' ruleBasicOp _ _ _ _ = - cannotSimplify + Skip -- | Remove the return values of a branch, that are not actually used -- after a branch. Standard dead code removal can remove the branch @@ -1236,8 +1231,9 @@ fb' = fb { bodyResult = pick fses } pat' = pick $ patternElements pat rettype' = pick rettype - in letBind_ (Pattern [] pat') $ If e1 tb' fb' $ IfAttr rettype' ifsort - | otherwise = cannotSimplify + in Simplify $ letBind_ (Pattern [] pat') $ If e1 tb' fb' $ IfAttr rettype' ifsort + + | otherwise = Skip -- Some helper functions diff --git a/src/Futhark/Optimise/TileLoops.hs b/src/Futhark/Optimise/TileLoops.hs --- a/src/Futhark/Optimise/TileLoops.hs +++ b/src/Futhark/Optimise/TileLoops.hs @@ -1,391 +1,916 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} --- | Perform a restricted form of loop tiling within kernel streams. --- We only tile primitive types, to avoid excessive local memory use. -module Futhark.Optimise.TileLoops - ( tileLoops ) - where - -import Control.Applicative -import Control.Monad.State -import Control.Monad.Reader -import qualified Data.Set as S -import qualified Data.Map.Strict as M -import Data.List -import Data.Maybe - -import Futhark.MonadFreshNames -import Futhark.Representation.Kernels - -import Futhark.Pass -import Futhark.Tools -import Futhark.Util (mapAccumLM) -import Futhark.Optimise.TileLoops.RegTiling3D - -tileLoops :: Pass Kernels Kernels -tileLoops = Pass "tile loops" "Tile stream loops inside kernels" $ - fmap Prog . mapM optimiseFunDef . progFunctions - -optimiseFunDef :: MonadFreshNames m => FunDef Kernels -> m (FunDef Kernels) -optimiseFunDef fundec = do - body' <- modifyNameSource $ runState $ - runReaderT m (scopeOfFParams (funDefParams fundec)) - return fundec { funDefBody = body' } - where m = optimiseBody $ funDefBody fundec - -type TileM = ReaderT (Scope Kernels) (State VNameSource) - -optimiseBody :: Body Kernels -> TileM (Body Kernels) -optimiseBody (Body () bnds res) = localScope (scopeOf bnds) $ - Body () <$> (mconcat <$> mapM optimiseStm (stmsToList bnds)) <*> pure res - -optimiseStm :: Stm Kernels -> TileM (Stms Kernels) -optimiseStm stmt@(Let pat aux (Op (HostOp old_kernel@(Kernel desc space ts body)))) = do - res3dtiling <- doRegTiling3D stmt - case res3dtiling of - Just (extra_bnds, stmt') -> return $ extra_bnds <> oneStm stmt' - Nothing -> do - (extra_bnds, space', body') <- tileInKernelBody mempty initial_variance space body - let new_kernel = Kernel desc space' ts body' - -- XXX: we should not change the type of the kernel (such as by - -- changing the number of groups being used for a kernel that - -- returns a result-per-group). - if kernelType old_kernel == kernelType new_kernel - then return $ extra_bnds <> oneStm (Let pat aux $ Op $ HostOp new_kernel) - else return $ oneStm $ Let pat aux $ Op $ HostOp old_kernel - where initial_variance = M.map mempty $ scopeOfKernelSpace space -optimiseStm (Let pat aux e) = - pure <$> (Let pat aux <$> mapExpM optimise e) - where optimise = identityMapper { mapOnBody = const optimiseBody } - -tileInKernelBody :: Names -> VarianceTable - -> KernelSpace -> KernelBody InKernel - -> TileM (Stms Kernels, KernelSpace, KernelBody InKernel) -tileInKernelBody branch_variant initial_variance initial_kspace (KernelBody () kstms kres) = do - (extra_bnds, kspace', kstms') <- - tileInStms branch_variant initial_variance initial_kspace kstms - return (extra_bnds, kspace', KernelBody () kstms' kres) - -tileInBody :: Names -> VarianceTable - -> KernelSpace -> Body InKernel - -> TileM (Stms Kernels, KernelSpace, Body InKernel) -tileInBody branch_variant initial_variance initial_kspace (Body () stms res) = do - (extra_bnds, kspace', stms') <- - tileInStms branch_variant initial_variance initial_kspace stms - return (extra_bnds, kspace', Body () stms' res) - -tileInStms :: Names -> VarianceTable - -> KernelSpace -> Stms InKernel - -> TileM (Stms Kernels, KernelSpace, Stms InKernel) -tileInStms branch_variant initial_variance initial_kspace kstms = do - ((kspace, extra_bndss), kstms') <- - mapAccumLM tileInKernelStatement (initial_kspace,mempty) $ stmsToList kstms - return (extra_bndss, kspace, stmsFromList kstms') - where variance = varianceInStms initial_variance kstms - - tileInKernelStatement (kspace, extra_bnds) - (Let pat attr (Op (GroupStream w max_chunk lam accs arrs))) - | max_chunk == w, - not $ null arrs, - chunk_size <- Var $ groupStreamChunkSize lam, - arr_chunk_params <- groupStreamArrParams lam, - maybe_1d_tiles <- - zipWith (is1dTileable branch_variant kspace variance chunk_size) arrs arr_chunk_params, - maybe_1_5d_tiles <- - zipWith (is1_5dTileable branch_variant kspace variance chunk_size) arrs arr_chunk_params, - Just mk_tilings <- - zipWithM (<|>) maybe_1d_tiles maybe_1_5d_tiles = do - - (kspaces, arr_chunk_params', tile_kstms) <- unzip3 <$> sequence mk_tilings - - let (kspace', kspace_bnds) = - case kspaces of - [] -> (kspace, mempty) - new_kspace : _ -> new_kspace - Body () lam_kstms lam_res <- syncAtEnd $ groupStreamLambdaBody lam - let lam_kstms' = mconcat tile_kstms <> lam_kstms - group_size = spaceGroupSize kspace - lam' = lam { groupStreamLambdaBody = Body () lam_kstms' lam_res - , groupStreamArrParams = arr_chunk_params' - } - - return ((kspace', extra_bnds <> kspace_bnds), - Let pat attr $ Op $ GroupStream w group_size lam' accs arrs) - - tileInKernelStatement (kspace, extra_bnds) - (Let pat attr (Op (GroupStream w max_chunk lam accs arrs))) - | w == max_chunk, - not $ null arrs, - FlatThreadSpace gspace <- spaceStructure kspace, - chunk_size <- Var $ groupStreamChunkSize lam, - arr_chunk_params <- groupStreamArrParams lam, - - Just mk_tilings <- - zipWithM (is2dTileable branch_variant kspace variance chunk_size) - arrs arr_chunk_params = do - - ((tile_size, tiled_group_size), tile_size_bnds) <- runBinder $ do - tile_size_key <- nameFromString . pretty <$> newVName "tile_size" - tile_size <- letSubExp "tile_size" $ Op $ GetSize tile_size_key SizeTile - tiled_group_size <- letSubExp "tiled_group_size" $ - BasicOp $ BinOp (Mul Int32) tile_size tile_size - return (tile_size, tiled_group_size) - - let (tiled_gspace,untiled_gspace) = splitAt 2 $ reverse gspace - -- Play with reversion to ensure we get increasing IDs for - -- ltids. This affects readability of generated code. - untiled_gspace' <- fmap reverse $ forM (reverse untiled_gspace) $ \(gtid,gdim) -> do - ltid <- newVName "ltid" - return (gtid,gdim, - ltid, constant (1::Int32)) - tiled_gspace' <- fmap reverse $ forM (reverse tiled_gspace) $ \(gtid,gdim) -> do - ltid <- newVName "ltid" - return (gtid,gdim, - ltid, tile_size) - let gspace' = reverse $ tiled_gspace' ++ untiled_gspace' - - -- We have to recalculate number of workgroups and - -- number of threads to fit the new workgroup size. - ((num_threads, num_groups), num_bnds) <- - runBinder $ sufficientGroups gspace' tiled_group_size - - let kspace' = kspace { spaceStructure = NestedThreadSpace gspace' - , spaceGroupSize = tiled_group_size - , spaceNumVirtGroups = num_groups - , spaceNumThreads = num_threads - , spaceNumGroups = num_groups - } - local_ids = map (\(_, _, ltid, _) -> ltid) gspace' - - (arr_chunk_params', tile_kstms) <- - fmap unzip $ forM mk_tilings $ \mk_tiling -> - mk_tiling tile_size local_ids - - Body () lam_kstms lam_res <- syncAtEnd $ groupStreamLambdaBody lam - let lam_kstms' = mconcat tile_kstms <> lam_kstms - lam' = lam { groupStreamLambdaBody = Body () lam_kstms' lam_res - , groupStreamArrParams = arr_chunk_params' - } - - return ((kspace', extra_bnds <> tile_size_bnds <> num_bnds), - Let pat attr $ Op $ GroupStream w tile_size lam' accs arrs) - - tileInKernelStatement (kspace, extra_bnds) - (Let pat attr (Op (GroupStream w maxchunk lam accs arrs))) = do - let branch_variant' = branch_variant <> - fromMaybe mempty (flip M.lookup variance =<< subExpVar w) - (bnds, kspace', lam') <- tileInStreamLambda branch_variant' variance kspace lam - return ((kspace', extra_bnds <> bnds), - Let pat attr $ Op $ GroupStream w maxchunk lam' accs arrs) - - tileInKernelStatement acc stm = - return (acc, stm) - -tileInStreamLambda :: Names -> VarianceTable -> KernelSpace -> GroupStreamLambda InKernel - -> TileM (Stms Kernels, KernelSpace, GroupStreamLambda InKernel) -tileInStreamLambda branch_variant variance kspace lam = do - (bnds, kspace', kbody') <- - tileInBody branch_variant variance' kspace $ groupStreamLambdaBody lam - return (bnds, kspace', lam { groupStreamLambdaBody = kbody' }) - where variance' = varianceInStms variance $ - bodyStms $ groupStreamLambdaBody lam - -is1dTileable :: MonadFreshNames m => - Names -> KernelSpace -> VarianceTable -> SubExp -> VName -> LParam InKernel - -> Maybe (m ((KernelSpace, Stms Kernels), - LParam InKernel, - Stms InKernel)) -is1dTileable branch_variant kspace variance block_size arr block_param = do - guard $ S.null $ M.findWithDefault mempty arr variance - guard $ S.null branch_variant - guard $ primType $ rowType $ paramType block_param - - return $ do - (outer_block_param, kstms) <- tile1d kspace block_size block_param - return ((kspace, mempty), outer_block_param, kstms) - -is1_5dTileable :: (MonadFreshNames m, HasScope Kernels m) => - Names -> KernelSpace -> VarianceTable - -> SubExp -> VName -> LParam InKernel - -> Maybe (m ((KernelSpace, Stms Kernels), - LParam InKernel, - Stms InKernel)) -is1_5dTileable branch_variant kspace variance block_size arr block_param = do - guard $ primType $ rowType $ paramType block_param - - (inner_gtid, inner_gdim) <- invariantToInnermostDimension - mk_structure <- - case spaceStructure kspace of - NestedThreadSpace{} -> Nothing - FlatThreadSpace gtids_and_gdims -> - return $ do - -- Force a functioning group size. XXX: not pretty. - let n_dims = length gtids_and_gdims - outer <- forM (take (n_dims-1) gtids_and_gdims) $ \(gtid, gdim) -> do - ltid <- newVName "ltid" - return (gtid, gdim, ltid, gdim) - - inner_ltid <- newVName "inner_ltid" - inner_ldim <- newVName "inner_ldim" - let compute_tiled_group_size = - mkLet [] [Ident inner_ldim $ Prim int32] $ - BasicOp $ BinOp (SMin Int32) (spaceGroupSize kspace) inner_gdim - structure = NestedThreadSpace $ outer ++ [(inner_gtid, inner_gdim, - inner_ltid, Var inner_ldim)] - ((num_threads, num_groups), num_bnds) <- runBinder $ do - threads_necessary <- - letSubExp "threads_necessary" =<< - foldBinOp (Mul Int32) - (constant (1::Int32)) (map snd gtids_and_gdims) - groups_necessary <- - letSubExp "groups_necessary" =<< - eDivRoundingUp Int32 (eSubExp threads_necessary) (eSubExp $ Var inner_ldim) - num_threads <- - letSubExp "num_threads" $ - BasicOp $ BinOp (Mul Int32) groups_necessary (Var inner_ldim) - return (num_threads, groups_necessary) - - let kspace' = kspace { spaceGroupSize = Var inner_ldim - , spaceNumGroups = num_groups - , spaceNumVirtGroups = num_groups - , spaceNumThreads = num_threads - , spaceStructure = structure - } - return (oneStm compute_tiled_group_size <> num_bnds, - kspace') - return $ do - (outer_block_param, kstms) <- tile1d kspace block_size block_param - (structure_bnds, kspace') <- mk_structure - return ((kspace', structure_bnds), outer_block_param, kstms) - where invariantToInnermostDimension :: Maybe (VName, SubExp) - invariantToInnermostDimension = - case reverse $ spaceDimensions kspace of - (i,d) : _ - | not $ i `S.member` M.findWithDefault mempty arr variance, - not $ i `S.member` branch_variant -> Just (i,d) - _ -> Nothing - -tile1d :: MonadFreshNames m => - KernelSpace - -> SubExp - -> LParam InKernel - -> m (LParam InKernel, Stms InKernel) -tile1d kspace block_size block_param = do - outer_block_param <- do - name <- newVName $ baseString (paramName block_param) ++ "_outer" - return block_param { paramName = name } - - let ltid = spaceLocalId kspace - read_elem_bnd <- do - name <- newVName $ baseString (paramName outer_block_param) ++ "_elem" - return $ - mkLet [] [Ident name $ rowType $ paramType outer_block_param] $ - BasicOp $ Index (paramName outer_block_param) [DimFix $ Var ltid] - - cid <- newVName "cid" - let block_cspace = combineSpace [(cid, block_size)] - block_pe = - PatElem (paramName block_param) $ paramType outer_block_param - write_block_stms = - [ Let (Pattern [] [block_pe]) (defAux ()) $ Op $ - Combine block_cspace [patElemType pe] [] $ - Body () (oneStm read_elem_bnd) [Var $ patElemName pe] - | pe <- patternElements $ stmPattern read_elem_bnd ] - - return (outer_block_param, stmsFromList write_block_stms) - -is2dTileable :: MonadFreshNames m => - Names -> KernelSpace -> VarianceTable -> SubExp -> VName -> LParam InKernel - -> Maybe (SubExp -> [VName] -> m (LParam InKernel, Stms InKernel)) -is2dTileable branch_variant kspace variance block_size arr block_param = do - guard $ primType $ rowType $ paramType block_param - - pt <- case rowType $ paramType block_param of - Prim pt -> return pt - _ -> Nothing - inner_perm <- invariantToOneOfTwoInnerDims - Just $ \tile_size local_is -> do - let num_outer = length local_is - 2 - perm = [0..num_outer-1] ++ map (+num_outer) inner_perm - invariant_i : variant_i : _ = reverse $ rearrangeShape perm local_is - (global_i,global_d):_ = rearrangeShape inner_perm $ drop num_outer $ spaceDimensions kspace - outer_block_param <- do - name <- newVName $ baseString (paramName block_param) ++ "_outer" - return block_param { paramName = name } - - elem_name <- newVName $ baseString (paramName outer_block_param) ++ "_elem" - let read_elem_bnd = mkLet [] [Ident elem_name $ Prim pt] $ - BasicOp $ Index (paramName outer_block_param) $ - fullSlice (paramType outer_block_param) [DimFix $ Var invariant_i] - - cids <- replicateM (length local_is - num_outer) $ newVName "cid" - let block_size_2d = Shape $ rearrangeShape inner_perm [tile_size, block_size] - block_cspace = combineSpace $ zip cids $ - rearrangeShape inner_perm [tile_size,block_size] - - block_name_2d <- newVName $ baseString (paramName block_param) ++ "_2d" - let block_pe = - PatElem block_name_2d $ - rowType (paramType outer_block_param) `arrayOfShape` block_size_2d - write_block_stm = - Let (Pattern [] [block_pe]) (defAux ()) $ - Op $ Combine block_cspace [Prim pt] [(global_i, global_d)] $ - Body () (oneStm read_elem_bnd) [Var elem_name] - - let index_block_kstms = - [mkLet [] [paramIdent block_param] $ - BasicOp $ Index block_name_2d $ - rearrangeShape inner_perm $ - fullSlice (rearrangeType inner_perm $ patElemType block_pe) - [DimFix $ Var variant_i]] - - return (outer_block_param, - oneStm write_block_stm <> stmsFromList index_block_kstms) - - where invariantToOneOfTwoInnerDims :: Maybe [Int] - invariantToOneOfTwoInnerDims = do - (j,_) : (i,_) : _ <- Just $ reverse $ spaceDimensions kspace - let variant_to = M.findWithDefault mempty arr variance - branch_invariant = not $ S.member j branch_variant || S.member i branch_variant - if branch_invariant && i `S.member` variant_to && not (j `S.member` variant_to) then - Just [0,1] - else if branch_invariant && j `S.member` variant_to && not (i `S.member` variant_to) then - Just [1,0] - else - Nothing - -syncAtEnd :: MonadFreshNames m => Body InKernel -> m (Body InKernel) -syncAtEnd (Body () stms res) = do - (res', stms') <- (`runBinderT` mempty) $ do - mapM_ addStm stms - map Var <$> letTupExp "sync" (Op $ Barrier res) - return $ Body () stms' res' - --- | The variance table keeps a mapping from a variable name --- (something produced by a 'Stm') to the kernel thread indices --- that name depends on. If a variable is not present in this table, --- that means it is bound outside the kernel (and so can be considered --- invariant to all dimensions). -type VarianceTable = M.Map VName Names - -varianceInStms :: VarianceTable -> Stms InKernel -> VarianceTable -varianceInStms = foldl varianceInStm - -varianceInStm :: VarianceTable -> Stm InKernel -> VarianceTable -varianceInStm variance bnd = - foldl' add variance $ patternNames $ stmPattern bnd - where add variance' v = M.insert v binding_variance variance' - look variance' v = S.insert v $ M.findWithDefault mempty v variance' - binding_variance = mconcat $ map (look variance) $ S.toList (freeIn bnd) - -sufficientGroups :: MonadBinder m => - [(VName, SubExp, VName, SubExp)] -> SubExp - -> m (SubExp, SubExp) -sufficientGroups gspace group_size = do - groups_in_dims <- forM gspace $ \(_, gd, _, ld) -> - letSubExp "groups_in_dim" =<< eDivRoundingUp Int32 (eSubExp gd) (eSubExp ld) - num_groups <- letSubExp "num_groups" =<< - foldBinOp (Mul Int32) (constant (1::Int32)) groups_in_dims - num_threads <- letSubExp "num_threads" $ - BasicOp $ BinOp (Mul Int32) num_groups group_size - return (num_threads, num_groups) +-- | Perform a restricted form of loop tiling within SegMaps. We only +-- tile primitive types, to avoid excessive local memory use. +module Futhark.Optimise.TileLoops + ( tileLoops ) + where + +import Control.Monad.State +import Control.Monad.Reader +import qualified Data.Sequence as Seq +import qualified Data.Map.Strict as M +import Data.List + +import Prelude hiding (quot) + +import Futhark.MonadFreshNames +import Futhark.Representation.Kernels +import Futhark.Transform.Rename +import Futhark.Pass +import Futhark.Tools + +tileLoops :: Pass Kernels Kernels +tileLoops = Pass "tile loops" "Tile stream loops inside kernels" $ + fmap Prog . mapM optimiseFunDef . progFunctions + +optimiseFunDef :: MonadFreshNames m => FunDef Kernels -> m (FunDef Kernels) +optimiseFunDef fundec = do + body' <- modifyNameSource $ runState $ + runReaderT m (scopeOfFParams (funDefParams fundec)) + return fundec { funDefBody = body' } + where m = optimiseBody $ funDefBody fundec + +type TileM = ReaderT (Scope Kernels) (State VNameSource) + +optimiseBody :: Body Kernels -> TileM (Body Kernels) +optimiseBody (Body () bnds res) = localScope (scopeOf bnds) $ + Body () <$> (mconcat <$> mapM optimiseStm (stmsToList bnds)) <*> pure res + +optimiseStm :: Stm Kernels -> TileM (Stms Kernels) +optimiseStm (Let pat aux (Op (SegOp (SegMap lvl@SegThread{} space ts kbody)))) = do + (host_stms, (lvl', space', kbody')) <- tileInKernelBody mempty initial_variance lvl space ts kbody + return $ host_stms <> + oneStm (Let pat aux $ Op $ SegOp $ SegMap lvl' space' ts kbody') + where initial_variance = M.map mempty $ scopeOfSegSpace space + +optimiseStm (Let pat aux e) = + pure <$> (Let pat aux <$> mapExpM optimise e) + where optimise = identityMapper { mapOnBody = \scope -> localScope scope . optimiseBody } + +tileInKernelBody :: Names -> VarianceTable + -> SegLevel -> SegSpace -> [Type] -> KernelBody Kernels + -> TileM (Stms Kernels, (SegLevel, SegSpace, KernelBody Kernels)) +tileInKernelBody branch_variant initial_variance lvl initial_kspace ts kbody + | Just kbody_res <- mapM isSimpleResult $ kernelBodyResult kbody = do + maybe_tiled <- + tileInBody branch_variant initial_variance lvl initial_kspace ts $ + Body () (kernelBodyStms kbody) kbody_res + case maybe_tiled of + Just (host_stms, tiling, tiledBody) -> do + (res', stms') <- + runBinder $ mapM (tilingTileReturns tiling) =<< tiledBody mempty + return (host_stms, (tilingLevel tiling, + tilingSpace tiling, + KernelBody () stms' res')) + Nothing -> + return (mempty, (lvl, initial_kspace, kbody)) + | otherwise = + return (mempty, (lvl, initial_kspace, kbody)) + where isSimpleResult (Returns se) = Just se + isSimpleResult _ = Nothing + +tileInBody :: Names -> VarianceTable + -> SegLevel -> SegSpace -> [Type] -> Body Kernels + -> TileM (Maybe (Stms Kernels, Tiling, TiledBody)) +tileInBody branch_variant initial_variance initial_lvl initial_space res_ts (Body () initial_kstms stms_res) = + descend mempty $ stmsToList initial_kstms + where + variance = varianceInStms initial_variance initial_kstms + + descend _ [] = + return Nothing + + descend prestms (stm_to_tile:poststms) + + -- 1D tiling of redomap. + | (gtid, kdim) : top_space_rev <- reverse $ unSegSpace initial_space, + Just (w, arrs, form) <- tileable stm_to_tile, + all (not . nameIn gtid . + flip (M.findWithDefault mempty) variance) arrs, + not $ gtid `nameIn` branch_variant, + (prestms', poststms') <- + preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms) = + + Just . injectPrelude initial_space variance prestms' (freeIn stm_to_tile) <$> + tileGeneric (tiling1d $ reverse top_space_rev) + initial_lvl res_ts (stmPattern stm_to_tile) + gtid kdim + w form (zip arrs $ repeat [0]) poststms' stms_res + + -- 2D tiling of redomap. + | (gtids, kdims) <- unzip $ unSegSpace initial_space, + Just (w, arrs, form) <- tileable stm_to_tile, + Just inner_perm <- mapM (invariantToOneOfTwoInnerDims branch_variant variance gtids) arrs, + gtid_y : gtid_x : top_gtids_rev <- reverse gtids, + kdim_y : kdim_x : top_kdims_rev <- reverse kdims, + (prestms', poststms') <- + preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms) = + + Just . injectPrelude initial_space variance prestms' (freeIn stm_to_tile) <$> + tileGeneric (tiling2d $ reverse $ zip top_gtids_rev top_kdims_rev) + initial_lvl res_ts (stmPattern stm_to_tile) + (gtid_x, gtid_y) (kdim_x, kdim_y) + w form (zip arrs inner_perm) poststms' stms_res + + -- Tiling inside for-loop. + | DoLoop [] merge (ForLoop i it bound []) loopbody <- stmExp stm_to_tile, + (prestms', poststms') <- + preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms)= do + + let branch_variant' = + branch_variant <> + mconcat (map (flip (M.findWithDefault mempty) variance) + (namesToList (freeIn bound))) + merge_params = map fst merge + + maybe_tiled <- + localScope (M.insert i (IndexInfo it) $ scopeOfFParams merge_params) $ + tileInBody branch_variant' variance initial_lvl initial_space + (map paramType merge_params) $ mkBody (bodyStms loopbody) (bodyResult loopbody) + + case maybe_tiled of + Nothing -> next + Just tiled -> + Just <$> tileDoLoop initial_space variance prestms' + (freeIn loopbody <> freeIn merge) tiled + res_ts (stmPattern stm_to_tile) (stmAux stm_to_tile) + merge i it bound poststms' stms_res + + | otherwise = next + + where next = localScope (scopeOf stm_to_tile) $ + descend (prestms <> oneStm stm_to_tile) poststms + +-- | Move statements from prelude to postlude if they are not used in +-- the tiled statement anyway. +preludeToPostlude :: VarianceTable + -> Stms Kernels -> Stm Kernels -> Stms Kernels + -> (Stms Kernels, Stms Kernels) +preludeToPostlude variance prelude stm_to_tile postlude = + (prelude_used, prelude_not_used <> postlude) + where used_in_tiled = freeIn stm_to_tile + + used_in_stm_variant = + (used_in_tiled<>) $ mconcat $ + map (flip (M.findWithDefault mempty) variance) $ + namesToList used_in_tiled + + used stm = any (`nameIn` used_in_stm_variant) $ + patternNames $ stmPattern stm + + (prelude_used, prelude_not_used) = + Seq.partition used prelude + +-- | Partition prelude statements preceding a tiled loop (or something +-- containing a tiled loop) into three categories: +-- +-- 1) Group-level statements that are invariant to the threads in the group. +-- +-- 2) Thread-variant statements that should be computed once with a segmap_thread_scalar. +-- +-- 3) Thread-variant statements that should be recomputed whenever +-- they are needed. +-- +-- The third category duplicates computation, so we only want to do it +-- when absolutely necessary. Currently, this is necessary for +-- results that are views of an array (slicing, rotate, etc), because +-- these cannot be efficiently represented by a scalar segmap (they'll +-- be manifested in memory). +partitionPrelude :: VarianceTable -> Stms Kernels -> Names + -> (Stms Kernels, Stms Kernels, Stms Kernels) +partitionPrelude variance prestms tiled_kdims = + (invariant_prestms, precomputed_variant_prestms, recomputed_variant_prestms) + where + invariantTo names stm = + case patternNames (stmPattern stm) of + [] -> True -- Does not matter. + v:_ -> not $ any (`nameIn` names) $ namesToList $ + M.findWithDefault mempty v variance + (invariant_prestms, variant_prestms) = + Seq.partition (invariantTo tiled_kdims) prestms + + mustBeInlinedExp (BasicOp (Index _ slice)) = not $ null $ sliceDims slice + mustBeInlinedExp (BasicOp Rotate{}) = True + mustBeInlinedExp (BasicOp Rearrange{}) = True + mustBeInlinedExp (BasicOp Reshape{}) = True + mustBeInlinedExp _ = False + mustBeInlined = mustBeInlinedExp . stmExp + + must_be_inlined = namesFromList $ concatMap (patternNames . stmPattern) $ + stmsToList $ Seq.filter mustBeInlined variant_prestms + recompute stm = + any (`nameIn` must_be_inlined) (patternNames (stmPattern stm)) || + not (invariantTo must_be_inlined stm) + (recomputed_variant_prestms, precomputed_variant_prestms) = + Seq.partition recompute variant_prestms + +injectPrelude :: SegSpace -> VarianceTable + -> Stms Kernels -> Names + -> (Stms Kernels, Tiling, TiledBody) + -> (Stms Kernels, Tiling, TiledBody) +injectPrelude initial_space variance prestms used (host_stms, tiling, tiledBody) = + (host_stms, tiling, tiledBody') + where tiled_kdims = namesFromList $ map fst $ + filter (`notElem` unSegSpace (tilingSpace tiling)) $ + unSegSpace initial_space + + tiledBody' privstms = do + let (invariant_prestms, + precomputed_variant_prestms, + recomputed_variant_prestms) = + partitionPrelude variance prestms tiled_kdims + + addStms invariant_prestms + + let live_set = namesToList $ liveSet precomputed_variant_prestms $ + used <> freeIn recomputed_variant_prestms + prelude_arrs <- inScopeOf precomputed_variant_prestms $ + doPrelude tiling precomputed_variant_prestms live_set + + let prelude_privstms = + PrivStms recomputed_variant_prestms $ + mkReadPreludeValues prelude_arrs live_set + + tiledBody (prelude_privstms <> privstms) + +tileDoLoop :: SegSpace -> VarianceTable + -> Stms Kernels -> Names + -> (Stms Kernels, Tiling, TiledBody) + -> [Type] -> Pattern Kernels -> StmAux (ExpAttr Kernels) + -> [(FParam Kernels, SubExp)] -> VName -> IntType -> SubExp + -> Stms Kernels -> Result + -> TileM (Stms Kernels, Tiling, TiledBody) +tileDoLoop initial_space variance prestms used_in_body (host_stms, tiling, tiledBody) res_ts pat aux merge i it bound poststms poststms_res = do + + let (invariant_prestms, + precomputed_variant_prestms, + recomputed_variant_prestms) = + partitionPrelude variance prestms tiled_kdims + + let (mergeparams, mergeinits) = unzip merge + + -- Expand the loop merge parameters to be arrays. + tileDim t = arrayOf t (tilingTileShape tiling) $ uniqueness t + + tiledBody' privstms = inScopeOf host_stms $ do + addStms invariant_prestms + + let live_set = namesToList $ liveSet precomputed_variant_prestms used_in_body + prelude_arrs <- inScopeOf precomputed_variant_prestms $ + doPrelude tiling precomputed_variant_prestms live_set + + mergeparams' <- forM mergeparams $ \(Param pname pt) -> + Param <$> newVName (baseString pname ++ "_group") <*> pure (tileDim pt) + + let merge_ts = map paramType mergeparams + + let inloop_privstms = + PrivStms recomputed_variant_prestms $ + mkReadPreludeValues prelude_arrs live_set + + mergeinit' <- + fmap (map Var) $ certifying (stmAuxCerts aux) $ + tilingSegMap tiling "tiled_loopinit" (scalarLevel tiling) $ + \in_bounds slice -> + fmap (map Var) $ protectOutOfBounds "loopinit" in_bounds merge_ts $ do + addPrivStms slice inloop_privstms + addPrivStms slice privstms + return mergeinits + + let merge' = zip mergeparams' mergeinit' + + let indexMergeParams slice = + localScope (scopeOfFParams mergeparams') $ + forM_ (zip mergeparams mergeparams') $ \(to, from) -> + letBindNames_ [paramName to] $ BasicOp $ Index (paramName from) $ + fullSlice (paramType from) slice + + loopbody' <- runBodyBinder $ resultBody . map Var <$> + tiledBody (privstms <> inloop_privstms <> PrivStms mempty indexMergeParams) + accs' <- letTupExp "tiled_inside_loop" $ + DoLoop [] merge' (ForLoop i it bound []) loopbody' + + postludeGeneric tiling privstms pat accs' poststms poststms_res res_ts + + return (host_stms, tiling, tiledBody') + + where tiled_kdims = namesFromList $ map fst $ + filter (`notElem` unSegSpace (tilingSpace tiling)) $ + unSegSpace initial_space + +doPrelude :: Tiling -> Stms Kernels -> [VName] -> Binder Kernels [VName] +doPrelude tiling prestms prestms_live = + -- Create a SegMap that takes care of the prelude for every thread. + tilingSegMap tiling "prelude" (scalarLevel tiling) $ + \in_bounds _slice -> do + ts <- mapM lookupType prestms_live + fmap (map Var) $ letTupExp "pre" =<< + eIf (toExp in_bounds) + (do addStms prestms + resultBodyM $ map Var prestms_live) + (eBody $ map eBlank ts) + +liveSet :: FreeIn a => Stms Kernels -> a -> Names +liveSet stms after = + namesFromList (concatMap (patternNames . stmPattern) stms) `namesIntersection` + freeIn after + +tileable :: Stm Kernels + -> Maybe (SubExp, [VName], + (Commutativity, Lambda Kernels, [SubExp], Lambda Kernels)) +tileable stm + | Op (OtherOp (Screma w form arrs)) <- stmExp stm, + Just (reds, map_lam) <- isRedomapSOAC form, + Reduce red_comm red_lam red_nes <- singleReduce reds, + lambdaReturnType map_lam == lambdaReturnType red_lam, -- No mapout arrays. + not $ null arrs, + all primType $ lambdaReturnType map_lam, + all (primType . paramType) $ lambdaParams map_lam = + Just (w, arrs, (red_comm, red_lam, red_nes, map_lam)) + | otherwise = + Nothing + +-- | Statements that we insert directly into every thread-private +-- SegMaps. This is for things that cannot efficiently be computed +-- once in advance in the prelude SegMap, primarily (exclusively?) +-- array slicing operations. +data PrivStms = PrivStms (Stms Kernels) ReadPrelude + +privStms :: Stms Kernels -> PrivStms +privStms stms = PrivStms stms $ const $ return () + +addPrivStms :: Slice SubExp -> PrivStms -> Binder Kernels () +addPrivStms local_slice (PrivStms stms readPrelude) = do + readPrelude local_slice + addStms stms + +instance Semigroup PrivStms where + PrivStms stms_x readPrelude_x <> PrivStms stms_y readPrelude_y = + PrivStms stms_z readPrelude_z + where stms_z = stms_x <> stms_y + readPrelude_z slice = readPrelude_x slice >> readPrelude_y slice + +instance Monoid PrivStms where + mempty = privStms mempty + +type ReadPrelude = Slice SubExp -> Binder Kernels () + +-- | Information about a loop that has been tiled inside a kernel, as +-- well as the kinds of changes that we would then like to perform on +-- the kernel. +data Tiling = + Tiling + { tilingSegMap :: String -> SegLevel + -> (PrimExp VName -> Slice SubExp -> Binder Kernels [SubExp]) + -> Binder Kernels [VName] + -- The boolean PrimExp indicates whether they are in-bounds. + + , tilingReadTile :: TileKind -> PrivStms + -> SubExp -> [(VName, [Int])] + -> Binder Kernels [VName] + + , tilingProcessTile :: PrivStms + -> Commutativity -> Lambda Kernels -> Lambda Kernels + -> [(VName, [Int])] -> [VName] + -> Binder Kernels [VName] + + , tilingProcessResidualTile :: PrivStms + -> Commutativity -> Lambda Kernels -> Lambda Kernels + -> SubExp -> [VName] -> SubExp + -> [(VName, [Int])] + -> Binder Kernels [VName] + + , tilingTileReturns :: VName -> Binder Kernels KernelResult + + , tilingSpace :: SegSpace + + , tilingTileShape :: Shape + + , tilingLevel :: SegLevel + + , tilingNumWholeTiles :: SubExp + } + +type DoTiling gtids kdims = + SegLevel -> gtids -> kdims -> SubExp -> Binder Kernels Tiling + +scalarLevel :: Tiling -> SegLevel +scalarLevel tiling = + SegThreadScalar (segNumGroups lvl) (segGroupSize lvl) SegNoVirt + where lvl = tilingLevel tiling + +protectOutOfBounds :: String -> PrimExp VName -> [Type] -> Binder Kernels [SubExp] + -> Binder Kernels [VName] +protectOutOfBounds desc in_bounds ts m = + letTupExp desc =<< eIf (toExp in_bounds) (resultBody <$> m) (eBody $ map eBlank ts) + +postludeGeneric :: Tiling -> PrivStms + -> Pattern Kernels -> [VName] + -> Stms Kernels -> Result -> [Type] + -> Binder Kernels [VName] +postludeGeneric tiling privstms pat accs' poststms poststms_res res_ts = + tilingSegMap tiling "thread_res" (scalarLevel tiling) $ \in_bounds slice -> do + -- Read our per-thread result from the tiled loop. + forM_ (zip (patternNames pat) accs') $ \(us, everyone) -> + letBindNames_ [us] $ BasicOp $ Index everyone slice + + if poststms == mempty + then return poststms_res + else fmap (map Var) $ protectOutOfBounds "postlude" in_bounds res_ts $ do + addPrivStms slice privstms + addStms poststms + return poststms_res + +type TiledBody = PrivStms -> Binder Kernels [VName] + +tileGeneric :: DoTiling gtids kdims + -> SegLevel + -> [Type] + -> Pattern Kernels + -> gtids + -> kdims + -> SubExp + -> (Commutativity, Lambda Kernels, [SubExp], Lambda Kernels) + -> [(VName, [Int])] + -> Stms Kernels -> Result + -> TileM (Stms Kernels, Tiling, TiledBody) +tileGeneric doTiling initial_lvl res_ts pat gtids kdims w form arrs_and_perms poststms poststms_res = do + + (tiling, tiling_stms) <- runBinder $ doTiling initial_lvl gtids kdims w + + return (tiling_stms, tiling, tiledBody tiling) + + where + (red_comm, red_lam, red_nes, map_lam) = form + + tiledBody :: Tiling -> PrivStms -> Binder Kernels [VName] + tiledBody tiling privstms = do + let num_whole_tiles = tilingNumWholeTiles tiling + tile_shape = tilingTileShape tiling + + -- We don't use a Replicate here, because we want to enforce a + -- scalar memory space. + mergeinits <- tilingSegMap tiling "mergeinit" (scalarLevel tiling) $ \in_bounds slice -> + -- Constant neutral elements (a common case) do not need protection from OOB. + if freeIn red_nes == mempty + then return red_nes + else fmap (map Var) $ protectOutOfBounds "neutral" in_bounds (lambdaReturnType red_lam) $ do + addPrivStms slice privstms + return red_nes + + merge <- forM (zip (lambdaParams red_lam) mergeinits) $ \(p, mergeinit) -> + (,) <$> + newParam (baseString (paramName p) ++ "_merge") + (paramType p `arrayOfShape` tile_shape `toDecl` Unique) <*> + pure (Var mergeinit) + + tile_id <- newVName "tile_id" + let loopform = ForLoop tile_id Int32 num_whole_tiles [] + loopbody <- renameBody <=< runBodyBinder $ inScopeOf loopform $ + localScope (scopeOfFParams $ map fst merge) $ do + + -- Collectively read a tile. + tile <- tilingReadTile tiling TileFull privstms (Var tile_id) arrs_and_perms + + -- Now each thread performs a traversal of the tile and + -- updates its accumulator. + resultBody . map Var <$> + tilingProcessTile tiling privstms + red_comm red_lam map_lam + (zip tile (map snd arrs_and_perms)) (map (paramName . fst) merge) + + accs <- letTupExp "accs" $ DoLoop [] merge loopform loopbody + + -- We possibly have to traverse a residual tile. + red_lam' <- renameLambda red_lam + map_lam' <- renameLambda map_lam + accs' <- tilingProcessResidualTile tiling privstms + red_comm red_lam' map_lam' + num_whole_tiles accs w arrs_and_perms + + -- Create a SegMap that takes care of the postlude for every thread. + postludeGeneric tiling privstms pat accs' poststms poststms_res res_ts + +data TileKind = TilePartial | TileFull + +mkReadPreludeValues :: [VName] -> [VName] -> ReadPrelude +mkReadPreludeValues prestms_live_arrs prestms_live slice = + fmap mconcat $ forM (zip prestms_live_arrs prestms_live) $ \(arr, v) -> do + arr_t <- lookupType arr + letBindNames_ [v] $ BasicOp $ Index arr $ fullSlice arr_t slice + +tileReturns :: [(VName, SubExp)] -> [(SubExp, SubExp)] -> VName -> Binder Kernels KernelResult +tileReturns dims_on_top dims arr = do + let unit_dims = replicate (length dims_on_top) (intConst Int32 1) + arr' <- if null dims_on_top then return arr + else do arr_t <- lookupType arr + let new_shape = unit_dims ++ arrayDims arr_t + letExp (baseString arr) $ BasicOp $ Reshape (map DimNew new_shape) arr + let tile_dims = zip (map snd dims_on_top) unit_dims ++ dims + return $ TileReturns tile_dims arr' + +segMap1D :: String + -> SegLevel + -> (VName -> Binder Kernels [SubExp]) + -> Binder Kernels [VName] +segMap1D desc lvl f = do + ltid <- newVName "ltid" + ltid_flat <- newVName "ltid_flat" + let space = SegSpace ltid_flat [(ltid, unCount $ segGroupSize lvl)] + + ((ts, res), stms) <- runBinder $ do + res <- f ltid + ts <- mapM subExpType res + return (ts, res) + Body _ stms' res' <- renameBody $ mkBody stms res + + letTupExp desc $ Op $ SegOp $ + SegMap lvl space ts $ KernelBody () stms' $ map Returns res' + +reconstructGtids1D :: Count GroupSize SubExp -> VName -> VName -> VName + -> Binder Kernels () +reconstructGtids1D group_size gtid gid ltid = + letBindNames_ [gtid] =<< + toExp (LeafExp gid int32 * + primExpFromSubExp int32 (unCount group_size) + + LeafExp ltid int32) + +readTile1D :: SubExp -> VName -> VName + -> Count NumGroups SubExp -> Count GroupSize SubExp + -> TileKind -> PrivStms + -> SubExp + -> [(VName, [Int])] + -> Binder Kernels [VName] +readTile1D + tile_size gid gtid num_groups group_size + kind privstms tile_id arrs_and_perms = + + segMap1D "full_tile" (SegThread num_groups group_size SegNoVirt) $ \ltid -> do + j <- letSubExp "j" =<< + toExp (primExpFromSubExp int32 tile_id * + primExpFromSubExp int32 tile_size + + LeafExp ltid int32) + + reconstructGtids1D group_size gtid gid ltid + addPrivStms [DimFix $ Var ltid] privstms + + let arrs = map fst arrs_and_perms + arr_ts <- mapM lookupType arrs + let tile_ts = map rowType arr_ts + w = arraysSize 0 arr_ts + + let readTileElem arr = + -- No need for fullSlice because we are tiling only prims. + letExp "tile_elem" $ BasicOp $ Index arr [DimFix j] + fmap (map Var) $ + case kind of + TilePartial -> + letTupExp "pre" =<< eIf (toExp $ primExpFromSubExp int32 j .<. + primExpFromSubExp int32 w) + (resultBody <$> mapM (fmap Var . readTileElem) arrs) + (eBody $ map eBlank tile_ts) + TileFull -> + mapM readTileElem arrs + +processTile1D :: VName -> VName -> SubExp -> SubExp + -> Count NumGroups SubExp -> Count GroupSize SubExp + -> PrivStms + -> Commutativity -> Lambda Kernels -> Lambda Kernels + -> [(VName, [Int])] -> [VName] + -> Binder Kernels [VName] +processTile1D + gid gtid kdim tile_size num_groups group_size + privstms + red_comm red_lam map_lam tiles_and_perm accs = do + + let tile = map fst tiles_and_perm + + segMap1D "acc" (SegThreadScalar num_groups group_size SegNoVirt) $ \ltid -> do + + reconstructGtids1D group_size gtid gid ltid + addPrivStms [DimFix $ Var ltid] privstms + + -- We replace the neutral elements with the accumulators (this is + -- OK because the parallel semantics are not used after this + -- point). + thread_accs <- forM accs $ \acc -> + letSubExp "acc" $ BasicOp $ Index acc [DimFix $ Var ltid] + let form' = redomapSOAC [Reduce red_comm red_lam thread_accs] map_lam + + fmap (map Var) $ + letTupExp "acc" =<< eIf (toExp $ LeafExp gtid int32 .<. primExpFromSubExp int32 kdim) + (eBody [pure $ Op $ OtherOp $ Screma tile_size form' tile]) + (resultBodyM thread_accs) + +processResidualTile1D :: VName -> VName -> SubExp -> SubExp + -> Count NumGroups SubExp -> Count GroupSize SubExp -> PrivStms + -> Commutativity -> Lambda Kernels -> Lambda Kernels + -> SubExp -> [VName] -> SubExp -> [(VName, [Int])] + -> Binder Kernels [VName] +processResidualTile1D + gid gtid kdim tile_size num_groups group_size privstms red_comm red_lam map_lam + num_whole_tiles accs w arrs_and_perms = do + -- The number of residual elements that are not covered by + -- the whole tiles. + residual_input <- letSubExp "residual_input" $ + BasicOp $ BinOp (SRem Int32) w tile_size + + letTupExp "acc_after_residual" =<< + eIf (toExp $ primExpFromSubExp int32 residual_input .==. 0) + (resultBodyM $ map Var accs) + (nonemptyTile residual_input) + + where + nonemptyTile residual_input = runBodyBinder $ do + -- Collectively construct a tile. Threads that are out-of-bounds + -- provide a blank dummy value. + full_tile <- readTile1D tile_size gid gtid num_groups group_size + TilePartial privstms num_whole_tiles arrs_and_perms + tile <- forM full_tile $ \tile -> + letExp "partial_tile" $ BasicOp $ Index tile + [DimSlice (intConst Int32 0) residual_input (intConst Int32 1)] + + -- Now each thread performs a traversal of the tile and + -- updates its accumulator. + resultBody . map Var <$> processTile1D + gid gtid kdim residual_input num_groups group_size privstms + red_comm red_lam map_lam (zip tile $ repeat [0]) accs + +tiling1d :: [(VName,SubExp)] -> DoTiling VName SubExp +tiling1d dims_on_top initial_lvl gtid kdim w = do + gid <- newVName "gid" + gid_flat <- newVName "gid_flat" + + (lvl, space) <- + if null dims_on_top + then return (SegGroup (segNumGroups initial_lvl) (segGroupSize initial_lvl) $ segVirt initial_lvl, + SegSpace gid_flat [(gid, unCount $ segNumGroups initial_lvl)]) + else do + group_size <- letSubExp "computed_group_size" $ + BasicOp $ BinOp (SMin Int32) (unCount (segGroupSize initial_lvl)) kdim + + -- How many groups we need to exhaust the innermost dimension. + ldim <- letSubExp "ldim" =<< + eDivRoundingUp Int32 (eSubExp kdim) (eSubExp group_size) + + num_groups <- letSubExp "computed_num_groups" =<< + foldBinOp (Mul Int32) ldim (map snd dims_on_top) + + return (SegGroup (Count num_groups) (Count group_size) SegNoVirt, + SegSpace gid_flat $ dims_on_top ++ [(gid, ldim)]) + let tile_size = unCount $ segGroupSize lvl + + -- Number of whole tiles that fit in the input. + num_whole_tiles <- letSubExp "num_whole_tiles" $ BasicOp $ BinOp (SQuot Int32) w tile_size + return Tiling + { tilingSegMap = \desc lvl' f -> segMap1D desc lvl' $ \ltid -> do + letBindNames_ [gtid] =<< + toExp (LeafExp gid int32 * primExpFromSubExp int32 tile_size + + LeafExp ltid int32) + f (LeafExp gtid int32 .<. primExpFromSubExp int32 kdim) + [DimFix $ Var ltid] + + , tilingReadTile = + readTile1D tile_size gid gtid (segNumGroups lvl) (segGroupSize lvl) + + , tilingProcessTile = + processTile1D gid gtid kdim tile_size (segNumGroups lvl) (segGroupSize lvl) + + , tilingProcessResidualTile = + processResidualTile1D gid gtid kdim tile_size (segNumGroups lvl) (segGroupSize lvl) + + , tilingTileReturns = tileReturns dims_on_top [(kdim, tile_size)] + + , tilingTileShape = Shape [tile_size] + , tilingNumWholeTiles = num_whole_tiles + , tilingLevel = lvl + , tilingSpace = space + } + +invariantToOneOfTwoInnerDims :: Names -> M.Map VName Names -> [VName] -> VName + -> Maybe [Int] +invariantToOneOfTwoInnerDims branch_variant variance dims arr = do + j : i : _ <- Just $ reverse dims + let variant_to = M.findWithDefault mempty arr variance + branch_invariant = not $ nameIn j branch_variant || nameIn i branch_variant + if branch_invariant && i `nameIn` variant_to && not (j `nameIn` variant_to) then + Just [0,1] + else if branch_invariant && j `nameIn` variant_to && not (i `nameIn` variant_to) then + Just [1,0] + else + Nothing + +segMap2D :: String + -> SegLevel -> (SubExp, SubExp) + -> ((VName, VName) -> Binder Kernels [SubExp]) + -> Binder Kernels [VName] +segMap2D desc lvl (dim_x, dim_y) f = do + ltid_x <- newVName "ltid_x" + ltid_y <- newVName "ltid_y" + ltid_flat <- newVName "ltid_flat" + let space = SegSpace ltid_flat [(ltid_x, dim_x), (ltid_y, dim_y)] + + ((ts, res), stms) <- runBinder $ do + res <- f (ltid_x, ltid_y) + ts <- mapM subExpType res + return (ts, res) + Body _ stms' res' <- renameBody $ mkBody stms res + + letTupExp desc $ Op $ SegOp $ + SegMap lvl space ts $ KernelBody () stms' $ map Returns res' + + +-- Reconstruct the original gtids from group and local IDs. +reconstructGtids2D :: SubExp -> (VName, VName) -> (VName, VName) -> (VName, VName) + -> Binder Kernels () +reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y) = do + -- Reconstruct the original gtids from gid_x/gid_y and ltid_x/ltid_y. + letBindNames_ [gtid_x] =<< + toExp (LeafExp gid_x int32 * primExpFromSubExp int32 tile_size + + LeafExp ltid_x int32) + letBindNames_ [gtid_y] =<< + toExp (LeafExp gid_y int32 * primExpFromSubExp int32 tile_size + + LeafExp ltid_y int32) + +readTile2D :: (SubExp, SubExp) -> (VName, VName) -> (VName, VName) -> SubExp + -> Count NumGroups SubExp -> Count GroupSize SubExp + -> TileKind -> PrivStms -> SubExp + -> [(VName, [Int])] + -> Binder Kernels [VName] +readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size num_groups group_size kind privstms tile_id arrs_and_perms = + segMap2D "full_tile" (SegThread num_groups group_size SegNoVirt) (tile_size, tile_size) $ \(ltid_x, ltid_y) -> do + i <- letSubExp "i" =<< + toExp (primExpFromSubExp int32 tile_id * + primExpFromSubExp int32 tile_size + + LeafExp ltid_x int32) + j <- letSubExp "j" =<< + toExp (primExpFromSubExp int32 tile_id * + primExpFromSubExp int32 tile_size + + LeafExp ltid_y int32) + + reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y) + addPrivStms [DimFix $ Var ltid_x, DimFix $ Var ltid_y] privstms + + let (arrs, perms) = unzip arrs_and_perms + arr_ts <- mapM lookupType arrs + let tile_ts = map rowType arr_ts + w = arraysSize 0 arr_ts + + let readTileElem arr perm = + -- No need for fullSlice because we are tiling only prims. + letExp "tile_elem" $ BasicOp $ Index arr + [DimFix $ last $ rearrangeShape perm [i,j]] + readTileElemIfInBounds (tile_t, arr, perm) = do + let idx = last $ rearrangeShape perm [i,j] + othercheck = last $ rearrangeShape perm + [ LeafExp gtid_y int32 .<. primExpFromSubExp int32 kdim_y + , LeafExp gtid_x int32 .<. primExpFromSubExp int32 kdim_x + ] + eIf (toExp $ + primExpFromSubExp int32 idx .<. primExpFromSubExp int32 w .&&. othercheck) + (eBody [return $ BasicOp $ Index arr [DimFix idx]]) + (eBody [eBlank tile_t]) + + fmap (map Var) $ + case kind of + TilePartial -> + mapM (letExp "pre" <=< readTileElemIfInBounds) (zip3 tile_ts arrs perms) + TileFull -> + zipWithM readTileElem arrs perms + +processTile2D :: (VName, VName) -> (VName, VName) -> (SubExp, SubExp) -> SubExp + -> Count NumGroups SubExp -> Count GroupSize SubExp + -> PrivStms + -> Commutativity -> Lambda Kernels -> Lambda Kernels + -> [(VName,[Int])] -> [VName] + -> Binder Kernels [VName] +processTile2D + (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size num_groups group_size + privstms red_comm red_lam map_lam tiles_and_perms accs = do + + -- Might be truncated in case of a partial tile. + actual_tile_size <- arraysSize 0 <$> mapM (lookupType . fst) tiles_and_perms + + segMap2D "acc" (SegThreadScalar num_groups group_size SegNoVirt) (tile_size, tile_size) $ \(ltid_x, ltid_y) -> do + reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y) + + addPrivStms [DimFix $ Var ltid_x, DimFix $ Var ltid_y] privstms + + -- We replace the neutral elements with the accumulators (this is + -- OK because the parallel semantics are not used after this + -- point). + thread_accs <- forM accs $ \acc -> + letSubExp "acc" $ BasicOp $ Index acc [DimFix $ Var ltid_x, DimFix $ Var ltid_y] + let form' = redomapSOAC [Reduce red_comm red_lam thread_accs] map_lam + + tiles' <- forM tiles_and_perms $ \(tile, perm) -> do + tile_t <- lookupType tile + letExp "tile" $ BasicOp $ Index tile $ sliceAt tile_t (head perm) + [DimFix $ Var $ head $ rearrangeShape perm [ltid_x, ltid_y]] + + fmap (map Var) $ + letTupExp "acc" =<< eIf (toExp $ + LeafExp gtid_x int32 .<. primExpFromSubExp int32 kdim_x .&&. + LeafExp gtid_y int32 .<. primExpFromSubExp int32 kdim_y) + (eBody [pure $ Op $ OtherOp $ Screma actual_tile_size form' tiles']) + (resultBodyM thread_accs) + +processResidualTile2D :: (VName, VName) -> (VName, VName) -> (SubExp, SubExp) -> SubExp + -> Count NumGroups SubExp -> Count GroupSize SubExp -> PrivStms + -> Commutativity -> Lambda Kernels -> Lambda Kernels + -> SubExp -> [VName] -> SubExp -> [(VName, [Int])] + -> Binder Kernels [VName] +processResidualTile2D + gids gtids kdims tile_size num_groups group_size privstms red_comm red_lam map_lam + num_whole_tiles accs w arrs_and_perms = do + -- The number of residual elements that are not covered by + -- the whole tiles. + residual_input <- letSubExp "residual_input" $ + BasicOp $ BinOp (SRem Int32) w tile_size + + letTupExp "acc_after_residual" =<< + eIf (toExp $ primExpFromSubExp int32 residual_input .==. 0) + (resultBodyM $ map Var accs) + (nonemptyTile residual_input) + + where + nonemptyTile residual_input = renameBody <=< runBodyBinder $ do + -- Collectively construct a tile. Threads that are out-of-bounds + -- provide a blank dummy value. + full_tile <- readTile2D kdims gtids gids tile_size num_groups group_size + TilePartial privstms num_whole_tiles arrs_and_perms + + tile <- forM full_tile $ \tile -> + letExp "partial_tile" $ BasicOp $ Index tile + [DimSlice (intConst Int32 0) residual_input (intConst Int32 1), + DimSlice (intConst Int32 0) residual_input (intConst Int32 1)] + + -- Now each thread performs a traversal of the tile and + -- updates its accumulator. + resultBody . map Var <$> + processTile2D gids gtids kdims tile_size num_groups group_size + privstms red_comm red_lam map_lam + (zip tile (map snd arrs_and_perms)) accs + +tiling2d :: [(VName,SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp) +tiling2d dims_on_top _initial_lvl (gtid_x, gtid_y) (kdim_x, kdim_y) w = do + gid_x <- newVName "gid_x" + gid_y <- newVName "gid_y" + + tile_size_key <- nameFromString . pretty <$> newVName "tile_size" + tile_size <- letSubExp "tile_size" $ Op $ GetSize tile_size_key SizeTile + group_size <- letSubExp "group_size" $ BasicOp $ BinOp (Mul Int32) tile_size tile_size + + num_groups_x <- letSubExp "num_groups_x" =<< + eDivRoundingUp Int32 (eSubExp kdim_x) (eSubExp tile_size) + num_groups_y <- letSubExp "num_groups_y" =<< + eDivRoundingUp Int32 (eSubExp kdim_y) (eSubExp tile_size) + + num_groups <- letSubExp "num_groups_top" =<< + foldBinOp (Mul Int32) num_groups_x + (num_groups_y : map snd dims_on_top) + + gid_flat <- newVName "gid_flat" + let lvl = SegGroup (Count num_groups) (Count group_size) SegNoVirt + space = SegSpace gid_flat $ + dims_on_top ++ [(gid_x, num_groups_x), (gid_y, num_groups_y)] + + -- Number of whole tiles that fit in the input. + num_whole_tiles <- letSubExp "num_whole_tiles" $ + BasicOp $ BinOp (SQuot Int32) w tile_size + return Tiling + { tilingSegMap = \desc lvl' f -> segMap2D desc lvl' (tile_size, tile_size) $ \(ltid_x, ltid_y) -> do + reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y) + f (LeafExp gtid_x int32 .<. primExpFromSubExp int32 kdim_x .&&. + LeafExp gtid_y int32 .<. primExpFromSubExp int32 kdim_y) + [DimFix $ Var ltid_x, DimFix $ Var ltid_y] + + , tilingReadTile = readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size (segNumGroups lvl) (segGroupSize lvl) + , tilingProcessTile = processTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size (segNumGroups lvl) (segGroupSize lvl) + , tilingProcessResidualTile = processResidualTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size (segNumGroups lvl) (segGroupSize lvl) + + , tilingTileReturns = tileReturns dims_on_top [(kdim_x, tile_size), (kdim_y, tile_size)] + + , tilingTileShape = Shape [tile_size, tile_size] + , tilingNumWholeTiles = num_whole_tiles + , tilingLevel = lvl + , tilingSpace = space + } + +-- | The variance table keeps a mapping from a variable name +-- (something produced by a 'Stm') to the kernel thread indices +-- that name depends on. If a variable is not present in this table, +-- that means it is bound outside the kernel (and so can be considered +-- invariant to all dimensions). +type VarianceTable = M.Map VName Names + +varianceInStms :: VarianceTable -> Stms Kernels -> VarianceTable +varianceInStms = foldl varianceInStm + +varianceInStm :: VarianceTable -> Stm Kernels -> VarianceTable +varianceInStm variance bnd = + foldl' add variance $ patternNames $ stmPattern bnd + where add variance' v = M.insert v binding_variance variance' + look variance' v = oneName v <> M.findWithDefault mempty v variance' + binding_variance = mconcat $ map (look variance) $ namesToList (freeIn bnd) diff --git a/src/Futhark/Optimise/TileLoops/RegTiling3D.hs b/src/Futhark/Optimise/TileLoops/RegTiling3D.hs deleted file mode 100644 --- a/src/Futhark/Optimise/TileLoops/RegTiling3D.hs +++ /dev/null @@ -1,736 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} --- | Perform a restricted form of register tiling corresponding to --- the following pattern: --- * a stream is perfectly nested inside a kernel with at least --- three parallel dimension (the perfectly nested restriction --- can be relaxed a bit); --- * all streamed arrays are one dimensional; --- * all streamed arrays are variant to exacly one of the three --- innermost parallel dimensions, and conversly for each of --- the three innermost parallel dimensions, there is at least --- one streamed array variant to it; --- * the stream's result is a tuple of scalar values, which are --- also the "thread-in-space" return of the kernel. --- Target code can be found in "tests/reg-tiling/reg-tiling-3d.fut". -module Futhark.Optimise.TileLoops.RegTiling3D - ( doRegTiling3D ) - where - -import Control.Monad.State -import Control.Monad.Reader -import qualified Data.Set as S -import qualified Data.Map.Strict as M -import Data.List -import Data.Maybe - -import Futhark.MonadFreshNames -import Futhark.Representation.Kernels -import Futhark.Tools -import Futhark.Transform.Substitute -import Futhark.Transform.Rename - -type TileM = ReaderT (Scope Kernels) (State VNameSource) -type VarianceTable = M.Map VName Names - -maxRegTile :: Int32 -maxRegTile = 30 - -mkRegTileSe :: Int32 -> SubExp -mkRegTileSe = constant - --- | Expects a kernel statement as argument. --- CONDITIONS for 3D tiling optimization to fire are: --- 1. a) The kernel body can be broken into --- scalar-code-1 ++ [GroupStream stmt] ++ scalar-code-2. --- b) The kernels has a "ThreadsReturn ThreadsInSpace" result, --- and obviously the result is variant to the 3rd dimension --- (counter from innermost to outermost) --- 2. For the GroupStream (morally StreamSeq): --- a) the arrays' outersize must equal the maximal chunk size --- b) the streamed arrays are one dimensional --- c) each of the array arguments of GroupStream are variant --- to exactly one of the three innermost-parallel dimension --- of the kernel. This condition can be relaxed by interchanging --- kernel dimensions whenever possible. --- 3. For scalar-code-1: --- a) each of the statements is a slice that produces one of the --- streamed arrays --- 4. For simplicity assume scalar-code-2 is empty! --- (To be extended later.) --- ASSUME the initial kernel is (as in tests/reg-tiling/reg-tiling-3d.fut): --- --- kernel map(num groups: num_groups, group size: group_size, --- num threads: num_threads, global TID -> global_tid, --- local TID -> local_tid, group ID -> group_id) --- (gtid_z < size_z, gtid_y < size_xy, --- gtid_x < size_xy) : {f32} { --- let {[size_com]f32 flags} = fss_6664[gtid_z, --- 0i32:+size_com*1i32] --- let {[size_com]f32 ass} = ass_6662[gtid_y, 0i32:+size_com*1i32] --- let {[size_com]f32 bss} = res_6687[gtid_x, 0i32:+size_com*1i32] --- let {f32 res_ker} = --- stream(size_com, size_com, --- fn (int chunk_size_out, int chunk_offset_6736, f32 acc_out, --- [chunk_size_out]f32 flags_chunk_out, --- [chunk_size_out]f32 ass_chunk_out, --- [chunk_size_out]f32 bss_chunk_out) => --- let {f32 res_out} = --- stream(chunk_size_out, 1i32, --- fn (int chunk_size_in, int i_6743, f32 acc_in, --- [chunk_size_in]f32 flags_chunk_in, --- [chunk_size_in]f32 ass_chunk_in, --- [chunk_size_in]f32 bss_chunk_in) => --- let {f32 f} = flags_chunk_in[0i32] --- let {f32 a} = ass_chunk_in[0i32] --- let {f32 b} = bss_chunk_in[0i32] --- let {bool cond} = lt32(f, 9.0f32) --- let {f32 tmp} = --- if cond --- then { --- let {f32 tmp1} = fmul32(a, b) --- in {tmp1} --- } else {0.0f32} --- let {f32 res_in} = fadd32(acc_in, tmp) --- in {res_in}, --- {acc_out}, --- flags_chunk_out, ass_chunk_out, bss_chunk_out) --- in {res_out}, --- {0.0f32}, --- flags, ass, bss) --- return {thread in space returns res_ker} --- } --- -doRegTiling3D :: Stm Kernels -> TileM (Maybe (Stms Kernels, Stm Kernels)) -doRegTiling3D (Let pat aux (Op (HostOp old_kernel))) - | Kernel kerhint space kertp (KernelBody () kstms kres) <- old_kernel, - FlatThreadSpace gspace <- spaceStructure space, - initial_variance <- M.map mempty $ scopeOfKernelSpace space, - variance <- varianceInStms initial_variance kstms, - local_tid <- spaceLocalId space, - (_,_) : (_,_) : (gidz,m_M) : _ <- reverse $ spaceDimensions space, - (code1, Just stream_stmt, code2) <- matchCodeStreamCode kstms, - Let pat_strm aux_strm (Op (GroupStream w w0 lam accs arrs)) <- stream_stmt, - not (null accs), - reg_tile <- maxRegTile `quot` fromIntegral (length accs), - reg_tile_se <- mkRegTileSe reg_tile, - w == w0, - arr_chunk_params <- groupStreamArrParams lam, - Just _ <- is3dTileable mempty space variance - arrs arr_chunk_params, - Just arr_tab0 <- foldl (processIndirections $ S.fromList arrs) - (Just M.empty) code1, - -- for simplicity, assume a single result, which is variant to - -- the outer parallel dimension (for sanity sake, it should be) - ker_res_nms <- mapMaybe retThreadInSpace kres, - length ker_res_nms == length kres, - Pattern [] ker_patels <- pat, - all primType kertp, - all (variantToOuterDim variance gidz) ker_res_nms = do - mm <- newVName "mm" - mask <- newVName "mask" - - -- let mm = gidz * regTile - let mm_stmt = mkInKerIntMulStmt mm (Var gidz) reg_tile_se - let mask_stm= mkLet [] [Ident mask $ Prim int32] $ BasicOp $ - BinOp (Shl Int32) - (Constant $ IntValue $ Int32Value 1 ) - (Constant $ IntValue $ Int32Value 31) - - -- process the z-variant arrays that need transposition; - -- these "manifest" statements should come before the kernel - (arr_tab,trnsp_tab) <- foldM (insertTranspose variance gidz) - (M.empty, M.empty) $ M.toList arr_tab0 - let manif_stms = map(\ (a_t, (a,i,tp)) -> - let perm = [i+1..arrayRank tp-1] ++ [0..i] - in mkLet [] [Ident a_t tp] $ - BasicOp $ Manifest perm a - ) $ M.toList trnsp_tab - - -- adjust the kernel space for 3d register tiling. - (space_stms, space_struct, tiled_group_size, num_threads, num_groups) <- - mkKerSpaceExtraStms reg_tile gspace - let kspace' = space { spaceStructure = space_struct - , spaceGroupSize = tiled_group_size - , spaceNumThreads = num_threads - , spaceNumGroups = num_groups - } - - -- most everything happans here! - mb_myloop <- translateStreamsToLoop (reg_tile,mask,gidz,m_M,mm,local_tid,tiled_group_size) - variance arr_tab w lam accs arrs $ - patternValueElements pat_strm - - -- ToDo: adjust the new kernel with - -- 1. in-place update return: for this you will need to `scratch` - -- the result array before the kernel - -- 2. adjust the range of gidz to `(m_M + TILE_REG -1)/ TILE_REG` - -- 3. transpose the array invariant to the third-inner dim - case mb_myloop of - Nothing -> return Nothing - Just (myloop, strm_res_inv, strm_res_var) -> do - -- make loop statement - loop_var_res <- forM strm_res_var $ \(PatElem nm attr) -> do - clone_patel_nms <- replicateM (fromIntegral reg_tile) $ newVName $ baseString nm - return $ map (`PatElem` attr) clone_patel_nms - let pat_loop = Pattern [] $ strm_res_inv ++ concat loop_var_res - let stm_loop = Let pat_loop aux_strm myloop - - -- get variant ker-results and corresponding pattern elements - let ker_var_res_patels = - filter (\(r,_) -> variantToOuterDim variance gidz r) $ - zip ker_res_nms ker_patels - (ker_var_res, ker_var_patels) = unzip ker_var_res_patels - (code2_var, code2_inv) = - partition (variantToOuterDim variance gidz . patElemName . - head . patternValueElements . stmPattern) code2 - - -- make the scratch statements for kernel results variant to the z-parallel dimension - scratch_nms_stms <- mapM mkScratchStm ker_var_patels - let (scratch_nms, scratch_stms) = unzip scratch_nms_stms - loop_var_nms_tr = transpose $ map (map patElemName) loop_var_res - - -- clone the statements in code2 variant to the z-parallel dimension, - -- by encapsulating them inside if-then-else in which the then-body - -- terminates with an in-place update corresponding to the result! - strm_var_nms = map patElemName strm_res_var - (ip_out_nms, unrolled_code) <- - foldM (cloneVarCode2 mm space strm_var_nms ker_var_res_patels code2_var) - (scratch_nms, []) $ zip [0..reg_tile-1] loop_var_nms_tr - - -- replace the `ThreadsInSpace` kernel return to an `InPlace` return - -- for the z-variant kernel results - let ker_res_ip_tp_tab = M.fromList $ zip ker_var_res ip_out_nms - (kres', kertp') = unzip $ - zipWith (\r tp -> case M.lookup r ker_res_ip_tp_tab of - Nothing -> (ThreadsReturn (Var r), tp) - Just (dims, arr, ivs) -> (WriteReturn dims arr ivs, tp) - ) ker_res_nms kertp - - -- finally, put everything together - kstms' = stmsFromList $ mask_stm : mm_stmt : stm_loop : (code2_inv ++ unrolled_code) - ker_body = KernelBody () kstms' kres' - new_ker = Op $ HostOp $ Kernel kerhint kspace' kertp' ker_body - extra_stms = space_stms <> stmsFromList (scratch_stms ++ manif_stms) - return $ Just (extra_stms, Let pat aux new_ker) - - where -- | Checks that the statement is a slice that produces one of the - -- streamed arrays. Also that the streamed array is one dimensional. - -- Accumulates the information in a table for later use. - processIndirections :: S.Set VName - -> Maybe (M.Map VName (VName, Slice SubExp, Type)) - -> Stm InKernel - -> Maybe (M.Map VName (VName, Slice SubExp, Type)) - processIndirections arrs acc (Let patt _ (BasicOp (Index arr_nm slc))) = - case (acc, patternValueElements patt) of - (Nothing, _) -> Nothing - (Just tab, [p]) -> do - let (p_nm, p_tp) = (patElemName p, patElemType p) - case (S.member p_nm arrs, p_tp) of - (True, Array _ (Shape [_]) _) -> - Just $ M.insert p_nm (arr_nm,slc,p_tp) tab - _ -> Nothing - (_, _) -> Nothing - processIndirections _ _ _ = Nothing - - -- | The second Map accumulator keeps tracks of the arrays that - -- are variant to the z-parallel dimension and need to be transposed; - -- the `Int` field refers to the index of the z-variant dimension, and - -- the `Type` field refers to the type of the original global array. - -- The first accumulator table is updated to refer to the transposed-array - -- name, whenever such a case is discovered; otherwise it just accumulates. - insertTranspose :: VarianceTable -> VName - -> (M.Map VName (VName, Slice SubExp, Type), M.Map VName (VName,Int,Type)) - -> (VName, (VName, Slice SubExp, Type)) - -> TileM (M.Map VName (VName, Slice SubExp, Type), M.Map VName (VName,Int,Type)) - insertTranspose variance gidz (tab, trnsp) (p_nm, (arr_nm,slc,p_tp)) = - case findIndex (variantSliceDim variance gidz) slc of - Nothing -> return (M.insert p_nm (arr_nm,slc,p_tp) tab, trnsp) - Just i -> do - arr_tp <- lookupType arr_nm - arr_tr_nm <- newVName $ baseString arr_nm ++ "_transp" - let tab' = M.insert p_nm (arr_tr_nm,slc,p_tp) tab - let trnsp' = M.insert arr_tr_nm (arr_nm, i, arr_tp) trnsp - return (tab', trnsp') - - variantSliceDim :: VarianceTable -> VName -> DimIndex SubExp -> Bool - variantSliceDim variance gidz (DimFix (Var vnm)) = variantToOuterDim variance gidz vnm - variantSliceDim _ _ _ = False - - mkInKerIntMulStmt :: VName -> SubExp -> SubExp -> Stm InKernel - mkInKerIntMulStmt res_nm0 op1_se op2_se = - mkLet [] [Ident res_nm0 $ Prim int32] $ - BasicOp $ BinOp (Mul Int32) op1_se op2_se - - retThreadInSpace (ThreadsReturn (Var r)) = Just r - retThreadInSpace _ = Nothing - -doRegTiling3D _ = return Nothing - -translateStreamsToLoop :: (Int32,VName,VName,SubExp,VName,VName,SubExp) -> - VarianceTable -> - M.Map VName (VName, Slice SubExp, Type) -> - SubExp -> GroupStreamLambda InKernel -> - [SubExp] -> [VName] -> [PatElem InKernel] - -> TileM (Maybe (Exp InKernel, [PatElem InKernel], [PatElem InKernel])) -translateStreamsToLoop (reg_tile, mask,gidz,m_M,mm,local_tid, group_size) variance - arr_tab w_o lam_o accs_o_p arrs_o_p strm_ress - | -- 1. We assume the inner stream (of chunk 1) is directly nested - -- inside the outer stream and also takes its arguments (array - -- and accumulators) from the outer stream (all checked). - -- Also all accumulators have primitive types (otherwise - -- they cannot be efficiently stored in registers anyway). - accs_o_f <- groupStreamAccParams lam_o, - arrs_o_f <- groupStreamArrParams lam_o, - [Let _ _ (Op (GroupStream _ ct1i32 lam_i accs_i_p arrs_i_p))] <- - stmsToList $ bodyStms $ groupStreamLambdaBody lam_o, - ct1i32 == (Constant $ IntValue $ Int32Value 1), - accs_i_f <- groupStreamAccParams lam_i, - arrs_i_f <- groupStreamArrParams lam_i, - and $ zipWith (==) (map subExpVar accs_i_p) (map (Just . paramName) accs_o_f), - and $ zipWith (==) arrs_i_p $ map paramName arrs_o_f, - all (primType . paramType) accs_o_f, - -- 2. The intent is to flatten the two streams into a loop, so - -- we reuse the index of the inner stream for the result-loop index, - -- and we will modify the body of the inner lambda `body_i` for the - -- result loop. - loop_ind_nm <- groupStreamChunkOffset lam_i, - body_i <- groupStreamLambdaBody lam_i, - -- 3. We transfer the slicing information (from sclar-code-1) to - -- the array-formal arguments of the inner stream. - arr_tab' <- foldl (\ tab (a_o_p, a_o_f, a_i_p, a_i_f) -> - case (paramName a_o_f == a_i_p, M.lookup a_o_p tab) of - (True, Just info) -> M.insert (paramName a_i_f) info tab - _ -> tab - ) arr_tab $ zip4 arrs_o_p arrs_o_f arrs_i_p arrs_i_f, - -- 4. We translate the inner stream's accumulator to a FParam, required for - -- mapping it as a result-loop variant variable. - accs_i_f' <- map translParamToFParam accs_i_f, - -- 5. We break the "loop" statements into two parts: - -- a) the ones invariant to the z parallel dimension `invar_out_stms`, - -- b) the ones variant to the z parallel dimension `var_out_stms`, and - -- c) the ones corresponding to indexing operations on variant arrays `var_ind_stms`. - (invar_out_stms, var_ind_stms, var_out_stms) <- - foldl (\ (acc_inv, acc_inds, acc_var) stmt -> - let nm = patElemName $ head $ patternValueElements $ stmPattern stmt - in if not $ variantToOuterDim variance gidz nm - then (stmt : acc_inv,acc_inds,acc_var) - else case stmt of - Let _ _ (BasicOp (Index arr_nm [DimFix _])) -> - case M.lookup arr_nm arr_tab' of - Just _ -> (acc_inv,stmt:acc_inds,acc_var) - Nothing -> (acc_inv,acc_inds,stmt:acc_var) - _ -> (acc_inv,acc_inds,stmt:acc_var) - ) ([],[],[]) $ reverse $ stmsToList $ bodyStms body_i, - -- 6. We check that the variables used in the index statements referring to - -- streamed arrays that are variant to the z parallel dimension (`var_ind_stms`) - -- depend only on variables defined in the invariant stms to the z parallel dimension. - var_nms <- concatMap (patternNames . stmPattern) var_out_stms, - null $ S.intersection (S.fromList var_nms) $ - S.unions (map freeIn var_ind_stms), - -- 7. We assume (check) for simplicity that all accumulator initializers - -- of the outer stream are invariant to the z parallel dimension. - loop_ini_vs <- subExpVars accs_o_p, - all (not . variantToOuterDim variance gidz) loop_ini_vs, - -- 8. We assume that all results of the inner-stream body are variables - -- (for simplicity); they should have been simplified anyways if not! - loop_res0 <- bodyResult body_i, - loop_res <- subExpVars loop_res0, - length loop_res == length loop_res0 = do - -- I. After all these conditions, we finally start by partitioning - -- the stream's accumulators and results into the ones that are - -- variant to the z-parallel dimension and the ones that are not. - let (loop_var_p_i_r, loop_inv_p_i_r) = - partition (\(_,_,r,_) -> variantToOuterDim variance gidz r) $ - zip4 accs_i_f' accs_o_p loop_res strm_ress - -- II. Transform the statements invariant to the z-parallel dimension - -- so that they perform indexing in the global arrays rather than - -- in the streamed arrays, i.e., eliminate the indirection. - inv_stms0 <- mapM (transfInvIndStm arr_tab' loop_ind_nm) invar_out_stms - let inv_stms = concat inv_stms0 - -- III. the index-statements variant to the z-parallel dimension are - -- transformed to combined regions. - m <- newVName "m" - ind_stms0 <- foldM (transfVarIndStm arr_tab' (reg_tile,loop_ind_nm,local_tid,group_size,m,m_M)) - (Just ([],M.empty)) $ reverse var_ind_stms - case ind_stms0 of - Nothing -> return Nothing - Just (ind_stms, subst_tab) -> do - -- IV. Add statement `let m = mm + local_tid` - -- Then perform the substitution `gidz -> m` on the combine regions. - let m_stmt = mkLet [] [Ident m $ Prim int32] $ - BasicOp $ BinOp (Add Int32) (Var mm) (Var local_tid) - tab_z_m_comb = M.insert gidz m M.empty - ind_stms' = m_stmt : map (substituteNames tab_z_m_comb) ind_stms - - -- V. We clone the variant statements regTile times and enclose - -- each one in a if-then-else testing whether `mm + local_id < m_M` - -- TODO: check that the statements do not involve In-Place updates! - let loop_var_p_i_r' = map (\(x,y,z,_)->(x,y,z)) loop_var_p_i_r - if_ress <- mapM (cloneVarStms subst_tab (mask,loop_ind_nm,mm,m_M,gidz) - loop_var_p_i_r' var_out_stms) [0..reg_tile-1] - -- VI. build the loop-variant vars/res/inis - let (if_stmt_clones0, var_ress_pars) = unzip if_ress - if_stmt_clones = concat if_stmt_clones0 - (_, var_ini, _, strm_var_res) = unzip4 loop_var_p_i_r - var_inis = concat $ replicate (fromIntegral reg_tile) var_ini - (var_ress, var_pars) = unzip $ concat var_ress_pars - (inv_pars, inv_inis, inv_ress, strm_inv_res) = unzip4 loop_inv_p_i_r - loop_form_acc = inv_pars ++ var_pars - loop_inis_acc = inv_inis ++ var_inis - loop_ress = inv_ress ++ var_ress - -- VII. Finally build the loop body and return it! - -- Insert an extra barrier at the begining of the loop; make - -- it dependent on the loop index so it cannot be hoisted! - ind_bar <- newVName "loop_ind" - let bar_stmt = mkLet [] [Ident loop_ind_nm $ Prim int32] $ Op (Barrier [Var ind_bar]) - stms_body_i' = bar_stmt : inv_stms ++ ind_stms' ++ if_stmt_clones - form = ForLoop ind_bar Int32 w_o [] - body_i' = Body (bodyAttr body_i) - (stmsFromList stms_body_i') $ - map Var loop_ress - myloop = DoLoop [] (zip loop_form_acc loop_inis_acc) form body_i' - free_in_body = freeIn body_i' - elim_vars = S.fromList $ arrs_i_p ++ arrs_o_p ++ - map paramName arrs_i_f ++ - map paramName accs_o_f - if null $ S.intersection free_in_body elim_vars - then return $ Just (myloop, strm_inv_res, strm_var_res) - else return Nothing -translateStreamsToLoop _ _ _ _ _ _ _ _ = return Nothing - --- | Clone the variant statements, by creating a new if-then-else --- big statement that cheks that `mm + i < m_M` for `i = 0...regTile-1` --- Return the if-then-else statement together with the result variables --- so that the body of the loop and the loop results and paramters can --- be constructed. --- In order to disallow hoisting from the loop we will generate: --- let zero = mask & loop_ind --- let mmpi = zero + mm + i -cloneVarStms :: M.Map VName (VName,Type) -> (VName, VName, VName, SubExp, VName) - -> [(FParam InKernel, SubExp, VName)] -> [Stm InKernel] - -> Int32 -> TileM ([Stm InKernel], [(VName,FParam InKernel)]) -cloneVarStms subst_tab (mask,loop_ind,mm,m_M,gidz) loop_info var_out_stms i = do - let (loop_par_origs, loop_inis, body_res_origs) = unzip3 loop_info - body_res_clones <- mapM (\x -> newVName $ baseString x ++ "_clone") body_res_origs - loop_par_nm_clones <- mapM (\x -> newVName $ baseString (paramName x) ++ "_clone") loop_par_origs - m <- newVName "m" - z <- newVName "zero" - ii<- newVName "unroll_ct" - let loop_par_clones = zipWith (\ p nm -> p { paramName = nm }) - loop_par_origs loop_par_nm_clones - res_types = map paramType loop_par_origs - i_se = Constant $ IntValue $ Int32Value i - - stmt_zero = mkLet [] [Ident z $ Prim int32] $ - BasicOp $ BinOp (And Int32) (Var mask) (Var loop_ind) - stmt_ii = mkLet [] [Ident ii $ Prim int32] $ - BasicOp $ BinOp (Add Int32) (Var z) i_se - m_stmt_other = - mkLet [] [Ident m $ Prim int32] $ - BasicOp $ BinOp (Add Int32) (Var mm) (Var ii) - read_sh_stms = - map (\ (scal,(sh_arr, el_tp)) -> - mkLet [] [Ident scal el_tp] $ - BasicOp $ Index sh_arr [DimFix i_se] - ) $ M.toList subst_tab - tab_z_m_other = foldl (\tab (old,new) -> M.insert (paramName old) new tab) - (M.insert gidz m M.empty) $ - zip loop_par_origs loop_par_nm_clones - var_out_stms' = map (substituteNames tab_z_m_other) $ - read_sh_stms ++ var_out_stms - cond_nm <- newVName "out3_inbounds" - -- if the statements are simple, i.e., "safe", then do not - -- encapsulate them in an if-then-else; this will result in - -- significant performance gains. - let simple = all simpleStm var_out_stms - let cond_stm = if simple - then mkLet [] [Ident cond_nm $ Prim Bool] $ - BasicOp $ SubExp (Constant $ BoolValue True) - else mkCondStmt m_M m cond_nm - -- TODO: we need to uniquely rename the then/else bodies! - then_body <- renameBody $ Body () (stmsFromList var_out_stms') (map Var body_res_origs) - let else_body = Body () mempty loop_inis - if_stmt = mkLet [] (zipWith Ident body_res_clones res_types) $ - If (Var cond_nm) then_body else_body $ - IfAttr (staticShapes res_types) IfFallback - -- we will substitute later the original loop formal-param names - -- with the newly created ones in the body - return ( [stmt_zero, stmt_ii, m_stmt_other, cond_stm, if_stmt] - , zip body_res_clones loop_par_clones ) - -mkCondStmt :: SubExp -> VName -> VName -> Stm InKernel -mkCondStmt m_M m cond_nm = - mkLet [] [Ident cond_nm $ Prim Bool] $ - BasicOp $ CmpOp (CmpSlt Int32) (Var m) m_M - -simpleStm :: Stm InKernel -> Bool -simpleStm (Let _ _ e) = safeExp e - -mkScratchStm :: PatElem Kernels -> TileM (([SubExp], VName, [([SubExp], SubExp)]), - Stm Kernels) -mkScratchStm ker_patel = do - let (unique_arr_tp, res_arr_nm0) = (patElemType ker_patel, patElemName ker_patel) - ptp = elemType unique_arr_tp - scrtch_arr_nm <- newVName $ baseString res_arr_nm0 ++ "_0" - let scratch_stm = mkLet [] [Ident scrtch_arr_nm unique_arr_tp] $ - BasicOp $ Scratch ptp $ arrayDims unique_arr_tp - return ((arrayDims unique_arr_tp, scrtch_arr_nm, []), scratch_stm) - --- | Arguments are: --- 1. @mm@ this is the length of z-parallel dimension divided by reg_tile --- 2. @space@: the kernel space --- 3. @strm_res_nms@: the z-variant results of the original stream --- 4. @keres_patels@: the kernel result names tupled with the corresponding --- pattern elements of the kernel statement. --- 5. @code2_var@: the z-variant statements of the code after the stream. --- 6. @ip_writes@: the "current" argument to a 'WriteReturn'. --- @unroll_code@: the current unrolled code. Both form a `foldM` accumulator. --- 7. @k@ the "current" clone number; --- @loop_res_nms@ the names of the loop result corresponding to the current clone. --- Result: --- 1. the argument for the current in-place update result, --- 2. a new if-statement is added to the unrolled-code accumulator which actually --- perform the in-place update. -cloneVarCode2 :: VName -> KernelSpace -> [VName] - -> [(VName, PatElem InKernel)] -> [Stm InKernel] - -> ([([SubExp], VName, [([SubExp], SubExp)])], - [Stm InKernel]) - -> (Int32, [VName]) - -> TileM ([([SubExp], VName, [([SubExp], SubExp)])], - [Stm InKernel]) -cloneVarCode2 mm space strm_res_nms keres_patels code2_var - (writes, unroll_code) (k, loop_res_nms) = do - let (ker_nms, pat_els) = unzip keres_patels - root_strs = map (baseString . patElemName) pat_els - ip_out_nms <- mapM (\s -> newVName $ s ++ "_out_" ++ pretty (k+1)) root_strs - m <- newVName "m" - -- make in-place update statements - let (gidx,_) : (gidy,_) : (gidz,m_M) : rev_outer_dims = reverse $ spaceDimensions space - (outer_dims, _) = unzip $ reverse rev_outer_dims - strip_dims = length $ outer_dims++[m,gidy,gidx] - ts = map (stripArray strip_dims . patElemType) pat_els - -- make if - cond_nm <- newVName "m_cond" - let i_se = Constant $ IntValue $ Int32Value k - m_stm = mkLet [] [Ident m $ Prim int32] $ - BasicOp $ BinOp (Add Int32) (Var mm) i_se - c_stm = mkCondStmt m_M m cond_nm - strm_loop_tab = M.fromList $ (gidz, m) : zip strm_res_nms loop_res_nms - then_body <- renameBody $ substituteNames strm_loop_tab $ - Body () (stmsFromList code2_var) $ map Var ker_nms - let else_body = Body () mempty $ map blank ts - if_stm = mkLet [] (zipWith Ident ip_out_nms ts) $ - If (Var cond_nm) then_body else_body $ - IfAttr (staticShapes ts) IfFallback - addWritePair (dims, arr, current) ker_nm = - (dims, arr, current ++ [(map Var $ outer_dims++[m,gidy,gidx], Var ker_nm)]) - return (zipWith addWritePair writes ip_out_nms, unroll_code ++ [m_stm, c_stm, if_stm]) - where blank (Prim t) = Constant $ blankPrimValue t - blank t = error $ "cloneVarCode2: cannot tile non-prim type " ++ pretty t - -helper3Stms :: VName -> SubExp -> SubExp -> Slice SubExp - -> VName -> Stm InKernel -> TileM [Stm InKernel] -helper3Stms loop_ind strd beg par_slc par_arr (Let ptt att _) = do - tmp1 <- newVName "tmp" - tmp2 <- newVName "ind" - let stmt1 = mkLet [] [Ident tmp1 $ Prim int32] $ - BasicOp $ BinOp (Mul Int32) (Var loop_ind) strd - stmt2 = mkLet [] [Ident tmp2 $ Prim int32] $ - BasicOp $ BinOp (Add Int32) beg (Var tmp1) - ndims = length par_slc - ind_exp = BasicOp (Index par_arr (take (ndims-1) par_slc ++ [DimFix $ Var tmp2])) - stmt3 = Let ptt att ind_exp - return [stmt1,stmt2,stmt3] - --- | Insert the necessary translations for a statement that is indexing --- in one of the streamed arrays, which is invariant to the z-parallel --- dimension. The index is necessarily `0` at this point, and we use `tab` --- to figure out to what global array does the streamed array actually --- refers to, and to compute the global index. -transfInvIndStm :: M.Map VName (VName, Slice SubExp, Type) - -> VName -> Stm InKernel - -> TileM [Stm InKernel] -transfInvIndStm tab loop_ind stm@(Let _ _ (BasicOp (Index arr_nm [DimFix _]))) - | Just (par_arr, par_slc@(_:_), _) <- M.lookup arr_nm tab, - DimSlice beg _ strd <- last par_slc = - helper3Stms loop_ind strd beg par_slc par_arr stm -transfInvIndStm _ _ stm = return [stm] - --- | Insert the necessary translations for a statement that is indexing --- inside one of the streamed arrays, which is variant to the outermost --- parallel dimension. -transfVarIndStm :: M.Map VName (VName, Slice SubExp, Type) - -> (Int32,VName,VName,SubExp,VName,SubExp) - -> Maybe ([Stm InKernel],M.Map VName (VName,Type)) - -> Stm InKernel - -> TileM (Maybe ([Stm InKernel],M.Map VName (VName,Type))) -transfVarIndStm tab (reg_tile,loop_ind,local_tid,group_size,m,m_M) acc - stm@(Let ptt _ (BasicOp (Index arr_nm [DimFix _]))) - | Just (tstms,stab) <- acc, - Just (par_arr, par_slc@(_:_), _) <- M.lookup arr_nm tab, - DimSlice beg _ strd <- last par_slc, - [pat_el] <- patternValueElements ptt, - el_tp <- patElemType pat_el, - pat_el_nm <- patElemName pat_el, - Prim _ <- el_tp = do - -- compute the index into the global array - stms3 <- helper3Stms loop_ind strd beg par_slc par_arr stm - let glb_ind_stms = stmsFromList stms3 - -- set up the combine part - sh_arr_1d <- newVName $ baseString par_arr ++ "_sh_1d" - cid <- newVName "cid" - let block_cspace = combineSpace [(cid,group_size)] - comb_exp = Op $ Combine block_cspace [el_tp] - [(local_tid, mkRegTileSe reg_tile), (m,m_M)] $ - Body () glb_ind_stms [Var pat_el_nm] - sh_arr_pe = PatElem sh_arr_1d $ - arrayOfShape el_tp $ Shape [group_size] - write_sh_arr_stmt = - Let (Pattern [] [sh_arr_pe]) (defAux ()) comb_exp - return $ Just (write_sh_arr_stmt:tstms, M.insert pat_el_nm (sh_arr_1d,el_tp) stab) -transfVarIndStm _ _ _ _ = return Nothing - --------------- ---- HELPES --- --------------- - --- | translates an LParam to an FParam -translParamToFParam :: LParam InKernel -> FParam InKernel -translParamToFParam = fmap (`toDecl` Nonunique) - --- | Tries to identified the following pattern: --- code folowed by a group stream followed by --- another code. -matchCodeStreamCode :: Stms InKernel -> - ([Stm InKernel], Maybe (Stm InKernel), [Stm InKernel]) -matchCodeStreamCode kstms = - foldl (\acc stmt -> - case (acc,stmt) of - ( (cd1,Nothing,cd2), Let _ _ (Op GroupStream{})) -> - (cd1, Just stmt, cd2) - ( (cd1, Nothing, cd2), _) -> (cd1++[stmt], Nothing, cd2) - ( (cd1,Just strm,cd2), _) -> (cd1,Just strm,cd2++[stmt]) - ) ([],Nothing,[]) (stmsToList kstms) - --- | Checks that all streamed arrays are variant to exacly one of --- the three innermost parallel dimensions, and conversly for --- each of the three innermost parallel dimensions, there is at --- least one streamed array variant to it. The result is the --- the number of the only variant parallel dimension for each array. -is3dTileable :: Names -> KernelSpace -> VarianceTable -> [VName] - -> [LParam InKernel] -> Maybe [Int] -is3dTileable branch_variant kspace variance arrs block_params = - let ok1 = all (primType . rowType . paramType) block_params - inner_perm0 = map variantOnlyToOneOfThreeInnerDims arrs - inner_perm = catMaybes inner_perm0 - ok2 = elem 0 inner_perm && elem 1 inner_perm && elem 2 inner_perm - ok3 = length inner_perm0 == length inner_perm - ok = ok1 && ok2 && ok3 - in if ok then Just inner_perm else Nothing - where variantOnlyToOneOfThreeInnerDims :: VName -> Maybe Int - variantOnlyToOneOfThreeInnerDims arr = do - (k,_) : (j,_) : (i,_) : _ <- Just $ reverse $ spaceDimensions kspace - let variant_to = M.findWithDefault mempty arr variance - branch_invariant = not $ S.member k branch_variant || - S.member j branch_variant || - S.member i branch_variant - if not branch_invariant - then Nothing - else if i `S.member` variant_to && not (j `S.member` variant_to) && not (k `S.member` variant_to) then Just 0 - else if not (i `S.member` variant_to) && j `S.member` variant_to && not (k `S.member` variant_to) then Just 1 - else if not (i `S.member` variant_to) && not (j `S.member` variant_to) && k `S.member` variant_to then Just 2 - else Nothing - -mkKerSpaceExtraStms :: Int32 -> [(VName, SubExp)] - -> TileM (Stms Kernels, SpaceStructure, SubExp, SubExp, SubExp) -mkKerSpaceExtraStms reg_tile gspace = do - dim_z_nm <- newVName "gidz_range" - tmp <- newVName "tmp" - let tmp_stm = mkLet [] [Ident tmp $ Prim int32] $ - BasicOp $ BinOp (Add Int32) m_M $ - Constant $ IntValue $ Int32Value (reg_tile-1) - rgz_stm = mkLet [] [Ident dim_z_nm $ Prim int32] $ - BasicOp $ BinOp (SQuot Int32) (Var tmp) $ - Constant $ IntValue $ Int32Value reg_tile - (gidx,sz_x) : (gidy,sz_y) : (gidz,m_M) : untiled_gspace = reverse gspace - - ((tile_size_x, tile_size_y, tiled_group_size), tile_size_bnds) <- runBinder $ do - tile_size_key <- nameFromString . pretty <$> newVName "tile_size" - tile_ct_size <- letSubExp "tile_size" $ Op $ GetSize tile_size_key SizeTile - tile_size_x <- letSubExp "tile_size_x" $ BasicOp $ - BinOp (SMin Int32) tile_ct_size sz_x - tile_size_y <- letSubExp "tile_size_y" $ BasicOp $ - BinOp (SMin Int32) tile_ct_size sz_y - tiled_group_size <- letSubExp "tiled_group_size" $ - BasicOp $ BinOp (Mul Int32) tile_size_x tile_size_y - return (tile_size_x, tile_size_y, tiled_group_size) - -- Play with reversion to ensure we get increasing IDs for - -- ltids. This affects readability of generated code. - untiled_gspace' <- fmap reverse $ forM (reverse untiled_gspace) $ \(gtid,gdim) -> do - ltid <- newVName "ltid" - return (gtid, gdim, ltid, constant (1::Int32)) - ltidz <- newVName "ltid" - let dim_z = (gidz, Var dim_z_nm, ltidz, constant (1::Int32)) - ltidy <- newVName "ltid" - let dim_y = (gidy, sz_y, ltidy, tile_size_y) - ltidx <- newVName "ltid" - let dim_x = (gidx, sz_x, ltidx, tile_size_x) - gspace' = reverse $ dim_x : dim_y : dim_z : untiled_gspace' - -- We have to recalculate number of workgroups and - -- number of threads to fit the new workgroup size. - ((num_threads, num_groups), num_bnds) <- - runBinder $ sufficientGroups gspace' tiled_group_size - - let extra_stms = oneStm tmp_stm <> oneStm rgz_stm <> tile_size_bnds <> num_bnds - return ( extra_stms, NestedThreadSpace gspace' - , tiled_group_size, num_threads, num_groups ) - - -variantToOuterDim :: VarianceTable -> VName -> VName -> Bool -variantToOuterDim variance gid_outer nm = - gid_outer == nm || gid_outer `S.member` M.findWithDefault mempty nm variance - -varianceInStms :: VarianceTable -> Stms InKernel -> VarianceTable -varianceInStms = foldl varianceInStm - -varianceInStm :: VarianceTable -> Stm InKernel -> VarianceTable -varianceInStm v0 bnd@(Let _ _ (Op (GroupStream _ _ lam accs arrs))) = - let v = defVarianceInStm v0 bnd - acc_lam_f = groupStreamAccParams lam - arr_lam_f = groupStreamArrParams lam - bdy_lam = groupStreamLambdaBody lam - stm_lam = bodyStms bdy_lam - - v' = foldl' (\vacc (v_a, v_f) -> - let vrc = S.insert v_a $ M.findWithDefault mempty v_a vacc - in M.insert v_f vrc vacc - ) v $ zip arrs $ map paramName arr_lam_f - v''= foldl' (\vacc (v_se, v_f) -> - case v_se of - Var v_a -> - let vrc = S.insert v_a $ M.findWithDefault mempty v_a vacc - in M.insert v_f vrc vacc - Constant _ -> vacc - ) v' $ zip accs $ map paramName acc_lam_f - in varianceInStms v'' stm_lam -varianceInStm variance bnd = defVarianceInStm variance bnd - -defVarianceInStm :: VarianceTable -> Stm InKernel -> VarianceTable -defVarianceInStm variance bnd = - foldl' add variance $ patternNames $ stmPattern bnd - where add variance' v = M.insert v binding_variance variance' - look variance' v = S.insert v $ M.findWithDefault mempty v variance' - binding_variance = mconcat $ map (look variance) $ S.toList (freeIn bnd) - -sufficientGroups :: MonadBinder m => - [(VName, SubExp, VName, SubExp)] -> SubExp - -> m (SubExp, SubExp) -sufficientGroups gspace group_size = do - groups_in_dims <- forM gspace $ \(_, gd, _, ld) -> - letSubExp "groups_in_dim" =<< eDivRoundingUp Int32 (eSubExp gd) (eSubExp ld) - num_groups <- letSubExp "num_groups" =<< - foldBinOp (Mul Int32) (constant (1::Int32)) groups_in_dims - num_threads <- letSubExp "num_threads" $ - BasicOp $ BinOp (Mul Int32) num_groups group_size - return (num_threads, num_groups) diff --git a/src/Futhark/Optimise/Unstream.hs b/src/Futhark/Optimise/Unstream.hs --- a/src/Futhark/Optimise/Unstream.hs +++ b/src/Futhark/Optimise/Unstream.hs @@ -1,27 +1,21 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} --- | Turn GroupStreams that operate on entire input or thread-variant --- sizes into do-loops, thus aiding subsequent optimisation. It is --- very important that this is run *after* any access-pattern-related --- optimisation, because this pass will destroy information. -module Futhark.Optimise.Unstream - ( unstream ) - where +-- | Sequentialise any remaining SOACs. It is very important that +-- this is run *after* any access-pattern-related optimisation, +-- because this pass will destroy information. +module Futhark.Optimise.Unstream (unstream) where import Control.Monad.State import Control.Monad.Reader -import qualified Data.Set as S -import qualified Data.Map as M -import Futhark.Representation.AST.Attributes.Aliases -import qualified Futhark.Analysis.Alias as Alias import Futhark.MonadFreshNames import Futhark.Representation.Kernels import Futhark.Pass import Futhark.Tools +import qualified Futhark.Transform.FirstOrderTransform as FOT unstream :: Pass Kernels Kernels -unstream = Pass "unstream" "Remove whole-array streams in kernels" $ +unstream = Pass "unstream" "sequentialise remaining SOACs" $ intraproceduralTransformation optimiseFunDef optimiseFunDef :: MonadFreshNames m => FunDef Kernels -> m (FunDef Kernels) @@ -38,57 +32,29 @@ localScope (scopeOf stms) $ Body () <$> (stmsFromList . concat <$> mapM optimiseStm (stmsToList stms)) <*> pure res +optimiseKernelBody :: KernelBody Kernels -> UnstreamM (KernelBody Kernels) +optimiseKernelBody (KernelBody () stms res) = + localScope (scopeOf stms) $ + KernelBody () <$> (stmsFromList . concat <$> mapM optimiseStm (stmsToList stms)) <*> pure res + +optimiseLambda :: Lambda Kernels -> UnstreamM (Lambda Kernels) +optimiseLambda lam = localScope (scopeOfLParams $ lambdaParams lam) $ do + body <- optimiseBody $ lambdaBody lam + return lam { lambdaBody = body} + optimiseStm :: Stm Kernels -> UnstreamM [Stm Kernels] -optimiseStm (Let pat aux (Op (HostOp op))) = do - inv <- S.fromList . M.keys <$> askScope - let mapper = identityKernelMapper { mapOnKernelKernelBody = onKernelBody } - onKernelBody kbody = do - stms' <- localScope (scopeOfKernelSpace (kernelSpace op)) $ - runBinder_ $ optimiseInKernelStms inv $ kernelBodyStms kbody - return kbody { kernelBodyStms = stms' } +optimiseStm (Let pat _ (Op (OtherOp soac))) = do + stms <- runBinder_ $ FOT.transformSOAC pat soac + fmap concat $ localScope (scopeOf stms) $ mapM optimiseStm $ stmsToList stms - op' <- mapKernelM mapper op - return [Let pat aux $ Op $ HostOp op'] +optimiseStm (Let pat aux (Op (SegOp op))) = + localScope (scopeOfSegSpace $ segSpace op) $ + pure <$> (Let pat aux . Op . SegOp <$> mapSegOpM optimise op) + where optimise = identitySegOpMapper { mapOnSegOpBody = optimiseKernelBody + , mapOnSegOpLambda = optimiseLambda + } optimiseStm (Let pat aux e) = pure <$> (Let pat aux <$> mapExpM optimise e) where optimise = identityMapper { mapOnBody = \scope -> localScope scope . optimiseBody } - -type Invariant = S.Set VName - -type InKernelM = Binder InKernel - -optimiseInKernelStms :: Invariant -> Stms InKernel -> InKernelM () -optimiseInKernelStms inv = mapM_ (optimiseInKernelStm inv) . stmsToList - -optimiseInKernelStm :: Invariant -> Stm InKernel -> InKernelM () -optimiseInKernelStm inv (Let pat aux (Op (GroupStream w max_chunk lam accs arrs))) - | max_chunk == w || maybe False (`S.notMember` inv) (subExpVar w) = do - let GroupStreamLambda chunk_size chunk_offset acc_params arr_params body = lam - letBindNames_ [chunk_size] $ BasicOp $ SubExp $ constant (1::Int32) - - loop_body <- insertStmsM $ do - forM_ (zip arr_params arrs) $ \(p,a) -> - letBindNames_ [paramName p] $ - BasicOp $ Index a $ fullSlice (paramType p) - [DimSlice (Var chunk_offset) (Var chunk_size) (constant (1::Int32))] - localScope (scopeOfLParams acc_params) $ optimiseInBody inv body - - -- Some accumulators may be updated in-place and must hence be unique. - let lam_consumed = consumedInBody $ Alias.analyseBody $ groupStreamLambdaBody lam - uniqueIfConsumed p | paramName p `S.member` lam_consumed = - fmap (`toDecl` Unique) p - | otherwise = fmap (`toDecl` Nonunique) p - merge = zip (map uniqueIfConsumed acc_params) accs - certifying (stmAuxCerts aux) $ - letBind_ pat $ DoLoop [] merge (ForLoop chunk_offset Int32 w []) loop_body -optimiseInKernelStm inv (Let pat aux e) = - addStm =<< (Let pat aux <$> mapExpM optimise e) - where optimise = identityMapper - { mapOnBody = \scope -> localScope scope . optimiseInBody inv } - -optimiseInBody :: Invariant -> Body InKernel -> InKernelM (Body InKernel) -optimiseInBody inv body = do - stms' <- collectStms_ $ optimiseInKernelStms inv $ bodyStms body - return body { bodyStms = stms' } diff --git a/src/Futhark/Pass/ExpandAllocations.hs b/src/Futhark/Pass/ExpandAllocations.hs --- a/src/Futhark/Pass/ExpandAllocations.hs +++ b/src/Futhark/Pass/ExpandAllocations.hs @@ -8,8 +8,8 @@ import Control.Monad.Except import Control.Monad.State import Control.Monad.Reader +import Control.Monad.Writer import qualified Data.Map.Strict as M -import qualified Data.Set as S import Data.Maybe import Data.List @@ -26,8 +26,9 @@ import qualified Futhark.Representation.Kernels as Kernels import Futhark.Representation.Kernels.Simplify as Kernels import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun -import Futhark.Pass.ExtractKernels.BlockedKernel (nonSegRed) +import Futhark.Pass.ExtractKernels.BlockedKernel (segThread, nonSegRed) import Futhark.Pass.ExplicitAllocations (explicitAllocationsInStms) +import Futhark.Transform.Rename (renameStm) import Futhark.Util.IntegralExp import Futhark.Util (mapAccumLM) @@ -64,7 +65,7 @@ where transform = identityMapper { mapOnBody = \scope -> localScope scope . transformBody } -nameInfoConv :: NameInfo ExplicitMemory -> NameInfo InKernel +nameInfoConv :: NameInfo ExplicitMemory -> NameInfo ExplicitMemory nameInfoConv (LetInfo mem_info) = LetInfo mem_info nameInfoConv (FParamInfo mem_info) = FParamInfo mem_info nameInfoConv (LParamInfo mem_info) = LParamInfo mem_info @@ -72,168 +73,178 @@ transformExp :: Exp ExplicitMemory -> ExpandM (Stms ExplicitMemory, Exp ExplicitMemory) -transformExp (Op (Inner (HostOp (Kernel desc kspace ts kbody)))) = do - let (kbody', allocs) = extractKernelBodyAllocations kbody - variantAlloc (Var v) = v `S.member` bound_in_kernel - variantAlloc _ = False - (variant_allocs, invariant_allocs) = M.partition (variantAlloc . fst) allocs - - (alloc_stms, alloc_offsets) <- - memoryRequirements kspace (kernelBodyStms kbody) variant_allocs invariant_allocs - - scope <- askScope - let scope' = scopeOfKernelSpace kspace <> M.map nameInfoConv scope - kbody'' <- either compilerLimitationS pure $ - runOffsetM scope' alloc_offsets $ offsetMemoryInKernelBody kbody' - - return (alloc_stms, - Op $ Inner $ HostOp $ Kernel desc kspace ts kbody'') - - where bound_in_kernel = - S.fromList $ M.keys $ scopeOfKernelSpace kspace <> - scopeOf (kernelBodyStms kbody) - -transformExp (Op (Inner (HostOp (SegMap kspace ts kbody)))) = do - (alloc_stms, (_, kbody')) <- transformScanRed kspace [] kbody +transformExp (Op (Inner (SegOp (SegMap lvl space ts kbody)))) = do + (alloc_stms, (_, kbody')) <- transformScanRed lvl space [] kbody return (alloc_stms, - Op $ Inner $ HostOp $ SegMap kspace ts kbody') + Op $ Inner $ SegOp $ SegMap lvl space ts kbody') -transformExp (Op (Inner (HostOp (SegRed kspace reds ts kbody)))) = do +transformExp (Op (Inner (SegOp (SegRed lvl space reds ts kbody)))) = do (alloc_stms, (lams, kbody')) <- - transformScanRed kspace (map segRedLambda reds) kbody + transformScanRed lvl space (map segRedLambda reds) kbody let reds' = zipWith (\red lam -> red { segRedLambda = lam }) reds lams return (alloc_stms, - Op $ Inner $ HostOp $ SegRed kspace reds' ts kbody') + Op $ Inner $ SegOp $ SegRed lvl space reds' ts kbody') -transformExp (Op (Inner (HostOp (SegScan kspace scan_op nes ts kbody)))) = do - (alloc_stms, (scan_op', kbody')) <- - transformScanRed kspace [scan_op] kbody +transformExp (Op (Inner (SegOp (SegScan lvl space scan_op nes ts kbody)))) = do + (alloc_stms, (scan_op', kbody')) <- transformScanRed lvl space [scan_op] kbody return (alloc_stms, - Op $ Inner $ HostOp $ SegScan kspace (head scan_op') nes ts kbody') + Op $ Inner $ SegOp $ SegScan lvl space (head scan_op') nes ts kbody') -transformExp (Op (Inner (HostOp (SegGenRed kspace ops ts kbody)))) = do - (alloc_stms, (lams, kbody')) <- - transformScanRed kspace (map genReduceOp ops) kbody - let ops' = zipWith (\red lam -> red { genReduceOp = lam }) ops lams +transformExp (Op (Inner (SegOp (SegGenRed lvl space ops ts kbody)))) = do + (alloc_stms, (lams', kbody')) <- transformScanRed lvl space lams kbody + let ops' = zipWith onOp ops lams' return (alloc_stms, - Op $ Inner $ HostOp $ SegGenRed kspace ops' ts kbody') + Op $ Inner $ SegOp $ SegGenRed lvl space ops' ts kbody') + where lams = map genReduceOp ops + onOp op lam = op { genReduceOp = lam } transformExp e = return (mempty, e) -transformScanRed :: KernelSpace - -> [Lambda InKernel] - -> KernelBody InKernel - -> ExpandM (Stms ExplicitMemory, ([Lambda InKernel], KernelBody InKernel)) -transformScanRed kspace ops kbody = do - let (kbody', kbody_allocs) = extractKernelBodyAllocations kbody - (ops', ops_allocs) = unzip $ map extractLambdaAllocations ops - variantAlloc (Var v) = v `S.member` bound_in_kernel +transformScanRed :: SegLevel -> SegSpace + -> [Lambda ExplicitMemory] + -> KernelBody ExplicitMemory + -> ExpandM (Stms ExplicitMemory, ([Lambda ExplicitMemory], KernelBody ExplicitMemory)) +transformScanRed lvl space ops kbody = do + bound_outside <- asks $ namesFromList . M.keys + let (kbody', kbody_allocs) = + extractKernelBodyAllocations (bound_outside<>bound_in_kernel) kbody + (ops', ops_allocs) = unzip $ map (extractLambdaAllocations bound_outside) ops + variantAlloc (Var v) = v `nameIn` bound_in_kernel variantAlloc _ = False allocs = kbody_allocs <> mconcat ops_allocs (variant_allocs, invariant_allocs) = M.partition (variantAlloc . fst) allocs - allocsForBody variant_allocs invariant_allocs kspace kbody' $ \alloc_stms kbody'' -> do + allocsForBody variant_allocs invariant_allocs lvl space kbody' $ \alloc_stms kbody'' -> do ops'' <- forM ops' $ \op' -> localScope (scopeOf op') $ offsetMemoryInLambda op' return (alloc_stms, (ops'', kbody'')) - where bound_in_kernel = - S.fromList $ map fst (spaceDimensions kspace) ++ - M.keys (scopeOfKernelSpace kspace <> - scopeOf (kernelBodyStms kbody)) + where bound_in_kernel = namesFromList $ M.keys $ scopeOfSegSpace space <> + scopeOf (kernelBodyStms kbody) allocsForBody :: M.Map VName (SubExp, Space) -> M.Map VName (SubExp, Space) - -> KernelSpace - -> KernelBody InKernel - -> (Stms ExplicitMemory -> KernelBody InKernel -> OffsetM b) + -> SegLevel -> SegSpace + -> KernelBody ExplicitMemory + -> (Stms ExplicitMemory -> KernelBody ExplicitMemory -> OffsetM b) -> ExpandM b -allocsForBody variant_allocs invariant_allocs kspace kbody' m = do - (alloc_stms, alloc_offsets) <- - memoryRequirements kspace (kernelBodyStms kbody') variant_allocs invariant_allocs +allocsForBody variant_allocs invariant_allocs lvl space kbody' m = do + (alloc_offsets, alloc_stms) <- + memoryRequirements lvl space + (kernelBodyStms kbody') variant_allocs invariant_allocs scope <- askScope - let scope' = scopeOfKernelSpace kspace <> M.map nameInfoConv scope + let scope' = scopeOfSegSpace space <> M.map nameInfoConv scope either compilerLimitationS pure $ runOffsetM scope' alloc_offsets $ do kbody'' <- offsetMemoryInKernelBody kbody' m alloc_stms kbody'' -memoryRequirements :: KernelSpace - -> Stms InKernel +memoryRequirements :: SegLevel -> SegSpace + -> Stms ExplicitMemory -> M.Map VName (SubExp, Space) -> M.Map VName (SubExp, Space) - -> ExpandM (Stms ExplicitMemory, RebaseMap) -memoryRequirements kspace kstms variant_allocs invariant_allocs = do - num_threads64 <- newVName "num_threads64" - let num_threads64_pat = Pattern [] [PatElem num_threads64 $ MemPrim int64] - num_threads64_bnd = Let num_threads64_pat (defAux ()) $ BasicOp $ - ConvOp (SExt Int32 Int64) (spaceNumThreads kspace) + -> ExpandM (RebaseMap, Stms ExplicitMemory) +memoryRequirements lvl space kstms variant_allocs invariant_allocs = do + ((num_threads, num_threads64), num_threads_stms) <- runBinder $ do + num_threads <- letSubExp "num_threads" $ BasicOp $ BinOp (Mul Int32) + (unCount $ segNumGroups lvl) (unCount $ segGroupSize lvl) + num_threads64 <- letSubExp "num_threads64" $ BasicOp $ ConvOp (SExt Int32 Int64) num_threads + return (num_threads, num_threads64) (invariant_alloc_stms, invariant_alloc_offsets) <- - expandedInvariantAllocations - (Var num_threads64, spaceNumGroups kspace, spaceGroupSize kspace) - (spaceGlobalId kspace, spaceGroupId kspace, spaceLocalId kspace) invariant_allocs + inScopeOf num_threads_stms $ expandedInvariantAllocations + (num_threads64, segNumGroups lvl, segGroupSize lvl) + space invariant_allocs (variant_alloc_stms, variant_alloc_offsets) <- - expandedVariantAllocations kspace kstms variant_allocs + inScopeOf num_threads_stms $ expandedVariantAllocations num_threads space kstms variant_allocs - let alloc_offsets = invariant_alloc_offsets <> variant_alloc_offsets - alloc_stms = invariant_alloc_stms <> variant_alloc_stms + return (invariant_alloc_offsets <> variant_alloc_offsets, + num_threads_stms <> invariant_alloc_stms <> variant_alloc_stms) - return (oneStm num_threads64_bnd <> alloc_stms, alloc_offsets) +-- | A description of allocations that have been extracted, and how +-- much memory (and which space) is needed. +type Extraction = M.Map VName (SubExp, Space) -- | Extract allocations from 'Thread' statements with -- 'extractThreadAllocations'. -extractKernelBodyAllocations :: KernelBody InKernel - -> (KernelBody InKernel, - M.Map VName (SubExp, Space)) -extractKernelBodyAllocations = extractGenericBodyAllocations kernelBodyStms $ +extractKernelBodyAllocations :: Names -> KernelBody ExplicitMemory + -> (KernelBody ExplicitMemory, + Extraction) +extractKernelBodyAllocations bound_outside = + extractGenericBodyAllocations bound_outside kernelBodyStms $ \stms kbody -> kbody { kernelBodyStms = stms } -extractBodyAllocations :: Body InKernel - -> (Body InKernel, - M.Map VName (SubExp, Space)) -extractBodyAllocations = extractGenericBodyAllocations bodyStms $ +extractBodyAllocations :: Names -> Body ExplicitMemory + -> (Body ExplicitMemory, Extraction) +extractBodyAllocations bound_outside = + extractGenericBodyAllocations bound_outside bodyStms $ \stms body -> body { bodyStms = stms } -extractLambdaAllocations :: Lambda InKernel - -> (Lambda InKernel, - M.Map VName (SubExp, Space)) -extractLambdaAllocations lam = (lam { lambdaBody = body' }, allocs) - where (body', allocs) = extractBodyAllocations $ lambdaBody lam +extractLambdaAllocations :: Names -> Lambda ExplicitMemory + -> (Lambda ExplicitMemory, Extraction) +extractLambdaAllocations bound_outside lam = (lam { lambdaBody = body' }, allocs) + where (body', allocs) = extractBodyAllocations bound_outside $ lambdaBody lam -extractGenericBodyAllocations :: (body -> Stms InKernel) - -> (Stms InKernel -> body -> body) +extractGenericBodyAllocations :: Names + -> (body -> Stms ExplicitMemory) + -> (Stms ExplicitMemory -> body -> body) -> body -> (body, - M.Map VName (SubExp, Space)) -extractGenericBodyAllocations get_stms set_stms body = - let (allocs, stms) = mapAccumL extract M.empty $ stmsToList $ get_stms body - in (set_stms (mconcat stms) body, allocs) - where extract allocs bnd = - let (bnds, body_allocs) = extractThreadAllocations $ oneStm bnd - in (allocs <> body_allocs, bnds) + Extraction) +extractGenericBodyAllocations bound_outside get_stms set_stms body = + let (stms, allocs) = runWriter $ fmap catMaybes $ + mapM (extractStmAllocations bound_outside) $ + stmsToList $ get_stms body + in (set_stms (stmsFromList stms) body, allocs) -extractThreadAllocations :: Stms InKernel - -> (Stms InKernel, M.Map VName (SubExp, Space)) -extractThreadAllocations bnds = - let (allocs, bnds') = mapAccumL isAlloc M.empty $ stmsToList bnds - in (stmsFromList $ catMaybes bnds', allocs) - where isAlloc allocs (Let (Pattern [] [patElem]) _ (Op (Alloc size space))) - | space `notElem` [Space "private", Space "local"] = - (M.insert (patElemName patElem) (size, space) allocs, - Nothing) +extractStmAllocations :: Names -> Stm ExplicitMemory + -> Writer Extraction (Maybe (Stm ExplicitMemory)) +extractStmAllocations bound_outside (Let (Pattern [] [patElem]) _ (Op (Alloc size space))) + | space `notElem` + [Space "private", Space "local"] ++ + map Space (M.keys allScalarMemory), + visibleOutside size = do + tell $ M.singleton (patElemName patElem) (size, space) + return Nothing - isAlloc allocs bnd = - (allocs, Just bnd) + where visibleOutside (Var v) = v `nameIn` bound_outside + visibleOutside Constant{} = True -expandedInvariantAllocations :: (SubExp,SubExp, SubExp) - -> (VName, VName, VName) - -> M.Map VName (SubExp, Space) +extractStmAllocations bound_outside stm = do + e <- mapExpM expMapper $ stmExp stm + return $ Just $ stm { stmExp = e } + where expMapper = identityMapper { mapOnBody = const onBody + , mapOnOp = onOp } + + onBody body = do + let (body', allocs) = extractBodyAllocations bound_outside body + tell allocs + return body' + + onOp (Inner (SegOp op)) = Inner . SegOp <$> mapSegOpM opMapper op + onOp op = return op + + opMapper = identitySegOpMapper { mapOnSegOpLambda = onLambda + , mapOnSegOpBody = onKernelBody + } + + onKernelBody body = do + let (body', allocs) = extractKernelBodyAllocations bound_outside body + tell allocs + return body' + + onLambda lam = do + body <- onBody $ lambdaBody lam + return lam { lambdaBody = body } + +expandedInvariantAllocations :: (SubExp, Count NumGroups SubExp, Count GroupSize SubExp) + -> SegSpace + -> Extraction -> ExpandM (Stms ExplicitMemory, RebaseMap) -expandedInvariantAllocations (num_threads64, num_groups, group_size) - (_thread_index, group_id, local_id) +expandedInvariantAllocations (num_threads64, Count num_groups, Count group_size) + segspace invariant_allocs = do -- We expand the invariant allocations by adding an inner dimension -- equal to the number of kernel threads. @@ -253,28 +264,29 @@ newBase (old_shape, _) = let num_dims = length old_shape - perm = [0, num_dims+1] ++ [1..num_dims] - root_ixfun = IxFun.iota (primExpFromSubExp int32 num_groups : old_shape - ++ [primExpFromSubExp int32 group_size]) + perm = num_dims : [0..num_dims-1] + root_ixfun = IxFun.iota (old_shape + ++ [primExpFromSubExp int32 num_groups * + primExpFromSubExp int32 group_size]) permuted_ixfun = IxFun.permute root_ixfun perm untouched d = DimSlice (fromInt32 0) d (fromInt32 1) offset_ixfun = IxFun.slice permuted_ixfun $ - [DimFix (LeafExp group_id int32), - DimFix (LeafExp local_id int32)] ++ + DimFix (LeafExp (segFlat segspace) int32) : map untouched old_shape in offset_ixfun -expandedVariantAllocations :: KernelSpace -> Stms InKernel - -> M.Map VName (SubExp, Space) +expandedVariantAllocations :: SubExp + -> SegSpace -> Stms ExplicitMemory + -> Extraction -> ExpandM (Stms ExplicitMemory, RebaseMap) -expandedVariantAllocations _ _ variant_allocs +expandedVariantAllocations _ _ _ variant_allocs | null variant_allocs = return (mempty, mempty) -expandedVariantAllocations kspace kstms variant_allocs = do +expandedVariantAllocations num_threads kspace kstms variant_allocs = do let sizes_to_blocks = removeCommonSizes variant_allocs variant_sizes = map fst sizes_to_blocks (slice_stms, offsets, size_sums) <- - sliceKernelSizes variant_sizes kspace kstms + sliceKernelSizes num_threads variant_sizes kspace kstms -- Note the recursive call to expand allocations inside the newly -- produced kernels. slice_stms_tmp <- ExplicitMemory.simplifyStms =<< explicitAllocationsInStms slice_stms @@ -296,8 +308,8 @@ return (Let allocpat (defAux ()) $ Op $ Alloc total_size space, M.singleton mem $ newBase offset) - num_threads = primExpFromSubExp int32 $ spaceNumThreads kspace - gtid = LeafExp (spaceGlobalId kspace) int32 + num_threads' = primExpFromSubExp int32 num_threads + gtid = LeafExp (segFlat kspace) int32 -- For the variant allocations, we add an inner dimension, -- which is then offset by a thread-specific amount. @@ -306,9 +318,9 @@ elems_per_thread = ConvOpExp (SExt Int64 Int32) (primExpFromSubExp int64 size_per_thread) `quot` pt_size - root_ixfun = IxFun.iota [elems_per_thread, num_threads] + root_ixfun = IxFun.iota [elems_per_thread, num_threads'] offset_ixfun = IxFun.slice root_ixfun - [DimSlice (fromInt32 0) num_threads (fromInt32 1), + [DimSlice (fromInt32 0) num_threads' (fromInt32 1), DimFix gtid] shapechange = if length old_shape == 1 then map DimCoercion old_shape @@ -319,13 +331,13 @@ type RebaseMap = M.Map VName (([PrimExp VName], PrimType) -> IxFun) -newtype OffsetM a = OffsetM (ReaderT (Scope InKernel) +newtype OffsetM a = OffsetM (ReaderT (Scope ExplicitMemory) (ReaderT RebaseMap (Either String)) a) deriving (Applicative, Functor, Monad, - HasScope InKernel, LocalScope InKernel, + HasScope ExplicitMemory, LocalScope ExplicitMemory, MonadError String) -runOffsetM :: Scope InKernel -> RebaseMap -> OffsetM a -> Either String a +runOffsetM :: Scope ExplicitMemory -> RebaseMap -> OffsetM a -> Either String a runOffsetM scope offsets (OffsetM m) = runReaderT (runReaderT m scope) offsets @@ -337,7 +349,7 @@ offsets <- askRebaseMap return $ ($ x) <$> M.lookup name offsets -offsetMemoryInKernelBody :: KernelBody InKernel -> OffsetM (KernelBody InKernel) +offsetMemoryInKernelBody :: KernelBody ExplicitMemory -> OffsetM (KernelBody ExplicitMemory) offsetMemoryInKernelBody kbody = do scope <- askScope stms' <- stmsFromList . snd <$> @@ -345,7 +357,7 @@ (stmsToList $ kernelBodyStms kbody) return kbody { kernelBodyStms = stms' } -offsetMemoryInBody :: Body InKernel -> OffsetM (Body InKernel) +offsetMemoryInBody :: Body ExplicitMemory -> OffsetM (Body ExplicitMemory) offsetMemoryInBody (Body attr stms res) = do scope <- askScope stms' <- stmsFromList . snd <$> @@ -353,7 +365,7 @@ (stmsToList stms) return $ Body attr stms' res -offsetMemoryInStm :: Stm InKernel -> OffsetM (Scope InKernel, Stm InKernel) +offsetMemoryInStm :: Stm ExplicitMemory -> OffsetM (Scope ExplicitMemory, Stm ExplicitMemory) offsetMemoryInStm (Let pat attr e) = do pat' <- offsetMemoryInPattern pat e' <- localScope (scopeOfPattern pat') $ offsetMemoryInExp e @@ -379,7 +391,7 @@ where inst Ext{} = Nothing inst (Free x) = return x -offsetMemoryInPattern :: Pattern InKernel -> OffsetM (Pattern InKernel) +offsetMemoryInPattern :: Pattern ExplicitMemory -> OffsetM (Pattern ExplicitMemory) offsetMemoryInPattern (Pattern ctx vals) = do mapM_ inspectCtx ctx Pattern ctx <$> mapM inspectVal vals @@ -418,12 +430,12 @@ IxFun.rebase (fmap (fmap Free) new_base') ixfun offsetMemoryInBodyReturns br = return br -offsetMemoryInLambda :: Lambda InKernel -> OffsetM (Lambda InKernel) -offsetMemoryInLambda lam = do +offsetMemoryInLambda :: Lambda ExplicitMemory -> OffsetM (Lambda ExplicitMemory) +offsetMemoryInLambda lam = inScopeOf lam $ do body <- offsetMemoryInBody $ lambdaBody lam return $ lam { lambdaBody = body } -offsetMemoryInExp :: Exp InKernel -> OffsetM (Exp InKernel) +offsetMemoryInExp :: Exp ExplicitMemory -> OffsetM (Exp ExplicitMemory) offsetMemoryInExp (DoLoop ctx val form body) = do let (ctxparams, ctxinit) = unzip ctx (valparams, valinit) = unzip val @@ -431,46 +443,31 @@ valparams' <- mapM offsetMemoryInParam valparams body' <- localScope (scopeOfFParams ctxparams' <> scopeOfFParams valparams' <> scopeOf form) (offsetMemoryInBody body) return $ DoLoop (zip ctxparams' ctxinit) (zip valparams' valinit) form body' -offsetMemoryInExp (Op (Inner (GroupStream w max_chunk lam accs arrs))) = do - lam_accs <- mapM offsetMemoryInParam $ groupStreamAccParams lam - lam_arrs <- mapM offsetMemoryInParam $ groupStreamArrParams lam - let lam' = lam { groupStreamAccParams = lam_accs - , groupStreamArrParams = lam_arrs - } - body <- localScope (scopeOf lam') $ offsetMemoryInBody $ groupStreamLambdaBody lam - let lam'' = lam' { groupStreamLambdaBody = body } - return $ Op $ Inner $ GroupStream w max_chunk lam'' accs arrs -offsetMemoryInExp (Op (Inner (GroupReduce w lam input))) = do - body <- localScope (scopeOf lam) $ offsetMemoryInBody $ lambdaBody lam - let lam' = lam { lambdaBody = body } - return $ Op $ Inner $ GroupReduce w lam' input -offsetMemoryInExp (Op (Inner (GroupScan w lam input))) = do - body <- localScope (scopeOf lam) $ offsetMemoryInBody $ lambdaBody lam - let lam' = lam { lambdaBody = body } - return $ Op $ Inner $ GroupScan w lam' input -offsetMemoryInExp (Op (Inner (GroupGenReduce w dests lam nes vals locks))) = do - lam_params <- mapM offsetMemoryInParam $ lambdaParams lam - let lam' = lam { lambdaParams = lam_params } - body <- localScope (scopeOf lam') $ offsetMemoryInBody $ lambdaBody lam - let lam'' = lam' { lambdaBody = body } - return $ Op $ Inner $ GroupGenReduce w dests lam'' nes vals locks -offsetMemoryInExp (Op (Inner (Combine cspace ts active body))) = - Op . Inner . Combine cspace ts active <$> offsetMemoryInBody body offsetMemoryInExp e = mapExpM recurse e where recurse = identityMapper { mapOnBody = \bscope -> localScope bscope . offsetMemoryInBody , mapOnBranchType = offsetMemoryInBodyReturns + , mapOnOp = onOp } + onOp (Inner (SegOp op)) = Inner . SegOp <$> mapSegOpM segOpMapper op + where segOpMapper = + identitySegOpMapper { mapOnSegOpBody = offsetMemoryInKernelBody + , mapOnSegOpLambda = offsetMemoryInLambda + } + onOp op = return op + ---- Slicing allocation sizes out of a kernel. -unAllocInKernelStms :: Stms InKernel - -> Either String (Stms Kernels.InKernel) -unAllocInKernelStms = unAllocStms False +unAllocKernelsStms :: Stms ExplicitMemory -> Either String (Stms Kernels.Kernels) +unAllocKernelsStms = unAllocStms False where unAllocBody (Body attr stms res) = Body attr <$> unAllocStms True stms <*> pure res + unAllocKernelBody (KernelBody attr stms res) = + KernelBody attr <$> unAllocStms True stms <*> pure res + unAllocStms nested = fmap (stmsFromList . catMaybes) . mapM (unAllocStm nested) . stmsToList @@ -480,28 +477,6 @@ unAllocStm _ (Let pat attr e) = Just <$> (Let <$> unAllocPattern pat <*> pure attr <*> mapExpM unAlloc' e) - unAllocKernelExp (Barrier se) = - return $ Barrier se - unAllocKernelExp (SplitSpace o w i elems_per_thread) = - return $ SplitSpace o w i elems_per_thread - unAllocKernelExp (Combine cspace ts active body) = - Combine cspace ts active <$> unAllocBody body - unAllocKernelExp (GroupReduce w lam input) = - GroupReduce w <$> unAllocLambda lam <*> pure input - unAllocKernelExp (GroupScan w lam input) = - GroupScan w <$> unAllocLambda lam <*> pure input - unAllocKernelExp (GroupStream w maxchunk lam accs arrs) = - GroupStream w maxchunk <$> unAllocStreamLambda lam <*> pure accs <*> pure arrs - unAllocKernelExp (GroupGenReduce w arrs op bucket vals locks) = - GroupGenReduce w arrs <$> unAllocLambda op <*> - pure bucket <*> pure vals <*> pure locks - - unAllocStreamLambda (GroupStreamLambda chunk_size chunk_offset - acc_params arr_params body) = - GroupStreamLambda chunk_size chunk_offset - (unParams acc_params) (unParams arr_params) <$> - unAllocBody body - unAllocLambda (Lambda params body ret) = Lambda (unParams params) <$> unAllocBody body <*> pure ret @@ -512,8 +487,20 @@ <*> maybe bad return (mapM (rephrasePatElem unAttr) val) where bad = Left $ "Cannot handle memory in pattern " ++ pretty pat - unAllocOp Alloc{} = Left "unhandled Op" - unAllocOp (Inner op) = unAllocKernelExp op + unAllocOp Alloc{} = Left "unAllocOp: unhandled Alloc" + unAllocOp (Inner OtherOp{}) = Left "unAllocOp: unhandled OtherOp" + unAllocOp (Inner (SplitSpace o w i elems_per_thread)) = + return $ SplitSpace o w i elems_per_thread + unAllocOp (Inner (GetSize name sclass)) = + return $ GetSize name sclass + unAllocOp (Inner (GetSizeMax sclass)) = + return $ GetSizeMax sclass + unAllocOp (Inner (CmpSizeLe name sclass x)) = + return $ CmpSizeLe name sclass x + unAllocOp (Inner (SegOp op)) = SegOp <$> mapSegOpM mapper op + where mapper = identitySegOpMapper { mapOnSegOpLambda = unAllocLambda + , mapOnSegOpBody = unAllocKernelBody + } unParam p = maybe bad return $ traverse unAttr p where bad = Left $ "Cannot handle memory-typed parameter '" ++ pretty p ++ "'" @@ -521,7 +508,6 @@ unT t = maybe bad return $ unAttr t where bad = Left $ "Cannot handle memory type '" ++ pretty t ++ "'" - unAlloc' :: Mapper InKernel Kernels.InKernel (Either String) unAlloc' = Mapper { mapOnBody = const unAllocBody , mapOnRetType = unT , mapOnBranchType = unT @@ -537,29 +523,27 @@ unAttr (MemArray pt shape u _) = Just $ Array pt shape u unAttr MemMem{} = Nothing -unAllocScope :: Scope ExplicitMemory -> Scope Kernels.InKernel +unAllocScope :: Scope ExplicitMemory -> Scope Kernels.Kernels unAllocScope = M.mapMaybe unInfo where unInfo (LetInfo attr) = LetInfo <$> unAttr attr unInfo (FParamInfo attr) = FParamInfo <$> unAttr attr unInfo (LParamInfo attr) = LParamInfo <$> unAttr attr unInfo (IndexInfo it) = Just $ IndexInfo it -removeCommonSizes :: M.Map VName (SubExp, Space) +removeCommonSizes :: Extraction -> [(SubExp, [(VName, Space)])] removeCommonSizes = M.toList . foldl' comb mempty . M.toList where comb m (mem, (size, space)) = M.insertWith (++) size [(mem, space)] m -sliceKernelSizes :: [SubExp] -> KernelSpace -> Stms InKernel +sliceKernelSizes :: SubExp -> [SubExp] -> SegSpace -> Stms ExplicitMemory -> ExpandM (Stms Kernels.Kernels, [VName], [VName]) -sliceKernelSizes sizes kspace kstms = do - kstms' <- either compilerLimitationS return $ unAllocInKernelStms kstms +sliceKernelSizes num_threads sizes space kstms = do + kstms' <- either compilerLimitationS return $ unAllocKernelsStms kstms let num_sizes = length sizes i64s = replicate num_sizes $ Prim int64 - inkernels_scope <- asks unAllocScope - - let kernels_scope = castScope inkernels_scope + kernels_scope <- asks unAllocScope - (max_lam, _) <- flip runBinderT inkernels_scope $ do + (max_lam, _) <- flip runBinderT kernels_scope $ do xs <- replicateM num_sizes $ newParam "x" (Prim int64) ys <- replicateM num_sizes $ newParam "y" (Prim int64) (zs, stms) <- localScope (scopeOfLParams $ xs ++ ys) $ collectStms $ @@ -569,7 +553,7 @@ flat_gtid_lparam <- Param <$> newVName "flat_gtid" <*> pure (Prim (IntType Int32)) - (size_lam', _) <- flip runBinderT inkernels_scope $ do + (size_lam', _) <- flip runBinderT kernels_scope $ do params <- replicateM num_sizes $ newParam "x" (Prim int64) (zs, stms) <- localScope (scopeOfLParams params <> scopeOfLParams [flat_gtid_lparam]) $ collectStms $ do @@ -577,7 +561,7 @@ -- Even though this SegRed is one-dimensional, we need to -- provide indexes corresponding to the original potentially -- multi-dimensional construct. - let (kspace_gtids, kspace_dims) = unzip $ spaceDimensions kspace + let (kspace_gtids, kspace_dims) = unzip $ unSegSpace space new_inds = unflattenIndex (map (primExpFromSubExp int32) kspace_dims) (primExpFromSubExp int32 $ Var $ paramName flat_gtid_lparam) @@ -586,23 +570,23 @@ mapM_ addStm kstms' return sizes - localScope (scopeOfKernelSpace kspace) $ - Kernels.simplifyLambda kspace -- XXX, is this the right KernelSpace? - (Lambda [flat_gtid_lparam] (Body () stms zs) i64s) [] + localScope (scopeOfSegSpace space) $ + Kernels.simplifyLambda (Lambda [flat_gtid_lparam] (Body () stms zs) i64s) [] ((maxes_per_thread, size_sums), slice_stms) <- flip runBinderT kernels_scope $ do num_threads_64 <- letSubExp "num_threads" $ - BasicOp $ ConvOp (SExt Int32 Int64) $ spaceNumThreads kspace + BasicOp $ ConvOp (SExt Int32 Int64) num_threads pat <- basicPattern [] <$> replicateM num_sizes (newIdent "max_per_thread" $ Prim int64) thread_space_iota <- letExp "thread_space_iota" $ BasicOp $ - Iota (spaceNumThreads kspace) (intConst Int32 0) (intConst Int32 1) Int32 + Iota num_threads (intConst Int32 0) (intConst Int32 1) Int32 let red_op = SegRedOp Commutative max_lam (replicate num_sizes $ intConst Int64 0) mempty - addStms =<< - nonSegRed pat (spaceNumThreads kspace) [red_op] size_lam' [thread_space_iota] + lvl <- segThread "segred" + addStms =<< mapM renameStm =<< + nonSegRed lvl pat num_threads [red_op] size_lam' [thread_space_iota] size_sums <- forM (patternNames pat) $ \threads_max -> letExp "size_sum" $ diff --git a/src/Futhark/Pass/ExplicitAllocations.hs b/src/Futhark/Pass/ExplicitAllocations.hs --- a/src/Futhark/Pass/ExplicitAllocations.hs +++ b/src/Futhark/Pass/ExplicitAllocations.hs @@ -23,11 +23,7 @@ import Futhark.Representation.Kernels import Futhark.Optimise.Simplify.Lore - (mkWiseBody, - mkWiseLetStm, - removeExpWisdom, - - removeScopeWisdom) + (mkWiseBody, mkWiseLetStm, removeExpWisdom, removeScopeWisdom) import Futhark.MonadFreshNames import Futhark.Representation.ExplicitMemory import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun @@ -38,9 +34,6 @@ import Futhark.Pass import Futhark.Util (splitFromEnd, takeLast) -type InInKernel = Futhark.Representation.Kernels.InKernel -type OutInKernel = Futhark.Representation.ExplicitMemory.InKernel - data AllocStm = SizeComputation VName (PrimExp VName) | Allocation VName SubExp Space | ArrayCopy VName VName @@ -55,9 +48,13 @@ bindAllocStm (ArrayCopy name src) = letBindNames_ [name] $ BasicOp $ Copy src +defaultExpHints :: (Monad m, Attributes lore) => Exp lore -> m [ExpHint] +defaultExpHints e = return $ replicate (expExtTypeSize e) NoHint + class (MonadFreshNames m, HasScope lore m, ExplicitMemorish lore) => Allocator lore m where addAllocStm :: AllocStm -> m () + askDefaultSpace :: m Space default addAllocStm :: (Allocable fromlore lore, Op lore ~ MemOp inner, @@ -84,7 +81,7 @@ return size expHints :: Exp lore -> m [ExpHint] - expHints e = return $ replicate (expExtTypeSize e) NoHint + expHints = defaultExpHints allocateMemory :: Allocator lore m => String -> SubExp -> Space -> m VName @@ -123,17 +120,18 @@ , aggressiveReuse :: Bool -- ^ Aggressively try to reuse memory in do-loops - -- should be True inside kernels, False outside. + , allocSpace :: Space + -- ^ When allocating memory, put it in this memory space. + -- This is primarily used to ensure that group-wide + -- statements store their results in local memory. , allocInOp :: Op fromlore -> AllocM fromlore tolore (Op tolore) + , envExpHints :: Exp tolore -> AllocM fromlore tolore [ExpHint] } boundDims :: ChunkMap -> AllocEnv fromlore tolore -> AllocEnv fromlore tolore boundDims m env = env { chunkMap = m <> chunkMap env } -boundDim :: VName -> SubExp -> AllocEnv fromlore tolore - -> AllocEnv fromlore tolore -boundDim name se = boundDims $ M.singleton name se - -- | Monad for adding allocations to an entire program. newtype AllocM fromlore tolore a = AllocM (BinderT tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a) @@ -162,30 +160,20 @@ collectStms (AllocM m) = AllocM $ collectBinderStms m certifying cs (AllocM m) = AllocM $ certifyingBinder cs m -instance Allocable fromlore OutInKernel => +instance Allocable fromlore ExplicitMemory => Allocator ExplicitMemory (AllocM fromlore ExplicitMemory) where - expHints = kernelExpHints - -instance Allocable fromlore OutInKernel => - Allocator OutInKernel (AllocM fromlore OutInKernel) where - expHints = inKernelExpHints + expHints e = do + f <- asks envExpHints + f e + askDefaultSpace = asks allocSpace runAllocM :: MonadFreshNames m => (Op fromlore -> AllocM fromlore tolore (Op tolore)) + -> (Exp tolore -> AllocM fromlore tolore [ExpHint]) -> AllocM fromlore tolore a -> m a -runAllocM handleOp (AllocM m) = +runAllocM handleOp hints (AllocM m) = fmap fst $ modifyNameSource $ runState $ runReaderT (runBinderT m mempty) env - where env = AllocEnv mempty False handleOp - -subAllocM :: (SameScope tolore1 tolore2, ExplicitMemorish tolore2) => - (Op fromlore1 -> AllocM fromlore1 tolore1 (Op tolore1)) -> Bool - -> AllocM fromlore1 tolore1 a - -> AllocM fromlore2 tolore2 a -subAllocM handleOp b (AllocM m) = do - scope <- castScope <$> askScope - chunks <- asks chunkMap - let env = AllocEnv chunks b handleOp - fmap fst $ modifyNameSource $ runState $ runReaderT (runBinderT m scope) env + where env = AllocEnv mempty False DefaultSpace handleOp hints -- | Monad for adding allocations to a single pattern. newtype PatAllocM lore a = PatAllocM (RWS @@ -201,10 +189,7 @@ instance Allocator ExplicitMemory (PatAllocM ExplicitMemory) where addAllocStm = tell . pure dimAllocationSize = return - -instance Allocator OutInKernel (PatAllocM OutInKernel) where - addAllocStm = tell . pure - dimAllocationSize = return + askDefaultSpace = return DefaultSpace runPatAllocM :: MonadFreshNames m => PatAllocM lore a -> Scope lore @@ -290,11 +275,11 @@ return $ PatElem (identName ident) summary MemArray bt _ u ret -> do - let space = case ret of - Just (ReturnsNewBlock mem_space _ _) -> mem_space - _ -> DefaultSpace - (mem,(ident',ixfun)) <- lift $ memForBindee ident - tell ([PatElem (identName mem) $ MemMem space], + space <- case ret of + Just (ReturnsNewBlock mem_space _ _) -> return mem_space + _ -> lift askDefaultSpace + (mem,(ident',ixfun)) <- lift $ memForBindee ident space + tell ([PatElem (identName mem) $ MemMem space], []) return $ PatElem (identName ident') $ MemArray bt shape u $ ArrayIn (identName mem) ixfun @@ -319,7 +304,7 @@ summaryForBindage (Mem space) _ = return $ MemMem space summaryForBindage t@(Array bt shape u) NoHint = do - m <- allocForArray t DefaultSpace + m <- allocForArray t =<< askDefaultSpace return $ directIndexFunction bt shape u m t summaryForBindage t (Hint ixfun space) = do let bt = elemType t @@ -330,11 +315,11 @@ return $ MemArray bt (arrayShape t) NoUniqueness $ ArrayIn m ixfun memForBindee :: (MonadFreshNames m) => - Ident + Ident -> Space -> m (Ident, (Ident, IxFun)) -memForBindee ident = do - mem <- newIdent memname $ Mem DefaultSpace +memForBindee ident space = do + mem <- newIdent memname (Mem space) return (mem, (ident, IxFun.iota $ map (primExpFromSubExp int32) $ arrayDims t)) where memname = baseString (identName ident) <> "_mem" @@ -403,13 +388,13 @@ if space /= Space "local" && reuse && u == Unique && - loopInvariantShape mergeparam && - IxFun.isLinear ixfun + loopInvariantShape mergeparam then return (mergeparam { paramAttr = MemArray bt shape Unique $ ArrayIn mem ixfun }, lift . ensureArrayIn (paramType mergeparam) mem ixfun) - else doDefault mergeparam space + else do def_space <- asks allocSpace + doDefault mergeparam def_space - allocInMergeParam (mergeparam, _) = doDefault mergeparam DefaultSpace + allocInMergeParam (mergeparam, _) = doDefault mergeparam =<< lift askDefaultSpace doDefault mergeparam space = do mergeparam' <- allocInFParam mergeparam space @@ -442,9 +427,10 @@ ensureDirectArray space_ok v = do (mem, ixfun) <- lookupArraySummary v Mem mem_space <- lookupType mem + default_space <- askDefaultSpace if IxFun.isDirect ixfun && maybe True (==mem_space) space_ok then return (mem, Var v) - else needCopy (fromMaybe DefaultSpace space_ok) + else needCopy (fromMaybe default_space space_ok) where needCopy space = -- We need to do a new allocation, copy 'v', and make a new -- binding for the size of the memory block. @@ -469,7 +455,8 @@ funcallArgs args = do (valargs, mem_and_size_args) <- runWriterT $ forM args $ \(arg,d) -> do t <- lift $ subExpType arg - arg' <- linearFuncallArg t DefaultSpace arg + space <- lift askDefaultSpace + arg' <- linearFuncallArg t space arg return (arg', d) return $ map (,Observe) mem_and_size_args <> valargs @@ -493,7 +480,7 @@ Stms Kernels -> m (Stms ExplicitMemory) explicitAllocationsInStms stms = do scope <- askScope - runAllocM handleHostOp $ localScope scope $ allocInStms stms return + runAllocM handleHostOp kernelExpHints $ localScope scope $ allocInStms stms return memoryInRetType :: [RetType Kernels] -> [RetType ExplicitMemory] memoryInRetType ts = evalState (mapM addAttr ts) $ startOfFreeIDRange ts @@ -512,102 +499,44 @@ allocInFun :: MonadFreshNames m => FunDef Kernels -> m (FunDef ExplicitMemory) allocInFun (FunDef entry fname rettype params fbody) = - runAllocM handleHostOp $ + runAllocM handleHostOp kernelExpHints $ allocInFParams (zip params $ repeat DefaultSpace) $ \params' -> do fbody' <- insertStmsM $ allocInFunBody (map (const $ Just DefaultSpace) rettype) fbody return $ FunDef entry fname (memoryInRetType rettype) params' fbody' -handleHostOp :: HostOp Kernels (Kernel InInKernel) - -> AllocM Kernels ExplicitMemory (MemOp (HostOp ExplicitMemory (Kernel OutInKernel))) +handleHostOp :: HostOp Kernels (SOAC Kernels) + -> AllocM Kernels ExplicitMemory (MemOp (HostOp ExplicitMemory ())) +handleHostOp (SplitSpace o w i elems_per_thread) = + return $ Inner $ SplitSpace o w i elems_per_thread handleHostOp (GetSize key size_class) = return $ Inner $ GetSize key size_class handleHostOp (GetSizeMax size_class) = return $ Inner $ GetSizeMax size_class handleHostOp (CmpSizeLe key size_class x) = return $ Inner $ CmpSizeLe key size_class x -handleHostOp (HostOp (Kernel desc space kernel_ts kbody)) = - subInKernel space $ - Inner . HostOp . Kernel desc space kernel_ts <$> - localScope (scopeOfKernelSpace space) (allocInKernelBody kbody) - -handleHostOp (HostOp (SegMap space ts body)) = do - body' <- subInKernel space $ - localScope (scopeOfKernelSpace space) $ allocInKernelBody body - return $ Inner $ HostOp $ SegMap space ts body' - -handleHostOp (HostOp (SegRed space reds ts body)) = do - body' <- subInKernel space $ - localScope (scopeOfKernelSpace space) $ allocInKernelBody body - reds' <- forM reds $ \(SegRedOp comm lam nes shape) -> do - lam' <- allocInSegRedLambda space lam - return $ SegRedOp comm lam' nes shape - return $ Inner $ HostOp $ SegRed space reds' ts body' - -handleHostOp (HostOp (SegScan space scan_op nes ts body)) = do - body' <- subInKernel space $ - localScope (scopeOfKernelSpace space) $ allocInKernelBody body - scan_op' <- allocInSegRedLambda space scan_op - return $ Inner $ HostOp $ SegScan space scan_op' nes ts body' - -handleHostOp (HostOp (SegGenRed space ops ts body)) = do - body' <- subInKernel space $ - localScope (scopeOfKernelSpace space) $ allocInKernelBody body - ops' <- forM ops $ \op -> do - lam <- allocInSegRedLambda space $ genReduceOp op - return op { genReduceOp = lam } - return $ Inner $ HostOp $ SegGenRed space ops' ts body' - -subInKernel :: KernelSpace -> AllocM InInKernel OutInKernel a - -> AllocM fromlore2 ExplicitMemory a -subInKernel space = subAllocM handleKernelExp True - where handleKernelExp (Barrier se) = - return $ Inner $ Barrier se - - handleKernelExp (SplitSpace o w i elems_per_thread) = - return $ Inner $ SplitSpace o w i elems_per_thread - - handleKernelExp (Combine cspace ts active body) = - Inner . Combine cspace ts active <$> allocInBodyNoDirect body - - handleKernelExp (GroupReduce w lam input) = do - summaries <- mapM lookupArraySummary arrs - lam' <- allocInReduceLambda space lam summaries - return $ Inner $ GroupReduce w lam' input - where arrs = map snd input - - handleKernelExp (GroupScan w lam input) = do - summaries <- mapM lookupArraySummary arrs - lam' <- allocInReduceLambda space lam summaries - return $ Inner $ GroupScan w lam' input - where arrs = map snd input - - handleKernelExp (GroupGenReduce w dests op bucket vs locks) = do - let (x_params, y_params) = splitAt (length vs) $ lambdaParams op - sliceDest dest = do - dest_t <- lookupType dest - sliceInfo dest $ fullSlice dest_t $ map DimFix bucket - x_params' <- zipWith Param (map paramName x_params) <$> - mapM sliceDest dests - y_params' <- zipWith Param (map paramName y_params) <$> - mapM subExpMemInfo vs +handleHostOp (OtherOp op) = + fail $ "Cannot allocate memory in SOAC: " ++ pretty op - op' <- allocInLambda (x_params'<>y_params') (lambdaBody op) (lambdaReturnType op) - return $ Inner $ GroupGenReduce w dests op' bucket vs locks +handleHostOp (SegOp op) = + Inner . SegOp <$> handleSegOp op - handleKernelExp (GroupStream w maxchunk lam accs arrs) = do - acc_summaries <- mapM accSummary accs - arr_summaries <- mapM lookupArraySummary arrs - lam' <- allocInGroupStreamLambda maxchunk lam acc_summaries arr_summaries - return $ Inner $ GroupStream w maxchunk lam' accs arrs - where accSummary (Constant v) = return $ MemPrim $ primValueType v - accSummary (Var v) = lookupMemInfo v +handleSegOp :: SegOp Kernels + -> AllocM Kernels ExplicitMemory (SegOp ExplicitMemory) +handleSegOp op = allocAtLevel (segLevel op) $ mapSegOpM mapper op + where scope = scopeOfSegSpace $ segSpace op + mapper = identitySegOpMapper + { mapOnSegOpBody = localScope scope . allocInKernelBody (segLevel op) + , mapOnSegOpLambda = allocInBinOpLambda (segLevel op) $ segSpace op + } -allocInBodyNoDirect :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - Body fromlore -> AllocM fromlore tolore (Body tolore) -allocInBodyNoDirect (Body _ bnds res) = - allocInStms bnds $ \bnds' -> - return $ Body () bnds' res +allocAtLevel :: SegLevel -> AllocM fromlore tlore a -> AllocM fromlore tlore a +allocAtLevel lvl = local $ \env -> env { allocSpace = space + , aggressiveReuse = True + } + where space = case lvl of SegThread{} -> DefaultSpace + SegThreadScalar{} -> DefaultSpace + SegGroup{} -> Space "local" bodyReturnMemCtx :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => SubExp -> AllocM fromlore tolore [SubExp] @@ -767,65 +696,34 @@ Mem space -> return (p { paramAttr = MemMem space }, a) -allocInReduceLambda :: KernelSpace - -> Lambda InInKernel - -> [(VName, IxFun)] - -> AllocM InInKernel OutInKernel (Lambda OutInKernel) -allocInReduceLambda space lam input_summaries = do - let (acc_params, arr_params) = - splitAt (length input_summaries) $ lambdaParams lam - this_index = LeafExp (spaceGlobalId space) int32 - other_index = this_index + primExpFromSubExp int32 (spaceNumThreads space) - - acc_params' <- - allocInReduceParameters this_index $ - zip acc_params input_summaries - arr_params' <- - allocInReduceParameters other_index $ - zip arr_params input_summaries - - allocInLambda (acc_params' ++ arr_params') - (lambdaBody lam) (lambdaReturnType lam) - -allocInReduceParameters :: PrimExp VName - -> [(LParam InInKernel, (VName, IxFun))] - -> AllocM InInKernel OutInKernel [LParam ExplicitMemory] -allocInReduceParameters my_id = mapM allocInReduceParameter - where allocInReduceParameter (p, (mem, ixfun)) = - case paramType p of - (Array bt shape u) -> - let ixfun' = IxFun.slice ixfun $ - fullSliceNum (IxFun.shape ixfun) [DimFix my_id] - in return p { paramAttr = MemArray bt shape u $ ArrayIn mem ixfun' } - Prim bt -> - return p { paramAttr = MemPrim bt } - Mem space -> - return p { paramAttr = MemMem space } - -allocInSegRedLambda :: KernelSpace -> Lambda InInKernel - -> AllocM Kernels ExplicitMemory (Lambda OutInKernel) -allocInSegRedLambda space lam = do +allocInBinOpLambda :: SegLevel -> SegSpace -> Lambda Kernels + -> AllocM Kernels ExplicitMemory (Lambda ExplicitMemory) +allocInBinOpLambda lvl (SegSpace flat _) lam = do + num_threads <- letSubExp "num_threads" $ BasicOp $ BinOp (Mul Int32) + (unCount (segNumGroups lvl)) (unCount (segGroupSize lvl)) let (acc_params, arr_params) = splitAt (length (lambdaParams lam) `div` 2) $ lambdaParams lam - this_index = LeafExp (spaceGlobalId space) int32 - other_index = this_index + primExpFromSubExp int32 (spaceNumThreads space) + index_x = LeafExp flat int32 + index_y = index_x + primExpFromSubExp int32 num_threads (acc_params', arr_params') <- - allocInSegRedParameters (spaceNumThreads space) this_index other_index acc_params arr_params + allocInBinOpParams num_threads index_x index_y acc_params arr_params - subInKernel space $ allocInLambda (acc_params' ++ arr_params') + local (\env -> env { envExpHints = inThreadExpHints }) $ + allocInLambda (acc_params' ++ arr_params') (lambdaBody lam) (lambdaReturnType lam) -allocInSegRedParameters :: SubExp - -> PrimExp VName -> PrimExp VName - -> [LParam InInKernel] - -> [LParam InInKernel] - -> AllocM Kernels ExplicitMemory ([LParam ExplicitMemory], [LParam ExplicitMemory]) -allocInSegRedParameters num_threads my_id other_id xs ys = unzip <$> zipWithM alloc xs ys +allocInBinOpParams :: SubExp + -> PrimExp VName -> PrimExp VName + -> [LParam Kernels] + -> [LParam Kernels] + -> AllocM Kernels ExplicitMemory ([LParam ExplicitMemory], [LParam ExplicitMemory]) +allocInBinOpParams num_threads my_id other_id xs ys = unzip <$> zipWithM alloc xs ys where alloc x y = case paramType x of Array bt shape u -> do - twice_num_threads <- letSubExp "twice_num_threads" $ - BasicOp $ BinOp (Mul Int32) num_threads $ intConst Int32 2 + twice_num_threads <- + letSubExp "twice_num_threads" $ + BasicOp $ BinOp (Mul Int32) num_threads $ intConst Int32 2 let t = paramType x `arrayOfRow` twice_num_threads mem <- allocForArray t DefaultSpace -- XXX: this iota ixfun is a bit inefficient; leading to uncoalesced access. @@ -844,97 +742,40 @@ return (x { paramAttr = MemMem space }, y { paramAttr = MemMem space }) -allocInChunkedParameters :: PrimExp VName - -> [(LParam InInKernel, (VName, IxFun))] - -> AllocM InInKernel OutInKernel [LParam OutInKernel] -allocInChunkedParameters offset = mapM allocInChunkedParameter - where allocInChunkedParameter (p, (mem, ixfun)) = - case paramType p of - Array bt shape u -> - let ixfun' = IxFun.offsetIndex ixfun offset - in return p { paramAttr = MemArray bt shape u $ ArrayIn mem ixfun' } - Prim bt -> - return p { paramAttr = MemPrim bt } - Mem space -> - return p { paramAttr = MemMem space } - -allocInLambda :: [LParam OutInKernel] -> Body InInKernel -> [Type] - -> AllocM InInKernel OutInKernel (Lambda OutInKernel) +allocInLambda :: [LParam ExplicitMemory] -> Body Kernels -> [Type] + -> AllocM Kernels ExplicitMemory (Lambda ExplicitMemory) allocInLambda params body rettype = do body' <- localScope (scopeOfLParams params) $ allocInStms (bodyStms body) $ \bnds' -> return $ Body () bnds' $ bodyResult body return $ Lambda params body' rettype -allocInKernelBody :: KernelBody InInKernel - -> AllocM InInKernel OutInKernel (KernelBody OutInKernel) -allocInKernelBody (KernelBody () stms res) = - allocInStms stms $ \stms' -> - return $ KernelBody () stms' res +allocInKernelBody :: SegLevel -> KernelBody Kernels + -> AllocM Kernels ExplicitMemory (KernelBody ExplicitMemory) +allocInKernelBody lvl (KernelBody () stms res) = + local f $ allocInStms stms $ \stms' -> return $ KernelBody () stms' res + where f = case lvl of SegThread{} -> inThread + SegThreadScalar{} -> inThread + SegGroup{} -> inGroup + inThread env = env { envExpHints = inThreadExpHints } + inGroup env = env { envExpHints = inGroupExpHints } class SizeSubst op where opSizeSubst :: PatternT attr -> op -> ChunkMap -instance SizeSubst op => SizeSubst (HostOp lore op) where - opSizeSubst pat (HostOp op) = opSizeSubst pat op - opSizeSubst _ _ = mempty - -instance SizeSubst (Kernel lore) where +instance SizeSubst (HostOp lore op) where + opSizeSubst (Pattern _ [size]) (SplitSpace _ _ _ elems_per_thread) = + M.singleton (patElemName size) elems_per_thread opSizeSubst _ _ = mempty instance SizeSubst op => SizeSubst (MemOp op) where opSizeSubst pat (Inner op) = opSizeSubst pat op opSizeSubst _ _ = mempty -instance SizeSubst (KernelExp lore) where - opSizeSubst (Pattern _ [size]) (SplitSpace _ _ _ elems_per_thread) = - M.singleton (patElemName size) elems_per_thread - opSizeSubst _ _ = mempty - sizeSubst :: SizeSubst (Op lore) => Stm lore -> ChunkMap sizeSubst (Let pat _ (Op op)) = opSizeSubst pat op sizeSubst _ = mempty -allocInGroupStreamLambda :: SubExp - -> GroupStreamLambda InInKernel - -> [MemBound NoUniqueness] - -> [(VName, IxFun)] - -> AllocM InInKernel OutInKernel (GroupStreamLambda OutInKernel) -allocInGroupStreamLambda maxchunk lam acc_summaries arr_summaries = do - let GroupStreamLambda block_size block_offset acc_params arr_params body = lam - - acc_params' <- - allocInAccParameters acc_params acc_summaries - arr_params' <- - allocInChunkedParameters (LeafExp block_offset int32) $ - zip arr_params arr_summaries - - body' <- localScope (M.insert block_size (IndexInfo Int32) $ - M.insert block_offset (IndexInfo Int32) $ - scopeOfLParams $ acc_params' ++ arr_params') $ - local (boundDim block_size maxchunk) $ do - body' <- allocInBodyNoDirect body - insertStmsM $ do - -- We copy the result of the body to whereever the accumulators are stored. - addStms (bodyStms body') - let maybeCopyResult r p = - case paramAttr p of - MemArray _ _ _ (ArrayIn mem ixfun) -> - ensureArrayIn (paramType p) mem ixfun r - _ -> - return r - resultBodyM =<< - zipWithM maybeCopyResult (bodyResult body') acc_params' - return $ - GroupStreamLambda block_size block_offset acc_params' arr_params' body' - -allocInAccParameters :: [LParam InInKernel] - -> [MemBound NoUniqueness] - -> AllocM InInKernel OutInKernel [LParam OutInKernel] -allocInAccParameters = zipWithM allocInAccParameter - where allocInAccParameter p attr = return p { paramAttr = attr } - - mkLetNamesB' :: (Op (Lore m) ~ MemOp inner, MonadBinder m, ExpAttr (Lore m) ~ (), Allocator (Lore m) (PatAllocM (Lore m))) => @@ -962,21 +803,11 @@ mkBodyB stms res = return $ Body () stms res mkLetNamesB = mkLetNamesB' () -instance BinderOps OutInKernel where - mkExpAttrB _ _ = return () - mkBodyB stms res = return $ Body () stms res - mkLetNamesB = mkLetNamesB' () - instance BinderOps (Engine.Wise ExplicitMemory) where mkExpAttrB pat e = return $ Engine.mkWiseExpAttr pat () e mkBodyB stms res = return $ Engine.mkWiseBody () stms res mkLetNamesB = mkLetNamesB'' -instance BinderOps (Engine.Wise OutInKernel) where - mkExpAttrB pat e = return $ Engine.mkWiseExpAttr pat () e - mkBodyB stms res = return $ Engine.mkWiseBody () stms res - mkLetNamesB = mkLetNamesB'' - simplifiable :: (Engine.SimplifiableLore lore, ExpAttr lore ~ (), BodyAttr lore ~ (), @@ -1016,8 +847,7 @@ data ExpHint = NoHint | Hint IxFun Space -kernelExpHints :: (Allocator lore m, Op lore ~ MemOp (HostOp lore (Kernel somelore))) => - Exp lore -> m [ExpHint] +kernelExpHints :: Allocator ExplicitMemory m => Exp ExplicitMemory -> m [ExpHint] kernelExpHints (BasicOp (Manifest perm v)) = do dims <- arrayDims <$> lookupType v let perm_inv = rearrangeInverse perm @@ -1026,14 +856,11 @@ perm_inv return [Hint ixfun DefaultSpace] -kernelExpHints (Op (Inner (HostOp (Kernel _ space ts kbody)))) = - zipWithM (mapResultHint space) ts $ kernelBodyResult kbody - -kernelExpHints (Op (Inner (HostOp (SegMap space ts body)))) = - zipWithM (mapResultHint space) ts $ kernelBodyResult body +kernelExpHints (Op (Inner (SegOp (SegMap lvl@SegThread{} space ts body)))) = + zipWithM (mapResultHint lvl space) ts $ kernelBodyResult body -kernelExpHints (Op (Inner (HostOp (SegRed space reds ts body)))) = - (map (const NoHint) red_res <>) <$> zipWithM (mapResultHint space) (drop num_reds ts) map_res +kernelExpHints (Op (Inner (SegOp (SegRed lvl@SegThread{} space reds ts body)))) = + (map (const NoHint) red_res <>) <$> zipWithM (mapResultHint lvl space) (drop num_reds ts) map_res where num_reds = segRedResults reds (red_res, map_res) = splitAt num_reds $ kernelBodyResult body @@ -1041,9 +868,10 @@ return $ replicate (expExtTypeSize e) NoHint mapResultHint :: Allocator lore m => - KernelSpace -> Type -> KernelResult -> m ExpHint -mapResultHint space = hint - where num_threads = spaceNumThreads space + SegLevel -> SegSpace -> Type -> KernelResult -> m ExpHint +mapResultHint lvl space = hint + where num_threads = primExpFromSubExp int32 (unCount $ segNumGroups lvl) * + primExpFromSubExp int32 (unCount $ segGroupSize lvl) -- Heuristic: do not rearrange for returned arrays that are -- sufficiently small. @@ -1051,19 +879,18 @@ coalesceReturnOfShape bs [Constant (IntValue (Int32Value d))] = bs * d > 4 coalesceReturnOfShape _ _ = True - hint t (ThreadsReturn _) + hint t (Returns _) | coalesceReturnOfShape (primByteSize (elemType t)) $ arrayDims t = do - let space_dims = map snd $ spaceDimensions space + let space_dims = segSpaceDims space t_dims <- mapM dimAllocationSize $ arrayDims t return $ Hint (innermost space_dims t_dims) DefaultSpace - hint t (ConcatReturns SplitStrided{} w _ _ _) = do + hint t (ConcatReturns SplitStrided{} w _ _) = do t_dims <- mapM dimAllocationSize $ arrayDims t return $ Hint (innermost [w] t_dims) DefaultSpace - -- TODO: Can we make hint for ConcatRetuns when it has an offset? - hint Prim{} (ConcatReturns SplitContiguous w elems_per_thread Nothing _) = do - let ixfun_base = IxFun.iota $ map (primExpFromSubExp int32) [num_threads,elems_per_thread] + hint Prim{} (ConcatReturns SplitContiguous w elems_per_thread _) = do + let ixfun_base = IxFun.iota [num_threads, primExpFromSubExp int32 elems_per_thread] ixfun_tr = IxFun.permute ixfun_base [1,0] ixfun = IxFun.reshape ixfun_tr $ map (DimNew . primExpFromSubExp int32) [w] return $ Hint ixfun DefaultSpace @@ -1082,16 +909,19 @@ ixfun_rearranged = IxFun.permute ixfun_base perm_inv in ixfun_rearranged -inKernelExpHints :: (Allocator lore m, Op lore ~ MemOp (KernelExp somelore)) => - Exp lore -> m [ExpHint] -inKernelExpHints (Op (Inner (Combine (CombineSpace scatter cspace) ts _ _))) = - fmap (replicate (sum ns) NoHint ++) $ forM (drop (sum ns*2) ts) $ \t -> do - alloc_dims <- mapM dimAllocationSize $ dims ++ arrayDims t - let ixfun = IxFun.iota $ map (primExpFromSubExp int32) alloc_dims - return $ Hint ixfun $ Space "local" - where dims = map snd cspace - (_, ns, _) = unzip3 scatter -inKernelExpHints e = +inGroupExpHints :: Allocator ExplicitMemory m => Exp ExplicitMemory -> m [ExpHint] +inGroupExpHints (Op (Inner (SegOp (SegMap SegThreadScalar{} space ts _)))) = return $ do + t <- ts + case t of + Prim pt -> + return $ Hint (IxFun.iota $ map (primExpFromSubExp int32) $ + segSpaceDims space ++ arrayDims t) $ Space $ scalarMemory pt + _ -> + return NoHint +inGroupExpHints e = return $ replicate (expExtTypeSize e) NoHint + +inThreadExpHints :: Allocator ExplicitMemory m => Exp ExplicitMemory -> m [ExpHint] +inThreadExpHints e = mapM maybePrivate =<< expExtType e where maybePrivate t | arrayRank t > 0, diff --git a/src/Futhark/Pass/ExtractKernels.hs b/src/Futhark/Pass/ExtractKernels.hs --- a/src/Futhark/Pass/ExtractKernels.hs +++ b/src/Futhark/Pass/ExtractKernels.hs @@ -2,1541 +2,722 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE LambdaCase #-} -{-# LANGUAGE TypeSynonymInstances #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE RankNTypes #-} --- | Kernel extraction. --- --- In the following, I will use the term "width" to denote the amount --- of immediate parallelism in a map - that is, the outer size of the --- array(s) being used as input. --- --- = Basic Idea --- --- If we have: --- --- @ --- map --- map(f) --- bnds_a... --- map(g) --- @ --- --- Then we want to distribute to: --- --- @ --- map --- map(f) --- map --- bnds_a --- map --- map(g) --- @ --- --- But for now only if --- --- (0) it can be done without creating irregular arrays. --- Specifically, the size of the arrays created by @map(f)@, by --- @map(g)@ and whatever is created by @bnds_a@ that is also used --- in @map(g)@, must be invariant to the outermost loop. --- --- (1) the maps are _balanced_. That is, the functions @f@ and @g@ --- must do the same amount of work for every iteration. --- --- The advantage is that the map-nests containing @map(f)@ and --- @map(g)@ can now be trivially flattened at no cost, thus exposing --- more parallelism. Note that the @bnds_a@ map constitutes array --- expansion, which requires additional storage. --- --- = Distributing Sequential Loops --- --- As a starting point, sequential loops are treated like scalar --- expressions. That is, not distributed. However, sometimes it can --- be worthwhile to distribute if they contain a map: --- --- @ --- map --- loop --- map --- map --- @ --- --- If we distribute the loop and interchange the outer map into the --- loop, we get this: --- --- @ --- loop --- map --- map --- map --- map --- @ --- --- Now more parallelism may be available. --- --- = Unbalanced Maps --- --- Unbalanced maps will as a rule be sequentialised, but sometimes, --- there is another way. Assume we find this: --- --- @ --- map --- map(f) --- map(g) --- map --- @ --- --- Presume that @map(f)@ is unbalanced. By the simple rule above, we --- would then fully sequentialise it, resulting in this: --- --- @ --- map --- loop --- map --- map --- @ --- --- == Balancing by Loop Interchange --- --- The above is not ideal, as we cannot flatten the @map-loop@ nest, --- and we are thus limited in the amount of parallelism available. --- --- But assume now that the width of @map(g)@ is invariant to the outer --- loop. Then if possible, we can interchange @map(f)@ and @map(g)@, --- sequentialise @map(f)@ and distribute, interchanging the outer --- parallel loop into the sequential loop: --- --- @ --- loop(f) --- map --- map(g) --- map --- map --- @ --- --- After flattening the two nests we can obtain more parallelism. --- --- When distributing a map, we also need to distribute everything that --- the map depends on - possibly as its own map. When distributing a --- set of scalar bindings, we will need to know which of the binding --- results are used afterwards. Hence, we will need to compute usage --- information. --- --- = Redomap --- --- Redomap can be handled much like map. Distributed loops are --- distributed as maps, with the parameters corresponding to the --- neutral elements added to their bodies. The remaining loop will --- remain a redomap. Example: --- --- @ --- redomap(op, --- fn (v) => --- map(f) --- map(g), --- e,a) --- @ --- --- distributes to --- --- @ --- let b = map(fn v => --- let acc = e --- map(f), --- a) --- redomap(op, --- fn (v,dist) => --- map(g), --- e,a,b) --- @ --- --- Note that there may be further kernel extraction opportunities --- inside the @map(f)@. The downside of this approach is that the --- intermediate array (@b@ above) must be written to main memory. An --- often better approach is to just turn the entire @redomap@ into a --- single kernel. --- -module Futhark.Pass.ExtractKernels - (extractKernels) - where - -import Control.Monad.RWS.Strict -import Control.Monad.Reader -import Control.Monad.Writer.Strict -import Control.Monad.Trans.Maybe -import qualified Data.Set as S -import Data.Maybe -import Data.List - -import Futhark.Representation.SOACS -import qualified Futhark.Representation.SOACS.SOAC as SOAC -import Futhark.Representation.SOACS.Simplify (simplifyStms, simpleSOACS) -import qualified Futhark.Representation.Kernels as Out -import Futhark.Representation.Kernels.Kernel -import Futhark.MonadFreshNames -import Futhark.Tools -import qualified Futhark.Transform.FirstOrderTransform as FOT -import qualified Futhark.Pass.ExtractKernels.Kernelise as Kernelise -import Futhark.Transform.Rename -import Futhark.Pass -import Futhark.Transform.CopyPropagate -import Futhark.Pass.ExtractKernels.Distribution -import Futhark.Pass.ExtractKernels.ISRWIM -import Futhark.Pass.ExtractKernels.BlockedKernel -import Futhark.Pass.ExtractKernels.Interchange -import Futhark.Pass.ExtractKernels.Intragroup -import Futhark.Util -import Futhark.Util.Log - -type KernelsStms = Out.Stms Out.Kernels -type InKernelStms = Out.Stms Out.InKernel -type InKernelLambda = Out.Lambda Out.InKernel - --- | Transform a program using SOACs to a program using explicit --- kernels, using the kernel extraction transformation. -extractKernels :: Pass SOACS Out.Kernels -extractKernels = - Pass { passName = "extract kernels" - , passDescription = "Perform kernel extraction" - , passFunction = fmap Prog . mapM transformFunDef . progFunctions - } - --- In order to generate more stable threshold names, we keep track of --- the numbers used for thresholds separately from the ordinary name --- source, -data State = State { stateNameSource :: VNameSource - , stateThresholdCounter :: Int - } - -newtype DistribM a = DistribM (RWS (Scope Out.Kernels) Log State a) - deriving (Functor, Applicative, Monad, - HasScope Out.Kernels, LocalScope Out.Kernels, - MonadState State, - MonadLogger) - -instance MonadFreshNames DistribM where - getNameSource = gets stateNameSource - putNameSource src = modify $ \s -> s { stateNameSource = src } - -runDistribM :: (MonadLogger m, MonadFreshNames m) => - DistribM a -> m a -runDistribM (DistribM m) = do - (x, msgs) <- modifyNameSource $ \src -> - let (x, s, msgs) = runRWS m mempty (State src 0) - in ((x, msgs), stateNameSource s) - addLog msgs - return x - -transformFunDef :: (MonadFreshNames m, MonadLogger m) => - FunDef -> m (Out.FunDef Out.Kernels) -transformFunDef (FunDef entry name rettype params body) = runDistribM $ do - body' <- localScope (scopeOfFParams params) $ - transformBody mempty body - return $ FunDef entry name rettype params body' - -transformBody :: KernelPath -> Body -> DistribM (Out.Body Out.Kernels) -transformBody path body = do bnds <- transformStms path $ stmsToList $ bodyStms body - return $ mkBody bnds $ bodyResult body - -transformStms :: KernelPath -> [Stm] -> DistribM KernelsStms -transformStms _ [] = - return mempty -transformStms path (bnd:bnds) = - sequentialisedUnbalancedStm bnd >>= \case - Nothing -> do - bnd' <- transformStm path bnd - inScopeOf bnd' $ - (bnd'<>) <$> transformStms path bnds - Just bnds' -> - transformStms path $ stmsToList bnds' <> bnds - -sequentialisedUnbalancedStm :: Stm -> DistribM (Maybe (Stms SOACS)) -sequentialisedUnbalancedStm (Let pat _ (Op soac@(Screma _ form _))) - | Just (_, lam2) <- isRedomapSOAC form, - unbalancedLambda lam2, lambdaContainsParallelism lam2 = do - types <- asksScope scopeForSOACs - Just . snd <$> runBinderT (FOT.transformSOAC pat soac) types -sequentialisedUnbalancedStm _ = - return Nothing - -scopeForSOACs :: Scope Out.Kernels -> Scope SOACS -scopeForSOACs = castScope - -scopeForKernels :: Scope SOACS -> Scope Out.Kernels -scopeForKernels = castScope - -transformStm :: KernelPath -> Stm -> DistribM KernelsStms - -transformStm path (Let pat aux (Op (CmpThreshold what s))) = do - ((r, _), stms) <- cmpSizeLe s (Out.SizeThreshold path) what - runBinder_ $ do - addStms stms - addStm $ Let pat aux $ BasicOp $ SubExp r - -transformStm path (Let pat aux (If c tb fb rt)) = do - tb' <- transformBody path tb - fb' <- transformBody path fb - return $ oneStm $ Let pat aux $ If c tb' fb' rt - -transformStm path (Let pat aux (DoLoop ctx val form body)) = - localScope (castScope (scopeOf form) <> - scopeOfFParams mergeparams) $ - oneStm . Let pat aux . DoLoop ctx val form' <$> transformBody path body - where mergeparams = map fst $ ctx ++ val - form' = case form of - WhileLoop cond -> - WhileLoop cond - ForLoop i it bound ps -> - ForLoop i it bound ps - -transformStm path (Let pat (StmAux cs _) (Op (Screma w form arrs))) - | Just lam <- isMapSOAC form = - distributeMap path $ MapLoop pat cs w lam arrs - -transformStm path (Let res_pat (StmAux cs _) (Op (Screma w form arrs))) - | Just (scan_lam, nes) <- isScanSOAC form, - Just do_iswim <- iswim res_pat w scan_lam $ zip nes arrs = do - types <- asksScope scopeForSOACs - transformStms path =<< (stmsToList . snd <$> runBinderT (certifying cs do_iswim) types) - - -- We are only willing to generate code for scanomaps that do not - -- involve array accumulators, and do not have parallelism in their - -- map function. Such cases will fall through to the - -- screma-splitting case, and produce an ordinary map and scan. - -- Hopefully, the scan then triggers the ISWIM case above (otherwise - -- we will still crash in code generation). - | Just (scan_lam, nes, map_lam) <- isScanomapSOAC form, - all primType $ lambdaReturnType scan_lam, - not $ lambdaContainsParallelism map_lam = do - scan_lam_sequential <- Kernelise.transformLambda scan_lam - map_lam_sequential <- Kernelise.transformLambda map_lam - segScan res_pat w w scan_lam_sequential map_lam_sequential nes arrs [] [] - -transformStm path (Let res_pat (StmAux cs _) (Op (Screma w form arrs))) - | Just [Reduce comm red_fun nes] <- isReduceSOAC form, - let comm' | commutativeLambda red_fun = Commutative - | otherwise = comm, - Just do_irwim <- irwim res_pat w comm' red_fun $ zip nes arrs = do - types <- asksScope scopeForSOACs - bnds <- fst <$> runBinderT (simplifyStms =<< collectStms_ (certifying cs do_irwim)) types - transformStms path $ stmsToList bnds - -transformStm path (Let pat (StmAux cs _) (Op (Screma w form arrs))) - | Just (reds, map_lam) <- isRedomapSOAC form = do - - let paralleliseOuter = runBinder_ $ do - red_ops <- forM reds $ \(Reduce comm red_lam nes) -> do - (red_lam', nes', shape) <- determineReduceOp red_lam nes - let comm' | commutativeLambda red_lam = Commutative - | otherwise = comm - return $ SegRedOp comm' red_lam' nes' shape - map_lam_sequential <- Kernelise.transformLambda map_lam - addStms =<< - (fmap (certify cs) <$> nonSegRed pat w red_ops map_lam_sequential arrs) - - outerParallelBody = - renameBody =<< - (mkBody <$> paralleliseOuter <*> pure (map Var (patternNames pat))) - - paralleliseInner path' = do - (mapbnd, redbnd) <- redomapToMapAndReduce pat (w, comm', red_lam, map_lam, nes, arrs) - transformStms path' [certify cs mapbnd, certify cs redbnd] - where comm' | commutativeLambda red_lam = Commutative - | otherwise = comm - (Reduce comm red_lam nes) = singleReduce reds - - innerParallelBody path' = - renameBody =<< - (mkBody <$> paralleliseInner path' <*> pure (map Var (patternNames pat))) - - if not $ lambdaContainsParallelism map_lam - then paralleliseOuter - else if incrementalFlattening then do - ((outer_suff, outer_suff_key), suff_stms) <- - sufficientParallelism "suff_outer_redomap" w path - - outer_stms <- outerParallelBody - inner_stms <- innerParallelBody ((outer_suff_key, False):path) - - (suff_stms<>) <$> kernelAlternatives pat inner_stms [(outer_suff, outer_stms)] - else paralleliseOuter - --- Streams can be handled in two different ways - either we --- sequentialise the body or we keep it parallel and distribute. -transformStm path (Let pat (StmAux cs _) (Op (Stream w (Parallel _ _ _ []) map_fun arrs))) = do - -- No reduction part. Remove the stream and leave the body - -- parallel. It will be distributed. - types <- asksScope scopeForSOACs - transformStms path =<< - (stmsToList . snd <$> runBinderT (certifying cs $ sequentialStreamWholeArray pat w [] map_fun arrs) types) - -transformStm path (Let pat aux@(StmAux cs _) (Op (Stream w (Parallel o comm red_fun nes) fold_fun arrs))) - | incrementalFlattening = do - ((outer_suff, outer_suff_key), suff_stms) <- - sufficientParallelism "suff_outer_stream" w path - - outer_stms <- outerParallelBody ((outer_suff_key, True) : path) - inner_stms <- innerParallelBody ((outer_suff_key, False) : path) - - (suff_stms<>) <$> kernelAlternatives pat inner_stms [(outer_suff, outer_stms)] - - | otherwise = paralleliseOuter path - - where - paralleliseOuter path' - | any (not . primType) $ lambdaReturnType red_fun = do - -- Split into a chunked map and a reduction, with the latter - -- further transformed. - fold_fun_sequential <- Kernelise.transformLambda fold_fun - - let (red_pat_elems, concat_pat_elems) = - splitAt (length nes) $ patternValueElements pat - red_pat = Pattern [] red_pat_elems - - ((num_threads, red_results), stms) <- - streamMap (map (baseString . patElemName) red_pat_elems) concat_pat_elems w - Noncommutative fold_fun_sequential nes arrs - - reduce_soac <- reduceSOAC [Reduce comm' red_fun nes] - - (stms<>) <$> - inScopeOf stms - (transformStm path' $ Let red_pat aux $ - Op (Screma num_threads reduce_soac red_results)) - - | otherwise = do - red_fun_sequential <- Kernelise.transformLambda red_fun - fold_fun_sequential <- Kernelise.transformLambda fold_fun - fmap (certify cs) <$> - streamRed pat w comm' red_fun_sequential fold_fun_sequential nes arrs - - outerParallelBody path' = - renameBody =<< - (mkBody <$> paralleliseOuter path' <*> pure (map Var (patternNames pat))) - - paralleliseInner path' = do - types <- asksScope scopeForSOACs - transformStms path' . fmap (certify cs) =<< - (stmsToList . snd <$> runBinderT (sequentialStreamWholeArray pat w nes fold_fun arrs) types) - - innerParallelBody path' = - renameBody =<< - (mkBody <$> paralleliseInner path' <*> pure (map Var (patternNames pat))) - - comm' | commutativeLambda red_fun, o /= InOrder = Commutative - | otherwise = comm - -transformStm path (Let pat (StmAux cs _) (Op (Screma w form arrs))) = do - -- This screma is too complicated for us to immediately do - -- anything, so split it up and try again. - scope <- asksScope scopeForSOACs - transformStms path . map (certify cs) . stmsToList . snd =<< - runBinderT (dissectScrema pat w form arrs) scope - -transformStm path (Let pat _ (Op (Stream w (Sequential nes) fold_fun arrs))) = do - -- Remove the stream and leave the body parallel. It will be - -- distributed. - types <- asksScope scopeForSOACs - transformStms path =<< - (stmsToList . snd <$> runBinderT (sequentialStreamWholeArray pat w nes fold_fun arrs) types) - -transformStm _ (Let pat (StmAux cs _) (Op (Scatter w lam ivs as))) = runBinder_ $ do - lam' <- Kernelise.transformLambda lam - write_i <- newVName "write_i" - let (as_ws, as_ns, as_vs) = unzip3 as - (i_res, v_res) = splitAt (sum as_ns) $ bodyResult $ lambdaBody lam' - kstms = bodyStms $ lambdaBody lam' - krets = do (a_w, a, is_vs) <- zip3 as_ws as_vs $ chunks as_ns $ zip i_res v_res - return $ WriteReturn [a_w] a [ ([i],v) | (i,v) <- is_vs ] - body = KernelBody () kstms krets - inputs = do (p, p_a) <- zip (lambdaParams lam') ivs - return $ KernelInput (paramName p) (paramType p) p_a [Var write_i] - (bnds, kernel) <- - mapKernel w (FlatThreadSpace [(write_i,w)]) inputs (map rowType $ patternTypes pat) body - certifying cs $ do - addStms bnds - letBind_ pat $ Op $ HostOp kernel - -transformStm _ (Let orig_pat (StmAux cs _) (Op (GenReduce w ops bucket_fun imgs))) = do - bfun' <- Kernelise.transformLambda bucket_fun - genReduceKernel orig_pat [] [] cs w ops bfun' imgs - -transformStm _ bnd = - runBinder_ $ FOT.transformStmRecursively bnd - -data MapLoop = MapLoop Pattern Certificates SubExp Lambda [VName] - -mapLoopStm :: MapLoop -> Stm -mapLoopStm (MapLoop pat cs w lam arrs) = Let pat (StmAux cs ()) $ Op $ Screma w (mapSOAC lam) arrs - -sufficientParallelism :: String -> SubExp -> KernelPath - -> DistribM ((SubExp, Name), Out.Stms Out.Kernels) -sufficientParallelism desc what path = cmpSizeLe desc (Out.SizeThreshold path) what - -distributeMap :: KernelPath -> MapLoop -> DistribM KernelsStms -distributeMap path (MapLoop pat cs w lam arrs) = do - types <- askScope - let loopnest = MapNesting pat cs w $ zip (lambdaParams lam) arrs - env path' = KernelEnv { kernelNest = - singleNesting (Nesting mempty loopnest) - , kernelScope = - scopeForKernels (scopeOf lam) <> types - , kernelPath = - path' - } - exploitInnerParallelism path' = do - (acc', postkernels) <- runKernelM (env path') $ - distribute =<< distributeMapBodyStms acc (bodyStms $ lambdaBody lam) - - -- There may be a few final targets remaining - these correspond to - -- arrays that are identity mapped, and must have statements - -- inserted here. - return $ postKernelsStms postkernels <> - identityStms (outerTarget $ kernelTargets acc') - - if not incrementalFlattening then exploitInnerParallelism path - else do - - let exploitOuterParallelism path' = do - soactypes <- asksScope scopeForSOACs - (seq_lam, _) <- runBinderT (Kernelise.transformLambda lam) soactypes - (acc', postkernels) <- runKernelM (env path') $ distribute $ - addStmsToKernel (bodyStms $ lambdaBody seq_lam) acc - -- As above, we deal with identity mappings. - return $ postKernelsStms postkernels <> - identityStms (outerTarget $ kernelTargets acc') - - distributeMap' id (newKernel loopnest) path exploitOuterParallelism exploitInnerParallelism pat w lam - where acc = KernelAcc { kernelTargets = singleTarget (pat, bodyResult $ lambdaBody lam) - , kernelStms = mempty - } - - params_to_arrs = zip (map paramName $ lambdaParams lam) arrs - identityStms (rem_pat, res) = - stmsFromList $ zipWith identityStm (patternValueElements rem_pat) res - identityStm pe (Var v) - | Just arr <- lookup v params_to_arrs = - Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ Copy arr - identityStm pe se = - Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ Replicate (Shape [w]) se - -distributeMap' :: (HasScope Out.Kernels m, MonadFreshNames m) => - (forall a. DistribM a -> m a) - -> KernelNest -> KernelPath - -> (KernelPath -> m (Out.Stms Out.Kernels)) - -> (KernelPath -> m (Out.Stms Out.Kernels)) - -> PatternT Type - -> SubExp - -> LambdaT SOACS - -> m (Out.Stms Out.Kernels) -distributeMap' distribM loopnest path mk_seq_stms mk_par_stms pat nest_w lam = do - let res = map Var $ patternNames pat - - types <- askScope - ((outer_suff, outer_suff_key), outer_suff_stms) <- - distribM $ sufficientParallelism "suff_outer_par" nest_w path - - intra <- if worthIntraGroup lam then - flip runReaderT types $ intraGroupParallelise loopnest lam - else return Nothing - seq_body <- renameBody =<< mkBody <$> - mk_seq_stms ((outer_suff_key, True) : path) <*> pure res - let seq_alts = [(outer_suff, seq_body) | worthSequentialising lam] - - case intra of - Nothing -> do - par_body <- renameBody =<< mkBody <$> - mk_par_stms ((outer_suff_key, False) : path) <*> pure res - - (outer_suff_stms<>) <$> kernelAlternatives pat par_body seq_alts - - Just ((_intra_min_par, intra_avail_par), group_size, intra_prelude, intra_stms) -> do - -- We must check that all intra-group parallelism fits in a group. - ((intra_ok, intra_suff_key), intra_suff_stms) <- do - - ((intra_suff, suff_key), check_suff_stms) <- - distribM $ sufficientParallelism "suff_intra_par" intra_avail_par $ - (outer_suff_key, False) : path - - runBinder $ do - - addStms intra_prelude - - max_group_size <- - letSubExp "max_group_size" $ Op $ Out.GetSizeMax Out.SizeGroup - fits <- letSubExp "fits" $ BasicOp $ - CmpOp (CmpSle Int32) group_size max_group_size - - addStms check_suff_stms - - intra_ok <- letSubExp "intra_suff_and_fits" $ BasicOp $ BinOp LogAnd fits intra_suff - return (intra_ok, suff_key) - - group_par_body <- renameBody $ mkBody intra_stms res - - par_body <- renameBody =<< mkBody <$> - mk_par_stms ([(outer_suff_key, False), - (intra_suff_key, False)] - ++ path) <*> pure res - - ((outer_suff_stms<>intra_suff_stms)<>) <$> - kernelAlternatives pat par_body (seq_alts ++ [(intra_ok, group_par_body)]) - -data KernelEnv = KernelEnv { kernelNest :: Nestings - , kernelScope :: Scope Out.Kernels - , kernelPath :: KernelPath - } - -data KernelAcc = KernelAcc { kernelTargets :: Targets - , kernelStms :: InKernelStms - } - -data KernelRes = KernelRes { accPostKernels :: PostKernels - , accLog :: Log - } - -instance Semigroup KernelRes where - KernelRes ks1 log1 <> KernelRes ks2 log2 = - KernelRes (ks1 <> ks2) (log1 <> log2) - -instance Monoid KernelRes where - mempty = KernelRes mempty mempty - -newtype PostKernel = PostKernel { unPostKernel :: KernelsStms } - -newtype PostKernels = PostKernels [PostKernel] - -instance Semigroup PostKernels where - PostKernels xs <> PostKernels ys = PostKernels $ ys ++ xs - -instance Monoid PostKernels where - mempty = PostKernels mempty - -postKernelsStms :: PostKernels -> KernelsStms -postKernelsStms (PostKernels kernels) = mconcat $ map unPostKernel kernels - -typeEnvFromKernelAcc :: KernelAcc -> Scope Out.Kernels -typeEnvFromKernelAcc = scopeOfPattern . fst . outerTarget . kernelTargets - -addStmsToKernel :: InKernelStms -> KernelAcc -> KernelAcc -addStmsToKernel stms acc = - acc { kernelStms = stms <> kernelStms acc } - -addStmToKernel :: (LocalScope Out.Kernels m, MonadFreshNames m) => - Stm -> KernelAcc -> m KernelAcc -addStmToKernel bnd acc = do - stms <- runBinder_ $ Kernelise.transformStm bnd - return acc { kernelStms = stms <> kernelStms acc } - -newtype KernelM a = KernelM (ReaderT KernelEnv (WriterT KernelRes DistribM) a) - deriving (Functor, Applicative, Monad, - MonadReader KernelEnv, - MonadWriter KernelRes) - -liftDistribM :: DistribM a -> KernelM a -liftDistribM m = do - scope <- askScope - KernelM $ lift $ lift $ localScope scope m - -instance MonadFreshNames KernelM where - getNameSource = KernelM $ lift getNameSource - putNameSource = KernelM . lift . putNameSource - -instance HasScope Out.Kernels KernelM where - askScope = asks kernelScope - -instance LocalScope Out.Kernels KernelM where - localScope types = local $ \env -> - env { kernelScope = types <> kernelScope env } - -instance MonadLogger KernelM where - addLog msgs = tell mempty { accLog = msgs } - -runKernelM :: KernelEnv -> KernelM a -> DistribM (a, PostKernels) -runKernelM env (KernelM m) = do - (x, res) <- runWriterT $ runReaderT m env - addLog $ accLog res - return (x, accPostKernels res) - -collectKernels :: KernelM a -> KernelM (a, PostKernels) -collectKernels m = pass $ do - (x, res) <- listen m - return ((x, accPostKernels res), - const res { accPostKernels = mempty }) - -collectKernels_ :: KernelM () -> KernelM PostKernels -collectKernels_ = fmap snd . collectKernels - -localPath :: KernelPath -> KernelM a -> KernelM a -localPath path = local $ \env -> env { kernelPath = path } - -addKernels :: PostKernels -> KernelM () -addKernels ks = tell $ mempty { accPostKernels = ks } - -addKernel :: KernelsStms -> KernelM () -addKernel bnds = addKernels $ PostKernels [PostKernel bnds] - -withStm :: Stm -> KernelM a -> KernelM a -withStm bnd = local $ \env -> - env { kernelScope = - scopeForKernels (scopeOf [bnd]) <> kernelScope env - , kernelNest = - letBindInInnerNesting provided $ - kernelNest env - } - where provided = S.fromList $ patternNames $ stmPattern bnd - -mapNesting :: Pattern -> Certificates -> SubExp -> Lambda -> [VName] - -> KernelM a - -> KernelM a -mapNesting pat cs w lam arrs = local $ \env -> - env { kernelNest = pushInnerNesting nest $ kernelNest env - , kernelScope = scopeForKernels (scopeOf lam) <> kernelScope env - } - where nest = Nesting mempty $ - MapNesting pat cs w $ - zip (lambdaParams lam) arrs - -inNesting :: KernelNest -> KernelM a -> KernelM a -inNesting (outer, nests) = local $ \env -> - env { kernelNest = (inner, nests') - , kernelScope = mconcat (map scopeOf $ outer : nests) <> kernelScope env - } - where (inner, nests') = - case reverse nests of - [] -> (asNesting outer, []) - (inner' : ns) -> (asNesting inner', map asNesting $ outer : reverse ns) - asNesting = Nesting mempty - -unbalancedLambda :: Lambda -> Bool -unbalancedLambda lam = - unbalancedBody - (S.fromList $ map paramName $ lambdaParams lam) $ - lambdaBody lam - - where subExpBound (Var i) bound = i `S.member` bound - subExpBound (Constant _) _ = False - - unbalancedBody bound body = - any (unbalancedStm (bound <> boundInBody body) . stmExp) $ - bodyStms body - - -- XXX - our notion of balancing is probably still too naive. - unbalancedStm bound (Op (Stream w _ _ _)) = - w `subExpBound` bound - unbalancedStm bound (Op (Screma w _ _)) = - w `subExpBound` bound - unbalancedStm _ Op{} = - False - unbalancedStm _ DoLoop{} = False - - unbalancedStm bound (If cond tbranch fbranch _) = - cond `subExpBound` bound && - (unbalancedBody bound tbranch || unbalancedBody bound fbranch) - - unbalancedStm _ (BasicOp _) = - False - unbalancedStm _ (Apply fname _ _ _) = - not $ isBuiltInFunction fname - -bodyContainsParallelism :: Body -> Bool -bodyContainsParallelism = any (isMap . stmExp) . bodyStms - where isMap Op{} = True - isMap _ = False - -lambdaContainsParallelism :: Lambda -> Bool -lambdaContainsParallelism = bodyContainsParallelism . lambdaBody - --- | Returns the sizes of nested parallelism. -nestedParallelism :: Body -> [SubExp] -nestedParallelism = concatMap (parallelism . stmExp) . bodyStms - where parallelism (Op (Scatter w _ _ _)) = [w] - parallelism (Op (Screma w _ _)) = [w] - parallelism (Op (Stream w Sequential{} lam _)) - | chunk_size_param : _ <- lambdaParams lam = - let update (Var v) | v == paramName chunk_size_param = w - update se = se - in map update $ nestedParallelism $ lambdaBody lam - parallelism (DoLoop _ _ _ body) = nestedParallelism body - parallelism _ = [] - --- | A lambda is worth sequentialising if it contains nested --- parallelism of an interesting kind. -worthSequentialising :: Lambda -> Bool -worthSequentialising lam = interesting $ lambdaBody lam - where interesting body = any (interesting' . stmExp) $ bodyStms body - interesting' (Op (Screma _ form@(ScremaForm _ _ lam') _)) - | isJust $ isMapSOAC form = worthSequentialising lam' - interesting' (Op Scatter{}) = False -- Basically a map. - interesting' (DoLoop _ _ _ body) = interesting body - interesting' (Op _) = True - interesting' _ = False - --- | Intra-group parallelism is worthwhile if the lambda contains --- non-map nested parallelism, or any nested parallelism inside a --- loop. -worthIntraGroup :: Lambda -> Bool -worthIntraGroup lam = interesting $ lambdaBody lam - where interesting body = not (null $ nestedParallelism body) && - not (onlyMaps $ bodyStms body) - onlyMaps = all $ isMapOrSeq . stmExp - isMapOrSeq (Op (Screma _ form@(ScremaForm _ _ lam') _)) - | isJust $ isMapSOAC form = not $ worthIntraGroup lam' - isMapOrSeq (Op Scatter{}) = True -- Basically a map. - isMapOrSeq (DoLoop _ _ _ body) = - null $ nestedParallelism body - isMapOrSeq (Op _) = False - isMapOrSeq _ = True - --- Enable if you want the cool new versioned code. Beware: may be --- slower in practice. Caveat emptor (and you are the emptor). -incrementalFlattening :: Bool -incrementalFlattening = isJust $ lookup "FUTHARK_INCREMENTAL_FLATTENING" unixEnvironment - -distributeInnerMap :: MapLoop -> KernelAcc - -> KernelM KernelAcc -distributeInnerMap maploop@(MapLoop pat cs w lam arrs) acc - | unbalancedLambda lam, lambdaContainsParallelism lam = - addStmToKernel (mapLoopStm maploop) acc - | not incrementalFlattening = - distributeNormally - | otherwise = - distributeSingleStm acc (mapLoopStm maploop) >>= \case - Just (post_kernels, res, nest, acc') - | Just (perm, _pat_unused) <- permutationAndMissing pat res -> do - addKernels post_kernels - multiVersion perm nest acc' - _ -> distributeNormally - where - lam_bnds = bodyStms $ lambdaBody lam - lam_res = bodyResult $ lambdaBody lam - - def_acc = KernelAcc { kernelTargets = pushInnerTarget - (pat, bodyResult $ lambdaBody lam) $ - kernelTargets acc - , kernelStms = mempty - } - - distributeNormally = - distribute =<< - leavingNesting maploop =<< - mapNesting pat cs w lam arrs - (distribute =<< distributeMapBodyStms def_acc lam_bnds) - - multiVersion perm nest acc' = do - -- The kernel can be distributed by itself, so now we can - -- decide whether to just sequentialise, or exploit inner - -- parallelism. - let map_nesting = MapNesting pat cs w $ zip (lambdaParams lam) arrs - -- Normally the permutation is for the output pattern, but - -- we can't really change that, so we change the result - -- order instead. - lam_res' = rearrangeShape (rearrangeInverse perm) lam_res - nest' = pushInnerKernelNesting (pat, lam_res') map_nesting nest - extra_scope = targetsScope $ kernelTargets acc' - - exploitInnerParallelism path' = - fmap postKernelsStms $ collectKernels_ $ localPath path' $ - localScope extra_scope $ inNesting nest' $ void $ - distribute =<< leavingNesting maploop =<< distribute =<< - distributeMapBodyStms def_acc lam_bnds - - -- XXX: we do not construct a new KernelPath when - -- sequentialising. This is only OK as long as further - -- versioning does not take place down that branch (it currently - -- does not). - (nestw_bnds, nestw, sequentialised_kernel) <- localScope extra_scope $ do - sequentialised_map_body <- - localScope (scopeOfLParams (lambdaParams lam)) $ runBinder_ $ - Kernelise.transformStms lam_bnds - let kbody = KernelBody () sequentialised_map_body $ - map ThreadsReturn lam_res' - constructKernel nest' kbody - - let outer_pat = loopNestingPattern $ fst nest - path <- asks kernelPath - addKernel =<< (nestw_bnds<>) <$> - localScope extra_scope (distributeMap' liftDistribM nest' path - (const $ return $ oneStm sequentialised_kernel) - exploitInnerParallelism - outer_pat nestw - lam { lambdaBody = (lambdaBody lam) { bodyResult = lam_res' }}) - - return acc' - -leavingNesting :: MapLoop -> KernelAcc -> KernelM KernelAcc -leavingNesting (MapLoop _ cs w lam arrs) acc = - case popInnerTarget $ kernelTargets acc of - Nothing -> - fail "The kernel targets list is unexpectedly small" - Just ((pat,res), newtargets) -> do - let acc' = acc { kernelTargets = newtargets } - if null $ kernelStms acc' - then return acc' - else do let kbody = Body () (kernelStms acc') res - used_in_body = freeIn kbody - (used_params, used_arrs) = - unzip $ - filter ((`S.member` used_in_body) . paramName . fst) $ - zip (lambdaParams lam) arrs - stms <- runBinder_ $ Kernelise.mapIsh pat cs w used_params kbody used_arrs - return $ addStmsToKernel stms acc' { kernelStms = mempty } - -distributeMapBodyStms :: KernelAcc -> Stms SOACS -> KernelM KernelAcc -distributeMapBodyStms orig_acc = onStms orig_acc . stmsToList - where - onStms acc [] = return acc - - onStms acc (Let pat (StmAux cs _) (Op (Stream w (Sequential accs) lam arrs)):stms) = do - types <- asksScope scopeForSOACs - stream_stms <- - snd <$> runBinderT (sequentialStreamWholeArray pat w accs lam arrs) types - stream_stms' <- - runReaderT (copyPropagateInStms simpleSOACS stream_stms) types - onStms acc $ stmsToList (fmap (certify cs) stream_stms') ++ stms - - onStms acc (stm:stms) = - -- It is important that stm is in scope if 'maybeDistributeStm' - -- wants to distribute, even if this causes the slightly silly - -- situation that stm is in scope of itself. - withStm stm $ maybeDistributeStm stm =<< onStms acc stms - -maybeDistributeStm :: Stm -> KernelAcc -> KernelM KernelAcc - -maybeDistributeStm bnd@(Let pat _ (Op (Screma w form arrs))) acc - | Just lam <- isMapSOAC form = - -- Only distribute inside the map if we can distribute everything - -- following the map. - distributeIfPossible acc >>= \case - Nothing -> addStmToKernel bnd acc - Just acc' -> distribute =<< distributeInnerMap (MapLoop pat (stmCerts bnd) w lam arrs) acc' - -maybeDistributeStm bnd@(Let pat _ (DoLoop [] val form@ForLoop{} body)) acc - | null (patternContextElements pat), bodyContainsParallelism body = - distributeSingleStm acc bnd >>= \case - Just (kernels, res, nest, acc') - | S.null $ freeIn form `S.intersection` boundInKernelNest nest, - Just (perm, pat_unused) <- permutationAndMissing pat res -> - -- We need to pretend pat_unused was used anyway, by adding - -- it to the kernel nest. - localScope (typeEnvFromKernelAcc acc') $ do - addKernels kernels - nest' <- expandKernelNest pat_unused nest - types <- asksScope scopeForSOACs - - bnds <- runReaderT - (interchangeLoops nest' (SeqLoop perm pat val form body)) types - path <- asks kernelPath - bnds' <- liftDistribM $ transformStms path $ stmsToList bnds - addKernel bnds' - return acc' - _ -> - addStmToKernel bnd acc - -maybeDistributeStm stm@(Let pat _ (If cond tbranch fbranch ret)) acc - | null (patternContextElements pat), - bodyContainsParallelism tbranch || bodyContainsParallelism fbranch || - any (not . primType) (ifReturns ret) = - distributeSingleStm acc stm >>= \case - Just (kernels, res, nest, acc') - | S.null $ (freeIn cond <> freeIn ret) `S.intersection` - boundInKernelNest nest, - Just (perm, pat_unused) <- permutationAndMissing pat res -> - -- We need to pretend pat_unused was used anyway, by adding - -- it to the kernel nest. - localScope (typeEnvFromKernelAcc acc') $ do - nest' <- expandKernelNest pat_unused nest - addKernels kernels - types <- asksScope scopeForSOACs - let branch = Branch perm pat cond tbranch fbranch ret - stms <- runReaderT (interchangeBranch nest' branch) types - - path <- asks kernelPath - stms' <- liftDistribM $ transformStms path $ stmsToList stms - addKernel stms' - return acc' - _ -> - addStmToKernel stm acc - -maybeDistributeStm (Let pat (StmAux cs _) (Op (Screma w form arrs))) acc - | Just [Reduce comm lam nes] <- isReduceSOAC form, - Just m <- irwim pat w comm lam $ zip nes arrs = do - types <- asksScope scopeForSOACs - (_, bnds) <- runBinderT (certifying cs m) types - distributeMapBodyStms acc bnds - --- Parallelise segmented scatters. -maybeDistributeStm bnd@(Let pat (StmAux cs _) (Op (Scatter w lam ivs as))) acc = - distributeSingleStm acc bnd >>= \case - Just (kernels, res, nest, acc') - | Just (perm, pat_unused) <- permutationAndMissing pat res -> - localScope (typeEnvFromKernelAcc acc') $ do - nest' <- expandKernelNest pat_unused nest - lam' <- Kernelise.transformLambda lam - addKernels kernels - addKernel =<< segmentedScatterKernel nest' perm pat cs w lam' ivs as - return acc' - _ -> - addStmToKernel bnd acc - --- Parallelise segmented GenReduce. -maybeDistributeStm bnd@(Let pat (StmAux cs _) (Op (GenReduce w ops lam as))) acc = - distributeSingleStm acc bnd >>= \case - Just (kernels, res, nest, acc') - | Just (perm, pat_unused) <- permutationAndMissing pat res -> - localScope (typeEnvFromKernelAcc acc') $ do - lam' <- Kernelise.transformLambda lam - nest' <- expandKernelNest pat_unused nest - addKernels kernels - addKernel =<< segmentedGenReduceKernel nest' perm cs w ops lam' as - return acc' - _ -> - addStmToKernel bnd acc - --- If the scan can be distributed by itself, we will turn it into a --- segmented scan. --- --- If the scan cannot be distributed by itself, it will be --- sequentialised in the default case for this function. -maybeDistributeStm bnd@(Let pat (StmAux cs _) (Op (Screma w form arrs))) acc - | Just (lam, nes, map_lam) <- isScanomapSOAC form = - distributeSingleStm acc bnd >>= \case - Just (kernels, res, nest, acc') - | Just (perm, pat_unused) <- permutationAndMissing pat res -> - -- We need to pretend pat_unused was used anyway, by adding - -- it to the kernel nest. - localScope (typeEnvFromKernelAcc acc') $ do - nest' <- expandKernelNest pat_unused nest - map_lam' <- Kernelise.transformLambda map_lam - lam' <- Kernelise.transformLambda lam - localScope (typeEnvFromKernelAcc acc') $ - segmentedScanomapKernel nest' perm w lam' map_lam' nes arrs >>= - kernelOrNot cs bnd acc kernels acc' - _ -> - addStmToKernel bnd acc - --- If the reduction can be distributed by itself, we will turn it into a --- segmented reduce. --- --- If the reduction cannot be distributed by itself, it will be --- sequentialised in the default case for this function. -maybeDistributeStm bnd@(Let pat (StmAux cs _) (Op (Screma w form arrs))) acc - | Just (reds, map_lam) <- isRedomapSOAC form, - Reduce comm lam nes <- singleReduce reds, - isIdentityLambda map_lam || incrementalFlattening = - distributeSingleStm acc bnd >>= \case - Just (kernels, res, nest, acc') - | Just (perm, pat_unused) <- permutationAndMissing pat res -> - -- We need to pretend pat_unused was used anyway, by adding - -- it to the kernel nest. - localScope (typeEnvFromKernelAcc acc') $ do - nest' <- expandKernelNest pat_unused nest - lam' <- Kernelise.transformLambda lam - map_lam' <- Kernelise.transformLambda map_lam - - let comm' | commutativeLambda lam = Commutative - | otherwise = comm - - regularSegmentedRedomapKernel nest' perm w comm' lam' map_lam' nes arrs >>= - kernelOrNot cs bnd acc kernels acc' - _ -> - addStmToKernel bnd acc - -maybeDistributeStm (Let pat (StmAux cs _) (Op (Screma w form arrs))) acc - | incrementalFlattening || isNothing (isRedomapSOAC form) = do - -- This with-loop is too complicated for us to immediately do - -- anything, so split it up and try again. - scope <- asksScope scopeForSOACs - distributeMapBodyStms acc . fmap (certify cs) . snd =<< - runBinderT (dissectScrema pat w form arrs) scope - -maybeDistributeStm (Let pat aux (BasicOp (Replicate (Shape (d:ds)) v))) acc - | [t] <- patternTypes pat = do - -- XXX: We need a temporary dummy binding to prevent an empty - -- map body. The kernel extractor does not like empty map - -- bodies. - tmp <- newVName "tmp" - let rowt = rowType t - newbnd = Let pat aux $ Op $ Screma d (mapSOAC lam) [] - tmpbnd = Let (Pattern [] [PatElem tmp rowt]) aux $ - BasicOp $ Replicate (Shape ds) v - lam = Lambda { lambdaReturnType = [rowt] - , lambdaParams = [] - , lambdaBody = mkBody (oneStm tmpbnd) [Var tmp] - } - maybeDistributeStm newbnd acc - -maybeDistributeStm bnd@(Let _ aux (BasicOp Copy{})) acc = - distributeSingleUnaryStm acc bnd $ \_ 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 - | not $ primType $ typeOf pe = - distributeSingleUnaryStm acc bnd $ \_ 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 - let r = length (snd nest) + 1 - perm' = [0..r-1] ++ map (+r) perm - -- We need to add a copy, because the original map nest - -- will have produced an array without aliases, and so must we. - arr' <- newVName $ baseString arr - arr_t <- lookupType arr - return $ stmsFromList - [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 - 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 - let rots' = map (const $ intConst Int32 0) (kernelNestWidths nest) ++ rots - return $ oneStm $ Let outerpat aux $ BasicOp $ Rotate rots' arr - --- XXX? This rule is present to avoid the case where an in-place --- update is distributed as its own kernel, as this would mean thread --- then writes the entire array that it updated. This is problematic --- because the in-place updates is O(1), but writing the array is --- O(n). It is OK if the in-place update is preceded, followed, or --- nested inside a sequential loop or similar, because that will --- probably be O(n) by itself. As a hack, we only distribute if there --- does not appear to be a loop following. The better solution is to --- depend on memory block merging for this optimisation, but it is not --- ready yet. -maybeDistributeStm (Let pat aux (BasicOp (Update arr [DimFix i] v))) acc - | [t] <- patternTypes pat, - arrayRank t == 1, - not $ any (amortises . stmExp) $ kernelStms acc = do - let w = arraySize 0 t - et = stripArray 1 t - lam = Lambda { lambdaParams = [] - , lambdaReturnType = [Prim int32, et] - , lambdaBody = mkBody mempty [i, v] } - maybeDistributeStm (Let pat aux $ Op $ Scatter (intConst Int32 1) lam [] [(w, 1, arr)]) acc - where amortises DoLoop{} = True - amortises Op{} = True - amortises _ = False - -maybeDistributeStm stm@(Let _ aux (BasicOp (Concat d x xs w))) acc = - distributeSingleStm acc stm >>= \case - Just (kernels, _, nest, acc') -> - localScope (typeEnvFromKernelAcc acc') $ - segmentedConcat nest >>= - kernelOrNot (stmAuxCerts aux) stm acc kernels acc' - _ -> - addStmToKernel stm acc - - where segmentedConcat nest = - isSegmentedOp nest [0] w mempty mempty [] (x:xs) $ - \pat _ _ _ _ (x':xs') _ -> - let d' = d + length (snd nest) + 1 - in addStm $ Let pat aux $ BasicOp $ Concat d' x' xs' w - -maybeDistributeStm bnd acc = - addStmToKernel bnd acc - -distributeSingleUnaryStm :: KernelAcc - -> Stm - -> (KernelNest -> Pattern -> VName -> KernelM (Stms Out.Kernels)) - -> KernelM KernelAcc -distributeSingleUnaryStm acc bnd f = - distributeSingleStm acc bnd >>= \case - Just (kernels, res, nest, acc') - | res == map Var (patternNames $ stmPattern bnd), - (outer, inners) <- nest, - [(arr_p, arr)] <- loopNestingParamsAndArrs outer, - boundInKernelNest nest `S.intersection` freeIn bnd - == S.singleton (paramName arr_p) -> do - addKernels kernels - let outerpat = loopNestingPattern $ fst nest - localScope (typeEnvFromKernelAcc acc') $ do - (arr', pre_stms) <- repeatMissing arr (outer:inners) - f_stms <- inScopeOf pre_stms $ f nest outerpat arr' - addKernel $ pre_stms <> f_stms - return acc' - _ -> addStmToKernel bnd acc - where -- | For an imperfectly mapped array, repeat the missing - -- dimensions to make it look like it was in fact perfectly - -- mapped. - repeatMissing arr inners = do - arr_t <- lookupType arr - let shapes = determineRepeats arr arr_t inners - if all (==Shape []) shapes then return (arr, mempty) - else do - let (outer_shapes, inner_shape) = repeatShapes shapes arr_t - arr_t' = repeatDims outer_shapes inner_shape arr_t - arr' <- newVName $ baseString arr - return (arr', oneStm $ Let (Pattern [] [PatElem arr' arr_t']) (defAux ()) $ - BasicOp $ Repeat outer_shapes inner_shape arr) - - determineRepeats arr arr_t nests - | (skipped, arr_nest:nests') <- break (hasInput arr) nests, - [(arr_p, _)] <- loopNestingParamsAndArrs arr_nest = - Shape (map loopNestingWidth skipped) : - determineRepeats (paramName arr_p) (rowType arr_t) nests' - | otherwise = - Shape (map loopNestingWidth nests) : replicate (arrayRank arr_t) (Shape []) - - hasInput arr nest - | [(_, arr')] <- loopNestingParamsAndArrs nest, arr' == arr = True - | otherwise = False - - -distribute :: KernelAcc -> KernelM KernelAcc -distribute acc = - fromMaybe acc <$> distributeIfPossible acc - -distributeIfPossible :: KernelAcc -> KernelM (Maybe KernelAcc) -distributeIfPossible acc = do - nest <- asks kernelNest - tryDistribute nest (kernelTargets acc) (kernelStms acc) >>= \case - Nothing -> return Nothing - Just (targets, kernel) -> do - addKernel kernel - return $ Just KernelAcc { kernelTargets = targets - , kernelStms = mempty - } - -distributeSingleStm :: KernelAcc -> Stm - -> KernelM (Maybe (PostKernels, Result, KernelNest, KernelAcc)) -distributeSingleStm acc bnd = do - nest <- asks kernelNest - tryDistribute nest (kernelTargets acc) (kernelStms acc) >>= \case - Nothing -> return Nothing - Just (targets, distributed_bnds) -> - tryDistributeStm nest targets bnd >>= \case - Nothing -> return Nothing - Just (res, targets', new_kernel_nest) -> - return $ Just (PostKernels [PostKernel distributed_bnds], - res, - new_kernel_nest, - KernelAcc { kernelTargets = targets' - , kernelStms = mempty - }) - -segmentedScatterKernel :: KernelNest - -> [Int] - -> Pattern - -> Certificates - -> SubExp - -> InKernelLambda - -> [VName] -> [(SubExp,Int,VName)] - -> KernelM KernelsStms -segmentedScatterKernel nest perm scatter_pat cs scatter_w lam ivs dests = do - -- We replicate some of the checking done by 'isSegmentedOp', but - -- things are different because a scatter is not a reduction or - -- scan. - -- - -- First, pretend that the scatter is also part of the nesting. The - -- KernelNest we produce here is technically not sensible, but it's - -- good enough for flatKernel to work. - let nest' = pushInnerKernelNesting (scatter_pat, bodyResult $ lambdaBody lam) - (MapNesting scatter_pat cs scatter_w $ zip (lambdaParams lam) ivs) nest - (nest_bnds, w, ispace, kernel_inps) <- flatKernel nest' - - let (as_ws, as_ns, as) = unzip3 dests - - -- The input/output arrays ('as') _must_ correspond to some kernel - -- input, or else the original nested scatter would have been - -- ill-typed. Find them. - as_inps <- mapM (findInput kernel_inps) as - - runBinder_ $ do - addStms nest_bnds - - let rts = concatMap (take 1) $ chunks as_ns $ - drop (sum as_ns) $ lambdaReturnType lam - (is,vs) = splitAt (sum as_ns) $ bodyResult $ lambdaBody lam - k_body = KernelBody () (bodyStms $ lambdaBody lam) $ - map (inPlaceReturn ispace) $ - zip3 as_ws as_inps $ chunks as_ns $ zip is vs - - (k_bnds, k) <- - mapKernel w (FlatThreadSpace ispace) kernel_inps rts k_body - - addStms k_bnds - - let pat = Pattern [] $ rearrangeShape perm $ - patternValueElements $ loopNestingPattern $ fst nest - - certifying cs $ letBind_ pat $ Op $ HostOp k - where findInput kernel_inps a = - maybe bad return $ find ((==a) . kernelInputName) kernel_inps - bad = fail "Ill-typed nested scatter encountered." - - inPlaceReturn ispace (aw, inp, is_vs) = - WriteReturn (init ws++[aw]) (kernelInputArray inp) - [ (map Var (init gtids)++[i], v) | (i,v) <- is_vs ] - where (gtids,ws) = unzip ispace - -segmentedGenReduceKernel :: KernelNest - -> [Int] - -> Certificates - -> SubExp - -> [SOAC.GenReduceOp SOACS] - -> InKernelLambda - -> [VName] - -> KernelM KernelsStms -segmentedGenReduceKernel nest perm cs genred_w ops lam arrs = do - -- We replicate some of the checking done by 'isSegmentedOp', but - -- things are different because a GenReduce is not a reduction or - -- scan. - (nest_stms, _, ispace, inputs) <- flatKernel nest - let orig_pat = Pattern [] $ rearrangeShape perm $ - patternValueElements $ loopNestingPattern $ fst nest - - -- The input/output arrays _must_ correspond to some kernel input, - -- or else the original nested GenReduce would have been ill-typed. - -- Find them. - ops' <- forM ops $ \(SOAC.GenReduceOp num_bins dests nes op) -> - SOAC.GenReduceOp num_bins - <$> mapM (fmap kernelInputArray . findInput inputs) dests - <*> pure nes - <*> pure op - liftDistribM $ (nest_stms<>) <$> - inScopeOf nest_stms - (genReduceKernel orig_pat ispace inputs cs genred_w ops' lam arrs) - where findInput kernel_inps a = - maybe bad return $ find ((==a) . kernelInputName) kernel_inps - bad = fail "Ill-typed nested GenReduce encountered." - -genReduceKernel :: Pattern -> [(VName, SubExp)] -> [KernelInput] - -> Certificates -> SubExp -> [SOAC.GenReduceOp SOACS] - -> InKernelLambda -> [VName] - -> DistribM KernelsStms -genReduceKernel orig_pat ispace inputs cs genred_w ops lam arrs = runBinder_ $ do - ops' <- forM ops $ \(SOAC.GenReduceOp num_bins dests nes op) -> do - (op', nes', shape) <- determineReduceOp op nes - return $ Out.GenReduceOp num_bins dests nes' shape op' - - let isDest = flip elem $ concatMap Out.genReduceDest ops' - inputs' = filter (not . isDest . kernelInputArray) inputs - - certifying cs $ - addStms =<< segGenRed orig_pat genred_w ispace inputs' ops' lam arrs - -determineReduceOp :: (MonadBinder m, Lore m ~ Out.Kernels) => - Lambda -> [SubExp] -> m (Out.Lambda Out.InKernel, [SubExp], Shape) -determineReduceOp lam nes = - -- FIXME? We are assuming that the accumulator is a replicate, and - -- we fish out its value in a gross way. - case mapM subExpVar nes of - Just ne_vs' -> do - let (shape, lam') = isVectorMap lam - nes' <- forM ne_vs' $ \ne_v -> do - ne_v_t <- lookupType ne_v - letSubExp "genred_ne" $ - BasicOp $ Index ne_v $ fullSlice ne_v_t $ - replicate (shapeRank shape) $ DimFix $ intConst Int32 0 - lam'' <- Kernelise.transformLambda lam' - return (lam'', nes', shape) - Nothing -> do - lam' <- Kernelise.transformLambda lam - return (lam', nes, mempty) - -isVectorMap :: Lambda -> (Shape, Lambda) -isVectorMap lam - | [Let (Pattern [] pes) _ (Op (Screma w form arrs))] <- - stmsToList $ bodyStms $ lambdaBody lam, - bodyResult (lambdaBody lam) == map (Var . patElemName) pes, - Just map_lam <- isMapSOAC form, - arrs == map paramName (lambdaParams lam) = - let (shape, lam') = isVectorMap map_lam - in (Shape [w] <> shape, lam') - | otherwise = (mempty, lam) - -segmentedScanomapKernel :: KernelNest - -> [Int] - -> SubExp - -> InKernelLambda -> InKernelLambda - -> [SubExp] -> [VName] - -> KernelM (Maybe KernelsStms) -segmentedScanomapKernel nest perm segment_size lam map_lam nes arrs = - isSegmentedOp nest perm segment_size (freeIn lam) (freeIn map_lam) nes arrs $ - \pat total_num_elements ispace inps nes' _ _ -> - addStms =<< segScan pat total_num_elements segment_size lam map_lam nes' arrs ispace inps - -regularSegmentedRedomapKernel :: KernelNest - -> [Int] - -> SubExp -> Commutativity - -> InKernelLambda -> InKernelLambda -> [SubExp] -> [VName] - -> KernelM (Maybe KernelsStms) -regularSegmentedRedomapKernel nest perm segment_size comm lam map_lam nes arrs = - isSegmentedOp nest perm segment_size (freeIn lam) (freeIn map_lam) nes arrs $ - \pat total_num_elements ispace inps nes' _ _ -> do - let red_op = SegRedOp comm lam nes' mempty - addStms =<< segRed pat total_num_elements segment_size [red_op] map_lam arrs ispace inps - -isSegmentedOp :: KernelNest - -> [Int] - -> SubExp - -> Names -> Names - -> [SubExp] -> [VName] - -> (Pattern - -> SubExp - -> [(VName, SubExp)] - -> [KernelInput] - -> [SubExp] -> [VName] -> [VName] - -> Binder Out.Kernels ()) - -> KernelM (Maybe KernelsStms) -isSegmentedOp nest perm segment_size free_in_op _free_in_fold_op nes arrs m = runMaybeT $ do - -- We must verify that array inputs to the operation are inputs to - -- the outermost loop nesting or free in the loop nest. Nothing - -- free in the op may be bound by the nest. Furthermore, the - -- neutral elements must be free in the loop nest. - -- - -- We must summarise any names from free_in_op that are bound in the - -- nest, and describe how to obtain them given segment indices. - - let bound_by_nest = boundInKernelNest nest - - (pre_bnds, nesting_size, ispace, kernel_inps) <- flatKernel nest - - unless (S.null $ free_in_op `S.intersection` bound_by_nest) $ - fail "Non-fold lambda uses nest-bound parameters." - - let indices = map fst ispace - - prepareNe (Var v) | v `S.member` bound_by_nest = - fail "Neutral element bound in nest" - prepareNe ne = return ne - - prepareArr arr = - case find ((==arr) . kernelInputName) kernel_inps of - Just inp - | kernelInputIndices inp == map Var indices -> - return $ return $ kernelInputArray inp - | not (kernelInputArray inp `S.member` bound_by_nest) -> - return $ replicateMissing ispace inp - Nothing | not (arr `S.member` bound_by_nest) -> - -- This input is something that is free inside - -- the loop nesting. We will have to replicate - -- it. - return $ - letExp (baseString arr ++ "_repd") - (BasicOp $ Replicate (Shape [nesting_size]) $ Var arr) - _ -> - fail "Input not free or outermost." - - nes' <- mapM prepareNe nes - - mk_arrs <- mapM prepareArr arrs - - lift $ runBinder_ $ do - addStms pre_bnds - - -- We must make sure all inputs are of size - -- segment_size*nesting_size. - total_num_elements <- - letSubExp "total_num_elements" $ BasicOp $ BinOp (Mul Int32) segment_size nesting_size - - let flatten arr = do - arr_shape <- arrayShape <$> lookupType arr - -- CHECKME: is the length the right thing here? We want to - -- reproduce the parameter type. - let reshape = reshapeOuter [DimNew total_num_elements] - (2+length (snd nest)) arr_shape - letExp (baseString arr ++ "_flat") $ - BasicOp $ Reshape reshape arr - - nested_arrs <- sequence mk_arrs - arrs' <- mapM flatten nested_arrs - - let pat = Pattern [] $ rearrangeShape perm $ - patternValueElements $ loopNestingPattern $ fst nest - - m pat total_num_elements ispace kernel_inps nes' nested_arrs arrs' - - where replicateMissing ispace inp = do - t <- lookupType $ kernelInputArray inp - let inp_is = kernelInputIndices inp - shapes = determineRepeats ispace inp_is - (outer_shapes, inner_shape) = repeatShapes shapes t - letExp "repeated" $ BasicOp $ - Repeat outer_shapes inner_shape $ kernelInputArray inp - - determineRepeats ispace (i:is) - | (skipped_ispace, ispace') <- span ((/=i) . Var . fst) ispace = - Shape (map snd skipped_ispace) : determineRepeats (drop 1 ispace') is - determineRepeats ispace _ = - [Shape $ map snd ispace] - -permutationAndMissing :: Pattern -> [SubExp] -> Maybe ([Int], [PatElem]) -permutationAndMissing pat res = do - let pes = patternValueElements pat - (_used,unused) = - partition ((`S.member` freeIn res) . patElemName) pes - res_expanded = res ++ map (Var . patElemName) unused - perm <- map (Var . patElemName) pes `isPermutationOf` res_expanded - return (perm, unused) - --- Add extra pattern elements to every kernel nesting level. -expandKernelNest :: MonadFreshNames m => - [PatElem] -> KernelNest -> m KernelNest -expandKernelNest pes (outer_nest, inner_nests) = do - let outer_size = loopNestingWidth outer_nest : - map loopNestingWidth inner_nests - inner_sizes = tails $ map loopNestingWidth inner_nests - outer_nest' <- expandWith outer_nest outer_size - inner_nests' <- zipWithM expandWith inner_nests inner_sizes - return (outer_nest', inner_nests') - where expandWith nest dims = do - pes' <- mapM (expandPatElemWith dims) pes - return nest { loopNestingPattern = - Pattern [] $ - patternElements (loopNestingPattern nest) <> pes' - } - - expandPatElemWith dims pe = do - name <- newVName $ baseString $ patElemName pe - return pe { patElemName = name - , patElemAttr = patElemType pe `arrayOfShape` Shape dims - } - -cmpSizeLe :: String -> Out.SizeClass -> SubExp - -> DistribM ((SubExp, Name), Out.Stms Out.Kernels) -cmpSizeLe desc size_class to_what = do - x <- gets stateThresholdCounter - modify $ \s -> s { stateThresholdCounter = x + 1} - let size_key = nameFromString $ desc ++ "_" ++ show x - runBinder $ do - cmp_res <- letSubExp desc $ Op $ CmpSizeLe size_key size_class to_what - return (cmp_res, size_key) - -kernelAlternatives :: (MonadFreshNames m, HasScope Out.Kernels m) => - Out.Pattern Out.Kernels - -> Out.Body Out.Kernels - -> [(SubExp, Out.Body Out.Kernels)] - -> m (Out.Stms Out.Kernels) -kernelAlternatives pat default_body [] = runBinder_ $ do - ses <- bodyBind default_body - forM_ (zip (patternNames pat) ses) $ \(name, se) -> - letBindNames_ [name] $ BasicOp $ SubExp se -kernelAlternatives pat default_body ((cond,alt):alts) = runBinder_ $ do - alts_pat <- fmap (Pattern []) $ forM (patternElements pat) $ \pe -> do - name <- newVName $ baseString $ patElemName pe - return pe { patElemName = name } - - alt_stms <- kernelAlternatives alts_pat default_body alts - let alt_body = mkBody alt_stms $ map Var $ patternValueNames alts_pat - - letBind_ pat $ If cond alt alt_body $ ifCommon $ patternTypes pat - -kernelOrNot :: Certificates -> Stm -> KernelAcc - -> PostKernels -> KernelAcc -> Maybe KernelsStms - -> KernelM KernelAcc -kernelOrNot cs bnd acc _ _ Nothing = - addStmToKernel (certify cs bnd) acc -kernelOrNot cs _ _ kernels acc' (Just bnds) = do - addKernels kernels - addKernel $ fmap (certify cs) bnds - return acc' +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE RankNTypes #-} +-- | Kernel extraction. +-- +-- In the following, I will use the term "width" to denote the amount +-- of immediate parallelism in a map - that is, the outer size of the +-- array(s) being used as input. +-- +-- = Basic Idea +-- +-- If we have: +-- +-- @ +-- map +-- map(f) +-- bnds_a... +-- map(g) +-- @ +-- +-- Then we want to distribute to: +-- +-- @ +-- map +-- map(f) +-- map +-- bnds_a +-- map +-- map(g) +-- @ +-- +-- But for now only if +-- +-- (0) it can be done without creating irregular arrays. +-- Specifically, the size of the arrays created by @map(f)@, by +-- @map(g)@ and whatever is created by @bnds_a@ that is also used +-- in @map(g)@, must be invariant to the outermost loop. +-- +-- (1) the maps are _balanced_. That is, the functions @f@ and @g@ +-- must do the same amount of work for every iteration. +-- +-- The advantage is that the map-nests containing @map(f)@ and +-- @map(g)@ can now be trivially flattened at no cost, thus exposing +-- more parallelism. Note that the @bnds_a@ map constitutes array +-- expansion, which requires additional storage. +-- +-- = Distributing Sequential Loops +-- +-- As a starting point, sequential loops are treated like scalar +-- expressions. That is, not distributed. However, sometimes it can +-- be worthwhile to distribute if they contain a map: +-- +-- @ +-- map +-- loop +-- map +-- map +-- @ +-- +-- If we distribute the loop and interchange the outer map into the +-- loop, we get this: +-- +-- @ +-- loop +-- map +-- map +-- map +-- map +-- @ +-- +-- Now more parallelism may be available. +-- +-- = Unbalanced Maps +-- +-- Unbalanced maps will as a rule be sequentialised, but sometimes, +-- there is another way. Assume we find this: +-- +-- @ +-- map +-- map(f) +-- map(g) +-- map +-- @ +-- +-- Presume that @map(f)@ is unbalanced. By the simple rule above, we +-- would then fully sequentialise it, resulting in this: +-- +-- @ +-- map +-- loop +-- map +-- map +-- @ +-- +-- == Balancing by Loop Interchange +-- +-- The above is not ideal, as we cannot flatten the @map-loop@ nest, +-- and we are thus limited in the amount of parallelism available. +-- +-- But assume now that the width of @map(g)@ is invariant to the outer +-- loop. Then if possible, we can interchange @map(f)@ and @map(g)@, +-- sequentialise @map(f)@ and distribute, interchanging the outer +-- parallel loop into the sequential loop: +-- +-- @ +-- loop(f) +-- map +-- map(g) +-- map +-- map +-- @ +-- +-- After flattening the two nests we can obtain more parallelism. +-- +-- When distributing a map, we also need to distribute everything that +-- the map depends on - possibly as its own map. When distributing a +-- set of scalar bindings, we will need to know which of the binding +-- results are used afterwards. Hence, we will need to compute usage +-- information. +-- +-- = Redomap +-- +-- Redomap can be handled much like map. Distributed loops are +-- distributed as maps, with the parameters corresponding to the +-- neutral elements added to their bodies. The remaining loop will +-- remain a redomap. Example: +-- +-- @ +-- redomap(op, +-- fn (v) => +-- map(f) +-- map(g), +-- e,a) +-- @ +-- +-- distributes to +-- +-- @ +-- let b = map(fn v => +-- let acc = e +-- map(f), +-- a) +-- redomap(op, +-- fn (v,dist) => +-- map(g), +-- e,a,b) +-- @ +-- +-- Note that there may be further kernel extraction opportunities +-- inside the @map(f)@. The downside of this approach is that the +-- intermediate array (@b@ above) must be written to main memory. An +-- often better approach is to just turn the entire @redomap@ into a +-- single kernel. +-- +module Futhark.Pass.ExtractKernels (extractKernels) where + +import Control.Monad.Identity +import Control.Monad.RWS.Strict +import Control.Monad.Reader +import Data.Maybe + +import Prelude hiding (log) + +import Futhark.Representation.SOACS +import Futhark.Representation.SOACS.Simplify (simplifyStms) +import qualified Futhark.Representation.Kernels as Out +import Futhark.Representation.Kernels.Kernel +import Futhark.MonadFreshNames +import Futhark.Tools +import qualified Futhark.Transform.FirstOrderTransform as FOT +import Futhark.Transform.Rename +import Futhark.Pass +import Futhark.Pass.ExtractKernels.Distribution +import Futhark.Pass.ExtractKernels.DistributeNests +import Futhark.Pass.ExtractKernels.ISRWIM +import Futhark.Pass.ExtractKernels.BlockedKernel +import Futhark.Pass.ExtractKernels.Intragroup +import Futhark.Util +import Futhark.Util.Log + +-- | Transform a program using SOACs to a program using explicit +-- kernels, using the kernel extraction transformation. +extractKernels :: Pass SOACS Out.Kernels +extractKernels = + Pass { passName = "extract kernels" + , passDescription = "Perform kernel extraction" + , passFunction = fmap Prog . mapM transformFunDef . progFunctions + } + +-- In order to generate more stable threshold names, we keep track of +-- the numbers used for thresholds separately from the ordinary name +-- source, +data State = State { stateNameSource :: VNameSource + , stateThresholdCounter :: Int + } + +newtype DistribM a = DistribM (RWS (Scope Out.Kernels) Log State a) + deriving (Functor, Applicative, Monad, + HasScope Out.Kernels, LocalScope Out.Kernels, + MonadState State, + MonadLogger) + +instance MonadFreshNames DistribM where + getNameSource = gets stateNameSource + putNameSource src = modify $ \s -> s { stateNameSource = src } + +runDistribM :: (MonadLogger m, MonadFreshNames m) => + DistribM a -> m a +runDistribM (DistribM m) = do + (x, msgs) <- modifyNameSource $ \src -> + let (x, s, msgs) = runRWS m mempty (State src 0) + in ((x, msgs), stateNameSource s) + addLog msgs + return x + +transformFunDef :: (MonadFreshNames m, MonadLogger m) => + FunDef SOACS -> m (Out.FunDef Out.Kernels) +transformFunDef (FunDef entry name rettype params body) = runDistribM $ do + body' <- localScope (scopeOfFParams params) $ + transformBody mempty body + return $ FunDef entry name rettype params body' + +transformBody :: KernelPath -> Body -> DistribM (Out.Body Out.Kernels) +transformBody path body = do bnds <- transformStms path $ stmsToList $ bodyStms body + return $ mkBody bnds $ bodyResult body + +transformStms :: KernelPath -> [Stm] -> DistribM KernelsStms +transformStms _ [] = + return mempty +transformStms path (bnd:bnds) = + sequentialisedUnbalancedStm bnd >>= \case + Nothing -> do + bnd' <- transformStm path bnd + inScopeOf bnd' $ + (bnd'<>) <$> transformStms path bnds + Just bnds' -> + transformStms path $ stmsToList bnds' <> bnds + +unbalancedLambda :: Lambda -> Bool +unbalancedLambda lam = + unbalancedBody + (namesFromList $ map paramName $ lambdaParams lam) $ + lambdaBody lam + + where subExpBound (Var i) bound = i `nameIn` bound + subExpBound (Constant _) _ = False + + unbalancedBody bound body = + any (unbalancedStm (bound <> boundInBody body) . stmExp) $ + bodyStms body + + -- XXX - our notion of balancing is probably still too naive. + unbalancedStm bound (Op (Stream w _ _ _)) = + w `subExpBound` bound + unbalancedStm bound (Op (Screma w _ _)) = + w `subExpBound` bound + unbalancedStm _ Op{} = + False + unbalancedStm _ DoLoop{} = False + + unbalancedStm bound (If cond tbranch fbranch _) = + cond `subExpBound` bound && + (unbalancedBody bound tbranch || unbalancedBody bound fbranch) + + unbalancedStm _ (BasicOp _) = + False + unbalancedStm _ (Apply fname _ _ _) = + not $ isBuiltInFunction fname + +sequentialisedUnbalancedStm :: Stm -> DistribM (Maybe (Stms SOACS)) +sequentialisedUnbalancedStm (Let pat _ (Op soac@(Screma _ form _))) + | Just (_, lam2) <- isRedomapSOAC form, + unbalancedLambda lam2, lambdaContainsParallelism lam2 = do + types <- asksScope scopeForSOACs + Just . snd <$> runBinderT (FOT.transformSOAC pat soac) types +sequentialisedUnbalancedStm _ = + return Nothing + +cmpSizeLe :: String -> Out.SizeClass -> [SubExp] + -> DistribM ((SubExp, Name), Out.Stms Out.Kernels) +cmpSizeLe desc size_class to_what = do + x <- gets stateThresholdCounter + modify $ \s -> s { stateThresholdCounter = x + 1} + let size_key = nameFromString $ desc ++ "_" ++ show x + runBinder $ do + to_what' <- letSubExp "comparatee" =<< + foldBinOp (Mul Int32) (intConst Int32 1) to_what + cmp_res <- letSubExp desc $ Op $ CmpSizeLe size_key size_class to_what' + return (cmp_res, size_key) + +kernelAlternatives :: (MonadFreshNames m, HasScope Out.Kernels m) => + Out.Pattern Out.Kernels + -> Out.Body Out.Kernels + -> [(SubExp, Out.Body Out.Kernels)] + -> m (Out.Stms Out.Kernels) +kernelAlternatives pat default_body [] = runBinder_ $ do + ses <- bodyBind default_body + forM_ (zip (patternNames pat) ses) $ \(name, se) -> + letBindNames_ [name] $ BasicOp $ SubExp se +kernelAlternatives pat default_body ((cond,alt):alts) = runBinder_ $ do + alts_pat <- fmap (Pattern []) $ forM (patternElements pat) $ \pe -> do + name <- newVName $ baseString $ patElemName pe + return pe { patElemName = name } + + alt_stms <- kernelAlternatives alts_pat default_body alts + let alt_body = mkBody alt_stms $ map Var $ patternValueNames alts_pat + + letBind_ pat $ If cond alt alt_body $ ifCommon $ patternTypes pat + +transformStm :: KernelPath -> Stm -> DistribM KernelsStms + +transformStm path (Let pat aux (Op (CmpThreshold what s))) = do + ((r, _), stms) <- cmpSizeLe s (Out.SizeThreshold path) [what] + runBinder_ $ do + addStms stms + addStm $ Let pat aux $ BasicOp $ SubExp r + +transformStm path (Let pat aux (If c tb fb rt)) = do + tb' <- transformBody path tb + fb' <- transformBody path fb + return $ oneStm $ Let pat aux $ If c tb' fb' rt + +transformStm path (Let pat aux (DoLoop ctx val form body)) = + localScope (castScope (scopeOf form) <> + scopeOfFParams mergeparams) $ + oneStm . Let pat aux . DoLoop ctx val form' <$> transformBody path body + where mergeparams = map fst $ ctx ++ val + form' = case form of + WhileLoop cond -> + WhileLoop cond + ForLoop i it bound ps -> + ForLoop i it bound ps + +transformStm path (Let pat (StmAux cs _) (Op (Screma w form arrs))) + | Just lam <- isMapSOAC form = + onMap path $ MapLoop pat cs w lam arrs + +transformStm path (Let res_pat (StmAux cs _) (Op (Screma w form arrs))) + | Just (scan_lam, nes) <- isScanSOAC form, + Just do_iswim <- iswim res_pat w scan_lam $ zip nes arrs = do + types <- asksScope scopeForSOACs + transformStms path =<< (stmsToList . snd <$> runBinderT (certifying cs do_iswim) types) + + -- We are only willing to generate code for scanomaps that do not + -- involve array accumulators, and do not have parallelism in their + -- map function. Such cases will fall through to the + -- screma-splitting case, and produce an ordinary map and scan. + -- Hopefully, the scan then triggers the ISWIM case above (otherwise + -- we will still crash in code generation). + | Just (scan_lam, nes, map_lam) <- isScanomapSOAC form, + all primType $ lambdaReturnType scan_lam, + not $ lambdaContainsParallelism map_lam = runBinder_ $ do + let scan_lam' = soacsLambdaToKernels scan_lam + map_lam' = soacsLambdaToKernels map_lam + lvl <- segThreadCapped [w] "segscan" $ NoRecommendation SegNoVirt + addStms =<< segScan lvl res_pat w scan_lam' map_lam' nes arrs [] [] + +transformStm path (Let res_pat (StmAux cs _) (Op (Screma w form arrs))) + | Just [Reduce comm red_fun nes] <- isReduceSOAC form, + let comm' | commutativeLambda red_fun = Commutative + | otherwise = comm, + Just do_irwim <- irwim res_pat w comm' red_fun $ zip nes arrs = do + types <- asksScope scopeForSOACs + bnds <- fst <$> runBinderT (simplifyStms =<< collectStms_ (certifying cs do_irwim)) types + transformStms path $ stmsToList bnds + +transformStm path (Let pat (StmAux cs _) (Op (Screma w form arrs))) + | Just (reds, map_lam) <- isRedomapSOAC form = do + + let paralleliseOuter = runBinder_ $ do + red_ops <- forM reds $ \(Reduce comm red_lam nes) -> do + (red_lam', nes', shape) <- determineReduceOp red_lam nes + let comm' | commutativeLambda red_lam' = Commutative + | otherwise = comm + return $ SegRedOp comm' red_lam' nes' shape + let map_lam_sequential = soacsLambdaToKernels map_lam + lvl <- segThreadCapped [w] "segred" $ NoRecommendation SegNoVirt + addStms =<< + (fmap (certify cs) <$> + nonSegRed lvl pat w red_ops map_lam_sequential arrs) + + outerParallelBody = + renameBody =<< + (mkBody <$> paralleliseOuter <*> pure (map Var (patternNames pat))) + + paralleliseInner path' = do + (mapbnd, redbnd) <- redomapToMapAndReduce pat (w, comm', red_lam, map_lam, nes, arrs) + transformStms path' [certify cs mapbnd, certify cs redbnd] + where comm' | commutativeLambda red_lam = Commutative + | otherwise = comm + (Reduce comm red_lam nes) = singleReduce reds + + innerParallelBody path' = + renameBody =<< + (mkBody <$> paralleliseInner path' <*> pure (map Var (patternNames pat))) + + if not $ lambdaContainsParallelism map_lam + then paralleliseOuter + else if incrementalFlattening then do + ((outer_suff, outer_suff_key), suff_stms) <- + sufficientParallelism "suff_outer_redomap" [w] path + + outer_stms <- outerParallelBody + inner_stms <- innerParallelBody ((outer_suff_key, False):path) + + (suff_stms<>) <$> kernelAlternatives pat inner_stms [(outer_suff, outer_stms)] + else paralleliseOuter + +-- Streams can be handled in two different ways - either we +-- sequentialise the body or we keep it parallel and distribute. +transformStm path (Let pat (StmAux cs _) (Op (Stream w (Parallel _ _ _ []) map_fun arrs))) = do + -- No reduction part. Remove the stream and leave the body + -- parallel. It will be distributed. + types <- asksScope scopeForSOACs + transformStms path =<< + (stmsToList . snd <$> runBinderT (certifying cs $ sequentialStreamWholeArray pat w [] map_fun arrs) types) + +transformStm path (Let pat aux@(StmAux cs _) (Op (Stream w (Parallel o comm red_fun nes) fold_fun arrs))) + | incrementalFlattening = do + ((outer_suff, outer_suff_key), suff_stms) <- + sufficientParallelism "suff_outer_stream" [w] path + + outer_stms <- outerParallelBody ((outer_suff_key, True) : path) + inner_stms <- innerParallelBody ((outer_suff_key, False) : path) + + (suff_stms<>) <$> kernelAlternatives pat inner_stms [(outer_suff, outer_stms)] + + | otherwise = paralleliseOuter path + + where + paralleliseOuter path' + | any (not . primType) $ lambdaReturnType red_fun = do + -- Split into a chunked map and a reduction, with the latter + -- further transformed. + let fold_fun' = soacsLambdaToKernels fold_fun + + let (red_pat_elems, concat_pat_elems) = + splitAt (length nes) $ patternValueElements pat + red_pat = Pattern [] red_pat_elems + + ((num_threads, red_results), stms) <- + streamMap (map (baseString . patElemName) red_pat_elems) concat_pat_elems w + Noncommutative fold_fun' nes arrs + + reduce_soac <- reduceSOAC [Reduce comm' red_fun nes] + + (stms<>) <$> + inScopeOf stms + (transformStm path' $ Let red_pat aux $ + Op (Screma num_threads reduce_soac red_results)) + + | otherwise = do + let red_fun_sequential = soacsLambdaToKernels red_fun + fold_fun_sequential = soacsLambdaToKernels fold_fun + fmap (certify cs) <$> + streamRed pat w comm' red_fun_sequential fold_fun_sequential nes arrs + + outerParallelBody path' = + renameBody =<< + (mkBody <$> paralleliseOuter path' <*> pure (map Var (patternNames pat))) + + paralleliseInner path' = do + types <- asksScope scopeForSOACs + transformStms path' . fmap (certify cs) =<< + (stmsToList . snd <$> runBinderT (sequentialStreamWholeArray pat w nes fold_fun arrs) types) + + innerParallelBody path' = + renameBody =<< + (mkBody <$> paralleliseInner path' <*> pure (map Var (patternNames pat))) + + comm' | commutativeLambda red_fun, o /= InOrder = Commutative + | otherwise = comm + +transformStm path (Let pat (StmAux cs _) (Op (Screma w form arrs))) = do + -- This screma is too complicated for us to immediately do + -- anything, so split it up and try again. + scope <- asksScope scopeForSOACs + transformStms path . map (certify cs) . stmsToList . snd =<< + runBinderT (dissectScrema pat w form arrs) scope + +transformStm path (Let pat _ (Op (Stream w (Sequential nes) fold_fun arrs))) = do + -- Remove the stream and leave the body parallel. It will be + -- distributed. + types <- asksScope scopeForSOACs + transformStms path =<< + (stmsToList . snd <$> + runBinderT (sequentialStreamWholeArray pat w nes fold_fun arrs) types) + +transformStm _ (Let pat (StmAux cs _) (Op (Scatter w lam ivs as))) = runBinder_ $ do + let lam' = soacsLambdaToKernels lam + write_i <- newVName "write_i" + let (as_ws, as_ns, as_vs) = unzip3 as + (i_res, v_res) = splitAt (sum as_ns) $ bodyResult $ lambdaBody lam' + kstms = bodyStms $ lambdaBody lam' + krets = do (a_w, a, is_vs) <- zip3 as_ws as_vs $ chunks as_ns $ zip i_res v_res + return $ WriteReturns [a_w] a [ ([i],v) | (i,v) <- is_vs ] + body = KernelBody () kstms krets + inputs = do (p, p_a) <- zip (lambdaParams lam') ivs + return $ KernelInput (paramName p) (paramType p) p_a [Var write_i] + (kernel, stms) <- + mapKernel segThreadCapped [(write_i,w)] inputs (map rowType $ patternTypes pat) body + certifying cs $ do + addStms stms + letBind_ pat $ Op $ SegOp kernel + +transformStm _ (Let orig_pat (StmAux cs _) (Op (GenReduce w ops bucket_fun imgs))) = do + let bfun' = soacsLambdaToKernels bucket_fun + genReduceKernel orig_pat [] [] cs w ops bfun' imgs + +transformStm _ bnd = + runBinder_ $ FOT.transformStmRecursively bnd + +sufficientParallelism :: String -> [SubExp] -> KernelPath + -> DistribM ((SubExp, Name), Out.Stms Out.Kernels) +sufficientParallelism desc ws path = cmpSizeLe desc (Out.SizeThreshold path) ws + +-- | Returns the sizes of nested parallelism. +nestedParallelism :: Body -> [SubExp] +nestedParallelism = concatMap (parallelism . stmExp) . bodyStms + where parallelism (Op (Scatter w _ _ _)) = [w] + parallelism (Op (Screma w _ _)) = [w] + parallelism (Op (Stream w Sequential{} lam _)) + | chunk_size_param : _ <- lambdaParams lam = + let update (Var v) | v == paramName chunk_size_param = w + update se = se + in map update $ nestedParallelism $ lambdaBody lam + parallelism (DoLoop _ _ _ body) = nestedParallelism body + parallelism _ = [] + +-- | Intra-group parallelism is worthwhile if the lambda contains +-- non-map nested parallelism, or any nested parallelism inside a +-- loop. +worthIntraGroup :: Lambda -> Bool +worthIntraGroup lam = interesting $ lambdaBody lam + where interesting body = not (null $ nestedParallelism body) && + not (onlyMaps $ bodyStms body) + onlyMaps = all $ isMapOrSeq . stmExp + isMapOrSeq (Op (Screma _ form@(ScremaForm _ _ lam') _)) + | isJust $ isMapSOAC form = not $ worthIntraGroup lam' + isMapOrSeq (Op Scatter{}) = True -- Basically a map. + isMapOrSeq (DoLoop _ _ _ body) = + null $ nestedParallelism body + isMapOrSeq (Op _) = False + isMapOrSeq _ = True + +-- | A lambda is worth sequentialising if it contains nested +-- parallelism of an interesting kind. +worthSequentialising :: Lambda -> Bool +worthSequentialising lam = interesting $ lambdaBody lam + where interesting body = any (interesting' . stmExp) $ bodyStms body + interesting' (Op (Screma _ form@(ScremaForm _ _ lam') _)) + | isJust $ isMapSOAC form = worthSequentialising lam' + interesting' (Op Scatter{}) = False -- Basically a map. + interesting' (DoLoop _ _ _ body) = interesting body + interesting' (Op _) = True + interesting' _ = False + + +onTopLevelStms :: KernelPath -> Stms SOACS -> DistNestT DistribM KernelsStms +onTopLevelStms path stms = do + scope <- askScope + lift $ localScope scope $ transformStms path $ stmsToList stms + +onMap :: KernelPath -> MapLoop -> DistribM KernelsStms +onMap path (MapLoop pat cs w lam arrs) = do + types <- askScope + let loopnest = MapNesting pat cs w $ zip (lambdaParams lam) arrs + env path' = DistEnv + { distNest = singleNesting (Nesting mempty loopnest) + , distScope = scopeOfPattern pat <> + scopeForKernels (scopeOf lam) <> + types + , distOnInnerMap = onInnerMap path' + , distOnTopLevelStms = onTopLevelStms path' + , distSegLevel = segThreadCapped + } + exploitInnerParallelism path' = + runDistNestT (env path') $ + distributeMapBodyStms acc (bodyStms $ lambdaBody lam) + + if not incrementalFlattening then exploitInnerParallelism path + else do + + let exploitOuterParallelism path' = do + let lam' = soacsLambdaToKernels lam + runDistNestT (env path') $ distribute $ + addStmsToKernel (bodyStms $ lambdaBody lam') acc + + onMap' (newKernel loopnest) path exploitOuterParallelism exploitInnerParallelism pat lam + where acc = DistAcc { distTargets = singleTarget (pat, bodyResult $ lambdaBody lam) + , distStms = mempty + } + +onMap' :: KernelNest -> KernelPath + -> (KernelPath -> DistribM (Out.Stms Out.Kernels)) + -> (KernelPath -> DistribM (Out.Stms Out.Kernels)) + -> Pattern + -> Lambda + -> DistribM (Out.Stms Out.Kernels) +onMap' loopnest path mk_seq_stms mk_par_stms pat lam = do + let nest_ws = kernelNestWidths loopnest + res = map Var $ patternNames pat + + types <- askScope + ((outer_suff, outer_suff_key), outer_suff_stms) <- + sufficientParallelism "suff_outer_par" nest_ws path + + intra <- if worthIntraGroup lam then + flip runReaderT types $ intraGroupParallelise loopnest lam + else return Nothing + seq_body <- renameBody =<< mkBody <$> + mk_seq_stms ((outer_suff_key, True) : path) <*> pure res + let seq_alts = [(outer_suff, seq_body) | worthSequentialising lam] + + case intra of + Nothing -> do + par_body <- renameBody =<< mkBody <$> + mk_par_stms ((outer_suff_key, False) : path) <*> pure res + + (outer_suff_stms<>) <$> kernelAlternatives pat par_body seq_alts + + Just ((_intra_min_par, intra_avail_par), group_size, log, intra_prelude, intra_stms) -> do + addLog log + -- We must check that all intra-group parallelism fits in a group. + ((intra_ok, intra_suff_key), intra_suff_stms) <- do + + ((intra_suff, suff_key), check_suff_stms) <- + sufficientParallelism "suff_intra_par" [intra_avail_par] $ + (outer_suff_key, False) : path + + runBinder $ do + + addStms intra_prelude + + max_group_size <- + letSubExp "max_group_size" $ Op $ Out.GetSizeMax Out.SizeGroup + fits <- letSubExp "fits" $ BasicOp $ + CmpOp (CmpSle Int32) group_size max_group_size + + addStms check_suff_stms + + intra_ok <- letSubExp "intra_suff_and_fits" $ BasicOp $ BinOp LogAnd fits intra_suff + return (intra_ok, suff_key) + + group_par_body <- renameBody $ mkBody intra_stms res + + par_body <- renameBody =<< mkBody <$> + mk_par_stms ([(outer_suff_key, False), + (intra_suff_key, False)] + ++ path) <*> pure res + + ((outer_suff_stms<>intra_suff_stms)<>) <$> + kernelAlternatives pat par_body (seq_alts ++ [(intra_ok, group_par_body)]) + +onInnerMap :: KernelPath -> MapLoop -> DistAcc -> DistNestT DistribM DistAcc +onInnerMap path maploop@(MapLoop pat cs w lam arrs) acc + | unbalancedLambda lam, lambdaContainsParallelism lam = + addStmToKernel (mapLoopStm maploop) acc + | not incrementalFlattening = + distributeMap maploop acc + | otherwise = + distributeSingleStm acc (mapLoopStm maploop) >>= \case + Just (post_kernels, res, nest, acc') + | Just (perm, _pat_unused) <- permutationAndMissing pat res -> do + addKernels post_kernels + multiVersion perm nest acc' + _ -> distributeMap maploop acc + + where + discardTargets acc' = + -- FIXME: work around bogus targets. + acc' { distTargets = singleTarget (mempty, mempty) } + + multiVersion perm nest acc' = do + -- The kernel can be distributed by itself, so now we can + -- decide whether to just sequentialise, or exploit inner + -- parallelism. + dist_env <- ask + let extra_scope = targetsScope $ distTargets acc' + scope <- (extra_scope<>) <$> askScope + + stms <- lift $ localScope scope $ do + let maploop' = MapLoop pat cs w lam arrs + + exploitInnerParallelism path' = do + let dist_env' = + dist_env { distOnTopLevelStms = onTopLevelStms path' + , distOnInnerMap = onInnerMap path' + } + runDistNestT dist_env' $ + inNesting nest $ localScope extra_scope $ + discardTargets <$> distributeMap maploop' acc { distStms = mempty } + + -- Normally the permutation is for the output pattern, but + -- we can't really change that, so we change the result + -- order instead. + let lam_res' = rearrangeShape perm $ bodyResult $ lambdaBody lam + lam' = lam { lambdaBody = (lambdaBody lam) { bodyResult = lam_res' } } + map_nesting = MapNesting pat cs w $ zip (lambdaParams lam) arrs + nest' = pushInnerKernelNesting (pat, lam_res') map_nesting nest + + -- XXX: we do not construct a new KernelPath when + -- sequentialising. This is only OK as long as further + -- versioning does not take place down that branch (it currently + -- does not). + (sequentialised_kernel, nestw_bnds) <- localScope extra_scope $ do + let sequentialised_lam = soacsLambdaToKernels lam' + constructKernel segThreadCapped nest' $ lambdaBody sequentialised_lam + + let outer_pat = loopNestingPattern $ fst nest + (nestw_bnds<>) <$> + onMap' nest' path + (const $ return $ oneStm sequentialised_kernel) + exploitInnerParallelism + outer_pat lam' + + addKernel stms + return acc' diff --git a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs --- a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs +++ b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs @@ -1,10 +1,13 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} module Futhark.Pass.ExtractKernels.BlockedKernel - ( segRed + ( MkSegLevel + , ThreadRecommendation(..) + , segRed , nonSegRed , segScan , segGenRed + , segMap , streamRed , streamMap @@ -13,20 +16,30 @@ , KernelInput(..) , readKernelInput - , newKernelSpace + , soacsLambdaToKernels + , soacsStmToKernels + , scopeForKernels + , scopeForSOACs + , getSize + , segThread + , segThreadCapped + , mkSegSpace ) where import Control.Monad import Control.Monad.Writer -import Data.Maybe +import Control.Monad.Identity import Data.List import Prelude hiding (quot) import Futhark.Analysis.PrimExp +import Futhark.Analysis.Rephrase import Futhark.Representation.AST +import Futhark.Representation.SOACS (SOACS) +import qualified Futhark.Representation.SOACS.SOAC as SOAC import Futhark.Representation.Kernels hiding (Prog, Body, Stm, Pattern, PatElem, BasicOp, Exp, Lambda, FunDef, FParam, LParam, RetType) @@ -40,11 +53,64 @@ size_key <- nameFromString . pretty <$> newVName desc letSubExp desc $ Op $ GetSize size_key size_class +numberOfGroups :: MonadBinder m => SubExp -> SubExp -> SubExp -> m (SubExp, SubExp) +numberOfGroups w group_size max_num_groups = do + -- If 'w' is small, we launch fewer groups than we normally would. + -- We don't want any idle groups. + w_div_group_size <- letSubExp "w_div_group_size" =<< + eDivRoundingUp Int64 (eSubExp w) (eSubExp group_size) + -- We also don't want zero groups. + num_groups_maybe_zero <- letSubExp "num_groups_maybe_zero" $ BasicOp $ BinOp (SMin Int64) + w_div_group_size max_num_groups + num_groups <- letSubExp "num_groups" $ + BasicOp $ BinOp (SMax Int64) (intConst Int64 1) + num_groups_maybe_zero + num_threads <- + letSubExp "num_threads" $ BasicOp $ BinOp (Mul Int64) num_groups group_size + return (num_groups, num_threads) + +segThread :: (MonadBinder m, Op (Lore m) ~ HostOp (Lore m) inner) => + String -> m SegLevel +segThread desc = + SegThread + <$> (Count <$> getSize (desc ++ "_num_groups") SizeNumGroups) + <*> (Count <$> getSize (desc ++ "_group_size") SizeGroup) + <*> pure SegVirt + +data ThreadRecommendation = ManyThreads | NoRecommendation SegVirt + +type MkSegLevel m = + [SubExp] -> String -> ThreadRecommendation -> BinderT Kernels m SegLevel + +-- | Like 'segThread', but cap the thread count to the input size. +-- This is more efficient for small kernels, e.g. summing a small +-- array. +segThreadCapped :: MonadFreshNames m => MkSegLevel m +segThreadCapped ws desc r = do + w <- letSubExp "nest_size" =<< foldBinOp (Mul Int32) (intConst Int32 1) ws + group_size <- getSize (desc ++ "_group_size") SizeGroup + + case r of + ManyThreads -> do + usable_groups <- letSubExp "segmap_usable_groups" =<< + eDivRoundingUp Int32 (eSubExp w) (eSubExp group_size) + return $ SegThread (Count usable_groups) (Count group_size) SegNoVirt + NoRecommendation v -> do + group_size_64 <- asIntS Int64 group_size + max_num_groups_64 <- asIntS Int64 =<< getSize (desc ++ "_max_num_groups") SizeNumGroups + w_64 <- asIntS Int64 w + (num_groups_64, _) <- numberOfGroups w_64 group_size_64 max_num_groups_64 + num_groups <- asIntS Int32 num_groups_64 + return $ SegThread (Count num_groups) (Count group_size) v + +mkSegSpace :: MonadFreshNames m => [(VName, SubExp)] -> m SegSpace +mkSegSpace dims = SegSpace <$> newVName "phys_tid" <*> pure dims + -- | Given a chunked fold lambda that takes its initial accumulator -- value as parameters, bind those parameters to the neutral element -- instead. kerneliseLambda :: MonadFreshNames m => - [SubExp] -> Lambda InKernel -> m (Lambda InKernel) + [SubExp] -> Lambda Kernels -> m (Lambda Kernels) kerneliseLambda nes lam = do thread_index <- newVName "thread_index" let thread_index_param = Param thread_index $ Prim int32 @@ -64,55 +130,67 @@ } prepareRedOrScan :: (MonadBinder m, Lore m ~ Kernels) => - SubExp -> SubExp - -> LambdaT InKernel + SubExp + -> Lambda Kernels -> [VName] -> [(VName, SubExp)] -> [KernelInput] - -> m (KernelSpace, KernelBody InKernel) -prepareRedOrScan total_num_elements w map_lam arrs ispace inps = do - (_, KernelSize num_groups group_size _ _ num_threads) <- blockedKernelSize =<< asIntS Int64 total_num_elements + -> m (SegSpace, KernelBody Kernels) +prepareRedOrScan w map_lam arrs ispace inps = do gtid <- newVName "gtid" - kspace <- newKernelSpace (num_groups, group_size, num_threads, num_groups) $ - FlatThreadSpace $ ispace ++ [(gtid, w)] - body <- fmap (uncurry (flip (KernelBody ()))) $ runBinder $ - localScope (scopeOfKernelSpace kspace) $ do - mapM_ (addStm <=< readKernelInput) inps + space <- mkSegSpace $ ispace ++ [(gtid, w)] + kbody <- fmap (uncurry (flip (KernelBody ()))) $ runBinder $ + localScope (scopeOfSegSpace space) $ do + mapM_ readKernelInput inps forM_ (zip (lambdaParams map_lam) arrs) $ \(p, arr) -> do arr_t <- lookupType arr letBindNames_ [paramName p] $ BasicOp $ Index arr $ fullSlice arr_t [DimFix $ Var gtid] - map ThreadsReturn <$> bodyBind (lambdaBody map_lam) + map Returns <$> bodyBind (lambdaBody map_lam) - return (kspace, body) + return (space, kbody) segRed :: (MonadFreshNames m, HasScope Kernels m) => - Pattern Kernels - -> SubExp + SegLevel + -> Pattern Kernels -> SubExp -- segment size - -> [SegRedOp InKernel] - -> Lambda InKernel + -> [SegRedOp Kernels] + -> Lambda Kernels -> [VName] -> [(VName, SubExp)] -- ispace = pair of (gtid, size) for the maps on "top" of this reduction -> [KernelInput] -- inps = inputs that can be looked up by using the gtids from ispace -> m (Stms Kernels) -segRed pat total_num_elements w ops map_lam arrs ispace inps = runBinder_ $ do - (kspace, kbody) <- prepareRedOrScan total_num_elements w map_lam arrs ispace inps - letBind_ pat $ Op $ HostOp $ - SegRed kspace ops (lambdaReturnType map_lam) kbody +segRed lvl pat w ops map_lam arrs ispace inps = runBinder_ $ do + (kspace, kbody) <- prepareRedOrScan w map_lam arrs ispace inps + letBind_ pat $ Op $ SegOp $ + SegRed lvl kspace ops (lambdaReturnType map_lam) kbody segScan :: (MonadFreshNames m, HasScope Kernels m) => - Pattern Kernels - -> SubExp + SegLevel + -> Pattern Kernels -> SubExp -- segment size - -> Lambda InKernel -> Lambda InKernel + -> Lambda Kernels -> Lambda Kernels -> [SubExp] -> [VName] -> [(VName, SubExp)] -- ispace = pair of (gtid, size) for the maps on "top" of this scan -> [KernelInput] -- inps = inputs that can be looked up by using the gtids from ispace -> m (Stms Kernels) -segScan pat total_num_elements w scan_lam map_lam nes arrs ispace inps = runBinder_ $ do - (kspace, kbody) <- prepareRedOrScan total_num_elements w map_lam arrs ispace inps - letBind_ pat $ Op $ HostOp $ - SegScan kspace scan_lam nes (lambdaReturnType map_lam) kbody +segScan lvl pat w scan_lam map_lam nes arrs ispace inps = runBinder_ $ do + (kspace, kbody) <- prepareRedOrScan w map_lam arrs ispace inps + letBind_ pat $ Op $ SegOp $ + SegScan lvl kspace scan_lam nes (lambdaReturnType map_lam) kbody +segMap :: (MonadFreshNames m, HasScope Kernels m) => + SegLevel + -> Pattern Kernels + -> SubExp -- segment size + -> Lambda Kernels + -> [VName] + -> [(VName, SubExp)] -- ispace = pair of (gtid, size) for the maps on "top" of this map + -> [KernelInput] -- inps = inputs that can be looked up by using the gtids from ispace + -> m (Stms Kernels) +segMap lvl pat w map_lam arrs ispace inps = runBinder_ $ do + (kspace, kbody) <- prepareRedOrScan w map_lam arrs ispace inps + letBind_ pat $ Op $ SegOp $ + SegMap lvl kspace (lambdaReturnType map_lam) kbody + dummyDim :: (MonadFreshNames m, MonadBinder m) => Pattern Kernels -> m (Pattern Kernels, [(VName, SubExp)], m ()) @@ -135,15 +213,16 @@ fullSlice from_t [DimFix $ intConst Int32 0]) nonSegRed :: (MonadFreshNames m, HasScope Kernels m) => - Pattern Kernels + SegLevel + -> Pattern Kernels -> SubExp - -> [SegRedOp InKernel] - -> Lambda InKernel + -> [SegRedOp Kernels] + -> Lambda Kernels -> [VName] -> m (Stms Kernels) -nonSegRed pat w ops map_lam arrs = runBinder_ $ do +nonSegRed lvl pat w ops map_lam arrs = runBinder_ $ do (pat', ispace, read_dummy) <- dummyDim pat - addStms =<< segRed pat' w w ops map_lam arrs ispace [] + addStms =<< segRed lvl pat' w ops map_lam arrs ispace [] read_dummy prepareStream :: (MonadBinder m, Lore m ~ Kernels) => @@ -151,12 +230,12 @@ -> [(VName, SubExp)] -> SubExp -> Commutativity - -> Lambda InKernel + -> Lambda Kernels -> [SubExp] -> [VName] - -> m (KernelSpace, [Type], KernelBody InKernel) + -> m (SubExp, SegSpace, [Type], KernelBody Kernels) prepareStream size ispace w comm fold_lam nes arrs = do - let (KernelSize num_groups group_size elems_per_thread _ num_threads) = size + let (KernelSize _ _ elems_per_thread _ num_threads) = size let (ordering, split_ordering) = case comm of Commutative -> (Disorder, SplitStrided num_threads) Noncommutative -> (InOrder, SplitContiguous) @@ -166,27 +245,26 @@ elems_per_thread_32 <- asIntS Int32 elems_per_thread gtid <- newVName "gtid" - kspace <- newKernelSpace (num_groups, group_size, num_threads, num_groups) $ - FlatThreadSpace $ ispace ++ [(gtid, num_threads)] + space <- mkSegSpace $ ispace ++ [(gtid, num_threads)] kbody <- fmap (uncurry (flip (KernelBody ()))) $ runBinder $ - localScope (scopeOfKernelSpace kspace) $ do + localScope (scopeOfSegSpace space) $ do (chunk_red_pes, chunk_map_pes) <- blockedPerThread gtid w size ordering fold_lam' (length nes) arrs let concatReturns pe = - ConcatReturns split_ordering w elems_per_thread_32 Nothing $ patElemName pe - return (map (ThreadsReturn . Var . patElemName) chunk_red_pes ++ + ConcatReturns split_ordering w elems_per_thread_32 $ patElemName pe + return (map (Returns . Var . patElemName) chunk_red_pes ++ map concatReturns chunk_map_pes) let (redout_ts, mapout_ts) = splitAt (length nes) $ lambdaReturnType fold_lam ts = redout_ts ++ map rowType mapout_ts - return (kspace, ts, kbody) + return (num_threads, space, ts, kbody) streamRed :: (MonadFreshNames m, HasScope Kernels m) => Pattern Kernels -> SubExp -> Commutativity - -> Lambda InKernel -> Lambda InKernel + -> Lambda Kernels -> Lambda Kernels -> [SubExp] -> [VName] -> m (Stms Kernels) streamRed pat w comm red_lam fold_lam nes arrs = runBinder_ $ do @@ -199,9 +277,10 @@ (redout_pat, ispace, read_dummy) <- dummyDim $ Pattern [] redout_pes let pat' = Pattern [] $ patternElements redout_pat ++ mapout_pes - (kspace, ts, kbody) <- prepareStream size ispace w comm fold_lam nes arrs + (_, kspace, ts, kbody) <- prepareStream size ispace w comm fold_lam nes arrs - letBind_ pat' $ Op $ HostOp $ SegRed kspace + lvl <- segThreadCapped [w] "stream_red" $ NoRecommendation SegNoVirt + letBind_ pat' $ Op $ SegOp $ SegRed lvl kspace [SegRedOp comm red_lam nes mempty] ts kbody read_dummy @@ -209,58 +288,56 @@ -- Similar to streamRed, but without the last reduction. streamMap :: (MonadFreshNames m, HasScope Kernels m) => [String] -> [PatElem Kernels] -> SubExp - -> Commutativity -> Lambda InKernel -> [SubExp] -> [VName] + -> Commutativity -> Lambda Kernels -> [SubExp] -> [VName] -> m ((SubExp, [VName]), Stms Kernels) streamMap out_desc mapout_pes w comm fold_lam nes arrs = runBinder $ do (_, size) <- blockedKernelSize =<< asIntS Int64 w - (kspace, ts, kbody) <- prepareStream size [] w comm fold_lam nes arrs + (threads, kspace, ts, kbody) <- prepareStream size [] w comm fold_lam nes arrs let redout_ts = take (length nes) ts redout_pes <- forM (zip out_desc redout_ts) $ \(desc, t) -> - PatElem <$> newVName desc <*> pure (t `arrayOfRow` spaceNumThreads kspace) + PatElem <$> newVName desc <*> pure (t `arrayOfRow` threads) let pat = Pattern [] $ redout_pes ++ mapout_pes - letBind_ pat $ Op $ HostOp $ SegMap kspace ts kbody + lvl <- segThreadCapped [w] "stream_map" $ NoRecommendation SegNoVirt + letBind_ pat $ Op $ SegOp $ SegMap lvl kspace ts kbody - return (spaceNumThreads kspace, map patElemName redout_pes) + return (threads, map patElemName redout_pes) segGenRed :: (MonadFreshNames m, HasScope Kernels m) => Pattern Kernels -> SubExp -> [(VName,SubExp)] -- ^ Segment indexes and sizes. -> [KernelInput] - -> [GenReduceOp InKernel] - -> Lambda InKernel -> [VName] + -> [GenReduceOp Kernels] + -> Lambda Kernels -> [VName] -> m (Stms Kernels) segGenRed pat arr_w ispace inps ops lam arrs = runBinder_ $ do - let (_, segment_sizes) = unzip ispace - arr_w_64 <- letSubExp "arr_w_64" =<< eConvOp (SExt Int32 Int64) (toExp arr_w) - segment_sizes_64 <- mapM (letSubExp "segment_size_64" <=< eConvOp (SExt Int32 Int64) . toExp) segment_sizes - total_w <- letSubExp "genreduce_elems" =<< foldBinOp (Mul Int64) arr_w_64 segment_sizes_64 - (_, KernelSize num_groups group_size _ _ num_threads) <- - blockedKernelSize total_w - gtid <- newVName "gtid" - kspace <- newKernelSpace (num_groups, group_size, num_threads, num_groups) $ - FlatThreadSpace $ ispace ++ [(gtid, arr_w)] + space <- mkSegSpace $ ispace ++ [(gtid, arr_w)] kbody <- fmap (uncurry (flip $ KernelBody ())) $ runBinder $ - localScope (scopeOfKernelSpace kspace) $ do - mapM_ (addStm <=< readKernelInput) inps + localScope (scopeOfSegSpace space) $ do + mapM_ readKernelInput inps forM_ (zip (lambdaParams lam) arrs) $ \(p, arr) -> do arr_t <- lookupType arr letBindNames_ [paramName p] $ BasicOp $ Index arr $ fullSlice arr_t [DimFix $ Var gtid] - map ThreadsReturn <$> bodyBind (lambdaBody lam) + map Returns <$> bodyBind (lambdaBody lam) - letBind_ pat $ Op $ HostOp $ SegGenRed kspace ops (lambdaReturnType lam) kbody + -- It is important not to launch unnecessarily many threads for + -- histograms, because it may mean we unnecessarily need to reduce + -- subhistograms as well. + lvl <- segThreadCapped (arr_w : map snd ispace) "seggenred" $ NoRecommendation SegNoVirt -blockedPerThread :: (MonadBinder m, Lore m ~ InKernel) => - VName -> SubExp -> KernelSize -> StreamOrd -> Lambda InKernel + letBind_ pat $ Op $ SegOp $ SegGenRed lvl space ops (lambdaReturnType lam) kbody + +blockedPerThread :: (MonadBinder m, Lore m ~ Kernels) => + VName -> SubExp -> KernelSize -> StreamOrd -> Lambda Kernels -> Int -> [VName] - -> m ([PatElem InKernel], [PatElem InKernel]) + -> m ([PatElem Kernels], [PatElem Kernels]) blockedPerThread thread_gtid w kernel_size ordering lam num_nonconcat arrs = do let (_, chunk_size, [], arr_params) = partitionChunkedKernelFoldParameters 0 $ lambdaParams lam @@ -296,7 +373,7 @@ return (chunk_red_pes, chunk_map_pes) -splitArrays :: (MonadBinder m, Lore m ~ InKernel) => +splitArrays :: (MonadBinder m, Lore m ~ Kernels) => VName -> [VName] -> SplitOrdering -> SubExp -> SubExp -> SubExp -> [VName] -> m () @@ -330,22 +407,6 @@ } deriving (Eq, Ord, Show) -numberOfGroups :: MonadBinder m => SubExp -> SubExp -> SubExp -> m (SubExp, SubExp) -numberOfGroups w group_size max_num_groups = do - -- If 'w' is small, we launch fewer groups than we normally would. - -- We don't want any idle groups. - w_div_group_size <- letSubExp "w_div_group_size" =<< - eDivRoundingUp Int64 (eSubExp w) (eSubExp group_size) - -- We also don't want zero groups. - num_groups_maybe_zero <- letSubExp "num_groups_maybe_zero" $ BasicOp $ BinOp (SMin Int64) - w_div_group_size max_num_groups - num_groups <- letSubExp "num_groups" $ - BasicOp $ BinOp (SMax Int64) (intConst Int64 1) - num_groups_maybe_zero - num_threads <- - letSubExp "num_threads" $ BasicOp $ BinOp (Mul Int64) num_groups group_size - return (num_groups, num_threads) - blockedKernelSize :: (MonadBinder m, Lore m ~ Kernels) => SubExp -> m (SubExp, KernelSize) blockedKernelSize w = do @@ -365,57 +426,35 @@ return (max_num_groups, KernelSize num_groups' group_size per_thread_elements w num_threads') -createsArrays :: KernelBody InKernel -> Bool -createsArrays = getAny . execWriter . mapM_ onStm . kernelBodyStms - where onStm stm = do - when (any (not . primType) $ patternTypes $ stmPattern stm) $ tell $ Any True - walkExpM walker $ stmExp stm - walker = identityWalker { walkOnBody = mapM_ onStm . bodyStms } - mapKernelSkeleton :: (HasScope Kernels m, MonadFreshNames m) => - SubExp -> SpaceStructure -> [KernelInput] -> Bool - -> m (KernelSpace, - Stms Kernels, - Stms InKernel) -mapKernelSkeleton w ispace inputs creates_arrays = do - ((group_size, num_threads, num_groups, virt_groups), ksize_bnds) <- runBinder $ - -- If the kernel creates arrays internally (meaning it will - -- require memory expansion), we want to truncate the amount of - -- threads. Otherwise, have at it! This is a bit of a hack - in - -- principle, we should make this decision later, when we have a - -- clearer idea of what is happening inside the kernel. - if not creates_arrays then do - group_size <- getSize "group_size" SizeGroup - num_groups <- letSubExp "num_groups" =<< eDivRoundingUp Int32 - (eSubExp w) (eSubExp group_size) - num_threads <- letSubExp "num_threads" $ - BasicOp $ BinOp (Mul Int32) num_groups group_size - return (group_size, num_threads, num_groups, num_groups) + [(VName, SubExp)] -> [KernelInput] + -> m (SegSpace, Stms Kernels) +mapKernelSkeleton ispace inputs = do + read_input_bnds <- runBinder_ $ mapM readKernelInput inputs - else do - (_, ksize) <- blockedKernelSize =<< asIntS Int64 w - virt_groups <- letSubExp "virt_groups" =<< eDivRoundingUp Int32 - (eSubExp w) (eSubExp (kernelWorkgroupSize ksize)) - return (kernelWorkgroupSize ksize, kernelNumThreads ksize, - kernelWorkgroups ksize, virt_groups) + space <- mkSegSpace ispace + return (space, read_input_bnds) - read_input_bnds <- stmsFromList <$> mapM readKernelInput inputs +mapKernel :: (HasScope Kernels m, MonadFreshNames m) => + MkSegLevel m + -> [(VName, SubExp)] -> [KernelInput] + -> [Type] -> KernelBody Kernels + -> m (SegOp Kernels, Stms Kernels) +mapKernel mk_lvl ispace inputs rts (KernelBody () kstms krets) = runBinderT' $ do + (space, read_input_stms) <- mapKernelSkeleton ispace inputs - let ksize = (num_groups, group_size, num_threads, virt_groups) + let kbody' = KernelBody () (read_input_stms <> kstms) krets - space <- newKernelSpace ksize ispace - return (space, ksize_bnds, read_input_bnds) + -- If the kernel creates arrays (meaning it will require memory + -- expansion), we want to truncate the amount of threads. + -- Otherwise, have at it! This is a bit of a hack - in principle, + -- we should make this decision later, when we have a clearer idea + -- of what is happening inside the kernel. + let r = if all primType rts then ManyThreads else NoRecommendation SegVirt -mapKernel :: (HasScope Kernels m, MonadFreshNames m) => - SubExp -> SpaceStructure -> [KernelInput] - -> [Type] -> KernelBody InKernel - -> m (Stms Kernels, Kernel InKernel) -mapKernel w ispace inputs rts kbody@(KernelBody () kstms krets) = do - (space, ksize_bnds, read_input_bnds) <- mapKernelSkeleton w ispace inputs $ - createsArrays kbody + lvl <- mk_lvl (map snd ispace) "segmap" r - let kbody' = KernelBody () (read_input_bnds <> kstms) krets - return (ksize_bnds, Kernel (KernelDebugHints "map" []) space rts kbody') + return $ SegMap lvl space rts kbody' data KernelInput = KernelInput { kernelInputName :: VName , kernelInputType :: Type @@ -424,24 +463,46 @@ } deriving (Show) -readKernelInput :: (HasScope scope m, Monad m) => - KernelInput -> m (Stm InKernel) +readKernelInput :: (MonadBinder m, Lore m ~ Kernels) => + KernelInput -> m () readKernelInput inp = do let pe = PatElem (kernelInputName inp) $ kernelInputType inp arr_t <- lookupType $ kernelInputArray inp - return $ Let (Pattern [] [pe]) (defAux ()) $ + letBind_ (Pattern [] [pe]) $ BasicOp $ Index (kernelInputArray inp) $ fullSlice arr_t $ map DimFix $ kernelInputIndices inp -newKernelSpace :: MonadFreshNames m => - (SubExp,SubExp,SubExp,SubExp) -> SpaceStructure -> m KernelSpace -newKernelSpace (num_groups, group_size, num_threads, virt_groups) ispace = - KernelSpace - <$> newVName "global_tid" - <*> newVName "local_tid" - <*> newVName "group_id" - <*> pure num_threads - <*> pure num_groups - <*> pure group_size - <*> pure virt_groups - <*> pure ispace +injectSOACS :: (Monad m, + SameScope from to, + ExpAttr from ~ ExpAttr to, + BodyAttr from ~ BodyAttr to, + RetType from ~ RetType to, + BranchType from ~ BranchType to, + Op from ~ SOAC from) => + (SOAC to -> Op to) -> Rephraser m from to +injectSOACS f = Rephraser { rephraseExpLore = return + , rephraseBodyLore = return + , rephraseLetBoundLore = return + , rephraseFParamLore = return + , rephraseLParamLore = return + , rephraseOp = fmap f . onSOAC + , rephraseRetType = return + , rephraseBranchType = return + } + where onSOAC = SOAC.mapSOACM mapper + mapper = SOAC.SOACMapper { SOAC.mapOnSOACSubExp = return + , SOAC.mapOnSOACVName = return + , SOAC.mapOnSOACLambda = rephraseLambda $ injectSOACS f + } + +soacsStmToKernels :: Stm SOACS -> Stm Kernels +soacsStmToKernels = runIdentity . rephraseStm (injectSOACS OtherOp) + +soacsLambdaToKernels :: Lambda SOACS -> Lambda Kernels +soacsLambdaToKernels = runIdentity . rephraseLambda (injectSOACS OtherOp) + +scopeForSOACs :: Scope Kernels -> Scope SOACS +scopeForSOACs = castScope + +scopeForKernels :: Scope SOACS -> Scope Kernels +scopeForKernels = castScope diff --git a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs @@ -0,0 +1,915 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE RankNTypes #-} +module Futhark.Pass.ExtractKernels.DistributeNests + ( KernelsStms + , MapLoop(..) + , mapLoopStm + + , bodyContainsParallelism + , lambdaContainsParallelism + , determineReduceOp + , incrementalFlattening + , genReduceKernel + + , DistEnv (..) + , DistAcc (..) + , runDistNestT + , DistNestT + + , distributeMap + + , distribute + , distributeSingleStm + , distributeMapBodyStms + , postKernelsStms + , addStmsToKernel + , addStmToKernel + , permutationAndMissing + , addKernels + , addKernel + , inNesting + ) +where + +import Control.Arrow (first) +import Control.Monad.Identity +import Control.Monad.RWS.Strict +import Control.Monad.Reader +import Control.Monad.Writer.Strict +import Control.Monad.Trans.Maybe +import Data.Maybe +import Data.List + +import Futhark.Representation.SOACS +import qualified Futhark.Representation.SOACS.SOAC as SOAC +import Futhark.Representation.SOACS.Simplify (simpleSOACS) +import qualified Futhark.Representation.Kernels as Out +import Futhark.Representation.Kernels.Kernel +import Futhark.MonadFreshNames +import Futhark.Tools +import Futhark.Transform.Rename +import Futhark.Transform.CopyPropagate +import Futhark.Pass.ExtractKernels.Distribution +import Futhark.Pass.ExtractKernels.ISRWIM +import Futhark.Pass.ExtractKernels.BlockedKernel hiding (segThread, segThreadCapped) +import Futhark.Pass.ExtractKernels.Interchange +import Futhark.Util +import Futhark.Util.Log + +data MapLoop = MapLoop Pattern Certificates SubExp Lambda [VName] + +mapLoopStm :: MapLoop -> Stm +mapLoopStm (MapLoop pat cs w lam arrs) = Let pat (StmAux cs ()) $ Op $ Screma w (mapSOAC lam) arrs + +type KernelsStms = Out.Stms Out.Kernels + +data DistEnv m = + DistEnv { distNest :: Nestings + , distScope :: Scope Out.Kernels + , distOnTopLevelStms :: Stms SOACS -> DistNestT m (Stms Out.Kernels) + , distOnInnerMap :: MapLoop -> DistAcc -> DistNestT m DistAcc + , distSegLevel :: MkSegLevel m + } + +data DistAcc = + DistAcc { distTargets :: Targets + , distStms :: KernelsStms + } + +data DistRes = + DistRes { accPostKernels :: PostKernels + , accLog :: Log + } + +instance Semigroup DistRes where + DistRes ks1 log1 <> DistRes ks2 log2 = + DistRes (ks1 <> ks2) (log1 <> log2) + +instance Monoid DistRes where + mempty = DistRes mempty mempty + +newtype PostKernel = PostKernel { unPostKernel :: KernelsStms } + +newtype PostKernels = PostKernels [PostKernel] + +instance Semigroup PostKernels where + PostKernels xs <> PostKernels ys = PostKernels $ ys ++ xs + +instance Monoid PostKernels where + mempty = PostKernels mempty + +postKernelsStms :: PostKernels -> KernelsStms +postKernelsStms (PostKernels kernels) = mconcat $ map unPostKernel kernels + +typeEnvFromDistAcc :: DistAcc -> Scope Out.Kernels +typeEnvFromDistAcc = scopeOfPattern . fst . outerTarget . distTargets + +addStmsToKernel :: KernelsStms -> DistAcc -> DistAcc +addStmsToKernel stms acc = + acc { distStms = stms <> distStms acc } + +addStmToKernel :: Monad m => Stm -> DistAcc -> m DistAcc +addStmToKernel stm acc = do + let stm' = soacsStmToKernels stm + return acc { distStms = oneStm stm' <> distStms acc } + +newtype DistNestT m a = DistNestT (ReaderT (DistEnv m) (WriterT DistRes m) a) + deriving (Functor, Applicative, Monad, + MonadReader (DistEnv m), + MonadWriter DistRes) + +instance MonadTrans DistNestT where + lift = DistNestT . lift . lift + +instance MonadFreshNames m => MonadFreshNames (DistNestT m) where + getNameSource = DistNestT $ lift getNameSource + putNameSource = DistNestT . lift . putNameSource + +instance Monad m => HasScope Out.Kernels (DistNestT m) where + askScope = asks distScope + +instance Monad m => LocalScope Out.Kernels (DistNestT m) where + localScope types = local $ \env -> + env { distScope = types <> distScope env } + +instance Monad m => MonadLogger (DistNestT m) where + addLog msgs = tell mempty { accLog = msgs } + +runDistNestT :: MonadLogger m => + DistEnv m -> DistNestT m DistAcc -> m (Out.Stms Out.Kernels) +runDistNestT env (DistNestT m) = do + (acc, res) <- runWriterT $ runReaderT m env + addLog $ accLog res + -- There may be a few final targets remaining - these correspond to + -- arrays that are identity mapped, and must have statements + -- inserted here. + return $ + postKernelsStms (accPostKernels res) <> + identityStms (outerTarget $ distTargets acc) + where outermost = nestingLoop $ + case distNest env of (nest, []) -> nest + (_, nest : _) -> nest + params_to_arrs = map (first paramName) $ + loopNestingParamsAndArrs outermost + + identityStms (rem_pat, res) = + stmsFromList $ zipWith identityStm (patternValueElements rem_pat) res + identityStm pe (Var v) + | Just arr <- lookup v params_to_arrs = + Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ Copy arr + identityStm pe se = + Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ + Replicate (Shape [loopNestingWidth outermost]) se + +addKernels :: Monad m => PostKernels -> DistNestT m () +addKernels ks = tell $ mempty { accPostKernels = ks } + +addKernel :: Monad m => KernelsStms -> DistNestT m () +addKernel bnds = addKernels $ PostKernels [PostKernel bnds] + +withStm :: Monad m => Stm -> DistNestT m a -> DistNestT m a +withStm stm = local $ \env -> + env { distScope = + scopeForKernels (scopeOf stm) <> distScope env + , distNest = + letBindInInnerNesting provided $ + distNest env + } + where provided = namesFromList $ patternNames $ stmPattern stm + +mapNesting :: Monad m => + Pattern -> Certificates -> SubExp -> Lambda -> [VName] + -> DistNestT m a + -> DistNestT m a +mapNesting pat cs w lam arrs = local $ \env -> + env { distNest = pushInnerNesting nest $ distNest env + , distScope = scopeForKernels (scopeOf lam) <> distScope env + } + where nest = Nesting mempty $ + MapNesting pat cs w $ + zip (lambdaParams lam) arrs + +inNesting :: Monad m => + KernelNest -> DistNestT m a -> DistNestT m a +inNesting (outer, nests) = local $ \env -> + env { distNest = (inner, nests') + , distScope = mconcat (map scopeOf $ outer : nests) <> distScope env + } + where (inner, nests') = + case reverse nests of + [] -> (asNesting outer, []) + (inner' : ns) -> (asNesting inner', map asNesting $ outer : reverse ns) + asNesting = Nesting mempty + +bodyContainsParallelism :: Body -> Bool +bodyContainsParallelism = any (isMap . stmExp) . bodyStms + where isMap Op{} = True + isMap _ = False + +lambdaContainsParallelism :: Lambda -> Bool +lambdaContainsParallelism = bodyContainsParallelism . lambdaBody + +-- Enable if you want the cool new versioned code. Beware: may be +-- slower in practice. Caveat emptor (and you are the emptor). +incrementalFlattening :: Bool +incrementalFlattening = isJust $ lookup "FUTHARK_INCREMENTAL_FLATTENING" unixEnvironment + +leavingNesting :: Monad m => MapLoop -> DistAcc -> DistNestT m DistAcc +leavingNesting (MapLoop _ cs w lam arrs) acc = + case popInnerTarget $ distTargets acc of + Nothing -> + fail "The kernel targets list is unexpectedly small" + Just ((pat,res), newtargets) -> do + let acc' = acc { distTargets = newtargets } + if null $ distStms acc' + then return acc' + else do let body = Body () (distStms acc') res + used_in_body = freeIn body + (used_params, used_arrs) = + unzip $ + filter ((`nameIn` used_in_body) . paramName . fst) $ + zip (lambdaParams lam) arrs + lam' = Lambda { lambdaParams = used_params + , lambdaBody = body + , lambdaReturnType = map rowType $ patternTypes pat + } + let stms = oneStm $ Let pat (StmAux cs ()) $ Op $ + OtherOp $ Screma w (mapSOAC lam') used_arrs + return $ addStmsToKernel stms acc' { distStms = mempty } + +distributeMapBodyStms :: MonadFreshNames m => DistAcc -> Stms SOACS -> DistNestT m DistAcc +distributeMapBodyStms orig_acc = distribute <=< onStms orig_acc . stmsToList + where + onStms acc [] = return acc + + onStms acc (Let pat (StmAux cs _) (Op (Stream w (Sequential accs) lam arrs)):stms) = do + types <- asksScope scopeForSOACs + stream_stms <- + snd <$> runBinderT (sequentialStreamWholeArray pat w accs lam arrs) types + stream_stms' <- + runReaderT (copyPropagateInStms simpleSOACS stream_stms) types + onStms acc $ stmsToList (fmap (certify cs) stream_stms') ++ stms + + onStms acc (stm:stms) = + -- It is important that stm is in scope if 'maybeDistributeStm' + -- wants to distribute, even if this causes the slightly silly + -- situation that stm is in scope of itself. + withStm stm $ maybeDistributeStm stm =<< onStms acc stms + +onInnerMap :: Monad m => MapLoop -> DistAcc -> DistNestT m DistAcc +onInnerMap loop acc = do + f <- asks distOnInnerMap + f loop acc + +onTopLevelStms :: Monad m => Stms SOACS -> DistNestT m () +onTopLevelStms stms = do + f <- asks distOnTopLevelStms + addKernel =<< f stms + +maybeDistributeStm :: MonadFreshNames m => Stm -> DistAcc -> DistNestT m DistAcc + +maybeDistributeStm bnd@(Let pat _ (Op (Screma w form arrs))) acc + | Just lam <- isMapSOAC form = + -- Only distribute inside the map if we can distribute everything + -- following the map. + distributeIfPossible acc >>= \case + Nothing -> addStmToKernel bnd acc + Just acc' -> distribute =<< onInnerMap (MapLoop pat (stmCerts bnd) w lam arrs) acc' + +maybeDistributeStm bnd@(Let pat _ (DoLoop [] val form@ForLoop{} body)) acc + | null (patternContextElements pat), bodyContainsParallelism body = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | not $ freeIn form `namesIntersect` boundInKernelNest nest, + Just (perm, pat_unused) <- permutationAndMissing pat res -> + -- We need to pretend pat_unused was used anyway, by adding + -- it to the kernel nest. + localScope (typeEnvFromDistAcc acc') $ do + addKernels kernels + nest' <- expandKernelNest pat_unused nest + types <- asksScope scopeForSOACs + + bnds <- runReaderT + (interchangeLoops nest' (SeqLoop perm pat val form body)) types + onTopLevelStms bnds + return acc' + _ -> + addStmToKernel bnd acc + +maybeDistributeStm stm@(Let pat _ (If cond tbranch fbranch ret)) acc + | null (patternContextElements pat), + bodyContainsParallelism tbranch || bodyContainsParallelism fbranch || + any (not . primType) (ifReturns ret) = + distributeSingleStm acc stm >>= \case + Just (kernels, res, nest, acc') + | not $ + (freeIn cond <> freeIn ret) `namesIntersect` boundInKernelNest nest, + Just (perm, pat_unused) <- permutationAndMissing pat res -> + -- We need to pretend pat_unused was used anyway, by adding + -- it to the kernel nest. + localScope (typeEnvFromDistAcc acc') $ do + nest' <- expandKernelNest pat_unused nest + addKernels kernels + types <- asksScope scopeForSOACs + let branch = Branch perm pat cond tbranch fbranch ret + stms <- runReaderT (interchangeBranch nest' branch) types + onTopLevelStms stms + return acc' + _ -> + addStmToKernel stm acc + +maybeDistributeStm (Let pat (StmAux cs _) (Op (Screma w form arrs))) acc + | Just [Reduce comm lam nes] <- isReduceSOAC form, + Just m <- irwim pat w comm lam $ zip nes arrs = do + types <- asksScope scopeForSOACs + (_, bnds) <- runBinderT (certifying cs m) types + distributeMapBodyStms acc bnds + +-- Parallelise segmented scatters. +maybeDistributeStm bnd@(Let pat (StmAux cs _) (Op (Scatter w lam ivs as))) acc = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | Just (perm, pat_unused) <- permutationAndMissing pat res -> + localScope (typeEnvFromDistAcc acc') $ do + nest' <- expandKernelNest pat_unused nest + let lam' = soacsLambdaToKernels lam + addKernels kernels + addKernel =<< segmentedScatterKernel nest' perm pat cs w lam' ivs as + return acc' + _ -> + addStmToKernel bnd acc + +-- Parallelise segmented GenReduce. +maybeDistributeStm bnd@(Let pat (StmAux cs _) (Op (GenReduce w ops lam as))) acc = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | Just (perm, pat_unused) <- permutationAndMissing pat res -> + localScope (typeEnvFromDistAcc acc') $ do + let lam' = soacsLambdaToKernels lam + nest' <- expandKernelNest pat_unused nest + addKernels kernels + addKernel =<< segmentedGenReduceKernel nest' perm cs w ops lam' as + return acc' + _ -> + addStmToKernel bnd acc + +-- If the scan can be distributed by itself, we will turn it into a +-- segmented scan. +-- +-- If the scan cannot be distributed by itself, it will be +-- sequentialised in the default case for this function. +maybeDistributeStm bnd@(Let pat (StmAux cs _) (Op (Screma w form arrs))) acc + | Just (lam, nes, map_lam) <- isScanomapSOAC form = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | Just (perm, pat_unused) <- permutationAndMissing pat res -> + -- We need to pretend pat_unused was used anyway, by adding + -- it to the kernel nest. + localScope (typeEnvFromDistAcc acc') $ do + nest' <- expandKernelNest pat_unused nest + let map_lam' = soacsLambdaToKernels map_lam + lam' = soacsLambdaToKernels lam + localScope (typeEnvFromDistAcc acc') $ + segmentedScanomapKernel nest' perm w lam' map_lam' nes arrs >>= + kernelOrNot cs bnd acc kernels acc' + _ -> + addStmToKernel bnd acc + +-- If the reduction can be distributed by itself, we will turn it into a +-- segmented reduce. +-- +-- If the reduction cannot be distributed by itself, it will be +-- sequentialised in the default case for this function. +maybeDistributeStm bnd@(Let pat (StmAux cs _) (Op (Screma w form arrs))) acc + | Just (reds, map_lam) <- isRedomapSOAC form, + Reduce comm lam nes <- singleReduce reds, + isIdentityLambda map_lam || incrementalFlattening = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | Just (perm, pat_unused) <- permutationAndMissing pat res -> + -- We need to pretend pat_unused was used anyway, by adding + -- it to the kernel nest. + localScope (typeEnvFromDistAcc acc') $ do + nest' <- expandKernelNest pat_unused nest + let lam' = soacsLambdaToKernels lam + map_lam' = soacsLambdaToKernels map_lam + + let comm' | commutativeLambda lam = Commutative + | otherwise = comm + + regularSegmentedRedomapKernel nest' perm w comm' lam' map_lam' nes arrs >>= + kernelOrNot cs bnd acc kernels acc' + _ -> + addStmToKernel bnd acc + +maybeDistributeStm (Let pat (StmAux cs _) (Op (Screma w form arrs))) acc + | incrementalFlattening || isNothing (isRedomapSOAC form) = do + -- This with-loop is too complicated for us to immediately do + -- anything, so split it up and try again. + scope <- asksScope scopeForSOACs + distributeMapBodyStms acc . fmap (certify cs) . snd =<< + runBinderT (dissectScrema pat w form arrs) scope + +maybeDistributeStm (Let pat aux (BasicOp (Replicate (Shape (d:ds)) v))) acc + | [t] <- patternTypes pat = do + -- XXX: We need a temporary dummy binding to prevent an empty + -- map body. The kernel extractor does not like empty map + -- bodies. + tmp <- newVName "tmp" + let rowt = rowType t + newbnd = Let pat aux $ Op $ Screma d (mapSOAC lam) [] + tmpbnd = Let (Pattern [] [PatElem tmp rowt]) aux $ + BasicOp $ Replicate (Shape ds) v + lam = Lambda { lambdaReturnType = [rowt] + , lambdaParams = [] + , lambdaBody = mkBody (oneStm tmpbnd) [Var tmp] + } + maybeDistributeStm newbnd acc + +maybeDistributeStm bnd@(Let _ aux (BasicOp Copy{})) acc = + distributeSingleUnaryStm acc bnd $ \_ 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 + | not $ primType $ typeOf pe = + distributeSingleUnaryStm acc bnd $ \_ 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 + let r = length (snd nest) + 1 + perm' = [0..r-1] ++ map (+r) perm + -- We need to add a copy, because the original map nest + -- will have produced an array without aliases, and so must we. + arr' <- newVName $ baseString arr + arr_t <- lookupType arr + return $ stmsFromList + [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 + 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 + let rots' = map (const $ intConst Int32 0) (kernelNestWidths nest) ++ rots + return $ oneStm $ Let outerpat aux $ BasicOp $ Rotate rots' arr + +-- XXX? This rule is present to avoid the case where an in-place +-- update is distributed as its own kernel, as this would mean thread +-- then writes the entire array that it updated. This is problematic +-- because the in-place updates is O(1), but writing the array is +-- O(n). It is OK if the in-place update is preceded, followed, or +-- nested inside a sequential loop or similar, because that will +-- probably be O(n) by itself. As a hack, we only distribute if there +-- does not appear to be a loop following. The better solution is to +-- depend on memory block merging for this optimisation, but it is not +-- ready yet. +maybeDistributeStm (Let pat aux (BasicOp (Update arr [DimFix i] v))) acc + | [t] <- patternTypes pat, + arrayRank t == 1, + not $ any (amortises . stmExp) $ distStms acc = do + let w = arraySize 0 t + et = stripArray 1 t + lam = Lambda { lambdaParams = [] + , lambdaReturnType = [Prim int32, et] + , lambdaBody = mkBody mempty [i, v] } + maybeDistributeStm (Let pat aux $ Op $ Scatter (intConst Int32 1) lam [] [(w, 1, arr)]) acc + where amortises DoLoop{} = True + amortises Op{} = True + amortises _ = False + +maybeDistributeStm stm@(Let _ aux (BasicOp (Concat d x xs w))) acc = + distributeSingleStm acc stm >>= \case + Just (kernels, _, nest, acc') -> + localScope (typeEnvFromDistAcc acc') $ + segmentedConcat nest >>= + kernelOrNot (stmAuxCerts aux) stm acc kernels acc' + _ -> + addStmToKernel stm acc + + where segmentedConcat nest = + isSegmentedOp nest [0] w mempty mempty [] (x:xs) $ + \pat _ _ _ (x':xs') _ -> + let d' = d + length (snd nest) + 1 + in addStm $ Let pat aux $ BasicOp $ Concat d' x' xs' w + +maybeDistributeStm bnd acc = + addStmToKernel bnd acc + +distributeSingleUnaryStm :: MonadFreshNames m => + DistAcc -> Stm + -> (KernelNest -> Pattern -> VName -> DistNestT m (Stms Out.Kernels)) + -> DistNestT m DistAcc +distributeSingleUnaryStm acc bnd f = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | res == map Var (patternNames $ stmPattern bnd), + (outer, inners) <- nest, + [(arr_p, arr)] <- loopNestingParamsAndArrs outer, + boundInKernelNest nest `namesIntersection` freeIn bnd + == oneName (paramName arr_p) -> do + addKernels kernels + let outerpat = loopNestingPattern $ fst nest + localScope (typeEnvFromDistAcc acc') $ do + (arr', pre_stms) <- repeatMissing arr (outer:inners) + f_stms <- inScopeOf pre_stms $ f nest outerpat arr' + addKernel $ pre_stms <> f_stms + return acc' + _ -> addStmToKernel bnd acc + where -- | For an imperfectly mapped array, repeat the missing + -- dimensions to make it look like it was in fact perfectly + -- mapped. + repeatMissing arr inners = do + arr_t <- lookupType arr + let shapes = determineRepeats arr arr_t inners + if all (==Shape []) shapes then return (arr, mempty) + else do + let (outer_shapes, inner_shape) = repeatShapes shapes arr_t + arr_t' = repeatDims outer_shapes inner_shape arr_t + arr' <- newVName $ baseString arr + return (arr', oneStm $ Let (Pattern [] [PatElem arr' arr_t']) (defAux ()) $ + BasicOp $ Repeat outer_shapes inner_shape arr) + + determineRepeats arr arr_t nests + | (skipped, arr_nest:nests') <- break (hasInput arr) nests, + [(arr_p, _)] <- loopNestingParamsAndArrs arr_nest = + Shape (map loopNestingWidth skipped) : + determineRepeats (paramName arr_p) (rowType arr_t) nests' + | otherwise = + Shape (map loopNestingWidth nests) : replicate (arrayRank arr_t) (Shape []) + + hasInput arr nest + | [(_, arr')] <- loopNestingParamsAndArrs nest, arr' == arr = True + | otherwise = False + + +distribute :: MonadFreshNames m => DistAcc -> DistNestT m DistAcc +distribute acc = + fromMaybe acc <$> distributeIfPossible acc + +mkSegLevel :: MonadFreshNames m => DistNestT m (MkSegLevel (DistNestT m)) +mkSegLevel = do + mk_lvl <- asks distSegLevel + return $ \w desc r -> do + scope <- askScope + (lvl, stms) <- lift $ lift $ runBinderT (mk_lvl w desc r) scope + addStms stms + return lvl + +distributeIfPossible :: MonadFreshNames m => DistAcc -> DistNestT m (Maybe DistAcc) +distributeIfPossible acc = do + nest <- asks distNest + mk_lvl <- mkSegLevel + tryDistribute mk_lvl nest (distTargets acc) (distStms acc) >>= \case + Nothing -> return Nothing + Just (targets, kernel) -> do + addKernel kernel + return $ Just DistAcc { distTargets = targets + , distStms = mempty + } + +distributeSingleStm :: MonadFreshNames m => + DistAcc -> Stm + -> DistNestT m (Maybe (PostKernels, Result, KernelNest, DistAcc)) +distributeSingleStm acc bnd = do + nest <- asks distNest + mk_lvl <- mkSegLevel + tryDistribute mk_lvl nest (distTargets acc) (distStms acc) >>= \case + Nothing -> return Nothing + Just (targets, distributed_bnds) -> + tryDistributeStm nest targets bnd >>= \case + Nothing -> return Nothing + Just (res, targets', new_kernel_nest) -> + return $ Just (PostKernels [PostKernel distributed_bnds], + res, + new_kernel_nest, + DistAcc { distTargets = targets' + , distStms = mempty + }) + +segmentedScatterKernel :: MonadFreshNames m => + KernelNest + -> [Int] + -> Pattern + -> Certificates + -> SubExp + -> Out.Lambda Out.Kernels + -> [VName] -> [(SubExp,Int,VName)] + -> DistNestT m KernelsStms +segmentedScatterKernel nest perm scatter_pat cs scatter_w lam ivs dests = do + -- We replicate some of the checking done by 'isSegmentedOp', but + -- things are different because a scatter is not a reduction or + -- scan. + -- + -- First, pretend that the scatter is also part of the nesting. The + -- KernelNest we produce here is technically not sensible, but it's + -- good enough for flatKernel to work. + let nest' = pushInnerKernelNesting (scatter_pat, bodyResult $ lambdaBody lam) + (MapNesting scatter_pat cs scatter_w $ zip (lambdaParams lam) ivs) nest + (ispace, kernel_inps) <- flatKernel nest' + + let (as_ws, as_ns, as) = unzip3 dests + + -- The input/output arrays ('as') _must_ correspond to some kernel + -- input, or else the original nested scatter would have been + -- ill-typed. Find them. + as_inps <- mapM (findInput kernel_inps) as + + mk_lvl <- mkSegLevel + + let rts = concatMap (take 1) $ chunks as_ns $ + drop (sum as_ns) $ lambdaReturnType lam + (is,vs) = splitAt (sum as_ns) $ bodyResult $ lambdaBody lam + k_body = KernelBody () (bodyStms $ lambdaBody lam) $ + map (inPlaceReturn ispace) $ + zip3 as_ws as_inps $ chunks as_ns $ zip is vs + + (k, k_bnds) <- mapKernel mk_lvl ispace kernel_inps rts k_body + + runBinder_ $ do + addStms k_bnds + + let pat = Pattern [] $ rearrangeShape perm $ + patternValueElements $ loopNestingPattern $ fst nest + + certifying cs $ letBind_ pat $ Op $ SegOp k + where findInput kernel_inps a = + maybe bad return $ find ((==a) . kernelInputName) kernel_inps + bad = fail "Ill-typed nested scatter encountered." + + inPlaceReturn ispace (aw, inp, is_vs) = + WriteReturns (init ws++[aw]) (kernelInputArray inp) + [ (map Var (init gtids)++[i], v) | (i,v) <- is_vs ] + where (gtids,ws) = unzip ispace + +segmentedGenReduceKernel :: MonadFreshNames m => + KernelNest + -> [Int] + -> Certificates + -> SubExp + -> [SOAC.GenReduceOp SOACS] + -> Out.Lambda Out.Kernels + -> [VName] + -> DistNestT m KernelsStms +segmentedGenReduceKernel nest perm cs genred_w ops lam arrs = do + -- We replicate some of the checking done by 'isSegmentedOp', but + -- things are different because a GenReduce is not a reduction or + -- scan. + (ispace, inputs) <- flatKernel nest + let orig_pat = Pattern [] $ rearrangeShape perm $ + patternValueElements $ loopNestingPattern $ fst nest + + -- The input/output arrays _must_ correspond to some kernel input, + -- or else the original nested GenReduce would have been ill-typed. + -- Find them. + ops' <- forM ops $ \(SOAC.GenReduceOp num_bins dests nes op) -> + SOAC.GenReduceOp num_bins + <$> mapM (fmap kernelInputArray . findInput inputs) dests + <*> pure nes + <*> pure op + runBinder_ $ addStms =<< + genReduceKernel orig_pat ispace inputs cs genred_w ops' lam arrs + where findInput kernel_inps a = + maybe bad return $ find ((==a) . kernelInputName) kernel_inps + bad = fail "Ill-typed nested GenReduce encountered." + +genReduceKernel :: (MonadFreshNames m, HasScope Out.Kernels m) => + Pattern -> [(VName, SubExp)] -> [KernelInput] + -> Certificates -> SubExp -> [SOAC.GenReduceOp SOACS] + -> Out.Lambda Out.Kernels -> [VName] + -> m KernelsStms +genReduceKernel orig_pat ispace inputs cs genred_w ops lam arrs = runBinder_ $ do + ops' <- forM ops $ \(SOAC.GenReduceOp num_bins dests nes op) -> do + (op', nes', shape) <- determineReduceOp op nes + return $ Out.GenReduceOp num_bins dests nes' shape op' + + let isDest = flip elem $ concatMap Out.genReduceDest ops' + inputs' = filter (not . isDest . kernelInputArray) inputs + + certifying cs $ + addStms =<< segGenRed orig_pat genred_w ispace inputs' ops' lam arrs + +determineReduceOp :: (MonadBinder m, Lore m ~ Out.Kernels) => + Lambda -> [SubExp] -> m (Out.Lambda Out.Kernels, [SubExp], Shape) +determineReduceOp lam nes = + -- FIXME? We are assuming that the accumulator is a replicate, and + -- we fish out its value in a gross way. + case mapM subExpVar nes of + Just ne_vs' -> do + let (shape, lam') = isVectorMap lam + nes' <- forM ne_vs' $ \ne_v -> do + ne_v_t <- lookupType ne_v + letSubExp "genred_ne" $ + BasicOp $ Index ne_v $ fullSlice ne_v_t $ + replicate (shapeRank shape) $ DimFix $ intConst Int32 0 + let lam'' = soacsLambdaToKernels lam' + return (lam'', nes', shape) + Nothing -> do + let lam' = soacsLambdaToKernels lam + return (lam', nes, mempty) + +isVectorMap :: Lambda -> (Shape, Lambda) +isVectorMap lam + | [Let (Pattern [] pes) _ (Op (Screma w form arrs))] <- + stmsToList $ bodyStms $ lambdaBody lam, + bodyResult (lambdaBody lam) == map (Var . patElemName) pes, + Just map_lam <- isMapSOAC form, + arrs == map paramName (lambdaParams lam) = + let (shape, lam') = isVectorMap map_lam + in (Shape [w] <> shape, lam') + | otherwise = (mempty, lam) + +segmentedScanomapKernel :: MonadFreshNames m => + KernelNest + -> [Int] + -> SubExp + -> Out.Lambda Out.Kernels -> Out.Lambda Out.Kernels + -> [SubExp] -> [VName] + -> DistNestT m (Maybe KernelsStms) +segmentedScanomapKernel nest perm segment_size lam map_lam nes arrs = do + mk_lvl <- asks distSegLevel + isSegmentedOp nest perm segment_size (freeIn lam) (freeIn map_lam) nes arrs $ + \pat ispace inps nes' _ _ -> do + lvl <- mk_lvl (segment_size : map snd ispace) "segscan" $ NoRecommendation SegNoVirt + addStms =<< traverse renameStm =<< + segScan lvl pat segment_size lam map_lam nes' arrs ispace inps + +regularSegmentedRedomapKernel :: MonadFreshNames m => + KernelNest + -> [Int] + -> SubExp -> Commutativity + -> Out.Lambda Out.Kernels -> Out.Lambda Out.Kernels + -> [SubExp] -> [VName] + -> DistNestT m (Maybe KernelsStms) +regularSegmentedRedomapKernel nest perm segment_size comm lam map_lam nes arrs = do + mk_lvl <- asks distSegLevel + isSegmentedOp nest perm segment_size (freeIn lam) (freeIn map_lam) nes arrs $ + \pat ispace inps nes' _ _ -> do + let red_op = SegRedOp comm lam nes' mempty + lvl <- mk_lvl (segment_size : map snd ispace) "segred" $ NoRecommendation SegNoVirt + addStms =<< traverse renameStm =<< + segRed lvl pat segment_size [red_op] map_lam arrs ispace inps + +isSegmentedOp :: MonadFreshNames m => + KernelNest + -> [Int] + -> SubExp + -> Names -> Names + -> [SubExp] -> [VName] + -> (Pattern + -> [(VName, SubExp)] + -> [KernelInput] + -> [SubExp] -> [VName] -> [VName] + -> BinderT Out.Kernels m ()) + -> DistNestT m (Maybe KernelsStms) +isSegmentedOp nest perm segment_size free_in_op _free_in_fold_op nes arrs m = runMaybeT $ do + -- We must verify that array inputs to the operation are inputs to + -- the outermost loop nesting or free in the loop nest. Nothing + -- free in the op may be bound by the nest. Furthermore, the + -- neutral elements must be free in the loop nest. + -- + -- We must summarise any names from free_in_op that are bound in the + -- nest, and describe how to obtain them given segment indices. + + let bound_by_nest = boundInKernelNest nest + + (ispace, kernel_inps) <- flatKernel nest + + unless (not $ free_in_op `namesIntersect` bound_by_nest) $ + fail "Non-fold lambda uses nest-bound parameters." + + let indices = map fst ispace + + prepareNe (Var v) | v `nameIn` bound_by_nest = + fail "Neutral element bound in nest" + prepareNe ne = return ne + + prepareArr arr = + case find ((==arr) . kernelInputName) kernel_inps of + Just inp + | kernelInputIndices inp == map Var indices -> + return $ return $ kernelInputArray inp + | not (kernelInputArray inp `nameIn` bound_by_nest) -> + return $ replicateMissing ispace inp + Nothing | not (arr `nameIn` bound_by_nest) -> + -- This input is something that is free inside + -- the loop nesting. We will have to replicate + -- it. + return $ + letExp (baseString arr ++ "_repd") + (BasicOp $ Replicate (Shape $ map snd ispace) $ Var arr) + _ -> + fail "Input not free or outermost." + + nes' <- mapM prepareNe nes + + mk_arrs <- mapM prepareArr arrs + scope <- lift askScope + + lift $ lift $ flip runBinderT_ scope $ do + -- We must make sure all inputs are of size + -- segment_size*nesting_size. + total_num_elements <- + letSubExp "total_num_elements" =<< + foldBinOp (Mul Int32) segment_size (map snd ispace) + + let flatten arr = do + arr_shape <- arrayShape <$> lookupType arr + -- CHECKME: is the length the right thing here? We want to + -- reproduce the parameter type. + let reshape = reshapeOuter [DimNew total_num_elements] + (2+length (snd nest)) arr_shape + letExp (baseString arr ++ "_flat") $ + BasicOp $ Reshape reshape arr + + nested_arrs <- sequence mk_arrs + arrs' <- mapM flatten nested_arrs + + let pat = Pattern [] $ rearrangeShape perm $ + patternValueElements $ loopNestingPattern $ fst nest + + m pat ispace kernel_inps nes' nested_arrs arrs' + + where replicateMissing ispace inp = do + t <- lookupType $ kernelInputArray inp + let inp_is = kernelInputIndices inp + shapes = determineRepeats ispace inp_is + (outer_shapes, inner_shape) = repeatShapes shapes t + letExp "repeated" $ BasicOp $ + Repeat outer_shapes inner_shape $ kernelInputArray inp + + determineRepeats ispace (i:is) + | (skipped_ispace, ispace') <- span ((/=i) . Var . fst) ispace = + Shape (map snd skipped_ispace) : determineRepeats (drop 1 ispace') is + determineRepeats ispace _ = + [Shape $ map snd ispace] + +permutationAndMissing :: Pattern -> [SubExp] -> Maybe ([Int], [PatElem]) +permutationAndMissing pat res = do + let pes = patternValueElements pat + (_used,unused) = + partition ((`nameIn` freeIn res) . patElemName) pes + res_expanded = res ++ map (Var . patElemName) unused + perm <- map (Var . patElemName) pes `isPermutationOf` res_expanded + return (perm, unused) + +-- Add extra pattern elements to every kernel nesting level. +expandKernelNest :: MonadFreshNames m => + [PatElem] -> KernelNest -> m KernelNest +expandKernelNest pes (outer_nest, inner_nests) = do + let outer_size = loopNestingWidth outer_nest : + map loopNestingWidth inner_nests + inner_sizes = tails $ map loopNestingWidth inner_nests + outer_nest' <- expandWith outer_nest outer_size + inner_nests' <- zipWithM expandWith inner_nests inner_sizes + return (outer_nest', inner_nests') + where expandWith nest dims = do + pes' <- mapM (expandPatElemWith dims) pes + return nest { loopNestingPattern = + Pattern [] $ + patternElements (loopNestingPattern nest) <> pes' + } + + expandPatElemWith dims pe = do + name <- newVName $ baseString $ patElemName pe + return pe { patElemName = name + , patElemAttr = patElemType pe `arrayOfShape` Shape dims + } + +kernelOrNot :: MonadFreshNames m => + Certificates -> Stm -> DistAcc + -> PostKernels -> DistAcc -> Maybe KernelsStms + -> DistNestT m DistAcc +kernelOrNot cs bnd acc _ _ Nothing = + addStmToKernel (certify cs bnd) acc +kernelOrNot cs _ _ kernels acc' (Just bnds) = do + addKernels kernels + addKernel $ fmap (certify cs) bnds + return acc' + +distributeMap :: MonadFreshNames m => MapLoop -> DistAcc -> DistNestT m DistAcc +distributeMap maploop@(MapLoop pat cs w lam arrs) acc = + distribute =<< + leavingNesting maploop =<< + mapNesting pat cs w lam arrs + (distribute =<< distributeMapBodyStms acc' lam_bnds) + + where acc' = DistAcc { distTargets = pushInnerTarget + (pat, bodyResult $ lambdaBody lam) $ + distTargets acc + , distStms = mempty + } + + lam_bnds = bodyStms $ lambdaBody lam diff --git a/src/Futhark/Pass/ExtractKernels/Distribution.hs b/src/Futhark/Pass/ExtractKernels/Distribution.hs --- a/src/Futhark/Pass/ExtractKernels/Distribution.hs +++ b/src/Futhark/Pass/ExtractKernels/Distribution.hs @@ -47,7 +47,6 @@ import Control.Monad.RWS.Strict import Control.Monad.Trans.Maybe import qualified Data.Map.Strict as M -import qualified Data.Set as S import Data.Foldable import Data.Maybe import Data.List @@ -58,7 +57,8 @@ import Futhark.Util import Futhark.Transform.Rename import Futhark.Util.Log -import Futhark.Pass.ExtractKernels.BlockedKernel (mapKernel, KernelInput(..)) +import Futhark.Pass.ExtractKernels.BlockedKernel + (mapKernel, KernelInput(..), readKernelInput, MkSegLevel) type Target = (Pattern Kernels, Result) @@ -97,7 +97,7 @@ where (pes', res') = unzip $ filter (used . fst) $ zip (patternElements pat) res pat' = Pattern [] pes' inner_used = freeIn $ snd inner_target - used pe = patElemName pe `S.member` inner_used + used pe = patElemName pe `nameIn` inner_used popInnerTarget :: Targets -> Maybe (Target, Targets) popInnerTarget (Targets t ts) = @@ -131,11 +131,11 @@ loopNestingParams = map fst . loopNestingParamsAndArrs instance FreeIn LoopNesting where - freeIn (MapNesting pat cs w params_and_arrs) = - freeIn pat <> - freeIn cs <> - freeIn w <> - freeIn params_and_arrs + freeIn' (MapNesting pat cs w params_and_arrs) = + freeIn' pat <> + freeIn' cs <> + freeIn' w <> + freeIn' params_and_arrs data Nesting = Nesting { nestingLetBound :: Names , nestingLoop :: LoopNesting @@ -166,7 +166,7 @@ -- | Both parameters and let-bound. boundInNesting :: Nesting -> Names boundInNesting nesting = - S.fromList (map paramName (loopNestingParams loop)) <> + namesFromList (map paramName (loopNestingParams loop)) <> nestingLetBound nesting where loop = nestingLoop nesting @@ -213,12 +213,12 @@ -- uses of corresponding parameters from innermost nesting. removeArraysFromNest :: [VName] -> KernelNest -> KernelNest removeArraysFromNest orig_arrs (outer, inners) = - let (arrs, outer') = remove (S.fromList orig_arrs) outer + let (arrs, outer') = remove (namesFromList orig_arrs) outer (_, inners') = mapAccumL remove arrs inners in (outer', inners') where remove arrs nest = - let (discard, keep) = partition ((`S.member` arrs) . snd) $ loopNestingParamsAndArrs nest - in (S.fromList (map (paramName . fst) discard) <> arrs, + let (discard, keep) = partition ((`nameIn` arrs) . snd) $ loopNestingParamsAndArrs nest + in (namesFromList (map (paramName . fst) discard) <> arrs, nest { loopNestingParamsAndArrs = keep }) newKernel :: LoopNesting -> KernelNest @@ -231,7 +231,7 @@ boundInKernelNest = mconcat . boundInKernelNests boundInKernelNests :: KernelNest -> [Names] -boundInKernelNests = map (S.fromList . +boundInKernelNests = map (namesFromList . map (paramName . fst) . loopNestingParamsAndArrs) . kernelNestLoops @@ -240,56 +240,47 @@ kernelNestWidths = map loopNestingWidth . kernelNestLoops constructKernel :: (MonadFreshNames m, LocalScope Kernels m) => - KernelNest -> KernelBody InKernel - -> m (Stms Kernels, SubExp, Stm Kernels) -constructKernel kernel_nest inner_body = do - (w_bnds, w, ispace, inps) <- flatKernel kernel_nest - let used_inps = filter inputIsUsed inps - cs = loopNestingCertificates first_nest + MkSegLevel m -> KernelNest -> Body Kernels + -> m (Stm Kernels, Stms Kernels) +constructKernel mk_lvl kernel_nest inner_body = runBinderT' $ do + (ispace, inps) <- flatKernel kernel_nest + let cs = loopNestingCertificates first_nest + ispace_scope = M.fromList $ zip (map fst ispace) $ repeat $ IndexInfo Int32 pat = loopNestingPattern first_nest rts = map (stripArray (length ispace)) $ patternTypes pat - (ksize_bnds, k) <- inScopeOf w_bnds $ - mapKernel w (FlatThreadSpace ispace) used_inps rts inner_body + inner_body' <- fmap (uncurry (flip (KernelBody ()))) $ runBinder $ + localScope ispace_scope $ do + mapM_ readKernelInput $ filter inputIsUsed inps + map Returns <$> bodyBind inner_body - let kbnds = w_bnds <> ksize_bnds - return (kbnds, - w, - Let pat (StmAux cs ()) $ Op $ HostOp k) + (segop, aux_stms) <- lift $ mapKernel mk_lvl ispace [] rts inner_body' + + addStms aux_stms + + return $ Let pat (StmAux cs ()) $ Op $ SegOp segop where first_nest = fst kernel_nest - inputIsUsed input = kernelInputName input `S.member` - freeIn inner_body + inputIsUsed input = kernelInputName input `nameIn` freeIn inner_body -- | Flatten a kernel nesting to: -- --- (0) Ancillary prologue bindings. --- --- (1) The total number of threads, equal to the product of all --- nesting widths, and equal to the product of the index space. --- --- (2) The index space. +-- (1) The index space. -- --- (3) The kernel inputs - note that some of these may be unused. +-- (2) The kernel inputs - note that some of these may be unused. flatKernel :: MonadFreshNames m => KernelNest - -> m (Stms Kernels, - SubExp, - [(VName, SubExp)], + -> m ([(VName, SubExp)], [KernelInput]) flatKernel (MapNesting _ _ nesting_w params_and_arrs, []) = do i <- newVName "gtid" let inps = [ KernelInput pname ptype arr [Var i] | (Param pname ptype, arr) <- params_and_arrs ] - return (mempty, nesting_w, [(i,nesting_w)], inps) + return ([(i,nesting_w)], inps) flatKernel (MapNesting _ _ nesting_w params_and_arrs, nest : nests) = do i <- newVName "gtid" - (w_bnds, w, ispace, inps) <- flatKernel (nest, nests) - - w' <- newVName "nesting_size" - let w_bnd = mkLet [] [Ident w' $ Prim int32] $ - BasicOp $ BinOp (Mul Int32) w nesting_w + (ispace, inps) <- flatKernel (nest, nests) let inps' = map fixupInput inps isParam inp = @@ -301,8 +292,7 @@ | otherwise = inp - return (w_bnds <> oneStm w_bnd, Var w', (i, nesting_w) : ispace, - extra_inps i <> inps') + return ((i, nesting_w) : ispace, extra_inps i <> inps') where extra_inps i = [ KernelInput pname ptype arr [Var i] | (Param pname ptype, arr) <- params_and_arrs ] @@ -322,12 +312,12 @@ distributionBodyFromStms :: Attributes lore => Targets -> Stms lore -> (DistributionBody, Result) distributionBodyFromStms (Targets (inner_pat, inner_res) targets) stms = - let bound_by_stms = S.fromList $ M.keys $ scopeOf stms + let bound_by_stms = namesFromList $ M.keys $ scopeOf stms (inner_pat', inner_res', inner_identity_map, inner_expand_target) = removeIdentityMappingGeneral bound_by_stms inner_pat inner_res in (DistributionBody { distributionTarget = Targets (inner_pat', inner_res') targets - , distributionFreeInBody = fold (fmap freeIn stms) `S.difference` bound_by_stms + , distributionFreeInBody = fold (fmap freeIn stms) `namesSubtract` bound_by_stms , distributionIdentityMap = inner_identity_map , distributionExpandTarget = inner_expand_target }, @@ -356,7 +346,7 @@ -- | Can something of this type be taken outside the nest? -- I.e. are none of its dimensions bound inside the nest. distributableType = - S.null . S.intersection bound_in_nest . freeIn . arrayDims + (==mempty) . namesIntersection bound_in_nest . freeIn . arrayDims distributeAtNesting :: (HasScope t m, MonadFreshNames m) => Nesting @@ -376,14 +366,14 @@ let nest'@(MapNesting _ cs w params_and_arrs) = removeUnusedNestingParts free_in_kernel nest (params,arrs) = unzip params_and_arrs - param_names = S.fromList $ map paramName params + param_names = namesFromList $ map paramName params free_in_kernel' = - (freeIn nest' <> free_in_kernel) `S.difference` param_names + (freeIn nest' <> free_in_kernel) `namesSubtract` param_names required_from_nest = - free_in_kernel' `S.intersection` nest_let_bound + free_in_kernel' `namesIntersection` nest_let_bound required_from_nest_idents <- - forM (S.toList required_from_nest) $ \name -> do + forM (namesToList required_from_nest) $ \name -> do t <- lift $ lookupType name return $ Ident name t @@ -413,14 +403,14 @@ (params++free_params, arrs++map identName free_arrs) actual_param_names = - S.fromList $ map paramName actual_params + namesFromList $ map paramName actual_params nest'' = removeUnusedNestingParts free_in_kernel $ MapNesting pat cs w $ zip actual_params actual_arrs free_in_kernel'' = - (freeIn nest'' <> free_in_kernel) `S.difference` actual_param_names + (freeIn nest'' <> free_in_kernel) `namesSubtract` actual_param_names unless (all (distributableType . paramType) $ loopNestingParams nest'') $ @@ -439,7 +429,7 @@ inner_nest (distributionInnerPattern distrib_body) (newKernel, - distributionFreeInBody distrib_body `S.intersection` bound_in_nest) + distributionFreeInBody distrib_body `namesIntersection` bound_in_nest) (distributionIdentityMap distrib_body) [] $ singleTarget . distributionExpandTarget distrib_body @@ -449,7 +439,7 @@ let (pat', res', identity_map, expand_target) = removeIdentityMappingFromNesting - (S.fromList $ patternNames $ loopNestingPattern outer) pat res + (namesFromList $ patternNames $ loopNestingPattern outer) pat res distributeAtNesting nest @@ -466,7 +456,7 @@ where (params,arrs) = unzip params_and_arrs (used_params, used_arrs) = unzip $ - filter ((`S.member` used) . paramName . fst) $ + filter ((`nameIn` used) . paramName . fst) $ zip params arrs removeIdentityMappingGeneral :: Names -> Pattern Kernels -> Result @@ -489,8 +479,8 @@ identity_map, expandTarget) where isIdentity (patElem, Var v) - | not (v `S.member` bound) = Left (patElem, v) - isIdentity x = Right x + | not (v `nameIn` bound) = Left (patElem, v) + isIdentity x = Right x removeIdentityMappingFromNesting :: Names -> Pattern Kernels -> Result -> (Pattern Kernels, @@ -503,17 +493,18 @@ in (pat', res', identity_map, expand_target) tryDistribute :: (MonadFreshNames m, LocalScope Kernels m, MonadLogger m) => - Nestings -> Targets -> Stms InKernel + MkSegLevel m -> Nestings -> Targets -> Stms Kernels -> m (Maybe (Targets, Stms Kernels)) -tryDistribute _ targets stms | null stms = +tryDistribute _ _ targets stms | null stms = -- No point in distributing an empty kernel. return $ Just (targets, mempty) -tryDistribute nest targets stms = +tryDistribute mk_lvl nest targets stms = createKernelNest nest dist_body >>= \case Just (targets', distributed) -> do - (w_bnds, _, kernel_bnd) <- localScope (targetsScope targets') $ - constructKernel distributed inner_body + (kernel_bnd, w_bnds) <- + localScope (targetsScope targets') $ + constructKernel mk_lvl distributed $ mkBody stms inner_body_res distributed' <- renameStm kernel_bnd logMsg $ "distributing\n" ++ unlines (map pretty $ stmsToList stms) ++ @@ -525,7 +516,6 @@ Nothing -> return Nothing where (dist_body, inner_body_res) = distributionBodyFromStms targets stms - inner_body = KernelBody () stms $ map ThreadsReturn inner_body_res tryDistributeStm :: (MonadFreshNames m, HasScope t m, Attributes lore) => Nestings -> Targets -> Stm lore diff --git a/src/Futhark/Pass/ExtractKernels/Interchange.hs b/src/Futhark/Pass/ExtractKernels/Interchange.hs --- a/src/Futhark/Pass/ExtractKernels/Interchange.hs +++ b/src/Futhark/Pass/ExtractKernels/Interchange.hs @@ -15,7 +15,6 @@ ) where import Control.Monad.RWS.Strict -import qualified Data.Set as S import Data.Maybe import Data.List @@ -75,7 +74,7 @@ where free_in_body = freeIn body copyOrRemoveParam (param, arr) - | not (paramName param `S.member` free_in_body) = + | not (paramName param `nameIn` free_in_body) = return Nothing | otherwise = return $ Just (param, arr) diff --git a/src/Futhark/Pass/ExtractKernels/Intragroup.hs b/src/Futhark/Pass/ExtractKernels/Intragroup.hs --- a/src/Futhark/Pass/ExtractKernels/Intragroup.hs +++ b/src/Futhark/Pass/ExtractKernels/Intragroup.hs @@ -1,4 +1,5 @@ {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE TypeFamilies #-} -- | Extract limited nested parallelism for execution inside -- individual kernel workgroups. @@ -6,21 +7,25 @@ (intraGroupParallelise) where +import Control.Monad.Identity import Control.Monad.RWS import Control.Monad.Trans.Maybe import qualified Data.Map.Strict as M import qualified Data.Set as S +import Prelude hiding (log) + import Futhark.Analysis.PrimExp.Convert import Futhark.Representation.SOACS import qualified Futhark.Representation.Kernels as Out import Futhark.Representation.Kernels.Kernel import Futhark.MonadFreshNames import Futhark.Tools -import Futhark.Analysis.DataDependencies -import qualified Futhark.Pass.ExtractKernels.Kernelise as Kernelise +import Futhark.Pass.ExtractKernels.DistributeNests import Futhark.Pass.ExtractKernels.Distribution import Futhark.Pass.ExtractKernels.BlockedKernel +import Futhark.Util (chunks) +import Futhark.Util.Log -- | Convert the statements inside a map nest to kernel statements, -- attempting to parallelise any remaining (top-level) parallel @@ -36,21 +41,27 @@ -- exploitable parallelism". intraGroupParallelise :: (MonadFreshNames m, LocalScope Out.Kernels m) => KernelNest -> Lambda - -> m (Maybe ((SubExp, SubExp), SubExp, + -> m (Maybe ((SubExp, SubExp), SubExp, Log, Out.Stms Out.Kernels, Out.Stms Out.Kernels)) intraGroupParallelise knest lam = runMaybeT $ do - (w_stms, w, ispace, inps) <- lift $ flatKernel knest - let num_groups = w - body = lambdaBody lam + (ispace, inps) <- lift $ flatKernel knest - ltid <- newVName "ltid" - let group_variant = S.fromList [ltid] - (wss_min, wss_avail, kbody) <- + (num_groups, w_stms) <- lift $ runBinder $ + letSubExp "intra_num_groups" =<< + foldBinOp (Mul Int32) (intConst Int32 1) (map snd ispace) + + let body = lambdaBody lam + + group_size <- newVName "computed_group_size" + let intra_lvl = SegThread (Count num_groups) (Count $ Var group_size) SegNoVirt + + (wss_min, wss_avail, log, kbody) <- lift $ localScope (scopeOfLParams $ lambdaParams lam) $ - intraGroupParalleliseBody (dataDependencies body) group_variant ltid body + intraGroupParalleliseBody intra_lvl body known_outside <- lift $ M.keys <$> askScope - unless (all (`elem` known_outside) $ freeIn $ wss_min ++ wss_avail) $ + unless (all (`elem` known_outside) $ namesToList $ freeIn $ + wss_min ++ wss_avail) $ fail "Irregular parallelism" ((intra_avail_par, kspace, read_input_stms), prelude_stms) <- lift $ runBinder $ do @@ -68,148 +79,130 @@ -- The group size is either the maximum of the minimum parallelism -- exploited, or the desired parallelism (bounded by the max group -- size) in case there is no minimum. - group_size <- letSubExp "computed_group_size" =<< - if null ws_min - then eBinOp (SMin Int32) - (eSubExp =<< letSubExp "max_group_size" (Op $ Out.GetSizeMax Out.SizeGroup)) - (eSubExp intra_avail_par) - else foldBinOp' (SMax Int32) ws_min + letBindNames_ [group_size] =<< + if null ws_min + then eBinOp (SMin Int32) + (eSubExp =<< letSubExp "max_group_size" (Op $ Out.GetSizeMax Out.SizeGroup)) + (eSubExp intra_avail_par) + else foldBinOp' (SMax Int32) ws_min - let inputIsUsed input = kernelInputName input `S.member` freeIn body + let inputIsUsed input = kernelInputName input `nameIn` freeIn body used_inps = filter inputIsUsed inps addStms w_stms - - num_threads <- letSubExp "num_threads" $ - BasicOp $ BinOp (Mul Int32) num_groups group_size - - let ksize = (num_groups, group_size, num_threads, num_groups) - - kspace <- newKernelSpace ksize $ FlatThreadSpace $ ispace ++ [(ltid,group_size)] - - read_input_stms <- mapM readKernelInput used_inps - - return (intra_avail_par, kspace, read_input_stms) + read_input_stms <- runBinder_ $ mapM readKernelInput used_inps + space <- mkSegSpace ispace + return (intra_avail_par, space, read_input_stms) - let kbody' = kbody { kernelBodyStms = stmsFromList read_input_stms <> kernelBodyStms kbody } + let kbody' = kbody { kernelBodyStms = read_input_stms <> kernelBodyStms kbody } - -- The kernel itself is producing a "flat" result of shape - -- [num_groups]. We must explicitly reshape it to match the shapes - -- of our enclosing map-nests. let nested_pat = loopNestingPattern first_nest - flatPatElem pat_elem = do - let t' = arrayOfRow (length ispace `stripArray` patElemType pat_elem) num_groups - name <- newVName $ baseString (patElemName pat_elem) ++ "_flat" - return $ PatElem name t' - flat_pat <- lift $ Pattern [] <$> mapM flatPatElem (patternValueElements nested_pat) - - let rts = map rowType $ patternTypes flat_pat - kstm = Let flat_pat (StmAux cs ()) $ Op $ HostOp $ - Kernel (KernelDebugHints "map_intra_group" []) kspace rts kbody' - reshapeStm nested_pe flat_pe = - Let (Pattern [] [nested_pe]) (StmAux cs ()) $ - BasicOp $ Reshape (map DimNew $ arrayDims $ patElemType nested_pe) $ - patElemName flat_pe - reshape_stms = zipWith reshapeStm (patternElements nested_pat) - (patternElements flat_pat) + rts = map (length ispace `stripArray`) $ patternTypes nested_pat + lvl = SegGroup (Count num_groups) (Count $ Var group_size) SegNoVirt + kstm = Let nested_pat (StmAux cs ()) $ Op $ SegOp $ SegMap lvl kspace rts kbody' let intra_min_par = intra_avail_par - return ((intra_min_par, intra_avail_par), spaceGroupSize kspace, - prelude_stms, oneStm kstm <> stmsFromList reshape_stms) + return ((intra_min_par, intra_avail_par), Var group_size, log, + prelude_stms, oneStm kstm) where first_nest = fst knest cs = loopNestingCertificates first_nest -data Env = Env { _localTID :: VName - , _dataDeps :: Dependencies - , _groupVariant :: Names +data Acc = Acc { accMinPar :: S.Set [SubExp] + , accAvailPar :: S.Set [SubExp] + , accLog :: Log } -type IntraGroupM = BinderT Out.InKernel (RWS Env (S.Set [SubExp], S.Set [SubExp]) VNameSource) +instance Semigroup Acc where + Acc min_x avail_x log_x <> Acc min_y avail_y log_y = + Acc (min_x <> min_y) (avail_x <> avail_y) (log_x <> log_y) +instance Monoid Acc where + mempty = Acc mempty mempty mempty + +type IntraGroupM = + BinderT Out.Kernels (RWS () Acc VNameSource) + +instance MonadLogger IntraGroupM where + addLog log = tell mempty { accLog = log } + runIntraGroupM :: (MonadFreshNames m, HasScope Out.Kernels m) => - Env -> IntraGroupM () -> m ([[SubExp]], [[SubExp]], Out.Stms Out.InKernel) -runIntraGroupM env m = do + IntraGroupM () -> m (Acc, Out.Stms Out.Kernels) +runIntraGroupM m = do scope <- castScope <$> askScope modifyNameSource $ \src -> - let (((), kstms), src', (ws_min, ws_avail)) = runRWS (runBinderT m scope) env src - in ((S.toList ws_min, S.toList ws_avail, kstms), src') + let (((), kstms), src', acc) = runRWS (runBinderT m scope) () src + in ((acc, kstms), src') parallelMin :: [SubExp] -> IntraGroupM () -parallelMin ws = tell (S.singleton ws, S.singleton ws) - -parallelAvail :: [SubExp] -> IntraGroupM () -parallelAvail ws = tell (mempty, S.singleton ws) +parallelMin ws = tell mempty { accMinPar = S.singleton ws + , accAvailPar = S.singleton ws + } -intraGroupBody :: Body -> IntraGroupM (Out.Body Out.InKernel) -intraGroupBody body = do - stms <- collectStms_ $ mapM_ intraGroupStm $ bodyStms body +intraGroupBody :: SegLevel -> Body -> IntraGroupM (Out.Body Out.Kernels) +intraGroupBody lvl body = do + stms <- collectStms_ $ mapM_ (intraGroupStm lvl) $ bodyStms body return $ mkBody stms $ bodyResult body -intraGroupStm :: Stm -> IntraGroupM () -intraGroupStm stm@(Let pat _ e) = do - Env ltid deps group_variant <- ask - let groupInvariant (Var v) = - S.null . S.intersection group_variant . - flip (M.findWithDefault mempty) deps $ v - groupInvariant Constant{} = True +intraGroupStm :: SegLevel -> Stm -> IntraGroupM () +intraGroupStm lvl stm@(Let pat aux e) = do + scope <- askScope + let lvl' = SegThread (segNumGroups lvl) (segGroupSize lvl) SegNoVirt case e of - -- Cosmin hack: previously, only for loops were supported, - -- and only if `groupInvariant bound` holds; - -- Let's see what can possibly go wrong if we - -- completely generalize this (?) DoLoop ctx val form loopbody -> - localScope (scopeOf form') $ - localScope (scopeOfFParams $ map fst $ ctx ++ val) $ do - loopbody' <- intraGroupBody loopbody - letBind_ pat $ DoLoop ctx val form' loopbody' - where form' = case form of - ForLoop i it bound inps -> ForLoop i it bound inps - WhileLoop cond -> WhileLoop cond - - If cond tbody fbody ifattr - | groupInvariant cond -> do - tbody' <- intraGroupBody tbody - fbody' <- intraGroupBody fbody - letBind_ pat $ If cond tbody' fbody' ifattr + localScope (scopeOf form') $ + localScope (scopeOfFParams $ map fst $ ctx ++ val) $ do + loopbody' <- intraGroupBody lvl loopbody + certifying (stmAuxCerts aux) $ + letBind_ pat $ DoLoop ctx val form' loopbody' + where form' = case form of + ForLoop i it bound inps -> ForLoop i it bound inps + WhileLoop cond -> WhileLoop cond - Op (Screma w form arrs) | Just fun <- isMapSOAC form -> do - body_stms <- collectStms_ $ do - forM_ (zip (lambdaParams fun) arrs) $ \(p, arr) -> do - arr_t <- lookupType arr - letBindNames [paramName p] $ BasicOp $ Index arr $ - fullSlice arr_t [DimFix $ Var ltid] - Kernelise.transformStms $ bodyStms $ lambdaBody fun - let comb_body = mkBody body_stms $ bodyResult $ lambdaBody fun - ctid <- newVName "ctid" - letBind_ pat $ Op $ - Out.Combine (Out.combineSpace [(ctid, w)]) (lambdaReturnType fun) [] comb_body - mapM_ (parallelMin . arrayDims) $ patternTypes pat - parallelMin [w] + If cond tbody fbody ifattr -> do + tbody' <- intraGroupBody lvl tbody + fbody' <- intraGroupBody lvl fbody + certifying (stmAuxCerts aux) $ + letBind_ pat $ If cond tbody' fbody' ifattr Op (Screma w form arrs) - | Just (scanfun, nes, foldfun) <- isScanomapSOAC form -> do - let (scan_pes, map_pes) = - splitAt (length nes) $ patternElements pat - scan_input <- procInput ltid (Pattern [] map_pes) w foldfun nes arrs + | Just lam <- isMapSOAC form -> do + let loopnest = MapNesting pat (stmAuxCerts aux) w $ zip (lambdaParams lam) arrs + env = DistEnv { distNest = + singleNesting $ Nesting mempty loopnest + , distScope = + scopeOfPattern pat <> + scopeForKernels (scopeOf lam) <> scope + , distOnInnerMap = + distributeMap + , distOnTopLevelStms = + lift . collectStms_ . intraGroupStms lvl + , distSegLevel = \minw _ _ -> do + lift $ parallelMin minw + return lvl + } + acc = DistAcc { distTargets = singleTarget (pat, bodyResult $ lambdaBody lam) + , distStms = mempty + } - scanfun' <- Kernelise.transformLambda scanfun + addStms =<< + runDistNestT env (distributeMapBodyStms acc (bodyStms $ lambdaBody lam)) - letBind_ (Pattern [] scan_pes) $ - Op $ Out.GroupScan w scanfun' $ zip nes scan_input + Op (Screma w form arrs) + | Just (scanfun, nes, mapfun) <- isScanomapSOAC form -> do + let scanfun' = soacsLambdaToKernels scanfun + mapfun' = soacsLambdaToKernels mapfun + certifying (stmAuxCerts aux) $ + addStms =<< segScan lvl' pat w scanfun' mapfun' nes arrs [] [] parallelMin [w] Op (Screma w form arrs) | Just (reds, map_lam) <- isRedomapSOAC form, - Reduce _ red_lam nes <- singleReduce reds -> do - let (red_pes, map_pes) = - splitAt (length nes) $ patternElements pat - red_input <- procInput ltid (Pattern [] map_pes) w map_lam nes arrs - - red_lam' <- Kernelise.transformLambda red_lam - - letBind_ (Pattern [] red_pes) $ - Op $ Out.GroupReduce w red_lam' $ zip nes red_input + Reduce comm red_lam nes <- singleReduce reds -> do + let red_lam' = soacsLambdaToKernels red_lam + map_lam' = soacsLambdaToKernels map_lam + certifying (stmAuxCerts aux) $ + addStms =<< segRed lvl' pat w [SegRedOp comm red_lam' nes mempty] map_lam' arrs [] [] parallelMin [w] Op (Stream w (Sequential accs) lam arrs) @@ -219,95 +212,44 @@ runBinderT (sequentialStreamWholeArray pat w accs lam arrs) types let replace (Var v) | v == paramName chunk_size_param = w replace se = se - replaceSets (x, y) = (S.map (map replace) x, S.map (map replace) y) - censor replaceSets $ mapM_ intraGroupStm stream_bnds + replaceSets (Acc x y log) = + Acc (S.map (map replace) x) (S.map (map replace) y) log + censor replaceSets $ mapM_ (intraGroupStm lvl) stream_bnds Op (Scatter w lam ivs dests) -> do - parallelMin [w] - ctid <- newVName "ctid" - let cspace = Out.CombineSpace dests [(ctid, w)] - body_stms <- collectStms_ $ do - forM_ (zip (lambdaParams lam) ivs) $ \(p, arr) -> do - arr_t <- lookupType arr - letBindNames [paramName p] $ BasicOp $ Index arr $ - fullSlice arr_t [DimFix $ Var ltid] -- ltid on purpose to enable hoisting. - Kernelise.transformStms $ bodyStms $ lambdaBody lam - let body = mkBody body_stms $ bodyResult $ lambdaBody lam - letBind_ pat $ Op $ Out.Combine cspace (lambdaReturnType lam) mempty body - - BasicOp (Update dest slice (Var v)) -> do - let ws = sliceDims slice - activeForDim w i = BasicOp $ CmpOp (CmpSlt Int32) i w - parallelMin ws - dest' <- letExp "update_inp" $ Op $ Out.Barrier [Var dest] - let new_inds = unflattenIndex (map (primExpFromSubExp int32) ws) - (primExpFromSubExp int32 $ Var ltid) - new_inds' <- mapM (letSubExp "i" <=< toExp) new_inds - active <- letSubExp "active" =<< - foldBinOp LogAnd (constant True) =<< - mapM (letSubExp "active") (zipWith activeForDim ws new_inds') - (active_res, active_stms) <- collectStms $ do - slice' <- - mapM (letSubExp "j" <=< toExp) $ - fixSlice (map (fmap $ primExpFromSubExp int32) slice) new_inds - letInPlace "update_res" dest' (map DimFix slice') $ - BasicOp $ Index v $ map DimFix new_inds' - sync <- letSubExp "update_res" =<< eIf (eSubExp active) - (pure $ mkBody active_stms [Var active_res]) - (pure $ mkBody mempty [Var dest']) - letBind_ pat $ Op $ Out.Barrier [sync] + write_i <- newVName "write_i" + space <- mkSegSpace [(write_i, w)] - BasicOp (Copy arr) -> do - arr_t <- lookupType arr - let w = arraySize 0 arr_t - ctid <- newVName "copy_ctid" - letBind_ pat . Op . Out.Combine (Out.combineSpace [(ctid, w)]) [rowType arr_t] [] <=< - localScope (M.singleton ctid $ IndexInfo Int32) $ - insertStmsM $ resultBodyM . pure <=< letSubExp "v" $ - BasicOp $ Index arr $ fullSlice arr_t [DimFix $ Var ctid] + let lam' = soacsLambdaToKernels lam + (dests_ws, dests_ns, dests_vs) = unzip3 dests + (i_res, v_res) = splitAt (sum dests_ns) $ bodyResult $ lambdaBody lam' + krets = do (a_w, a, is_vs) <- zip3 dests_ws dests_vs $ chunks dests_ns $ zip i_res v_res + return $ WriteReturns [a_w] a [ ([i],v) | (i,v) <- is_vs ] + inputs = do (p, p_a) <- zip (lambdaParams lam') ivs + return $ KernelInput (paramName p) (paramType p) p_a [Var write_i] - BasicOp (Replicate (Shape outer_ws) se) - | [inner_ws] <- map (drop (length outer_ws) . arrayDims) $ patternTypes pat -> do - let ws = outer_ws ++ inner_ws - new_inds' <- replicateM (length ws) $ newVName "new_local_index" - let inner_inds' = drop (length outer_ws) new_inds' - space = Out.combineSpace $ zip new_inds' ws - index = case se of Var v -> BasicOp $ Index v $ - map (DimFix . Var) inner_inds' - Constant{} -> BasicOp $ SubExp se - body <- runBodyBinder $ eBody [pure index] - letBind_ pat $ Op $ - Out.Combine space (map (Prim . elemType) $ patternTypes pat) [] body - mapM_ (parallelAvail . arrayDims) $ patternTypes pat + kstms <- runBinder_ $ localScope (scopeOfSegSpace space) $ do + mapM_ readKernelInput inputs + addStms $ bodyStms $ lambdaBody lam' - _ -> - Kernelise.transformStm stm + certifying (stmAuxCerts aux) $ do + let ts = map rowType $ patternTypes pat + body = KernelBody () kstms krets + letBind_ pat $ Op $ SegOp $ SegMap lvl' space ts body - where procInput :: VName - -> Out.Pattern Out.InKernel - -> SubExp -> Lambda -> [SubExp] -> [VName] - -> IntraGroupM [VName] - procInput ltid map_pat w map_fun nes arrs = do - fold_stms <- collectStms_ $ do - forM_ (zip (lambdaParams map_fun) arrs) $ \(p, arr) -> do - arr_t <- lookupType arr - letBindNames_ [paramName p] $ BasicOp $ Index arr $ - fullSlice arr_t [DimFix $ Var ltid] + parallelMin [w] - Kernelise.transformStms $ bodyStms $ lambdaBody map_fun - let fold_body = mkBody fold_stms $ bodyResult $ lambdaBody map_fun + _ -> + addStm $ soacsStmToKernels stm - op_inps <- replicateM (length nes) (newVName "op_input") - ctid <- newVName "ctid" - letBindNames_ (op_inps ++ patternNames map_pat) $ Op $ - Out.Combine (Out.combineSpace [(ctid, w)]) (lambdaReturnType map_fun) [] fold_body - return op_inps +intraGroupStms :: SegLevel -> Stms SOACS -> IntraGroupM () +intraGroupStms lvl = mapM_ (intraGroupStm lvl) intraGroupParalleliseBody :: (MonadFreshNames m, HasScope Out.Kernels m) => - Dependencies -> Names -> VName -> Body - -> m ([[SubExp]], [[SubExp]], Out.KernelBody Out.InKernel) -intraGroupParalleliseBody deps group_variant ltid body = do - (min_ws, avail_ws, kstms) <- runIntraGroupM (Env ltid deps group_variant) $ - mapM_ intraGroupStm $ bodyStms body - return (min_ws, avail_ws, - KernelBody () kstms $ map GroupsReturn $ bodyResult body) + SegLevel -> Body + -> m ([[SubExp]], [[SubExp]], Log, Out.KernelBody Out.Kernels) +intraGroupParalleliseBody lvl body = do + (Acc min_ws avail_ws log, kstms) <- + runIntraGroupM $ intraGroupStms lvl $ bodyStms body + return (S.toList min_ws, S.toList avail_ws, log, + KernelBody () kstms $ map Returns $ bodyResult body) diff --git a/src/Futhark/Pass/ExtractKernels/Kernelise.hs b/src/Futhark/Pass/ExtractKernels/Kernelise.hs deleted file mode 100644 --- a/src/Futhark/Pass/ExtractKernels/Kernelise.hs +++ /dev/null @@ -1,283 +0,0 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} --- | Sequentialise to kernel statements. -module Futhark.Pass.ExtractKernels.Kernelise - ( transformStm - , transformStms - , transformBody - , transformLambda - , mapIsh - - , groupStreamMapAccumL - ) - where - -import Control.Monad -import qualified Data.Set as S - -import qualified Futhark.Analysis.Alias as Alias -import qualified Futhark.Transform.FirstOrderTransform as FOT -import Futhark.Representation.SOACS -import qualified Futhark.Representation.Kernels as Out -import Futhark.MonadFreshNames -import Futhark.Representation.AST.Attributes.Aliases -import Futhark.Tools - -type Transformer m = (MonadBinder m, - Lore m ~ Out.InKernel, - LocalScope (Lore m) m) - -transformStms :: Transformer m => Stms SOACS -> m () -transformStms = mapM_ transformStm . stmsToList - -transformStm :: Transformer m => Stm -> m () - -transformStm (Let pat aux (Op (Screma w form arrs))) - -- No map-out part - | Just (reds, map_lam) <- isRedomapSOAC form, - Reduce _ red_lam nes <- singleReduce reds, - patternSize pat == length nes = do - - fold_lam <- composeLambda nilFn red_lam map_lam - - chunk_size <- newVName "chunk_size" - chunk_offset <- newVName "chunk_offset" - let arr_idents = drop (length nes) $ patternIdents pat - (fold_acc_params, fold_elem_params) = - splitAt (length nes) $ lambdaParams fold_lam - arr_chunk_params <- mapM (mkArrChunkParam $ Var chunk_size) fold_elem_params - - map_arr_params <- forM arr_idents $ \arr -> - newParam (baseString (identName arr) <> "_in") $ - setOuterSize (identType arr) (Var chunk_size) - - fold_acc_params' <- forM fold_acc_params $ \p -> - newParam (baseString $ paramName p) $ paramType p - - let param_scope = - scopeOfLParams $ fold_acc_params' ++ arr_chunk_params ++ map_arr_params - - redomap_pes <- forM (patternValueElements pat) $ \pe -> - PatElem <$> newVName (baseString $ patElemName pe) <*> pure (patElemType pe) - - redomap_kstms <- collectStms_ $ localScope param_scope $ do - fold_lam' <- transformLambda fold_lam - groupStreamMapAccumL redomap_pes (Var chunk_size) fold_lam' - (map (Var . paramName) fold_acc_params') (map paramName arr_chunk_params) - - let stream_kbody = Out.Body () redomap_kstms $ - map (Var . patElemName) redomap_pes - stream_lam = Out.GroupStreamLambda { Out.groupStreamChunkSize = chunk_size - , Out.groupStreamChunkOffset = chunk_offset - , Out.groupStreamAccParams = fold_acc_params' - , Out.groupStreamArrParams = arr_chunk_params - , Out.groupStreamLambdaBody = stream_kbody - } - - -- Tricky reverse logic: we have to copy all the initial - -- accumulators that were *not* consumed in the original lambda, as - -- a GroupStream will write to its accumulators. - let consumed = consumedByLambda $ Alias.analyseLambda fold_lam - nes' <- forM (zip fold_acc_params nes) $ \(p,e) -> - case e of - Var v | not $ paramName p `S.member` consumed, - not $ primType $ paramType p -> - letSubExp "groupstream_mapaccum_copy" $ BasicOp $ Copy v - _ -> return e - - addStm $ Let pat aux $ Op $ Out.GroupStream w w stream_lam nes' arrs - - where mkArrChunkParam chunk_size arr_param = - newParam (baseString (paramName arr_param) <> "_chunk") $ - arrayOfRow (paramType arr_param) chunk_size - -transformStm (Let pat aux (Op (Stream w (Sequential accs) fold_lam arrs))) = do - let ret = lambdaReturnType fold_lam - -- Sequential streams can be transformed easily to a GroupStream. - -- But we have to create accumulator parameters for mapout. - - chunk_offset <- newVName "streamseq_chunk_offset" - - let (chunk_size_param, fold_acc_params, arr_chunk_params) = - partitionChunkedFoldParameters (length accs) $ lambdaParams fold_lam - chunk_size = paramName chunk_size_param - map_arr_tps = map (`setOuterSize` w) $ drop (length accs) ret - - mapout_arrs <- resultArray map_arr_tps - outarr_params <- forM map_arr_tps $ \map_arr_t -> - Param <$> newVName "redomap_outarr" <*> pure map_arr_t - - lam_body <- localScope (castScope (scopeOf fold_lam) <> - scopeOfLParams outarr_params) $ insertStmsM $ do - res <- bodyBind =<< transformBody (lambdaBody fold_lam) - -- Some results are to be returned; others to be copied into the - -- map-out arrays. - let (acc_res, mapout_res) = splitAt (length accs) res - - mapout_res' <- forM (zip outarr_params mapout_res) $ \(p, r) -> - let slice = fullSlice (paramType p) - [DimSlice (Var chunk_offset) (Var chunk_size) (constant (1::Int32))] - in fmap Var $ letInPlace "mapout_res" (paramName p) slice $ BasicOp $ SubExp r - - return $ resultBody $ acc_res++mapout_res' - - let stream_lam = Out.GroupStreamLambda - { Out.groupStreamChunkSize = chunk_size - , Out.groupStreamChunkOffset = chunk_offset - , Out.groupStreamAccParams = fold_acc_params ++ outarr_params - , Out.groupStreamArrParams = arr_chunk_params - , Out.groupStreamLambdaBody = lam_body - } - - -- Only copy the accs that were not consumed in the original stream. - let consumed = consumedByLambda $ Alias.analyseLambda fold_lam - accs' <- forM (zip fold_acc_params accs) $ \(p, acc) -> - case acc of - Var v | not $ paramName p `S.member` consumed, - not $ primType $ paramType p -> - letSubExp "streamseq_acc_copy" $ BasicOp $ Copy v - _ -> return acc - - addStm $ Let pat aux $ Op $ - Out.GroupStream w w stream_lam (accs'++map Var mapout_arrs) arrs - -transformStm (Let pat aux (DoLoop [] val (ForLoop i Int32 bound []) body)) = do - dummy_chunk_size <- newVName "dummy_chunk_size" - body' <- localScope (scopeOfFParams (map fst val)) $ transformBody body - let lam = Out.GroupStreamLambda { Out.groupStreamChunkSize = dummy_chunk_size - , Out.groupStreamChunkOffset = i - , Out.groupStreamAccParams = map (fmap fromDecl . fst) val - , Out.groupStreamArrParams = [] - , Out.groupStreamLambdaBody = body' } - - -- Copy the initial merge parameters that were not unique in the - -- original stream. - accs' <- forM val $ \(p, initial) -> - case initial of - Var v | not $ unique $ paramDeclType p, - not $ primType $ paramDeclType p -> - letSubExp "streamseq_merge_copy" $ BasicOp $ Copy v - _ -> return initial - - addStm $ Let pat aux $ Op $ Out.GroupStream - bound (constant (1::Int32)) lam accs' [] - -transformStm (Let pat aux (If cond tb fb ts)) = do - tb' <- transformBody tb - fb' <- transformBody fb - addStm $ Let pat aux $ If cond tb' fb' ts - -transformStm bnd = - FOT.transformStmRecursively bnd - -transformBody :: Transformer m => Body -> m (Out.Body Out.InKernel) -transformBody (Body attr bnds res) = do - stms <- collectStms_ $ transformStms bnds - return $ Out.Body attr stms res - -transformLambda :: (MonadFreshNames m, - HasScope lore m, - SameScope lore Out.InKernel) => - Lambda -> m (Out.Lambda Out.InKernel) -transformLambda (Lambda params body rettype) = do - body' <- runBodyBinder $ - localScope (scopeOfLParams params) $ - transformBody body - return $ Lambda params body' rettype - -groupStreamMapAccumL :: Transformer m => - [Out.PatElem Out.InKernel] - -> SubExp - -> Out.Lambda Out.InKernel - -> [SubExp] - -> [VName] - -> m () -groupStreamMapAccumL pes w fold_lam accexps arrexps = do - let acc_num = length accexps - res_tps = lambdaReturnType fold_lam - map_arr_tps = drop acc_num res_tps - - let fold_lam' = fold_lam { lambdaParams = take acc_num $ lambdaParams fold_lam } - fold_lam_aliases = Alias.analyseLambda fold_lam' - - mapout_arrs <- resultArray [ arrayOf t (Shape [w]) NoUniqueness - | t <- map_arr_tps ] - - (merge, i, redomap_loop) <- - FOT.doLoopMapAccumL' w fold_lam_aliases accexps [] mapout_arrs - - -- HACK: we manually inject the indexing here. - dummy_chunk_size <- newVName "groupstream_mapaccum_dummy_chunk_size" - let arr_params = drop acc_num $ lambdaParams fold_lam - arr_params_chunked <- forM arr_params $ \arr_param -> - newParam (baseString (paramName arr_param) <> "_chunked") $ - paramType arr_param `arrayOfRow` Var dummy_chunk_size - let index_bnds = do - (p, arr, arr_t) <- zip3 arr_params (map paramName arr_params_chunked) - (map paramType arr_params_chunked) - return $ mkLet [] [paramIdent p] $ - BasicOp $ Index arr $ fullSlice arr_t [DimFix $ constant (0::Int32)] - - let redomap_kbody = stmsFromList index_bnds `insertStms` redomap_loop - acc_params = map (fmap fromDecl . fst) merge - stream_lam = Out.GroupStreamLambda { Out.groupStreamChunkSize = dummy_chunk_size - , Out.groupStreamChunkOffset = i - , Out.groupStreamAccParams = acc_params - , Out.groupStreamArrParams = arr_params_chunked - , Out.groupStreamLambdaBody = redomap_kbody - } - - letBind_ (Pattern [] pes) $ Op $ - Out.GroupStream w (constant (1::Int32)) stream_lam (accexps++map Var mapout_arrs) arrexps - -resultArray :: MonadBinder m => [Type] -> m [VName] -resultArray = mapM oneArray - where oneArray t = letExp "result" $ BasicOp $ Scratch (elemType t) (arrayDims t) - -mapIsh :: Transformer m => - Pattern - -> Certificates - -> SubExp - -> [LParam] - -> Out.Body Out.InKernel - -> [VName] - -> m () -mapIsh pat cs w params (Out.Body () kstms kres) arrs = do - i <- newVName "i" - - outarrs <- resultArray $ patternTypes pat - - outarr_params <- forM (patternElements pat) $ \pe -> - newParam (baseString (patElemName pe) <> "_out") $ - patElemType pe - - dummy_chunk_size <- newVName "dummy_chunk_size" - params_chunked <- forM params $ \param -> - newParam (baseString (paramName param) <> "_chunked") $ - paramType param `arrayOfRow` Var dummy_chunk_size - - (outarr_params_new, write_elems) <- - fmap unzip $ forM (zip outarr_params kres) $ \(outarr_param, se) -> do - outarr_param_new <- newParam' (<>"_new") outarr_param - return (outarr_param_new, - mkLet [] [paramIdent outarr_param_new] $ BasicOp $ - Update (paramName outarr_param) - (fullSlice (paramType outarr_param) [DimFix $ Var i]) se) - - let index_stms = do - (p, arr, arr_t) <- zip3 params (map paramName params_chunked) $ - map paramType params_chunked - return $ mkLet [] [paramIdent p] $ - BasicOp $ Index arr $ fullSlice arr_t [DimFix $ constant (0::Int32)] - kbody' = Out.Body () (stmsFromList index_stms <> kstms <> stmsFromList write_elems) $ - map (Var . paramName) outarr_params_new - - let stream_lam = Out.GroupStreamLambda { Out.groupStreamChunkSize = dummy_chunk_size - , Out.groupStreamChunkOffset = i - , Out.groupStreamAccParams = outarr_params - , Out.groupStreamArrParams = params_chunked - , Out.groupStreamLambdaBody = kbody' - } - certifying cs $ addStm $ Let pat (StmAux cs ()) $ - Op $ Out.GroupStream w (constant (1::Int32)) stream_lam (map Var outarrs) arrs diff --git a/src/Futhark/Pass/FirstOrderTransform.hs b/src/Futhark/Pass/FirstOrderTransform.hs --- a/src/Futhark/Pass/FirstOrderTransform.hs +++ b/src/Futhark/Pass/FirstOrderTransform.hs @@ -1,6 +1,5 @@ module Futhark.Pass.FirstOrderTransform - ( firstOrderTransform - ) + ( firstOrderTransform ) where import Futhark.Transform.FirstOrderTransform (transformFunDef) @@ -9,7 +8,8 @@ import Futhark.Pass firstOrderTransform :: Pass SOACS Kernels -firstOrderTransform = Pass - "first order transform" - "Transform all second-order array combinators to for-loops." $ - intraproceduralTransformation transformFunDef +firstOrderTransform = + Pass + "first order transform" + "Transform all second-order array combinators to for-loops." $ + intraproceduralTransformation transformFunDef diff --git a/src/Futhark/Pass/KernelBabysitting.hs b/src/Futhark/Pass/KernelBabysitting.hs --- a/src/Futhark/Pass/KernelBabysitting.hs +++ b/src/Futhark/Pass/KernelBabysitting.hs @@ -10,7 +10,6 @@ import Control.Arrow (first) import Control.Monad.State.Strict import qualified Data.Map.Strict as M -import qualified Data.Set as S import Data.Foldable import Data.List import Data.Maybe @@ -57,10 +56,7 @@ Just (Let _ _ (BasicOp (Rearrange perm _))) -> Just $ Just $ rearrangeInverse perm Just (Let _ _ (BasicOp (Reshape _ arr))) -> nonlinearInMemory arr m Just (Let _ _ (BasicOp (Manifest perm _))) -> Just $ Just perm - Just (Let pat _ (Op (HostOp (Kernel _ _ ts _)))) -> - nonlinear =<< find ((==name) . patElemName . fst) - (zip (patternElements pat) ts) - Just (Let pat _ (Op (HostOp (SegMap _ ts _)))) -> + Just (Let pat _ (Op (SegOp (SegMap _ _ ts _)))) -> nonlinear =<< find ((==name) . patElemName . fst) (zip (patternElements pat) ts) _ -> Nothing @@ -71,16 +67,15 @@ | otherwise = Nothing transformStm :: ExpMap -> Stm Kernels -> BabysitM ExpMap - -transformStm expmap (Let pat aux (Op (HostOp op))) = do - let mapper = identityKernelMapper { mapOnKernelKernelBody = - transformKernelBody expmap (kernelSpace op) - } - op' <- mapKernelM mapper op - let stm' = Let pat aux $ Op $ HostOp op' +transformStm expmap (Let pat aux (Op (SegOp op))) = do + let mapper = identitySegOpMapper + { mapOnSegOpBody = + transformKernelBody expmap (segLevel op) (segSpace op) + } + op' <- mapSegOpM mapper op + let stm' = Let pat aux $ Op $ SegOp op' addStm stm' return $ M.fromList [ (name, stm') | name <- patternNames pat ] <> expmap - transformStm expmap (Let pat aux e) = do e' <- mapExpM (transform expmap) e let bnd' = Let pat aux e' @@ -91,45 +86,41 @@ transform expmap = identityMapper { mapOnBody = \scope -> localScope scope . transformBody expmap } -transformKernelBody :: ExpMap -> KernelSpace -> KernelBody InKernel - -> BabysitM (KernelBody InKernel) -transformKernelBody expmap space kbody = do +transformKernelBody :: ExpMap -> SegLevel -> SegSpace -> KernelBody Kernels + -> BabysitM (KernelBody Kernels) +transformKernelBody expmap lvl space kbody = do -- Go spelunking for accesses to arrays that are defined outside the -- kernel body and where the indices are kernel thread indices. scope <- askScope - let thread_gids = map fst $ spaceDimensions space - thread_local = S.fromList $ spaceGlobalId space : spaceLocalId space : thread_gids - free_ker_vars = freeIn kbody `S.difference` getKerVariantIds space + let thread_gids = map fst $ unSegSpace space + thread_local = namesFromList $ segFlat space : thread_gids + free_ker_vars = freeIn kbody `namesSubtract` getKerVariantIds space + num_threads <- letSubExp "num_threads" $ BasicOp $ BinOp (Mul Int32) + (unCount $ segNumGroups lvl) (unCount $ segGroupSize lvl) evalStateT (traverseKernelBodyArrayIndexes free_ker_vars thread_local - (castScope scope <> scopeOfKernelSpace space) - (ensureCoalescedAccess expmap (spaceDimensions space) num_threads) + (scope <> scopeOfSegSpace space) + (ensureCoalescedAccess expmap (unSegSpace space) num_threads) kbody) mempty - where num_threads = spaceNumThreads space - getKerVariantIds (KernelSpace glb_id loc_id grp_id _ _ _ _ (FlatThreadSpace strct)) = - let (gids, _) = unzip strct - in S.fromList $ [glb_id, loc_id, grp_id] ++ gids - getKerVariantIds (KernelSpace glb_id loc_id grp_id _ _ _ _ (NestedThreadSpace strct)) = - let (gids, _, lids, _) = unzip4 strct - in S.fromList $ [glb_id, loc_id, grp_id] ++ gids ++ lids + where getKerVariantIds = namesFromList . M.keys . scopeOfSegSpace type ArrayIndexTransform m = Names -> (VName -> Bool) -> -- thread local? (VName -> SubExp -> Bool)-> -- variant to a certain gid (given as first param)? (SubExp -> Maybe SubExp) -> -- split substitution? - Scope InKernel -> -- type environment + Scope Kernels -> -- type environment VName -> Slice SubExp -> m (Maybe (VName, Slice SubExp)) traverseKernelBodyArrayIndexes :: (Applicative f, Monad f) => Names -> Names - -> Scope InKernel + -> Scope Kernels -> ArrayIndexTransform f - -> KernelBody InKernel - -> f (KernelBody InKernel) + -> KernelBody Kernels + -> f (KernelBody Kernels) traverseKernelBodyArrayIndexes free_ker_vars thread_variant outer_scope f (KernelBody () kstms kres) = KernelBody () . stmsFromList <$> mapM (onStm (varianceInStms mempty kstms, @@ -141,11 +132,6 @@ onBody (variance, szsubst, scope') (lambdaBody lam) where scope' = scope <> scopeOfLParams (lambdaParams lam) - onStreamLambda (variance, szsubst, scope) lam = - (\body' -> lam { groupStreamLambdaBody = body' }) <$> - onBody (variance, szsubst, scope') (groupStreamLambdaBody lam) - where scope' = scope <> scopeOf lam - onBody (variance, szsubst, scope) (Body battr stms bres) = do stms' <- stmsFromList <$> mapM (onStm (variance', szsubst', scope')) (stmsToList stms) Body battr stms' <$> pure bres @@ -161,13 +147,12 @@ BasicOp $ Index arr' is' isGidVariant gid (Var v) = - gid == v || S.member gid (M.findWithDefault (S.singleton v) v variance) + gid == v || nameIn gid (M.findWithDefault (oneName v) v variance) isGidVariant _ _ = False isThreadLocal v = - not $ S.null $ - thread_variant `S.intersection` - M.findWithDefault (S.singleton v) v variance + thread_variant `namesIntersect` + M.findWithDefault (oneName v) v variance sizeSubst (Constant v) = Just $ Constant v sizeSubst (Var v) @@ -178,15 +163,12 @@ onStm (variance, szsubst, scope) (Let pat attr e) = Let pat attr <$> mapExpM (mapper (variance, szsubst, scope)) e - mapper ctx = identityMapper { mapOnBody = const (onBody ctx) - , mapOnOp = onOp ctx - } + onOp ctx (OtherOp soac) = + OtherOp <$> mapSOACM identitySOACMapper{ mapOnSOACLambda = onLambda ctx } soac + onOp _ op = return op - onOp ctx (GroupReduce w lam input) = - GroupReduce w <$> onLambda ctx lam <*> pure input - onOp ctx (GroupStream w maxchunk lam accs arrs) = - GroupStream w maxchunk <$> onStreamLambda ctx lam <*> pure accs <*> pure arrs - onOp _ stm = pure stm + mapper ctx = identityMapper { mapOnBody = const (onBody ctx) + , mapOnOp = onOp ctx } mkSizeSubsts = fold . fmap mkStmSizeSubst where mkStmSizeSubst (Let (Pattern [] [pe]) _ (Op (SplitSpace _ _ _ elems_per_i))) = @@ -233,7 +215,7 @@ not $ null thread_gids, inner_gid <- last thread_gids, length slice >= length perm, - slice' <- map (\i -> slice !! i) perm, + slice' <- map (slice !!) perm, DimFix inner_ind <- last slice', not $ null thread_gids, isGidVariant inner_gid inner_ind -> @@ -252,7 +234,7 @@ not $ tooSmallSlice (primByteSize (elemType t)) rem_slice, is /= map Var (take (length is) thread_gids) || length is == length thread_gids, not (null thread_gids || null is), - not ( S.member (last thread_gids) (S.union (freeIn is) (freeIn rem_slice)) ) -> + not (last thread_gids `nameIn` (freeIn is <> freeIn rem_slice)) -> return Nothing -- We are not fully indexing the array, and the indices are not @@ -263,7 +245,7 @@ not $ null rem_slice, not $ tooSmallSlice (primByteSize (elemType t)) rem_slice, is /= map Var (take (length is) thread_gids) || length is == length thread_gids, - any isThreadLocal (S.toList $ freeIn is) -> do + any isThreadLocal (namesToList $ freeIn is) -> do let perm = coalescingPermutation (length is) $ arrayRank t replace =<< lift (rearrangeInput (nonlinearInMemory arr expmap) perm arr) @@ -333,7 +315,7 @@ -- (because access is likely to be already coalesced) | any isCt is = Nothing - | any (`S.member` free_ker_vars) (mapMaybe mbVarId is) = + | any (`nameIn` free_ker_vars) (mapMaybe mbVarId is) = Nothing | not (null tgids), not (null is), @@ -454,12 +436,12 @@ type VarianceTable = M.Map VName Names -varianceInStms :: VarianceTable -> Stms InKernel -> VarianceTable +varianceInStms :: VarianceTable -> Stms Kernels -> VarianceTable varianceInStms t = foldl varianceInStm t . stmsToList -varianceInStm :: VarianceTable -> Stm InKernel -> VarianceTable +varianceInStm :: VarianceTable -> Stm Kernels -> VarianceTable varianceInStm variance bnd = foldl' add variance $ patternNames $ stmPattern bnd where add variance' v = M.insert v binding_variance variance' - look variance' v = S.insert v $ M.findWithDefault mempty v variance' - binding_variance = mconcat $ map (look variance) $ S.toList (freeIn bnd) + look variance' v = oneName v <> M.findWithDefault mempty v variance' + binding_variance = mconcat $ map (look variance) $ namesToList (freeIn bnd) diff --git a/src/Futhark/Pass/ResolveAssertions.hs b/src/Futhark/Pass/ResolveAssertions.hs --- a/src/Futhark/Pass/ResolveAssertions.hs +++ b/src/Futhark/Pass/ResolveAssertions.hs @@ -35,7 +35,7 @@ where rulebook = Simplify.soacRules <> ruleBook [ RuleBasicOp simplifyScalExp ] [] simplifyScalExp :: BinderOps lore => TopDownRuleBasicOp lore -simplifyScalExp vtable pat _ e = do +simplifyScalExp vtable pat _ e = Simplify $ do res <- SE.toScalExp (`ST.lookupScalExp` vtable) $ BasicOp e case res of -- If the sufficient condition is 'True', then it statically succeeds. diff --git a/src/Futhark/Passes.hs b/src/Futhark/Passes.hs --- a/src/Futhark/Passes.hs +++ b/src/Futhark/Passes.hs @@ -1,5 +1,4 @@ {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE OverloadedStrings #-} -- | Optimisation pipelines. module Futhark.Passes ( standardPipeline @@ -30,13 +29,11 @@ import Futhark.Representation.ExplicitMemory (ExplicitMemory) import Futhark.Representation.Kernels (Kernels) import Futhark.Representation.SOACS (SOACS) -import Futhark.Util standardPipeline :: Pipeline SOACS SOACS standardPipeline = passes [ simplifySOACS , inlineAndRemoveDeadFunctions - , simplifySOACS , performCSE True , simplifySOACS -- We run fusion twice @@ -50,47 +47,31 @@ , removeDeadFunctions ] --- Do we use in-place lowering? Currently enabled by default. Disable by --- setting the environment variable IN_PLACE_LOWERING=0. -usesInPlaceLowering :: Bool -usesInPlaceLowering = - isEnvVarSet "IN_PLACE_LOWERING" True - -inPlaceLoweringMaybe :: Pipeline Kernels Kernels -inPlaceLoweringMaybe = - if usesInPlaceLowering - then onePass inPlaceLowering - else passes [] - kernelsPipeline :: Pipeline SOACS Kernels kernelsPipeline = standardPipeline >>> onePass extractKernels >>> passes [ simplifyKernels , babysitKernels - , simplifyKernels , tileLoops , unstream , simplifyKernels - , performCSE True - , simplifyKernels - ] >>> - inPlaceLoweringMaybe + , inPlaceLowering + ] sequentialPipeline :: Pipeline SOACS Kernels sequentialPipeline = standardPipeline >>> onePass firstOrderTransform >>> passes [ simplifyKernels - ] >>> - inPlaceLoweringMaybe + , inPlaceLowering + ] sequentialCpuPipeline :: Pipeline SOACS ExplicitMemory sequentialCpuPipeline = sequentialPipeline >>> onePass explicitAllocations >>> - passes [ simplifyExplicitMemory - , performCSE False + passes [ performCSE False , simplifyExplicitMemory , doubleBuffer , simplifyExplicitMemory diff --git a/src/Futhark/Pkg/Solve.hs b/src/Futhark/Pkg/Solve.hs --- a/src/Futhark/Pkg/Solve.hs +++ b/src/Futhark/Pkg/Solve.hs @@ -41,7 +41,7 @@ depRoots :: PkgRevDeps -> S.Set PkgPath depRoots (PkgRevDeps m) = S.fromList $ M.keys m --- | Construct a 'BuildList' from a 'RoughBuildList'. This involves +-- Construct a 'BuildList' from a 'RoughBuildList'. This involves -- pruning all packages that cannot be reached from the root. buildList :: S.Set PkgPath -> RoughBuildList -> BuildList buildList roots (RoughBuildList pkgs) = diff --git a/src/Futhark/Pkg/Types.hs b/src/Futhark/Pkg/Types.hs --- a/src/Futhark/Pkg/Types.hs +++ b/src/Futhark/Pkg/Types.hs @@ -200,7 +200,7 @@ addRequiredToManifest new_r pm = let (old, requires') = mapAccumL add Nothing $ commented $ manifestRequire pm in (if isJust old - then pm { manifestRequire = const requires' <$> manifestRequire pm } + then pm { manifestRequire = requires' <$ manifestRequire pm } else pm { manifestRequire = (++[Right new_r]) <$> manifestRequire pm }, old) where add acc (Left c) = (acc, Left c) diff --git a/src/Futhark/Representation/AST/Attributes/Aliases.hs b/src/Futhark/Representation/AST/Attributes/Aliases.hs --- a/src/Futhark/Representation/AST/Attributes/Aliases.hs +++ b/src/Futhark/Representation/AST/Attributes/Aliases.hs @@ -20,12 +20,12 @@ import Control.Arrow (first) import Data.Monoid ((<>)) -import qualified Data.Set as S import Futhark.Representation.AST.Attributes (IsOp) import Futhark.Representation.AST.Syntax import Futhark.Representation.AST.Attributes.Patterns import Futhark.Representation.AST.Attributes.Types +import Futhark.Representation.AST.Attributes.Names class (Annotations lore, AliasedOp (Op lore), AliasesOf (LetAttr lore)) => Aliased lore where @@ -33,7 +33,7 @@ consumedInBody :: Body lore -> Names vnameAliases :: VName -> Names -vnameAliases = S.singleton +vnameAliases = oneName subExpAliases :: SubExp -> Names subExpAliases Constant{} = mempty @@ -77,9 +77,8 @@ ifAliases :: ([Names], Names) -> ([Names], Names) -> [Names] ifAliases (als1,cons1) (als2,cons2) = - map (S.filter notConsumed) $ zipWith mappend als1 als2 - where notConsumed = not . (`S.member` cons) - cons = cons1 <> cons2 + map (`namesSubtract` cons) $ zipWith mappend als1 als2 + where cons = cons1 <> cons2 funcallAliases :: [(SubExp, Diet)] -> [TypeBase shape Uniqueness] -> [Names] funcallAliases args t = @@ -94,10 +93,10 @@ (bodyAliases fb, consumedInBody fb) expAliases (BasicOp op) = primOpAliases op expAliases (DoLoop ctxmerge valmerge _ loopbody) = - map (`S.difference` merge_names) val_aliases + map (`namesSubtract` merge_names) val_aliases where (_ctx_aliases, val_aliases) = splitAt (length ctxmerge) $ bodyAliases loopbody - merge_names = S.fromList $ map (paramName . fst) $ ctxmerge ++ valmerge + merge_names = namesFromList $ map (paramName . fst) $ ctxmerge ++ valmerge expAliases (Apply _ args t _) = funcallAliases args $ retTypeValues t expAliases (Op op) = opAliases op @@ -131,7 +130,7 @@ consumedInExp (DoLoop _ merge _ _) = mconcat (map (subExpAliases . snd) $ filter (unique . paramDeclType . fst) merge) -consumedInExp (BasicOp (Update src _ _)) = S.singleton src +consumedInExp (BasicOp (Update src _ _)) = oneName src consumedInExp (Op op) = consumedInOp op consumedInExp _ = mempty diff --git a/src/Futhark/Representation/AST/Attributes/Names.hs b/src/Futhark/Representation/AST/Attributes/Names.hs --- a/src/Futhark/Representation/AST/Attributes/Names.hs +++ b/src/Futhark/Representation/AST/Attributes/Names.hs @@ -1,47 +1,130 @@ -{-# LANGUAGE TypeSynonymInstances, FlexibleInstances, FlexibleContexts, UndecidableInstances #-} +{-# LANGUAGE FlexibleInstances, FlexibleContexts, UndecidableInstances #-} -- | Facilities for determining which names are used in some syntactic -- construct. The most important interface is the 'FreeIn' class and -- its instances, but for reasons related to the Haskell type system, -- some constructs have specialised functions. module Futhark.Representation.AST.Attributes.Names - ( - -- * Class - FreeIn (..) - , Names - -- * Specialised Functions - , freeInStmsAndRes - -- * Bound Names - , boundInBody - , boundByStm - , boundByStms - , boundByLambda - - , FreeAttr(..) + ( -- * Free names + Names + , nameIn + , oneName + , namesFromList + , namesToList + , namesIntersection + , namesIntersect + , namesSubtract + , mapNames + -- * Class + , FreeIn (..) + , freeIn + -- * Specialised Functions + , freeInStmsAndRes + -- * Bound Names + , boundInBody + , boundByStm + , boundByStms + , boundByLambda + -- * Efficient computation + , FreeAttr(..) + , FV + , fvBind + , fvName + , fvNames ) where -import Control.Monad.Writer +import Control.Monad.State.Strict +import qualified Data.IntMap.Strict as IM import qualified Data.Map.Strict as M -import qualified Data.Set as S import Data.Foldable import Futhark.Representation.AST.Syntax import Futhark.Representation.AST.Traversals import Futhark.Representation.AST.Attributes.Patterns import Futhark.Representation.AST.Attributes.Scope +import Futhark.Util.Pretty +-- | A set of names. +newtype Names = Names { unNames :: IM.IntMap VName } + deriving (Eq, Show) + +instance Semigroup Names where + vs1 <> vs2 = Names $ unNames vs1 <> unNames vs2 + +instance Monoid Names where + mempty = Names mempty + +instance Pretty Names where + ppr = ppr . namesToList + +-- | Does the set of names contain this name? +nameIn :: VName -> Names -> Bool +nameIn v (Names vs) = baseTag v `IM.member` vs + +-- | Construct a name set from a list. Slow. +namesFromList :: [VName] -> Names +namesFromList vs = Names $ IM.fromList $ zip (map baseTag vs) vs + +-- | Turn a name set into a list of names. Slow. +namesToList :: Names -> [VName] +namesToList = IM.elems . unNames + +-- | Construct a name set from a single name. +oneName :: VName -> Names +oneName v = Names $ IM.singleton (baseTag v) v + +-- | The intersection of two name sets. +namesIntersection :: Names -> Names -> Names +namesIntersection (Names vs1) (Names vs2) = Names $ IM.intersection vs1 vs2 + +-- | Do the two name sets intersect? +namesIntersect :: Names -> Names -> Bool +namesIntersect vs1 vs2 = not $ IM.disjoint (unNames vs1) (unNames vs2) + +-- | Subtract the latter name set from the former. +namesSubtract :: Names -> Names -> Names +namesSubtract (Names vs1) (Names vs2) = Names $ IM.difference vs1 vs2 + +-- | Map over the names in a set. +mapNames :: (VName -> VName) -> Names -> Names +mapNames f vs = namesFromList $ map f $ namesToList vs + +-- | A computation to build a free variable set. +newtype FV = FV { unFV :: Names } +-- Right now the variable set is just stored explicitly, without the +-- fancy functional representation that GHC uses. Turns out it's +-- faster this way. + +instance Monoid FV where + mempty = FV mempty + +instance Semigroup FV where + FV fv1 <> FV fv2 = FV $ fv1 <> fv2 + +-- | Consider a variable to be bound in the given 'FV' computation. +fvBind :: Names -> FV -> FV +fvBind vs (FV fv) = FV $ fv `namesSubtract` vs + +-- | Take note of a variable reference. +fvName :: VName -> FV +fvName v = FV $ oneName v + +-- | Take note of a set of variable references. +fvNames :: Names -> FV +fvNames = FV + freeWalker :: (FreeAttr (ExpAttr lore), FreeAttr (BodyAttr lore), FreeIn (FParamAttr lore), FreeIn (LParamAttr lore), FreeIn (LetAttr lore), FreeIn (Op lore)) => - Walker lore (Writer Names) + Walker lore (State FV) freeWalker = identityWalker { - walkOnSubExp = tell . freeIn - , walkOnBody = tell . freeIn - , walkOnVName = tell . S.singleton - , walkOnOp = tell . freeIn + walkOnSubExp = modify . (<>) . freeIn' + , walkOnBody = modify . (<>) . freeIn' + , walkOnVName = modify . (<>) . fvName + , walkOnOp = modify . (<>) . freeIn' } -- | Return the set of variable names that are free in the given @@ -53,30 +136,37 @@ FreeIn (FParamAttr lore), FreeAttr (BodyAttr lore), FreeAttr (ExpAttr lore)) => - Stms lore -> Result -> Names + Stms lore -> Result -> FV freeInStmsAndRes stms res = - (freeIn res `mappend` fold (fmap freeIn stms)) - `S.difference` boundByStms stms + fvBind (boundByStms stms) $ fold (fmap freeIn' stms) <> freeIn' res -- | A class indicating that we can obtain free variable information -- from values of this type. class FreeIn a where - freeIn :: a -> Names + freeIn' :: a -> FV + freeIn' = fvNames . freeIn +-- | The free variables of some syntactic construct. +freeIn :: FreeIn a => a -> Names +freeIn = unFV . freeIn' + +instance FreeIn FV where + freeIn' = id + instance FreeIn () where - freeIn () = mempty + freeIn' () = mempty instance FreeIn Int where - freeIn = const mempty + freeIn' = const mempty instance (FreeIn a, FreeIn b) => FreeIn (a,b) where - freeIn (a,b) = freeIn a <> freeIn b + freeIn' (a,b) = freeIn' a <> freeIn' b instance (FreeIn a, FreeIn b, FreeIn c) => FreeIn (a,b,c) where - freeIn (a,b,c) = freeIn a <> freeIn b <> freeIn c + freeIn' (a,b,c) = freeIn' a <> freeIn' b <> freeIn' c instance FreeIn a => FreeIn [a] where - freeIn = fold . fmap freeIn + freeIn' = fold . fmap freeIn' instance (FreeAttr (ExpAttr lore), FreeAttr (BodyAttr lore), @@ -84,12 +174,9 @@ FreeIn (LParamAttr lore), FreeIn (LetAttr lore), FreeIn (Op lore)) => FreeIn (Lambda lore) where - freeIn (Lambda params body rettype) = - S.filter (`notElem` paramnames) $ inRet <> inParams <> inBody - where inRet = mconcat $ map freeIn rettype - inParams = mconcat $ map freeIn params - inBody = freeIn body - paramnames = map paramName params + freeIn' (Lambda params body rettype) = + fvBind (namesFromList $ map paramName params) $ + freeIn' rettype <> freeIn' params <> freeIn' body instance (FreeAttr (ExpAttr lore), FreeAttr (BodyAttr lore), @@ -97,8 +184,8 @@ FreeIn (LParamAttr lore), FreeIn (LetAttr lore), FreeIn (Op lore)) => FreeIn (Body lore) where - freeIn (Body attr stms res) = - precomputed attr $ freeIn attr <> freeInStmsAndRes stms res + freeIn' (Body attr stms res) = + precomputed attr $ freeIn' attr <> freeInStmsAndRes stms res instance (FreeAttr (ExpAttr lore), FreeAttr (BodyAttr lore), @@ -106,16 +193,16 @@ FreeIn (LParamAttr lore), FreeIn (LetAttr lore), FreeIn (Op lore)) => FreeIn (Exp lore) where - freeIn (DoLoop ctxmerge valmerge form loopbody) = + freeIn' (DoLoop ctxmerge valmerge form loopbody) = let (ctxparams, ctxinits) = unzip ctxmerge (valparams, valinits) = unzip valmerge - bound_here = S.fromList $ M.keys $ + bound_here = namesFromList $ M.keys $ scopeOf form <> scopeOfFParams (ctxparams ++ valparams) - in (freeIn (ctxinits ++ valinits) <> freeIn form <> - freeIn (ctxparams ++ valparams) <> freeIn loopbody) - `S.difference` bound_here - freeIn e = execWriter $ walkExpM freeWalker e + in fvBind bound_here $ + freeIn' (ctxinits ++ valinits) <> freeIn' form <> + freeIn' (ctxparams ++ valparams) <> freeIn' loopbody + freeIn' e = execState (walkExpM freeWalker e) mempty instance (FreeAttr (ExpAttr lore), FreeAttr (BodyAttr lore), @@ -123,78 +210,78 @@ FreeIn (LParamAttr lore), FreeIn (LetAttr lore), FreeIn (Op lore)) => FreeIn (Stm lore) where - freeIn (Let pat (StmAux cs attr) e) = - freeIn cs <> precomputed attr (freeIn attr <> freeIn e <> freeIn pat) + freeIn' (Let pat (StmAux cs attr) e) = + freeIn' cs <> precomputed attr (freeIn' attr <> freeIn' e <> freeIn' pat) instance FreeIn (Stm lore) => FreeIn (Stms lore) where - freeIn = fold . fmap freeIn + freeIn' = fold . fmap freeIn' instance FreeIn Names where - freeIn = id + freeIn' = fvNames instance FreeIn Bool where - freeIn _ = mempty + freeIn' _ = mempty instance FreeIn a => FreeIn (Maybe a) where - freeIn = maybe mempty freeIn + freeIn' = maybe mempty freeIn' instance FreeIn VName where - freeIn = S.singleton + freeIn' = fvName instance FreeIn Ident where - freeIn = freeIn . identType + freeIn' = freeIn' . identType instance FreeIn SubExp where - freeIn (Var v) = freeIn v - freeIn Constant{} = mempty + freeIn' (Var v) = freeIn' v + freeIn' Constant{} = mempty instance FreeIn d => FreeIn (ShapeBase d) where - freeIn = mconcat . map freeIn . shapeDims + freeIn' = freeIn' . shapeDims instance FreeIn d => FreeIn (Ext d) where - freeIn (Free x) = freeIn x - freeIn (Ext _) = mempty + freeIn' (Free x) = freeIn' x + freeIn' (Ext _) = mempty instance FreeIn shape => FreeIn (TypeBase shape u) where - freeIn (Array _ shape _) = freeIn shape - freeIn (Mem _) = mempty - freeIn (Prim _) = mempty + freeIn' (Array _ shape _) = freeIn' shape + freeIn' (Mem _) = mempty + freeIn' (Prim _) = mempty -instance FreeIn attr => FreeIn (ParamT attr) where - freeIn (Param _ attr) = freeIn attr +instance FreeIn attr => FreeIn (Param attr) where + freeIn' (Param _ attr) = freeIn' attr instance FreeIn attr => FreeIn (PatElemT attr) where - freeIn (PatElem _ attr) = freeIn attr + freeIn' (PatElem _ attr) = freeIn' attr instance FreeIn (LParamAttr lore) => FreeIn (LoopForm lore) where - freeIn (ForLoop _ _ bound loop_vars) = freeIn bound <> freeIn loop_vars - freeIn (WhileLoop cond) = freeIn cond + freeIn' (ForLoop _ _ bound loop_vars) = freeIn' bound <> freeIn' loop_vars + freeIn' (WhileLoop cond) = freeIn' cond instance FreeIn d => FreeIn (DimChange d) where - freeIn = Data.Foldable.foldMap freeIn + freeIn' = Data.Foldable.foldMap freeIn' instance FreeIn d => FreeIn (DimIndex d) where - freeIn = Data.Foldable.foldMap freeIn + freeIn' = Data.Foldable.foldMap freeIn' instance FreeIn attr => FreeIn (PatternT attr) where - freeIn (Pattern context values) = - mconcat (map freeIn $ context ++ values) `S.difference` bound_here - where bound_here = S.fromList $ map patElemName $ context ++ values + freeIn' (Pattern context values) = + fvBind bound_here $ freeIn' $ context ++ values + where bound_here = namesFromList $ map patElemName $ context ++ values instance FreeIn Certificates where - freeIn (Certificates cs) = freeIn cs + freeIn' (Certificates cs) = freeIn' cs instance FreeIn attr => FreeIn (StmAux attr) where - freeIn (StmAux cs attr) = freeIn cs <> freeIn attr + freeIn' (StmAux cs attr) = freeIn' cs <> freeIn' attr instance FreeIn a => FreeIn (IfAttr a) where - freeIn (IfAttr r _) = freeIn r + freeIn' (IfAttr r _) = freeIn' r -- | Either return precomputed free names stored in the attribute, or -- the freshly computed names. Relies on lazy evaluation to avoid the -- work. class FreeIn attr => FreeAttr attr where - precomputed :: attr -> Names -> Names + precomputed :: attr -> FV -> FV precomputed _ = id instance FreeAttr () where @@ -216,7 +303,7 @@ -- | The names bound by a binding. boundByStm :: Stm lore -> Names -boundByStm = S.fromList . patternNames . stmPattern +boundByStm = namesFromList . patternNames . stmPattern -- | The names bound by the bindings. boundByStms :: Stms lore -> Names diff --git a/src/Futhark/Representation/AST/Attributes/Patterns.hs b/src/Futhark/Representation/AST/Attributes/Patterns.hs --- a/src/Futhark/Representation/AST/Attributes/Patterns.hs +++ b/src/Futhark/Representation/AST/Attributes/Patterns.hs @@ -33,15 +33,15 @@ (existentialiseExtTypes, staticShapes, Typed(..), DeclTyped(..)) -- | The 'Type' of a parameter. -paramType :: Typed attr => ParamT attr -> Type +paramType :: Typed attr => Param attr -> Type paramType = typeOf -- | The 'DeclType' of a parameter. -paramDeclType :: DeclTyped attr => ParamT attr -> DeclType +paramDeclType :: DeclTyped attr => Param attr -> DeclType paramDeclType = declTypeOf -- | An 'Ident' corresponding to a parameter. -paramIdent :: Typed attr => ParamT attr -> Ident +paramIdent :: Typed attr => Param attr -> Ident paramIdent param = Ident (paramName param) (typeOf param) -- | An 'Ident' corresponding to a pattern element. diff --git a/src/Futhark/Representation/AST/Attributes/Ranges.hs b/src/Futhark/Representation/AST/Attributes/Ranges.hs --- a/src/Futhark/Representation/AST/Attributes/Ranges.hs +++ b/src/Futhark/Representation/AST/Attributes/Ranges.hs @@ -2,7 +2,6 @@ {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeSynonymInstances #-} {-# LANGUAGE FlexibleInstances #-} -- | Utility declarations for performing range analysis. module Futhark.Representation.AST.Attributes.Ranges @@ -24,7 +23,6 @@ where import Data.Monoid ((<>)) -import qualified Data.Set as S import qualified Data.Map.Strict as M import Futhark.Representation.AST.Attributes @@ -62,10 +60,10 @@ rename = substituteRename instance FreeIn KnownBound where - freeIn (VarBound v) = freeIn v - freeIn (MinimumBound b1 b2) = freeIn b1 <> freeIn b2 - freeIn (MaximumBound b1 b2) = freeIn b1 <> freeIn b2 - freeIn (ScalarBound e) = freeIn e + freeIn' (VarBound v) = freeIn' v + freeIn' (MinimumBound b1 b2) = freeIn' b1 <> freeIn' b2 + freeIn' (MaximumBound b1 b2) = freeIn' b1 <> freeIn' b2 + freeIn' (ScalarBound e) = freeIn' e instance FreeAttr KnownBound where precomputed _ = id @@ -233,7 +231,7 @@ expRanges (DoLoop ctxmerge valmerge (ForLoop i Int32 iterations _) body) = zipWith returnedRange valmerge $ rangesOf body where bound_in_loop = - S.fromList $ i : map (paramName . fst) (ctxmerge++valmerge) ++ + namesFromList $ i : map (paramName . fst) (ctxmerge++valmerge) ++ concatMap (patternNames . stmPattern) (bodyStms body) returnedRange mergeparam (lower, upper) = @@ -245,7 +243,7 @@ Just bound' <- boundToScalExp bound, let se_diff = AS.simplify (SE.SMinus (SE.Id (paramName param) $ IntType Int32) bound') M.empty, - S.null $ S.intersection bound_in_loop $ freeIn se_diff = + namesIntersect bound_in_loop $ freeIn se_diff = Just $ ScalarBound $ SE.SPlus (SE.subExpToScalExp mergeinit $ IntType Int32) $ SE.STimes se_diff $ SE.MaxMin False [SE.subExpToScalExp iterations $ IntType Int32, 0] diff --git a/src/Futhark/Representation/AST/Attributes/Scope.hs b/src/Futhark/Representation/AST/Attributes/Scope.hs --- a/src/Futhark/Representation/AST/Attributes/Scope.hs +++ b/src/Futhark/Representation/AST/Attributes/Scope.hs @@ -4,7 +4,6 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE StandaloneDeriving #-} -- | This module defines the concept of a type environment as a @@ -169,12 +168,12 @@ scopeOfPatElem (PatElem name attr) = M.singleton name $ LetInfo attr scopeOfLParams :: LParamAttr lore ~ attr => - [ParamT attr] -> Scope lore + [Param attr] -> Scope lore scopeOfLParams = M.fromList . map f where f param = (paramName param, LParamInfo $ paramAttr param) scopeOfFParams :: FParamAttr lore ~ attr => - [ParamT attr] -> Scope lore + [Param attr] -> Scope lore scopeOfFParams = M.fromList . map f where f param = (paramName param, FParamInfo $ paramAttr param) diff --git a/src/Futhark/Representation/AST/Attributes/TypeOf.hs b/src/Futhark/Representation/AST/Attributes/TypeOf.hs --- a/src/Futhark/Representation/AST/Attributes/TypeOf.hs +++ b/src/Futhark/Representation/AST/Attributes/TypeOf.hs @@ -23,7 +23,6 @@ , bodyExtType , primOpType , mapType - , subExpShapeContext -- * Return type , module Futhark.Representation.AST.RetType @@ -155,13 +154,6 @@ where bndscope = scopeOf stms boundInLet (Let pat _ _) = S.fromList $ patternNames pat bound = S.toList $ fold $ fmap boundInLet stms - --- | Given the return type of a function and the subexpressions --- returned by that function, return the size context. -subExpShapeContext :: HasScope t m => - [TypeBase ExtShape u] -> [SubExp] -> m [SubExp] -subExpShapeContext rettype ses = - extractShapeContext rettype <$> traverse (fmap arrayDims . subExpType) ses -- | Given the context and value merge parameters of a Futhark @loop@, -- produce the return type. diff --git a/src/Futhark/Representation/AST/Attributes/Types.hs b/src/Futhark/Representation/AST/Attributes/Types.hs --- a/src/Futhark/Representation/AST/Attributes/Types.hs +++ b/src/Futhark/Representation/AST/Attributes/Types.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE FlexibleContexts, FlexibleInstances, TypeSynonymInstances #-} +{-# LANGUAGE FlexibleContexts, FlexibleInstances #-} -- | Functions for inspecting and constructing various types. module Futhark.Representation.AST.Attributes.Types ( @@ -47,9 +47,7 @@ , extractShapeContext , shapeContext - , shapeContextSize , hasStaticShape - , hasStaticShapes , generaliseExtTypes , existentialiseExtTypes , shapeMapping @@ -367,10 +365,6 @@ where ext (Ext x) = Just x ext (Free _) = Nothing --- | The size of the set that would be returned by 'shapeContext'. -shapeContextSize :: [ExtType] -> Int -shapeContextSize = S.size . shapeContext - -- | If all dimensions of the given 'RetType' are statically known, -- return the corresponding list of 'Type'. hasStaticShape :: ExtType -> Maybe Type @@ -380,9 +374,6 @@ Array bt <$> (Shape <$> mapM isFree shape) <*> pure u where isFree (Free s) = Just s isFree (Ext _) = Nothing - -hasStaticShapes :: [ExtType] -> Maybe [Type] -hasStaticShapes = mapM hasStaticShape -- | Given two lists of 'ExtType's of the same length, return a list -- of 'ExtType's that is a subtype (as per 'isSubtypeOf') of the two diff --git a/src/Futhark/Representation/AST/Pretty.hs b/src/Futhark/Representation/AST/Pretty.hs --- a/src/Futhark/Representation/AST/Pretty.hs +++ b/src/Futhark/Representation/AST/Pretty.hs @@ -31,7 +31,7 @@ instance PrettyAnnot (PatElemT (TypeBase shape u)) where ppAnnot = const Nothing -instance PrettyAnnot (ParamT (TypeBase shape u)) where +instance PrettyAnnot (Param (TypeBase shape u)) where ppAnnot = const Nothing instance PrettyAnnot () where @@ -41,8 +41,8 @@ class (Annotations lore, Pretty (RetType lore), Pretty (BranchType lore), - Pretty (ParamT (FParamAttr lore)), - Pretty (ParamT (LParamAttr lore)), + Pretty (Param (FParamAttr lore)), + Pretty (Param (LParamAttr lore)), Pretty (PatElemT (LetAttr lore)), PrettyAnnot (PatElem lore), PrettyAnnot (FParam lore), @@ -137,15 +137,15 @@ instance Pretty (PatElemT Type) where ppr (PatElem name t) = ppr t <+> ppr name -instance Pretty (ParamT b) => Pretty (ParamT (a,b)) where +instance Pretty (Param b) => Pretty (Param (a,b)) where ppr = ppr . fmap snd -instance Pretty (ParamT DeclType) where +instance Pretty (Param DeclType) where ppr (Param name t) = ppr t <+> ppr name -instance Pretty (ParamT Type) where +instance Pretty (Param Type) where ppr (Param name t) = ppr t <+> ppr name @@ -158,13 +158,15 @@ (True, Nothing) -> equals e' (_, Just ann) -> equals (ann e') (False, Nothing) -> equals <+/> e' - where e' = ppr cs <> ppr e + where e' | linebreak = ppr cs ppr e + | otherwise = ppr cs <> ppr e linebreak = case e of DoLoop{} -> True Op{} -> True If{} -> True BasicOp ArrayLit{} -> False - _ -> False + BasicOp Assert{} -> True + _ -> cs /= mempty instance Pretty (BasicOp lore) where ppr (SubExp se) = ppr se @@ -253,7 +255,7 @@ ppr (Lambda [] _ []) = text "nilFn" ppr (Lambda params body rettype) = annot (mapMaybe ppAnnot params) $ - text "fn" <+> ppTuple' rettype <+> + text "fn" <+> ppTuple' rettype <+/> parens (commasep (map ppr params)) <+> text "=>" indent 2 (ppr body) diff --git a/src/Futhark/Representation/AST/Syntax.hs b/src/Futhark/Representation/AST/Syntax.hs --- a/src/Futhark/Representation/AST/Syntax.hs +++ b/src/Futhark/Representation/AST/Syntax.hs @@ -48,15 +48,13 @@ , Lambda -- * Definitions - , ParamT (..) + , Param (..) , FParam , LParam - , FunDefT (..) - , FunDef + , FunDef (..) , EntryPoint , EntryPointType(..) - , ProgT(..) - , Prog + , Prog(..) -- * Utils , oneStm @@ -337,23 +335,23 @@ -- | Type alias for namespacing reasons. type Lambda = LambdaT -type FParam lore = ParamT (FParamAttr lore) +type FParam lore = Param (FParamAttr lore) -type LParam lore = ParamT (LParamAttr lore) +type LParam lore = Param (LParamAttr lore) -- | Function Declarations -data FunDefT lore = FunDef { funDefEntryPoint :: Maybe EntryPoint - -- ^ Contains a value if this function is - -- an entry point. - , funDefName :: Name - , funDefRetType :: [RetType lore] - , funDefParams :: [FParam lore] - , funDefBody :: BodyT lore - } +data FunDef lore = FunDef { funDefEntryPoint :: Maybe EntryPoint + -- ^ Contains a value if this function is + -- an entry point. + , funDefName :: Name + , funDefRetType :: [RetType lore] + , funDefParams :: [FParam lore] + , funDefBody :: BodyT lore + } -deriving instance Annotations lore => Eq (FunDefT lore) -deriving instance Annotations lore => Show (FunDefT lore) -deriving instance Annotations lore => Ord (FunDefT lore) +deriving instance Annotations lore => Eq (FunDef lore) +deriving instance Annotations lore => Show (FunDef lore) +deriving instance Annotations lore => Ord (FunDef lore) -- | Information about the parameters and return value of an entry -- point. The first element is for parameters, the second for return @@ -373,12 +371,6 @@ -- ^ Maps directly. deriving (Eq, Show, Ord) --- | Type alias for namespace reasons. -type FunDef = FunDefT - -- | An entire Futhark program. -newtype ProgT lore = Prog { progFunctions :: [FunDef lore] } - deriving (Eq, Ord, Show) - --- | Type alias for namespace reasons. -type Prog = ProgT +newtype Prog lore = Prog { progFunctions :: [FunDef lore] } + deriving (Eq, Ord, Show) diff --git a/src/Futhark/Representation/AST/Syntax/Core.hs b/src/Futhark/Representation/AST/Syntax/Core.hs --- a/src/Futhark/Representation/AST/Syntax/Core.hs +++ b/src/Futhark/Representation/AST/Syntax/Core.hs @@ -38,8 +38,7 @@ , Ident (..) , Certificates(..) , SubExp(..) - , ParamT (..) - , Param + , Param (..) , DimIndex (..) , Slice , dimFix @@ -48,16 +47,12 @@ , unitSlice , fixSlice , PatElemT (..) - - -- * Miscellaneous - , Names ) where import Control.Monad.State import Data.Maybe import Data.Monoid ((<>)) import Data.String -import qualified Data.Set as S import qualified Data.Map.Strict as M import Data.Traversable @@ -227,25 +222,22 @@ | Var VName deriving (Show, Eq, Ord) --- | A function parameter. -data ParamT attr = Param - { paramName :: VName - -- ^ Name of the parameter. - , paramAttr :: attr - -- ^ Function parameter attribute. - } - deriving (Ord, Show, Eq) - --- | A type alias for namespace control. -type Param = ParamT +-- | A function or lambda parameter. +data Param attr = Param + { paramName :: VName + -- ^ Name of the parameter. + , paramAttr :: attr + -- ^ Function parameter attribute. + } + deriving (Ord, Show, Eq) -instance Foldable ParamT where +instance Foldable Param where foldMap = foldMapDefault -instance Functor ParamT where +instance Functor Param where fmap = fmapDefault -instance Traversable ParamT where +instance Traversable Param where traverse f (Param name attr) = Param name <$> f attr -- | How to index a single dimension of an array. @@ -314,9 +306,6 @@ instance Functor PatElemT where fmap f (PatElem name attr) = PatElem name (f attr) - --- | A set of names. -type Names = S.Set VName -- | An error message is a list of error parts, which are concatenated -- to form the final message. diff --git a/src/Futhark/Representation/AST/Traversals.hs b/src/Futhark/Representation/AST/Traversals.hs --- a/src/Futhark/Representation/AST/Traversals.hs +++ b/src/Futhark/Representation/AST/Traversals.hs @@ -42,6 +42,7 @@ import Control.Monad import Control.Monad.Identity import qualified Data.Traversable +import Data.Foldable (traverse_) import Data.Monoid ((<>)) import Futhark.Representation.AST.Syntax @@ -85,7 +86,7 @@ mapExpM tv (BasicOp (SubExp se)) = BasicOp <$> (SubExp <$> mapOnSubExp tv se) mapExpM tv (BasicOp (ArrayLit els rowt)) = - BasicOp <$> (pure ArrayLit <*> mapM (mapOnSubExp tv) els <*> + BasicOp <$> (ArrayLit <$> mapM (mapOnSubExp tv) els <*> mapOnType (mapOnSubExp tv) rowt) mapExpM tv (BasicOp (BinOp bop x y)) = BasicOp <$> (BinOp bop <$> mapOnSubExp tv x <*> mapOnSubExp tv y) @@ -109,7 +110,7 @@ BasicOp <$> (Update <$> mapOnVName tv arr <*> mapM (traverse (mapOnSubExp tv)) slice <*> mapOnSubExp tv se) mapExpM tv (BasicOp (Iota n x s et)) = - BasicOp <$> (pure Iota <*> mapOnSubExp tv n <*> mapOnSubExp tv x <*> mapOnSubExp tv s <*> pure et) + BasicOp <$> (Iota <$> mapOnSubExp tv n <*> mapOnSubExp tv x <*> mapOnSubExp tv s <*> pure et) mapExpM tv (BasicOp (Replicate shape vexp)) = BasicOp <$> (Replicate <$> mapOnShape tv shape <*> mapOnSubExp tv vexp) mapExpM tv (BasicOp (Repeat shapes innershape v)) = @@ -130,7 +131,7 @@ mapOnVName tv x <*> mapM (mapOnVName tv) ys <*> mapOnSubExp tv size) mapExpM tv (BasicOp (Copy e)) = - BasicOp <$> (pure Copy <*> mapOnVName tv e) + BasicOp <$> (Copy <$> mapOnVName tv e) mapExpM tv (BasicOp (Manifest perm e)) = BasicOp <$> (Manifest perm <$> mapOnVName tv e) mapExpM tv (BasicOp (Assert e msg loc)) = @@ -211,20 +212,80 @@ , walkOnOp = const $ return () } -walkMapper :: Monad m => Walker lore m -> Mapper lore lore m -walkMapper f = Mapper { - mapOnSubExp = wrap walkOnSubExp - , mapOnBody = const $ wrap walkOnBody - , mapOnVName = wrap walkOnVName - , mapOnRetType = wrap walkOnRetType - , mapOnBranchType = wrap walkOnBranchType - , mapOnFParam = wrap walkOnFParam - , mapOnLParam = wrap walkOnLParam - , mapOnOp = wrap walkOnOp - } - where wrap op k = op f k >> return k +walkOnShape :: Monad m => Walker lore m -> Shape -> m () +walkOnShape tv (Shape ds) = mapM_ (walkOnSubExp tv) ds --- | As 'walkBodyM', but for expressions. +walkOnType :: Monad m => + (SubExp -> m ()) -> Type -> m () +walkOnType _ Prim{} = return () +walkOnType _ Mem{} = return () +walkOnType f (Array _ shape _) = mapM_ f $ shapeDims shape + +walkOnLoopForm :: Monad m => Walker lore m -> LoopForm lore -> m () +walkOnLoopForm tv (ForLoop i _ bound loop_vars) = + walkOnVName tv i >> walkOnSubExp tv bound >> + mapM_ (walkOnLParam tv) loop_lparams >> mapM_ (walkOnVName tv) loop_arrs + where (loop_lparams,loop_arrs) = unzip loop_vars +walkOnLoopForm tv (WhileLoop cond) = + walkOnVName tv cond + +-- | As 'mapExpM', but do not construct a result AST. walkExpM :: Monad m => Walker lore m -> Exp lore -> m () -walkExpM f = void . mapExpM m - where m = walkMapper f +walkExpM tv (BasicOp (SubExp se)) = + walkOnSubExp tv se +walkExpM tv (BasicOp (ArrayLit els rowt)) = + mapM_ (walkOnSubExp tv) els >> walkOnType (walkOnSubExp tv) rowt +walkExpM tv (BasicOp (BinOp _ x y)) = + walkOnSubExp tv x >> walkOnSubExp tv y +walkExpM tv (BasicOp (CmpOp _ x y)) = + walkOnSubExp tv x >> walkOnSubExp tv y +walkExpM tv (BasicOp (ConvOp _ x)) = + walkOnSubExp tv x +walkExpM tv (BasicOp (UnOp _ x)) = + walkOnSubExp tv x +walkExpM tv (If c texp fexp (IfAttr ts _)) = + walkOnSubExp tv c >> walkOnBody tv texp >> + walkOnBody tv fexp >> mapM_ (walkOnBranchType tv) ts +walkExpM tv (Apply _ args ret _) = + mapM_ (walkOnSubExp tv . fst) args >> mapM_ (walkOnRetType tv) ret +walkExpM tv (BasicOp (Index arr slice)) = + walkOnVName tv arr >> mapM_ (traverse_ (walkOnSubExp tv)) slice +walkExpM tv (BasicOp (Update arr slice se)) = + walkOnVName tv arr >> + mapM_ (traverse_ (walkOnSubExp tv)) slice >> + walkOnSubExp tv se +walkExpM tv (BasicOp (Iota n x s _)) = + walkOnSubExp tv n >> walkOnSubExp tv x >> walkOnSubExp tv s +walkExpM tv (BasicOp (Replicate shape vexp)) = + walkOnShape tv shape >> walkOnSubExp tv vexp +walkExpM tv (BasicOp (Repeat shapes innershape v)) = + mapM_ (walkOnShape tv) shapes >> walkOnShape tv innershape >> walkOnVName tv v +walkExpM tv (BasicOp (Scratch _ shape)) = + mapM_ (walkOnSubExp tv) shape +walkExpM tv (BasicOp (Reshape shape arrexp)) = + mapM_ (traverse_ (walkOnSubExp tv)) shape >> walkOnVName tv arrexp +walkExpM tv (BasicOp (Rearrange _ e)) = + walkOnVName tv e +walkExpM tv (BasicOp (Rotate es e)) = + mapM_ (walkOnSubExp tv) es >> walkOnVName tv e +walkExpM tv (BasicOp (Concat _ x ys size)) = + walkOnVName tv x >> mapM_ (walkOnVName tv) ys >> walkOnSubExp tv size +walkExpM tv (BasicOp (Copy e)) = + walkOnVName tv e +walkExpM tv (BasicOp (Manifest _ e)) = + walkOnVName tv e +walkExpM tv (BasicOp (Assert e msg _)) = + walkOnSubExp tv e >> traverse_ (walkOnSubExp tv) msg +walkExpM tv (BasicOp (Opaque e)) = + walkOnSubExp tv e +walkExpM tv (DoLoop ctxmerge valmerge form loopbody) = do + mapM_ (walkOnFParam tv) ctxparams + mapM_ (walkOnFParam tv) valparams + walkOnLoopForm tv form + mapM_ (walkOnSubExp tv) ctxinits + mapM_ (walkOnSubExp tv) valinits + walkOnBody tv loopbody + where (ctxparams,ctxinits) = unzip ctxmerge + (valparams,valinits) = unzip valmerge +walkExpM tv (Op op) = + walkOnOp tv op diff --git a/src/Futhark/Representation/Aliases.hs b/src/Futhark/Representation/Aliases.hs --- a/src/Futhark/Representation/Aliases.hs +++ b/src/Futhark/Representation/Aliases.hs @@ -45,7 +45,6 @@ import Data.Maybe import Data.Monoid ((<>)) import qualified Data.Map.Strict as M -import qualified Data.Set as S import Futhark.Representation.AST.Syntax import Futhark.Representation.AST.Attributes @@ -85,10 +84,10 @@ substituteNames substs (Names' names) = Names' $ substituteNames substs names instance FreeIn Names' where - freeIn = const mempty + freeIn' = const mempty instance PP.Pretty Names' where - ppr = PP.commasep . map PP.ppr . S.toList . unNames + ppr = PP.commasep . map PP.ppr . namesToList . unNames -- | The aliases of the let-bound variable. type VarAliases = Names' @@ -162,7 +161,7 @@ bodyAliases body _ -> Nothing - expAttr = case S.toList $ unNames consumed of + expAttr = case namesToList $ unNames consumed of [] -> Nothing als -> Just $ PP.oneLine $ PP.text "-- Consumes " <> PP.commasep (map PP.ppr als) @@ -171,19 +170,17 @@ maybeComment [] = Nothing maybeComment cs = Just $ PP.folddoc (PP.) cs -aliasComment :: (PP.Pretty a, PP.Pretty b) => - a -> S.Set b -> Maybe PP.Doc +aliasComment :: PP.Pretty a => a -> Names -> Maybe PP.Doc aliasComment name als = - case S.toList als of + case namesToList als of [] -> Nothing als' -> Just $ PP.oneLine $ PP.text "-- " <> PP.ppr name <> PP.text " aliases " <> PP.commasep (map PP.ppr als') -resultAliasComment :: (PP.Pretty a, PP.Pretty b) => - a -> S.Set b -> Maybe PP.Doc +resultAliasComment :: PP.Pretty a => a -> Names -> Maybe PP.Doc resultAliasComment name als = - case S.toList als of + case namesToList als of [] -> Nothing als' -> Just $ PP.oneLine $ PP.text "-- Result of " <> PP.ppr name <> PP.text " aliases " <> @@ -273,15 +270,15 @@ mkContextAliases pat (DoLoop ctxmerge valmerge _ body) = let ctx = map fst ctxmerge init_als = zip mergenames $ map (subExpAliases . snd) $ ctxmerge ++ valmerge - expand als = als <> S.unions (mapMaybe (`lookup` init_als) (S.toList als)) + expand als = als <> mconcat (mapMaybe (`lookup` init_als) (namesToList als)) merge_als = zip mergenames $ - map ((`S.difference` mergenames_set) . expand) $ + map ((`namesSubtract` mergenames_set) . expand) $ bodyAliases body in if length ctx == length (patternContextElements pat) then map (fromMaybe mempty . flip lookup merge_als . paramName) ctx else map (const mempty) $ patternContextElements pat where mergenames = map (paramName . fst) $ ctxmerge ++ valmerge - mergenames_set = S.fromList mergenames + mergenames_set = namesFromList mergenames mkContextAliases pat (If _ tbranch fbranch _) = take (length $ patternContextNames pat) $ zipWith (<>) (bodyAliases tbranch) (bodyAliases fbranch) @@ -299,10 +296,9 @@ -- bound in bnds. let (aliases, consumed) = mkStmsAliases bnds res boundNames = - fold $ fmap (S.fromList . patternNames . stmPattern) bnds - bound = (`S.member` boundNames) - aliases' = map (S.filter (not . bound)) aliases - consumed' = S.filter (not . bound) consumed + fold $ fmap (namesFromList . patternNames . stmPattern) bnds + aliases' = map (`namesSubtract` boundNames) aliases + consumed' = consumed `namesSubtract` boundNames in (map Names' aliases', Names' consumed') mkStmsAliases :: Aliased lore => @@ -315,7 +311,7 @@ delve (aliasmap, consumed) (bnd:bnds') = delve (trackAliases (aliasmap, consumed) bnd) bnds' aliasClosure aliasmap names = - names `S.union` mconcat (map look $ S.toList names) + names <> mconcat (map look $ namesToList names) where look k = M.findWithDefault mempty k aliasmap -- | Everything consumed in the given bindings and result (even transitively). @@ -336,7 +332,7 @@ consumed' = consumed <> addAliasesOfAliases (consumedInStm bnd) in (aliasmap', consumed') where addAliasesOfAliases names = names <> aliasesOfAliases names - aliasesOfAliases = mconcat . map look . S.toList + aliasesOfAliases = mconcat . map look . namesToList look k = M.findWithDefault mempty k aliasmap mkAliasedLetStm :: (Attributes lore, CanBeAliased (Op lore)) => diff --git a/src/Futhark/Representation/ExplicitMemory.hs b/src/Futhark/Representation/ExplicitMemory.hs --- a/src/Futhark/Representation/ExplicitMemory.hs +++ b/src/Futhark/Representation/ExplicitMemory.hs @@ -56,7 +56,6 @@ module Futhark.Representation.ExplicitMemory ( -- * The Lore definition ExplicitMemory - , InKernel , MemOp (..) , MemInfo (..) , MemBound @@ -81,13 +80,15 @@ , ixFunMatchesInnerShape , existentialiseIxFun + , scalarMemory + , allScalarMemory + -- * Module re-exports , module Futhark.Representation.AST.Attributes , module Futhark.Representation.AST.Traversals , module Futhark.Representation.AST.Pretty , module Futhark.Representation.AST.Syntax , module Futhark.Representation.Kernels.Kernel - , module Futhark.Representation.Kernels.KernelExp , module Futhark.Analysis.PrimExp.Convert ) where @@ -104,7 +105,6 @@ import Futhark.Analysis.Metrics import Futhark.Representation.AST.Syntax import Futhark.Representation.Kernels.Kernel -import Futhark.Representation.Kernels.KernelExp import Futhark.Representation.AST.Attributes import Futhark.Representation.AST.Attributes.Aliases import Futhark.Representation.AST.Traversals @@ -127,7 +127,6 @@ -- | A lore containing explicit memory information. data ExplicitMemory -data InKernel type ExplicitMemorish lore = (SameScope lore ExplicitMemory, RetType lore ~ FunReturns, @@ -151,8 +150,8 @@ deriving (Eq, Ord, Show) instance FreeIn inner => FreeIn (MemOp inner) where - freeIn (Alloc size _) = freeIn size - freeIn (Inner k) = freeIn k + freeIn' (Alloc size _) = freeIn' size + freeIn' (Inner k) = freeIn' k instance TypedOp inner => TypedOp (MemOp inner) where opType (Alloc _ space) = pure [Mem space] @@ -225,15 +224,7 @@ type LParamAttr ExplicitMemory = MemInfo SubExp NoUniqueness MemBind type RetType ExplicitMemory = FunReturns type BranchType ExplicitMemory = BodyReturns - type Op ExplicitMemory = MemOp (HostOp ExplicitMemory (Kernel InKernel)) - -instance Annotations InKernel where - type LetAttr InKernel = MemInfo SubExp NoUniqueness MemBind - type FParamAttr InKernel = MemInfo SubExp Uniqueness MemBind - type LParamAttr InKernel = MemInfo SubExp NoUniqueness MemBind - type RetType InKernel = FunReturns - type BranchType InKernel = BodyReturns - type Op InKernel = MemOp (KernelExp InKernel) + type Op ExplicitMemory = MemOp (HostOp ExplicitMemory ()) -- | The index function representation used for memory annotations. type IxFun = IxFun.IxFun (PrimExp VName) @@ -289,9 +280,9 @@ declTypeOf (MemArray bt shape u _) = Array bt shape u instance (FreeIn d, FreeIn ret) => FreeIn (MemInfo d u ret) where - freeIn (MemArray _ shape _ ret) = freeIn shape <> freeIn ret - freeIn MemMem{} = mempty - freeIn MemPrim{} = mempty + freeIn' (MemArray _ shape _ ret) = freeIn' shape <> freeIn' ret + freeIn' MemMem{} = mempty + freeIn' MemPrim{} = mempty instance (Substitute d, Substitute ret) => Substitute (MemInfo d u ret) where substituteNames subst (MemArray bt shape u ret) = @@ -371,7 +362,7 @@ PP.text "@" <> PP.ppr mem <> PP.text "->" <> PP.ppr ixfun instance FreeIn MemBind where - freeIn (ArrayIn mem ixfun) = freeIn mem <> freeIn ixfun + freeIn' (ArrayIn mem ixfun) = freeIn' mem <> freeIn' ixfun -- | A description of the memory properties of an array being returned -- by an operation. @@ -433,8 +424,8 @@ Space s -> PP.text $ "@" ++ s instance FreeIn MemReturn where - freeIn (ReturnsInBlock v ixfun) = freeIn v <> freeIn ixfun - freeIn _ = mempty + freeIn' (ReturnsInBlock v ixfun) = freeIn' v <> freeIn' ixfun + freeIn' _ = mempty instance Engine.Simplifiable MemReturn where simplify (ReturnsNewBlock space i ixfun) = @@ -495,12 +486,11 @@ bodyReturnsToExpReturns = noUniquenessReturns . maybeReturns instance TC.CheckableOp ExplicitMemory where - checkOp (Alloc size _) = TC.require [Prim int64] size - checkOp (Inner op) = typeCheckHostOp (TC.subCheck . typeCheckKernel) op - -instance TC.CheckableOp InKernel where - checkOp (Alloc size _) = TC.require [Prim int64] size - checkOp (Inner k) = TC.subCheck $ typeCheckKernelExp k + checkOp = typeCheckExplicitMemoryOp Nothing + where typeCheckExplicitMemoryOp _ (Alloc size _) = + TC.require [Prim int64] size + typeCheckExplicitMemoryOp lvl (Inner op) = + typeCheckHostOp (typeCheckExplicitMemoryOp . Just) lvl (const $ return ()) op instance TC.Checkable ExplicitMemory where checkFParamLore = checkMemInfo @@ -512,16 +502,6 @@ matchReturnType = matchFunctionReturnType matchBranchType = matchBranchReturnType -instance TC.Checkable InKernel where - checkFParamLore = checkMemInfo - checkLParamLore = checkMemInfo - checkLetBoundLore = checkMemInfo - checkRetType = mapM_ TC.checkExtType . retTypeValues - primFParam name t = return $ Param name (MemPrim t) - matchPattern = matchPatternToExp - matchReturnType = matchFunctionReturnType - matchBranchType = matchBranchReturnType - matchFunctionReturnType :: ExplicitMemorish lore => [FunReturns] -> Result -> TC.TypeM lore () matchFunctionReturnType rettype result = do @@ -785,9 +765,6 @@ instance Attributes ExplicitMemory where expTypesFromPattern = return . map snd . snd . bodyReturnsFromPattern -instance Attributes InKernel where - expTypesFromPattern = return . map snd . snd . bodyReturnsFromPattern - bodyReturnsFromPattern :: PatternT (MemBound NoUniqueness) -> ([(VName,BodyReturns)], [(VName,BodyReturns)]) bodyReturnsFromPattern pat = @@ -817,11 +794,10 @@ instance (PP.Pretty u, PP.Pretty r) => PrettyAnnot (PatElemT (MemInfo SubExp u r)) where ppAnnot = bindeeAnnot patElemName patElemAttr -instance (PP.Pretty u, PP.Pretty r) => PrettyAnnot (ParamT (MemInfo SubExp u r)) where +instance (PP.Pretty u, PP.Pretty r) => PrettyAnnot (Param (MemInfo SubExp u r)) where ppAnnot = bindeeAnnot paramName paramAttr instance PrettyLore ExplicitMemory where -instance PrettyLore InKernel where bindeeAnnot :: (PP.Pretty u, PP.Pretty r) => (a -> VName) -> (a -> MemInfo SubExp u r) @@ -985,52 +961,28 @@ Op lore -> m [ExpReturns] opReturns op = extReturns <$> opType op -instance OpReturns ExplicitMemory where - opReturns (Alloc _ space) = - return [MemMem space] - opReturns (Inner (HostOp k@(Kernel _ _ _ body))) = - zipWithM correct (kernelBodyResult body) =<< (extReturns <$> opType k) - where correct (WriteReturn _ arr _) _ = varReturns arr - correct _ ret = return ret - opReturns (Inner (HostOp (SegGenRed _ ops _ _))) = - concat <$> mapM (mapM varReturns . genReduceDest) ops - opReturns k = - extReturns <$> opType k +segOpReturns :: (Monad m, HasScope ExplicitMemory m) => + SegOp ExplicitMemory -> m [ExpReturns] +segOpReturns k@(SegMap _ _ _ kbody) = + kernelBodyReturns kbody =<< (extReturns <$> opType k) +segOpReturns k@(SegRed _ _ _ _ kbody) = + kernelBodyReturns kbody =<< (extReturns <$> opType k) +segOpReturns k@(SegScan _ _ _ _ _ kbody) = + kernelBodyReturns kbody =<< (extReturns <$> opType k) +segOpReturns (SegGenRed _ _ ops _ _) = + concat <$> mapM (mapM varReturns . genReduceDest) ops -instance OpReturns InKernel where +kernelBodyReturns :: (HasScope ExplicitMemory m, Monad m) => + KernelBody ExplicitMemory -> [ExpReturns] -> m [ExpReturns] +kernelBodyReturns = zipWithM correct . kernelBodyResult + where correct (WriteReturns _ arr _) _ = varReturns arr + correct _ ret = return ret + +instance OpReturns ExplicitMemory where opReturns (Alloc _ space) = return [MemMem space] - - opReturns (Inner (GroupStream _ _ lam _ _)) = - forM (groupStreamAccParams lam) $ \param -> - case paramAttr param of - MemPrim bt -> - return $ MemPrim bt - MemArray et shape _ (ArrayIn mem ixfun) -> - return $ MemArray et (Shape $ map Free $ shapeDims shape) NoUniqueness $ - Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun - MemMem space -> - return $ MemMem space - - opReturns (Inner (GroupScan _ _ input)) = - mapM varReturns arrs - where arrs = map snd input - - opReturns (Inner (GroupGenReduce _ dests _ _ _ _)) = - mapM varReturns dests - - opReturns (Inner (Barrier res)) = mapM f res - where f (Var v) = varReturns v - f (Constant v) = return $ MemPrim $ primValueType v - - opReturns (Inner (Combine (CombineSpace scatter cspace) ts _ _)) = - (++) <$> mapM varReturns as <*> - pure (extReturns $ staticShapes $ map (`arrayOfShape` shape) $ drop (sum ns*2) ts) - where (_, ns, as) = unzip3 scatter - shape = Shape $ map snd cspace - - opReturns k = - extReturns <$> opType k + opReturns (Inner (SegOp op)) = segOpReturns op + opReturns k = extReturns <$> opType k applyFunReturns :: Typed attr => [FunReturns] @@ -1083,3 +1035,12 @@ ShapeBase num -> IxFun.IxFun num -> Bool ixFunMatchesInnerShape shape ixfun = drop 1 (IxFun.shape ixfun) == drop 1 (shapeDims shape) + +-- | Construct the scalar memory space corresponding to a given primitive type. +scalarMemory :: PrimType -> SpaceId +scalarMemory = ("scalar_"++) . pretty + +-- | A mapping from all scalar memory spaces to the 'PrimType' they +-- store. +allScalarMemory :: M.Map SpaceId PrimType +allScalarMemory = M.fromList $ zip (map scalarMemory allPrimTypes) allPrimTypes diff --git a/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs b/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs --- a/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs +++ b/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs @@ -135,11 +135,10 @@ instance FreeIn num => FreeIn (LMAD num) where - freeIn = foldMap freeIn + freeIn' = foldMap freeIn' instance FreeIn num => FreeIn (IxFun num) where - freeIn = foldMap freeIn - + freeIn' = foldMap freeIn' instance Functor LMAD where fmap f = runIdentity . traverse (return . f) @@ -334,7 +333,7 @@ else if n /= -1 then n + 1 else n ) 0 inds - in if d == -1 then [] else [p - d] + in [p - d | d /= -1] harmlessRotation' :: (Eq num, IntegralExp num) => LMADDim num -> DimIndex num -> Bool diff --git a/src/Futhark/Representation/ExplicitMemory/Simplify.hs b/src/Futhark/Representation/ExplicitMemory/Simplify.hs --- a/src/Futhark/Representation/ExplicitMemory/Simplify.hs +++ b/src/Futhark/Representation/ExplicitMemory/Simplify.hs @@ -10,7 +10,6 @@ where import Control.Monad -import qualified Data.Set as S import Data.List import qualified Futhark.Representation.AST.Syntax as AST @@ -18,8 +17,7 @@ hiding (Prog, BasicOp, Exp, Body, Stm, Pattern, PatElem, Lambda, FunDef, FParam, LParam, RetType) import Futhark.Representation.ExplicitMemory -import Futhark.Representation.Kernels.Simplify - (simplifyKernelOp, simplifyKernelExp) +import Futhark.Representation.Kernels.Simplify (simplifyKernelOp) import Futhark.Pass.ExplicitAllocations (simplifiable, arraySizeInBytesExp) import qualified Futhark.Analysis.SymbolTable as ST @@ -34,74 +32,62 @@ import Futhark.Util simpleExplicitMemory :: Simplify.SimpleOps ExplicitMemory -simpleExplicitMemory = simplifiable (simplifyKernelOp simpleInKernel inKernelEnv) - -simpleInKernel :: KernelSpace -> Simplify.SimpleOps InKernel -simpleInKernel = simplifiable . simplifyKernelExp +simpleExplicitMemory = simplifiable $ simplifyKernelOp $ const $ return ((), mempty) simplifyExplicitMemory :: Prog ExplicitMemory -> PassM (Prog ExplicitMemory) simplifyExplicitMemory = Simplify.simplifyProg simpleExplicitMemory callKernelRules - blockers { Engine.blockHoistBranch = isAlloc } + blockers { Engine.blockHoistBranch = blockAllocs } + where blockAllocs vtable _ (Let _ _ (Op Alloc{})) = + not $ ST.simplifyMemory vtable + blockAllocs _ _ _ = False simplifyStms :: (HasScope ExplicitMemory m, MonadFreshNames m) => Stms ExplicitMemory -> m (Stms ExplicitMemory) simplifyStms = Simplify.simplifyStms simpleExplicitMemory callKernelRules blockers -isAlloc :: Op lore ~ MemOp op => Engine.BlockPred lore -isAlloc _ (Let _ _ (Op Alloc{})) = True -isAlloc _ _ = False - isResultAlloc :: Op lore ~ MemOp op => Engine.BlockPred lore -isResultAlloc usage (Let (AST.Pattern [] [bindee]) _ (Op Alloc{})) = +isResultAlloc _ usage (Let (AST.Pattern [] [bindee]) _ (Op Alloc{})) = UT.isInResult (patElemName bindee) usage -isResultAlloc _ _ = False +isResultAlloc _ _ _ = False -- | Getting the roots of what to hoist, for now only variable -- names that represent array and memory-block sizes. getShapeNames :: (ExplicitMemorish lore, Op lore ~ MemOp op) => - Stm (Wise lore) -> S.Set VName + Stm (Wise lore) -> Names getShapeNames stm = let ts = map patElemType $ patternElements $ stmPattern stm in freeIn (concatMap arrayDims ts) <> case stmExp stm of Op (Alloc size _) -> freeIn size _ -> mempty -isAlloc0 :: Op lore ~ MemOp op => AST.Stm lore -> Bool -isAlloc0 (Let _ _ (Op Alloc{})) = True -isAlloc0 _ = False - -inKernelEnv :: Engine.Env InKernel -inKernelEnv = Engine.emptyEnv inKernelRules blockers +isAlloc :: Op lore ~ MemOp op => Engine.BlockPred lore +isAlloc _ _ (Let _ _ (Op Alloc{})) = True +isAlloc _ _ _ = False -blockers :: (ExplicitMemorish lore, Op lore ~ MemOp op) => - Simplify.HoistBlockers lore +blockers :: Simplify.HoistBlockers ExplicitMemory blockers = Engine.noExtraHoistBlockers { Engine.blockHoistPar = isAlloc , Engine.blockHoistSeq = isResultAlloc , Engine.getArraySizes = getShapeNames - , Engine.isAllocation = isAlloc0 + , Engine.isAllocation = isAlloc mempty mempty } callKernelRules :: RuleBook (Wise ExplicitMemory) callKernelRules = standardRules <> ruleBook [RuleBasicOp copyCopyToCopy, - RuleBasicOp removeIdentityCopy] [] - -inKernelRules :: RuleBook (Wise InKernel) -inKernelRules = standardRules <> - ruleBook [RuleBasicOp copyCopyToCopy, - RuleBasicOp removeIdentityCopy, - RuleIf unExistentialiseMemory] [] + RuleBasicOp removeIdentityCopy, + RuleIf unExistentialiseMemory] [] -- | If a branch is returning some existential memory, but the size of -- the array is not existential, then we can create a block of the -- proper size and always return there. -unExistentialiseMemory :: TopDownRuleIf (Wise InKernel) -unExistentialiseMemory _ pat _ (cond, tbranch, fbranch, ifattr) - | fixable <- foldl hasConcretisableMemory mempty $ patternElements pat, - not $ null fixable = do +unExistentialiseMemory :: TopDownRuleIf (Wise ExplicitMemory) +unExistentialiseMemory vtable pat _ (cond, tbranch, fbranch, ifattr) + | ST.simplifyMemory vtable, + fixable <- foldl hasConcretisableMemory mempty $ patternElements pat, + not $ null fixable = Simplify $ do -- Create non-existential memory blocks big enough to hold the -- arrays. @@ -133,7 +119,7 @@ tbranch' <- updateBody tbranch fbranch' <- updateBody fbranch letBind_ pat $ If cond tbranch' fbranch' ifattr - where onlyUsedIn name here = not $ any ((name `S.member`) . freeIn) $ + where onlyUsedIn name here = not $ any ((name `nameIn`) . freeIn) $ filter ((/=here) . patElemName) $ patternValueElements pat knownSize Constant{} = True @@ -153,7 +139,7 @@ (pat_elem, mem, space) : fixable | otherwise = fixable -unExistentialiseMemory _ _ _ _ = cannotSimplify +unExistentialiseMemory _ _ _ _ = Skip -- | If we are copying something that is itself a copy, just copy the -- original one instead. @@ -174,16 +160,16 @@ src_space == dest_space, dest_ixfun == src_ixfun = - certifying v1_cs $ letBind_ pat $ BasicOp $ Copy v2 + Simplify $ certifying v1_cs $ letBind_ pat $ BasicOp $ Copy v2 copyCopyToCopy vtable pat _ (Copy v0) | Just (BasicOp (Rearrange perm v1), v0_cs) <- ST.lookupExp v0 vtable, - Just (BasicOp (Copy v2), v1_cs) <- ST.lookupExp v1 vtable = do + Just (BasicOp (Copy v2), v1_cs) <- ST.lookupExp v1 vtable = Simplify $ do v0' <- certifying (v0_cs<>v1_cs) $ letExp "rearrange_v0" $ BasicOp $ Rearrange perm v2 letBind_ pat $ BasicOp $ Copy v0' -copyCopyToCopy _ _ _ _ = cannotSimplify +copyCopyToCopy _ _ _ _ = Skip -- | If the destination of a copy is the same as the source, just -- remove it. @@ -195,6 +181,6 @@ Just (_, MemArray _ _ _ (ArrayIn src_mem src_ixfun)) <- ST.entryLetBoundAttr =<< ST.lookup v vtable, dest_mem == src_mem, dest_ixfun == src_ixfun = - letBind_ pat $ BasicOp $ SubExp $ Var v + Simplify $ letBind_ pat $ BasicOp $ SubExp $ Var v -removeIdentityCopy _ _ _ _ = cannotSimplify +removeIdentityCopy _ _ _ _ = Skip diff --git a/src/Futhark/Representation/Kernels.hs b/src/Futhark/Representation/Kernels.hs --- a/src/Futhark/Representation/Kernels.hs +++ b/src/Futhark/Representation/Kernels.hs @@ -4,25 +4,24 @@ module Futhark.Representation.Kernels ( -- * The Lore definition Kernels - , InKernel -- * Module re-exports , module Futhark.Representation.AST.Attributes , module Futhark.Representation.AST.Traversals , module Futhark.Representation.AST.Pretty , module Futhark.Representation.AST.Syntax , module Futhark.Representation.Kernels.Kernel - , module Futhark.Representation.Kernels.KernelExp , module Futhark.Representation.Kernels.Sizes + , module Futhark.Representation.SOACS.SOAC ) where import Futhark.Representation.AST.Syntax import Futhark.Representation.Kernels.Kernel -import Futhark.Representation.Kernels.KernelExp import Futhark.Representation.Kernels.Sizes import Futhark.Representation.AST.Attributes import Futhark.Representation.AST.Traversals import Futhark.Representation.AST.Pretty +import Futhark.Representation.SOACS.SOAC hiding (GenReduceOp(..)) import Futhark.Binder import Futhark.Construct import qualified Futhark.TypeCheck as TypeCheck @@ -34,27 +33,17 @@ data Kernels instance Annotations Kernels where - type Op Kernels = HostOp Kernels (Kernel InKernel) + type Op Kernels = HostOp Kernels (SOAC Kernels) instance Attributes Kernels where expTypesFromPattern = return . expExtTypesFromPattern -data InKernel -instance Annotations InKernel where - type Op InKernel = KernelExp InKernel -instance Attributes InKernel where - expTypesFromPattern = return . expExtTypesFromPattern -instance PrettyLore InKernel where - instance TypeCheck.CheckableOp Kernels where - checkOp = typeCheckHostOp $ TypeCheck.subCheck . typeCheckKernel - -instance TypeCheck.CheckableOp InKernel where - checkOp = TypeCheck.subCheck . typeCheckKernelExp + checkOp = typeCheckKernelsOp Nothing + where typeCheckKernelsOp lvl = + typeCheckHostOp (typeCheckKernelsOp . Just) lvl typeCheckSOAC instance TypeCheck.Checkable Kernels where -instance TypeCheck.Checkable InKernel where - instance Bindable Kernels where mkBody = Body () mkExpPat ctx val _ = basicPattern ctx val @@ -62,17 +51,6 @@ mkLetNames = simpleMkLetNames instance BinderOps Kernels where - mkExpAttrB = bindableMkExpAttrB - mkBodyB = bindableMkBodyB - mkLetNamesB = bindableMkLetNamesB - -instance Bindable InKernel where - mkBody = Body () - mkExpPat ctx val _ = basicPattern ctx val - mkExpAttr _ _ = () - mkLetNames = simpleMkLetNames - -instance BinderOps InKernel where mkExpAttrB = bindableMkExpAttrB mkBodyB = bindableMkBodyB mkLetNamesB = bindableMkLetNamesB diff --git a/src/Futhark/Representation/Kernels/Kernel.hs b/src/Futhark/Representation/Kernels/Kernel.hs --- a/src/Futhark/Representation/Kernels/Kernel.hs +++ b/src/Futhark/Representation/Kernels/Kernel.hs @@ -8,969 +8,1023 @@ {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} module Futhark.Representation.Kernels.Kernel - ( Kernel(..) - , kernelType - , kernelSpace - , KernelDebugHints(..) - , GenReduceOp(..) - , SegRedOp(..) - , segRedResults - , KernelBody(..) - , KernelSpace(..) - , spaceDimensions - , SpaceStructure(..) - , scopeOfKernelSpace - , KernelResult(..) - , kernelResultSubExp - , KernelPath - - , chunkedKernelNonconcatOutputs - - , typeCheckKernel - - -- * Generic traversal - , KernelMapper(..) - , identityKernelMapper - , mapKernelM - , KernelWalker(..) - , identityKernelWalker - , walkKernelM - - -- * Host operations - , HostOp(..) - , typeCheckHostOp - ) - where - -import Control.Arrow (first) -import Control.Monad.Writer hiding (mapM_) -import Control.Monad.Identity hiding (mapM_) -import qualified Data.Set as S -import qualified Data.Map.Strict as M -import Data.Foldable -import Data.List - -import Futhark.Representation.AST -import qualified Futhark.Analysis.Alias as Alias -import qualified Futhark.Analysis.SymbolTable as ST -import Futhark.Analysis.PrimExp.Convert -import qualified Futhark.Util.Pretty as PP -import Futhark.Util.Pretty - ((), (<+>), ppr, commasep, Pretty, parens, text) -import Futhark.Transform.Substitute -import Futhark.Transform.Rename -import Futhark.Optimise.Simplify.Lore -import Futhark.Representation.Ranges - (Ranges, removeLambdaRanges, removeBodyRanges, mkBodyRanges) -import Futhark.Representation.AST.Attributes.Ranges -import Futhark.Representation.AST.Attributes.Aliases -import Futhark.Representation.Aliases - (Aliases, removeLambdaAliases, removeBodyAliases, removeStmAliases) -import Futhark.Representation.Kernels.KernelExp (SplitOrdering(..)) -import Futhark.Representation.Kernels.Sizes -import qualified Futhark.TypeCheck as TC -import Futhark.Analysis.Metrics -import Futhark.Tools (partitionChunkedKernelLambdaParameters) -import qualified Futhark.Analysis.Range as Range -import Futhark.Util (maybeNth) - --- | Some information about what goes into a kernel, and where it came --- from. Has no semantic meaning; only used for debugging generated --- code. -data KernelDebugHints = - KernelDebugHints { kernelName :: String - , kernelHints :: [(String, SubExp)] - -- ^ A mapping from a description to some - -- PrimType value. - } - deriving (Eq, Show, Ord) - -data GenReduceOp lore = - GenReduceOp { genReduceWidth :: SubExp - , genReduceDest :: [VName] - , genReduceNeutral :: [SubExp] - , genReduceShape :: Shape - -- ^ In case this operator is semantically a - -- vectorised operator (corresponding to a perfect map - -- nest in the SOACS representation), these are the - -- logical "dimensions". This is used to generate - -- more efficient code. - , genReduceOp :: LambdaT lore - } - deriving (Eq, Ord, Show) - -data SegRedOp lore = - SegRedOp { segRedComm :: Commutativity - , segRedLambda :: Lambda lore - , segRedNeutral :: [SubExp] - , segRedShape :: Shape - -- ^ In case this operator is semantically a vectorised - -- operator (corresponding to a perfect map nest in the - -- SOACS representation), these are the logical - -- "dimensions". This is used to generate more efficient - -- code. - } - deriving (Eq, Ord, Show) - --- | How many reduction results are produced by these 'SegRedOp's? -segRedResults :: [SegRedOp lore] -> Int -segRedResults = sum . map (length . segRedNeutral) - -data Kernel lore - = Kernel KernelDebugHints KernelSpace [Type] (KernelBody lore) - | SegMap KernelSpace [Type] (KernelBody lore) - | SegRed KernelSpace [SegRedOp lore] [Type] (KernelBody lore) - -- ^ The KernelSpace must always have at least two dimensions, - -- implying that the result of a SegRed is always an array. - | SegScan KernelSpace (Lambda lore) [SubExp] [Type] (KernelBody lore) - | SegGenRed KernelSpace [GenReduceOp lore] [Type] (KernelBody lore) - deriving (Eq, Show, Ord) - -kernelSpace :: Kernel lore -> KernelSpace -kernelSpace (Kernel _ kspace _ _) = kspace -kernelSpace (SegMap kspace _ _) = kspace -kernelSpace (SegRed kspace _ _ _) = kspace -kernelSpace (SegScan kspace _ _ _ _) = kspace -kernelSpace (SegGenRed kspace _ _ _) = kspace - -data KernelSpace = KernelSpace { spaceGlobalId :: VName - , spaceLocalId :: VName - , spaceGroupId :: VName - , spaceNumThreads :: SubExp - , spaceNumGroups :: SubExp - , spaceGroupSize :: SubExp -- flat group size - , spaceNumVirtGroups :: SubExp - -- How many groups should we pretend - -- exist? - , spaceStructure :: SpaceStructure - -- TODO: document what this spaceStructure is - -- used for - } - deriving (Eq, Show, Ord) --- ^ first three bound in the kernel, the rest are params to kernel - --- | Indices computed for each thread (or group) inside the kernel. --- This is an arbitrary-dimensional space that is generated from the --- flat GPU thread space. -data SpaceStructure = FlatThreadSpace - [(VName, SubExp)] -- gtids and dim sizes - | NestedThreadSpace - [(VName, -- gtid - SubExp, -- global dim size - VName, -- ltid - SubExp -- local dim sizes - )] - deriving (Eq, Show, Ord) - --- | Global thread IDs and their upper bound. -spaceDimensions :: KernelSpace -> [(VName, SubExp)] -spaceDimensions = structureDimensions . spaceStructure - where structureDimensions (FlatThreadSpace dims) = dims - structureDimensions (NestedThreadSpace dims) = - let (gtids, gdim_sizes, _, _) = unzip4 dims - in zip gtids gdim_sizes - --- | The body of a 'Kernel'. -data KernelBody lore = KernelBody { kernelBodyLore :: BodyAttr lore - , kernelBodyStms :: Stms lore - , kernelBodyResult :: [KernelResult] - } - -deriving instance Annotations lore => Ord (KernelBody lore) -deriving instance Annotations lore => Show (KernelBody lore) -deriving instance Annotations lore => Eq (KernelBody lore) - -data KernelResult = ThreadsReturn SubExp - -- ^ Each thread in the kernel space (which must - -- be non-empty) returns this. - | GroupsReturn SubExp - -- ^ Each group returns this. - | WriteReturn - [SubExp] -- Size of array. Must match number of dims. - VName -- Which array - [([SubExp], SubExp)] - -- Arbitrary number of index/value pairs. - | ConcatReturns - SplitOrdering -- Permuted? - SubExp -- The final size. - SubExp -- Per-thread (max) chunk size. - (Maybe SubExp) -- Optional precalculated offset. - VName -- Chunk by this thread. - deriving (Eq, Show, Ord) - -kernelResultSubExp :: KernelResult -> SubExp -kernelResultSubExp (ThreadsReturn se) = se -kernelResultSubExp (GroupsReturn se) = se -kernelResultSubExp (WriteReturn _ arr _) = Var arr -kernelResultSubExp (ConcatReturns _ _ _ _ v) = Var v - --- | Like 'Mapper', but just for 'Kernel's. -data KernelMapper flore tlore m = KernelMapper { - mapOnKernelSubExp :: SubExp -> m SubExp - , mapOnKernelLambda :: Lambda flore -> m (Lambda tlore) - , mapOnKernelBody :: Body flore -> m (Body tlore) - , mapOnKernelVName :: VName -> m VName - , mapOnKernelLParam :: LParam flore -> m (LParam tlore) - , mapOnKernelKernelBody :: KernelBody flore -> m (KernelBody tlore) - } - --- | A mapper that simply returns the 'Kernel' verbatim. -identityKernelMapper :: Monad m => KernelMapper lore lore m -identityKernelMapper = KernelMapper { mapOnKernelSubExp = return - , mapOnKernelLambda = return - , mapOnKernelBody = return - , mapOnKernelVName = return - , mapOnKernelLParam = return - , mapOnKernelKernelBody = return - } - --- | Map a monadic action across the immediate children of a --- Kernel. The mapping does not descend recursively into subexpressions --- and is done left-to-right. -mapKernelM :: (Applicative m, Monad m) => - KernelMapper flore tlore m -> Kernel flore -> m (Kernel tlore) -mapKernelM tv (SegMap space ts body) = - SegMap - <$> mapOnKernelSpace tv space - <*> mapM (mapOnType $ mapOnKernelSubExp tv) ts - <*> mapOnKernelKernelBody tv body -mapKernelM tv (SegRed space reds ts body) = - SegRed - <$> mapOnKernelSpace tv space - <*> mapM onSegOp reds - <*> mapM (mapOnType $ mapOnKernelSubExp tv) ts - <*> mapOnKernelKernelBody tv body - where onSegOp (SegRedOp comm red_op nes shape) = - SegRedOp comm - <$> mapOnKernelLambda tv red_op - <*> mapM (mapOnKernelSubExp tv) nes - <*> (Shape <$> mapM (mapOnKernelSubExp tv) (shapeDims shape)) -mapKernelM tv (SegScan space scan_op nes ts body) = - SegScan - <$> mapOnKernelSpace tv space - <*> mapOnKernelLambda tv scan_op - <*> mapM (mapOnKernelSubExp tv) nes - <*> mapM (mapOnType $ mapOnKernelSubExp tv) ts - <*> mapOnKernelKernelBody tv body -mapKernelM tv (SegGenRed space ops ts body) = - SegGenRed - <$> mapOnKernelSpace tv space - <*> mapM onGenRedOp ops - <*> mapM (mapOnType $ mapOnKernelSubExp tv) ts - <*> mapOnKernelKernelBody tv body - where onGenRedOp (GenReduceOp w arrs nes shape op) = - GenReduceOp <$> mapOnKernelSubExp tv w - <*> mapM (mapOnKernelVName tv) arrs - <*> mapM (mapOnKernelSubExp tv) nes - <*> (Shape <$> mapM (mapOnKernelSubExp tv) (shapeDims shape)) - <*> mapOnKernelLambda tv op -mapKernelM tv (Kernel desc space ts kernel_body) = - Kernel <$> mapOnKernelDebugHints desc <*> - mapOnKernelSpace tv space <*> - mapM (mapOnKernelType tv) ts <*> - mapOnKernelKernelBody tv kernel_body - where mapOnKernelDebugHints (KernelDebugHints name kvs) = - KernelDebugHints name <$> - (zip (map fst kvs) <$> mapM (mapOnKernelSubExp tv . snd) kvs) - -mapOnKernelSpace :: Monad f => - KernelMapper flore tlore f -> KernelSpace -> f KernelSpace -mapOnKernelSpace tv (KernelSpace gtid ltid gid num_threads num_groups group_size virt_groups structure) = - KernelSpace gtid ltid gid -- all in binding position - <$> mapOnKernelSubExp tv num_threads - <*> mapOnKernelSubExp tv num_groups - <*> mapOnKernelSubExp tv group_size - <*> mapOnKernelSubExp tv virt_groups - <*> mapOnKernelStructure structure - where mapOnKernelStructure (FlatThreadSpace dims) = - FlatThreadSpace <$> (zip gtids <$> mapM (mapOnKernelSubExp tv) gdim_sizes) - where (gtids, gdim_sizes) = unzip dims - mapOnKernelStructure (NestedThreadSpace dims) = - NestedThreadSpace <$> (zip4 gtids - <$> mapM (mapOnKernelSubExp tv) gdim_sizes - <*> pure ltids - <*> mapM (mapOnKernelSubExp tv) ldim_sizes) - where (gtids, gdim_sizes, ltids, ldim_sizes) = unzip4 dims - -mapOnKernelType :: Monad m => - KernelMapper flore tlore m -> Type -> m Type -mapOnKernelType _tv (Prim pt) = pure $ Prim pt -mapOnKernelType tv (Array pt shape u) = Array pt <$> f shape <*> pure u - where f (Shape dims) = Shape <$> mapM (mapOnKernelSubExp tv) dims -mapOnKernelType _tv (Mem s) = pure $ Mem s - -instance (Attributes lore, FreeIn (LParamAttr lore)) => - FreeIn (Kernel lore) where - freeIn e = execWriter $ mapKernelM free e - where walk f x = tell (f x) >> return x - free = KernelMapper { mapOnKernelSubExp = walk freeIn - , mapOnKernelLambda = walk freeIn - , mapOnKernelBody = walk freeIn - , mapOnKernelVName = walk freeIn - , mapOnKernelLParam = walk freeIn - , mapOnKernelKernelBody = walk freeIn - } - --- | Like 'Walker', but just for 'Kernel's. -data KernelWalker lore m = KernelWalker { - walkOnKernelSubExp :: SubExp -> m () - , walkOnKernelLambda :: Lambda lore -> m () - , walkOnKernelBody :: Body lore -> m () - , walkOnKernelVName :: VName -> m () - , walkOnKernelLParam :: LParam lore -> m () - , walkOnKernelKernelBody :: KernelBody lore -> m () - } - --- | A no-op traversal. -identityKernelWalker :: Monad m => KernelWalker lore m -identityKernelWalker = KernelWalker { - walkOnKernelSubExp = const $ return () - , walkOnKernelLambda = const $ return () - , walkOnKernelBody = const $ return () - , walkOnKernelVName = const $ return () - , walkOnKernelLParam = const $ return () - , walkOnKernelKernelBody = const $ return () - } - -walkKernelMapper :: forall lore m. Monad m => - KernelWalker lore m -> KernelMapper lore lore m -walkKernelMapper f = KernelMapper { - mapOnKernelSubExp = wrap walkOnKernelSubExp - , mapOnKernelLambda = wrap walkOnKernelLambda - , mapOnKernelBody = wrap walkOnKernelBody - , mapOnKernelVName = wrap walkOnKernelVName - , mapOnKernelLParam = wrap walkOnKernelLParam - , mapOnKernelKernelBody = wrap walkOnKernelKernelBody - } - where wrap :: (KernelWalker lore m -> a -> m ()) -> a -> m a - wrap op k = op f k >> return k - --- | As 'mapKernelM', but ignoring the results. -walkKernelM :: Monad m => KernelWalker lore m -> Kernel lore -> m () -walkKernelM f = void . mapKernelM m - where m = walkKernelMapper f - -instance FreeIn KernelResult where - freeIn (GroupsReturn what) = freeIn what - freeIn (ThreadsReturn what) = freeIn what - freeIn (WriteReturn rws arr res) = freeIn rws <> freeIn arr <> freeIn res - freeIn (ConcatReturns o w per_thread_elems moffset v) = - freeIn o <> freeIn w <> freeIn per_thread_elems <> freeIn moffset <> freeIn v - -instance Attributes lore => FreeIn (KernelBody lore) where - freeIn (KernelBody attr stms res) = - (freeIn attr <> free_in_stms <> free_in_res) `S.difference` bound_in_stms - where free_in_stms = fold $ fmap freeIn stms - free_in_res = freeIn res - bound_in_stms = fold $ fmap boundByStm stms - -instance Attributes lore => Substitute (KernelBody lore) where - substituteNames subst (KernelBody attr stms res) = - KernelBody - (substituteNames subst attr) - (substituteNames subst stms) - (substituteNames subst res) - -instance Substitute KernelResult where - substituteNames subst (GroupsReturn se) = - GroupsReturn $ substituteNames subst se - substituteNames subst (ThreadsReturn se) = - ThreadsReturn $ substituteNames subst se - substituteNames subst (WriteReturn rws arr res) = - WriteReturn - (substituteNames subst rws) (substituteNames subst arr) - (substituteNames subst res) - substituteNames subst (ConcatReturns o w per_thread_elems moffset v) = - ConcatReturns - (substituteNames subst o) - (substituteNames subst w) - (substituteNames subst per_thread_elems) - (substituteNames subst moffset) - (substituteNames subst v) - -instance Substitute KernelSpace where - substituteNames subst (KernelSpace gtid ltid gid num_threads num_groups group_size virt_groups structure) = - KernelSpace (substituteNames subst gtid) - (substituteNames subst ltid) - (substituteNames subst gid) - (substituteNames subst num_threads) - (substituteNames subst num_groups) - (substituteNames subst group_size) - (substituteNames subst virt_groups) - (substituteNames subst structure) - -instance Substitute SpaceStructure where - substituteNames subst (FlatThreadSpace dims) = - FlatThreadSpace (map (substituteNames subst) dims) - substituteNames subst (NestedThreadSpace dims) = - NestedThreadSpace (map (substituteNames subst) dims) - -instance Attributes lore => Substitute (Kernel lore) where - substituteNames subst (Kernel desc space ts kbody) = - Kernel desc - (substituteNames subst space) - (substituteNames subst ts) - (substituteNames subst kbody) - substituteNames subst k = runIdentity $ mapKernelM substitute k - where substitute = - KernelMapper { mapOnKernelSubExp = return . substituteNames subst - , mapOnKernelLambda = return . substituteNames subst - , mapOnKernelBody = return . substituteNames subst - , mapOnKernelVName = return . substituteNames subst - , mapOnKernelLParam = return . substituteNames subst - , mapOnKernelKernelBody = return . substituteNames subst - } - -instance Attributes lore => Rename (KernelBody lore) where - rename (KernelBody attr stms res) = do - attr' <- rename attr - renamingStms stms $ \stms' -> - KernelBody attr' stms' <$> rename res - -instance Rename KernelResult where - rename = substituteRename - -scopeOfKernelSpace :: KernelSpace -> Scope lore -scopeOfKernelSpace (KernelSpace gtid ltid gid _ _ _ _ structure) = - M.fromList $ zip ([gtid, ltid, gid] ++ structure') $ repeat $ IndexInfo Int32 - where structure' = case structure of - FlatThreadSpace dims -> map fst dims - NestedThreadSpace dims -> - let (gtids, _, ltids, _) = unzip4 dims - in gtids ++ ltids - -instance Attributes lore => Rename (Kernel lore) where - rename = mapKernelM renamer - where renamer = KernelMapper rename rename rename rename rename rename - - -kernelResultShape :: KernelSpace -> Type -> KernelResult -> Type -kernelResultShape _ t (WriteReturn rws _ _) = - t `arrayOfShape` Shape rws -kernelResultShape space t (GroupsReturn _) = - t `arrayOfRow` spaceNumGroups space -kernelResultShape space t (ThreadsReturn _) = - foldr (flip arrayOfRow . snd) t $ spaceDimensions space -kernelResultShape _ t (ConcatReturns _ w _ _ _) = - t `arrayOfRow` w - -kernelType :: Kernel lore -> [Type] -kernelType (Kernel _ space ts body) = - zipWith (kernelResultShape space) ts $ kernelBodyResult body - -kernelType (SegMap space ts body) = - zipWith (kernelResultShape space) ts $ kernelBodyResult body - -kernelType (SegRed space reds ts body) = - red_ts ++ - zipWith (kernelResultShape space) map_ts - (drop (length red_ts) $ kernelBodyResult body) - where map_ts = drop (length red_ts) ts - segment_dims = init $ map snd $ spaceDimensions space - red_ts = do - op <- reds - let shape = Shape segment_dims <> segRedShape op - map (`arrayOfShape` shape) (lambdaReturnType $ segRedLambda op) - -kernelType (SegScan space _ _ ts _) = - map (`arrayOfShape` Shape dims) ts - where dims = map snd $ spaceDimensions space - -kernelType (SegGenRed space ops _ _) = do - op <- ops - let shape = Shape (segment_dims <> [genReduceWidth op]) <> genReduceShape op - map (`arrayOfShape` shape) (lambdaReturnType $ genReduceOp op) - where dims = map snd $ spaceDimensions space - segment_dims = init dims - -chunkedKernelNonconcatOutputs :: Lambda lore -> Int -chunkedKernelNonconcatOutputs fun = - length $ takeWhile (not . outerSizeIsChunk) $ lambdaReturnType fun - where outerSizeIsChunk = (==Var (paramName chunk)) . arraySize 0 - (_, chunk, _) = partitionChunkedKernelLambdaParameters $ lambdaParams fun - -instance TypedOp (Kernel lore) where - opType = pure . staticShapes . kernelType - -instance (Attributes lore, Aliased lore) => AliasedOp (Kernel lore) where - opAliases = map (const mempty) . kernelType - - consumedInOp (Kernel _ _ _ kbody) = - consumedInKernelBody kbody <> - mconcat (map consumedByReturn (kernelBodyResult kbody)) - where consumedByReturn (WriteReturn _ a _) = S.singleton a - consumedByReturn _ = mempty - consumedInOp (SegGenRed _ ops _ kbody) = - S.fromList (concatMap genReduceDest ops) <> - consumedInKernelBody kbody - consumedInOp (SegMap _ _ kbody) = - consumedInKernelBody kbody - consumedInOp (SegRed _ _ _ kbody) = - consumedInKernelBody kbody - consumedInOp (SegScan _ _ _ _ kbody) = - consumedInKernelBody kbody - -aliasAnalyseKernelBody :: (Attributes lore, - CanBeAliased (Op lore)) => - KernelBody lore - -> KernelBody (Aliases lore) -aliasAnalyseKernelBody (KernelBody attr stms res) = - let Body attr' stms' _ = Alias.analyseBody $ Body attr stms [] - in KernelBody attr' stms' res - -instance (Attributes lore, - Attributes (Aliases lore), - CanBeAliased (Op lore)) => CanBeAliased (Kernel lore) where - type OpWithAliases (Kernel lore) = Kernel (Aliases lore) - - addOpAliases = runIdentity . mapKernelM alias - where alias = KernelMapper return (return . Alias.analyseLambda) - (return . Alias.analyseBody) return return - (return . aliasAnalyseKernelBody) - - removeOpAliases = runIdentity . mapKernelM remove - where remove = KernelMapper return (return . removeLambdaAliases) - (return . removeBodyAliases) return return - (return . removeKernelBodyAliases) - removeKernelBodyAliases :: KernelBody (Aliases lore) - -> KernelBody lore - removeKernelBodyAliases (KernelBody (_, attr) stms res) = - KernelBody attr (fmap removeStmAliases stms) res - -instance Attributes lore => IsOp (Kernel lore) where - safeOp _ = True - cheapOp Kernel{} = False - cheapOp _ = True - -instance Ranged inner => RangedOp (Kernel inner) where - opRanges op = replicate (length $ kernelType op) unknownRange - -instance (Attributes lore, CanBeRanged (Op lore)) => CanBeRanged (Kernel lore) where - type OpWithRanges (Kernel lore) = Kernel (Ranges lore) - - removeOpRanges = runIdentity . mapKernelM remove - where remove = KernelMapper return (return . removeLambdaRanges) - (return . removeBodyRanges) return return - (return . removeKernelBodyRanges) - removeKernelBodyRanges = error "removeKernelBodyRanges" - addOpRanges = Range.runRangeM . mapKernelM add - where add = KernelMapper return Range.analyseLambda - Range.analyseBody return return addKernelBodyRanges - addKernelBodyRanges (KernelBody attr stms res) = - Range.analyseStms stms $ \stms' -> do - let attr' = (mkBodyRanges stms $ map kernelResultSubExp res, attr) - return $ KernelBody attr' stms' res - -instance (Attributes lore, CanBeWise (Op lore)) => CanBeWise (Kernel lore) where - type OpWithWisdom (Kernel lore) = Kernel (Wise lore) - - removeOpWisdom = runIdentity . mapKernelM remove - where remove = KernelMapper return - (return . removeLambdaWisdom) - (return . removeBodyWisdom) - return return - (return . removeKernelBodyWisdom) - removeKernelBodyWisdom :: KernelBody (Wise lore) - -> KernelBody lore - removeKernelBodyWisdom (KernelBody attr stms res) = - let Body attr' stms' _ = removeBodyWisdom $ Body attr stms [] - in KernelBody attr' stms' res - -instance Attributes lore => ST.IndexOp (Kernel lore) where - indexOp vtable k (Kernel _ space _ kbody) is = do - ThreadsReturn se <- maybeNth k $ kernelBodyResult kbody - let (gtids, _) = unzip $ spaceDimensions space - guard $ length gtids == length is - let prim_table = M.fromList $ zip gtids $ zip is $ repeat mempty - prim_table' = foldl expandPrimExpTable prim_table $ kernelBodyStms kbody - case se of - Var v -> M.lookup v prim_table' - _ -> Nothing - where expandPrimExpTable table stm - | [v] <- patternNames $ stmPattern stm, - Just (pe,cs) <- - runWriterT $ primExpFromExp (asPrimExp table) $ stmExp stm = - M.insert v (pe, stmCerts stm <> cs) table - | otherwise = - table - - asPrimExp table v - | Just (e,cs) <- M.lookup v table = tell cs >> return e - | Just (Prim pt) <- ST.lookupType v vtable = - return $ LeafExp v pt - | otherwise = lift Nothing - - indexOp _ _ _ _ = Nothing - -consumedInKernelBody :: Aliased lore => - KernelBody lore -> Names -consumedInKernelBody (KernelBody attr stms _) = - consumedInBody $ Body attr stms [] - -typeCheckKernel :: TC.Checkable lore => Kernel (Aliases lore) -> TC.TypeM lore () - -typeCheckKernel (SegMap space ts kbody) = do - checkSpace space - mapM_ TC.checkType ts - TC.binding (scopeOfKernelSpace space) $ checkKernelBody ts kbody - -typeCheckKernel (SegRed space reds ts body) = - checkScanRed space reds' ts body - where reds' = zip3 - (map segRedLambda reds) - (map segRedNeutral reds) - (map segRedShape reds) - -typeCheckKernel (SegScan space scan_op nes ts body) = - checkScanRed space [(scan_op, nes, mempty)] ts body - -typeCheckKernel (SegGenRed space ops ts body) = do - checkSpace space - mapM_ TC.checkType ts - - TC.binding (scopeOfKernelSpace space) $ do - nes_ts <- forM ops $ \(GenReduceOp dest_w dests nes shape op) -> do - TC.require [Prim int32] dest_w - nes' <- mapM TC.checkArg nes - mapM_ (TC.require [Prim int32]) $ shapeDims shape - - -- Operator type must match the type of neutral elements. - let stripVecDims = stripArray $ shapeRank shape - TC.checkLambda op $ map (TC.noArgAliases . first stripVecDims) $ nes' ++ nes' - let nes_t = map TC.argType nes' - unless (nes_t == lambdaReturnType op) $ - TC.bad $ TC.TypeError $ "SegGenRed operator has return type " ++ - prettyTuple (lambdaReturnType op) ++ " but neutral element has type " ++ - prettyTuple nes_t - - -- Arrays must have proper type. - let dest_shape = Shape (segment_dims <> [dest_w]) <> shape - forM_ (zip nes_t dests) $ \(t, dest) -> do - TC.requireI [t `arrayOfShape` dest_shape] dest - TC.consume =<< TC.lookupAliases dest - - return $ map (`arrayOfShape` shape) nes_t - - checkKernelBody ts body - - -- Return type of bucket function must be an index for each - -- operation followed by the values to write. - let bucket_ret_t = replicate (length ops) (Prim int32) ++ concat nes_ts - unless (bucket_ret_t == ts) $ - TC.bad $ TC.TypeError $ "SegGenRed body has return type " ++ - prettyTuple ts ++ " but should have type " ++ - prettyTuple bucket_ret_t - - where segment_dims = init $ map snd $ spaceDimensions space - -typeCheckKernel (Kernel _ space kts kbody) = do - checkSpace space - mapM_ TC.checkType kts - mapM_ (TC.require [Prim int32] . snd) $ spaceDimensions space - - TC.binding (scopeOfKernelSpace space) $ - checkKernelBody kts kbody - -checkKernelBody :: TC.Checkable lore => - [Type] -> KernelBody (Aliases lore) -> TC.TypeM lore () -checkKernelBody ts (KernelBody (_, attr) stms kres) = do - TC.checkBodyLore attr - TC.checkStms stms $ do - unless (length ts == length kres) $ - TC.bad $ TC.TypeError $ "Kernel return type is " ++ prettyTuple ts ++ - ", but body returns " ++ show (length kres) ++ " values." - zipWithM_ checkKernelResult kres ts - - where checkKernelResult (GroupsReturn what) t = - TC.require [t] what - checkKernelResult (ThreadsReturn what) t = - TC.require [t] what - checkKernelResult (WriteReturn rws arr res) t = do - mapM_ (TC.require [Prim int32]) rws - arr_t <- lookupType arr - forM_ res $ \(is, e) -> do - mapM_ (TC.require [Prim int32]) is - TC.require [t] e - unless (arr_t == t `arrayOfShape` Shape rws) $ - TC.bad $ TC.TypeError $ "WriteReturn returning " ++ - pretty e ++ " of type " ++ pretty t ++ ", shape=" ++ pretty rws ++ - ", but destination array has type " ++ pretty arr_t - TC.consume =<< TC.lookupAliases arr - checkKernelResult (ConcatReturns o w per_thread_elems moffset v) t = do - case o of - SplitContiguous -> return () - SplitStrided stride -> TC.require [Prim int32] stride - TC.require [Prim int32] w - TC.require [Prim int32] per_thread_elems - mapM_ (TC.require [Prim int32]) moffset - vt <- lookupType v - unless (vt == t `arrayOfRow` arraySize 0 vt) $ - TC.bad $ TC.TypeError $ "Invalid type for ConcatReturns " ++ pretty v - -checkScanRed :: TC.Checkable lore => - KernelSpace - -> [(Lambda (Aliases lore), [SubExp], Shape)] - -> [Type] - -> KernelBody (Aliases lore) - -> TC.TypeM lore () -checkScanRed space ops ts kbody = do - checkSpace space - mapM_ TC.checkType ts - - TC.binding (scopeOfKernelSpace space) $ do - ne_ts <- forM ops $ \(lam, nes, shape) -> do - mapM_ (TC.require [Prim int32]) $ shapeDims shape - nes' <- mapM TC.checkArg nes - - -- Operator type must match the type of neutral elements. - let stripVecDims = stripArray $ shapeRank shape - TC.checkLambda lam $ map (TC.noArgAliases . first stripVecDims) $ nes' ++ nes' - let nes_t = map TC.argType nes' - - unless (lambdaReturnType lam == nes_t) $ - TC.bad $ TC.TypeError "wrong type for operator or neutral elements." - - return $ map (`arrayOfShape` shape) nes_t - - let expecting = concat ne_ts - got = take (length expecting) ts - unless (expecting == got) $ - TC.bad $ TC.TypeError $ - "Wrong return for body (does not match neutral elements; expected " ++ - pretty expecting ++ "; found " ++ - pretty got ++ ")" - - checkKernelBody ts kbody - -checkSpace :: TC.Checkable lore => KernelSpace -> TC.TypeM lore () -checkSpace (KernelSpace _ _ _ num_threads num_groups group_size virt_groups structure) = do - mapM_ (TC.require [Prim int32]) [num_threads,num_groups,group_size,virt_groups] - case structure of - FlatThreadSpace dims -> - mapM_ (TC.require [Prim int32] . snd) dims - NestedThreadSpace dims -> - let (_, gdim_sizes, _, ldim_sizes) = unzip4 dims - in mapM_ (TC.require [Prim int32]) $ gdim_sizes ++ ldim_sizes - -instance OpMetrics (Op lore) => OpMetrics (Kernel lore) where - opMetrics (Kernel _ _ _ kbody) = - inside "Kernel" $ kernelBodyMetrics kbody - opMetrics (SegMap _ _ body) = - inside "SegMap" $ kernelBodyMetrics body - opMetrics (SegRed _ reds _ body) = - inside "SegRed" $ do mapM_ (lambdaMetrics . segRedLambda) reds - kernelBodyMetrics body - opMetrics (SegScan _ scan_op _ _ body) = - inside "SegScan" $ lambdaMetrics scan_op >> kernelBodyMetrics body - opMetrics (SegGenRed _ ops _ body) = - inside "SegGenRed" $ do mapM_ (lambdaMetrics . genReduceOp) ops - kernelBodyMetrics body - -kernelBodyMetrics :: OpMetrics (Op lore) => KernelBody lore -> MetricsM () -kernelBodyMetrics = mapM_ bindingMetrics . kernelBodyStms - -instance PrettyLore lore => PP.Pretty (Kernel lore) where - ppr (Kernel desc space ts body) = - text "kernel" <+> text (kernelName desc) <> - PP.align (ppr space) <+> - PP.colon <+> ppTuple' ts <+> PP.nestedBlock "{" "}" (ppr body) - - ppr (SegMap space ts body) = - text "segmap" <> - PP.align (ppr space) <+> PP.colon <+> ppTuple' ts <+> - PP.nestedBlock "{" "}" (ppr body) - - ppr (SegRed space reds ts body) = - text "segred" <> - PP.parens (PP.braces (mconcat $ intersperse (PP.comma <> PP.line) $ map ppOp reds)) - PP.align (ppr space) <+> PP.colon <+> ppTuple' ts <+> - PP.nestedBlock "{" "}" (ppr body) - where ppOp (SegRedOp comm lam nes shape) = - PP.braces (PP.commasep $ map ppr nes) <> PP.comma - ppr shape <> PP.comma - comm' <> ppr lam - where comm' = case comm of Commutative -> text "commutative " - Noncommutative -> mempty - - - ppr (SegScan space scan_op nes ts body) = - text "segscan" <> PP.parens (ppr scan_op <> PP.comma - PP.braces (PP.commasep $ map ppr nes)) - PP.align (ppr space) <+> PP.colon <+> ppTuple' ts <+> - PP.nestedBlock "{" "}" (ppr body) - - ppr (SegGenRed space ops ts body) = - text "seggenred" <> - PP.parens (PP.braces (mconcat $ intersperse (PP.comma <> PP.line) $ map ppOp ops)) - PP.align (ppr space) <+> PP.colon <+> ppTuple' ts <+> - PP.nestedBlock "{" "}" (ppr body) - where ppOp (GenReduceOp w dests nes shape op) = - ppr w <> PP.comma - PP.braces (PP.commasep $ map ppr dests) <> PP.comma - PP.braces (PP.commasep $ map ppr nes) <> PP.comma - ppr shape <> PP.comma - ppr op - -instance Pretty KernelSpace where - ppr (KernelSpace f_gtid f_ltid gid num_threads num_groups group_size virt_groups structure) = - parens (commasep [text "num groups:" <+> ppr num_groups, - text "group size:" <+> ppr group_size, - text "virt_num_groups:" <+> ppr virt_groups, - text "num threads:" <+> ppr num_threads, - text "global TID ->" <+> ppr f_gtid, - text "local TID ->" <+> ppr f_ltid, - text "group ID ->" <+> ppr gid]) structure' - where structure' = - case structure of - FlatThreadSpace dims -> flat dims - NestedThreadSpace space -> - parens (commasep $ do - (gtid,gd,ltid,ld) <- space - return $ ppr (gtid,ltid) <+> "<" <+> ppr (gd,ld)) - flat dims = parens $ commasep $ do - (i,d) <- dims - return $ ppr i <+> "<" <+> ppr d - -instance PrettyLore lore => Pretty (KernelBody lore) where - ppr (KernelBody _ stms res) = - PP.stack (map ppr (stmsToList stms)) - text "return" <+> PP.braces (PP.commasep $ map ppr res) - -instance Pretty KernelResult where - ppr (GroupsReturn what) = - text "group returns" <+> ppr what - ppr (ThreadsReturn what) = - text "thread returns" <+> ppr what - ppr (WriteReturn rws arr res) = - ppr arr <+> text "with" <+> PP.apply (map ppRes res) - where ppRes (is, e) = - PP.brackets (PP.commasep $ zipWith f is rws) <+> text "<-" <+> ppr e - f i rw = ppr i <+> text "<" <+> ppr rw - ppr (ConcatReturns o w per_thread_elems offset v) = - text "concat" <> suff <> - parens (commasep [ppr w, ppr per_thread_elems] <> offset_text) <+> - ppr v - where suff = case o of SplitContiguous -> mempty - SplitStrided stride -> text "Strided" <> parens (ppr stride) - offset_text = case offset of Nothing -> "" - Just se -> "," <+> "offset=" <> ppr se - ---- Host operations - --- | A host-level operation; parameterised by what else it can do. -data HostOp lore inner - = GetSize Name SizeClass - -- ^ Produce some runtime-configurable size. - | GetSizeMax SizeClass - -- ^ The maximum size of some class. - | CmpSizeLe Name SizeClass SubExp - -- ^ Compare size (likely a threshold) with some Int32 value. - | HostOp inner - -- ^ The arbitrary operation. - deriving (Eq, Ord, Show) - -instance Substitute inner => Substitute (HostOp lore inner) where - substituteNames substs (HostOp op) = - HostOp $ substituteNames substs op - substituteNames substs (CmpSizeLe name sclass x) = - CmpSizeLe name sclass $ substituteNames substs x - substituteNames _ x = x - -instance Rename inner => Rename (HostOp lore inner) where - rename (HostOp op) = HostOp <$> rename op - rename (CmpSizeLe name sclass x) = CmpSizeLe name sclass <$> rename x - rename x = pure x - -instance IsOp inner => IsOp (HostOp lore inner) where - safeOp (HostOp op) = safeOp op - safeOp _ = True - cheapOp (HostOp op) = cheapOp op - cheapOp _ = True - -instance TypedOp inner => TypedOp (HostOp lore inner) where - opType GetSize{} = pure [Prim int32] - opType GetSizeMax{} = pure [Prim int32] - opType CmpSizeLe{} = pure [Prim Bool] - opType (HostOp op) = opType op - -instance AliasedOp inner => AliasedOp (HostOp lore inner) where - opAliases (HostOp op) = opAliases op - opAliases _ = [mempty] - - consumedInOp (HostOp op) = consumedInOp op - consumedInOp _ = mempty - -instance RangedOp inner => RangedOp (HostOp lore inner) where - opRanges (HostOp op) = opRanges op - opRanges _ = [unknownRange] - -instance FreeIn inner => FreeIn (HostOp lore inner) where - freeIn (HostOp op) = freeIn op - freeIn (CmpSizeLe _ _ x) = freeIn x - freeIn _ = mempty - -instance CanBeAliased inner => CanBeAliased (HostOp lore inner) where - type OpWithAliases (HostOp lore inner) = HostOp (Aliases lore) (OpWithAliases inner) - - addOpAliases (HostOp op) = HostOp $ addOpAliases op - addOpAliases (GetSize name sclass) = GetSize name sclass - addOpAliases (GetSizeMax sclass) = GetSizeMax sclass - addOpAliases (CmpSizeLe name sclass x) = CmpSizeLe name sclass x - - removeOpAliases (HostOp op) = HostOp $ removeOpAliases op - removeOpAliases (GetSize name sclass) = GetSize name sclass - removeOpAliases (GetSizeMax sclass) = GetSizeMax sclass - removeOpAliases (CmpSizeLe name sclass x) = CmpSizeLe name sclass x - -instance CanBeRanged inner => CanBeRanged (HostOp lore inner) where - type OpWithRanges (HostOp lore inner) = HostOp (Ranges lore) (OpWithRanges inner) - - addOpRanges (HostOp op) = HostOp $ addOpRanges op - addOpRanges (GetSize name sclass) = GetSize name sclass - addOpRanges (GetSizeMax sclass) = GetSizeMax sclass - addOpRanges (CmpSizeLe name sclass x) = CmpSizeLe name sclass x - - removeOpRanges (HostOp op) = HostOp $ removeOpRanges op - removeOpRanges (GetSize name sclass) = GetSize name sclass - removeOpRanges (GetSizeMax sclass) = GetSizeMax sclass - removeOpRanges (CmpSizeLe name sclass x) = CmpSizeLe name sclass x - -instance CanBeWise inner => CanBeWise (HostOp lore inner) where - type OpWithWisdom (HostOp lore inner) = HostOp (Wise lore) (OpWithWisdom inner) - - removeOpWisdom (HostOp op) = HostOp $ removeOpWisdom op - removeOpWisdom (GetSize name sclass) = GetSize name sclass - removeOpWisdom (GetSizeMax sclass) = GetSizeMax sclass - removeOpWisdom (CmpSizeLe name sclass x) = CmpSizeLe name sclass x - -instance ST.IndexOp op => ST.IndexOp (HostOp lore op) where - indexOp vtable k (HostOp op) is = ST.indexOp vtable k op is - indexOp _ _ _ _ = Nothing - -instance PP.Pretty inner => PP.Pretty (HostOp lore inner) where - ppr (GetSize name size_class) = - text "get_size" <> parens (commasep [ppr name, ppr size_class]) - - ppr (GetSizeMax size_class) = - text "get_size_max" <> parens (ppr size_class) - - ppr (CmpSizeLe name size_class x) = - text "get_size" <> parens (commasep [ppr name, ppr size_class]) <+> - text "<=" <+> ppr x - - ppr (HostOp op) = ppr op - -instance OpMetrics inner => OpMetrics (HostOp lore inner) where - opMetrics GetSize{} = seen "GetSize" - opMetrics GetSizeMax{} = seen "GetSizeMax" - opMetrics CmpSizeLe{} = seen "CmpSizeLe" - opMetrics (HostOp op) = opMetrics op - -typeCheckHostOp :: TC.Checkable lore => - (inner -> TC.TypeM lore ()) - -> HostOp (Aliases lore) inner - -> TC.TypeM lore () -typeCheckHostOp _ GetSize{} = return () -typeCheckHostOp _ GetSizeMax{} = return () -typeCheckHostOp _ (CmpSizeLe _ _ x) = TC.require [Prim int32] x -typeCheckHostOp f (HostOp op) = f op + ( GenReduceOp(..) + , SegRedOp(..) + , segRedResults + , KernelBody(..) + , KernelResult(..) + , kernelResultSubExp + , SplitOrdering(..) + + -- * Segmented operations + , SegOp(..) + , SegLevel(..) + , SegVirt(..) + , segLevel + , segSpace + , typeCheckSegOp + , SegSpace(..) + , scopeOfSegSpace + , segSpaceDims + + -- ** Generic traversal + , SegOpMapper(..) + , identitySegOpMapper + , mapSegOpM + , SegOpWalker(..) + , identitySegOpWalker + , walkSegOpM + + -- * Host operations + , HostOp(..) + , typeCheckHostOp + + -- * Reexports + , module Futhark.Representation.Kernels.Sizes + ) + where + +import Control.Arrow (first) +import Control.Monad.State.Strict +import Control.Monad.Writer hiding (mapM_) +import Control.Monad.Identity hiding (mapM_) +import qualified Data.Map.Strict as M +import Data.Foldable +import Data.List + +import Futhark.Representation.AST +import qualified Futhark.Analysis.Alias as Alias +import qualified Futhark.Analysis.ScalExp as SE +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.Analysis.PrimExp.Convert +import qualified Futhark.Util.Pretty as PP +import Futhark.Util.Pretty + ((), (<+>), ppr, commasep, Pretty, parens, text) +import Futhark.Transform.Substitute +import Futhark.Transform.Rename +import Futhark.Optimise.Simplify.Lore +import Futhark.Representation.Ranges + (Ranges, removeLambdaRanges, removeStmRanges, mkBodyRanges) +import Futhark.Representation.AST.Attributes.Ranges +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Representation.Aliases + (Aliases, removeLambdaAliases, removeStmAliases) +import Futhark.Representation.Kernels.Sizes +import qualified Futhark.TypeCheck as TC +import Futhark.Analysis.Metrics +import qualified Futhark.Analysis.Range as Range +import Futhark.Util (maybeNth) + +-- | How an array is split into chunks. +data SplitOrdering = SplitContiguous + | SplitStrided SubExp + deriving (Eq, Ord, Show) + +instance FreeIn SplitOrdering where + freeIn' SplitContiguous = mempty + freeIn' (SplitStrided stride) = freeIn' stride + +instance Substitute SplitOrdering where + substituteNames _ SplitContiguous = + SplitContiguous + substituteNames subst (SplitStrided stride) = + SplitStrided $ substituteNames subst stride + +instance Rename SplitOrdering where + rename SplitContiguous = + pure SplitContiguous + rename (SplitStrided stride) = + SplitStrided <$> rename stride + +data GenReduceOp lore = + GenReduceOp { genReduceWidth :: SubExp + , genReduceDest :: [VName] + , genReduceNeutral :: [SubExp] + , genReduceShape :: Shape + -- ^ In case this operator is semantically a + -- vectorised operator (corresponding to a perfect map + -- nest in the SOACS representation), these are the + -- logical "dimensions". This is used to generate + -- more efficient code. + , genReduceOp :: LambdaT lore + } + deriving (Eq, Ord, Show) + +data SegRedOp lore = + SegRedOp { segRedComm :: Commutativity + , segRedLambda :: Lambda lore + , segRedNeutral :: [SubExp] + , segRedShape :: Shape + -- ^ In case this operator is semantically a vectorised + -- operator (corresponding to a perfect map nest in the + -- SOACS representation), these are the logical + -- "dimensions". This is used to generate more efficient + -- code. + } + deriving (Eq, Ord, Show) + +-- | How many reduction results are produced by these 'SegRedOp's? +segRedResults :: [SegRedOp lore] -> Int +segRedResults = sum . map (length . segRedNeutral) +-- | The body of a 'Kernel'. +data KernelBody lore = KernelBody { kernelBodyLore :: BodyAttr lore + , kernelBodyStms :: Stms lore + , kernelBodyResult :: [KernelResult] + } + +deriving instance Annotations lore => Ord (KernelBody lore) +deriving instance Annotations lore => Show (KernelBody lore) +deriving instance Annotations lore => Eq (KernelBody lore) + +data KernelResult = Returns SubExp + -- ^ Each "worker" in the kernel returns this. + -- Whether this is a result-per-thread or a + -- result-per-group depends on the 'SegLevel'. + | WriteReturns + [SubExp] -- Size of array. Must match number of dims. + VName -- Which array + [([SubExp], SubExp)] + -- Arbitrary number of index/value pairs. + | ConcatReturns + SplitOrdering -- Permuted? + SubExp -- The final size. + SubExp -- Per-thread/group (max) chunk size. + VName -- Chunk by this worker. + | TileReturns + [(SubExp, SubExp)] -- Total/tile for each dimension + VName -- Tile written by this worker. + -- The TileReturns must not expect more than one + -- result to be written per physical thread. + deriving (Eq, Show, Ord) + +kernelResultSubExp :: KernelResult -> SubExp +kernelResultSubExp (Returns se) = se +kernelResultSubExp (WriteReturns _ arr _) = Var arr +kernelResultSubExp (ConcatReturns _ _ _ v) = Var v +kernelResultSubExp (TileReturns _ v) = Var v + +instance FreeIn KernelResult where + freeIn' (Returns what) = freeIn' what + freeIn' (WriteReturns rws arr res) = freeIn' rws <> freeIn' arr <> freeIn' res + freeIn' (ConcatReturns o w per_thread_elems v) = + freeIn' o <> freeIn' w <> freeIn' per_thread_elems <> freeIn' v + freeIn' (TileReturns dims v) = + freeIn' dims <> freeIn' v + +instance Attributes lore => FreeIn (KernelBody lore) where + freeIn' (KernelBody attr stms res) = + fvBind bound_in_stms $ freeIn' attr <> freeIn' stms <> freeIn' res + where bound_in_stms = fold $ fmap boundByStm stms + +instance Attributes lore => Substitute (KernelBody lore) where + substituteNames subst (KernelBody attr stms res) = + KernelBody + (substituteNames subst attr) + (substituteNames subst stms) + (substituteNames subst res) + +instance Substitute KernelResult where + substituteNames subst (Returns se) = + Returns $ substituteNames subst se + substituteNames subst (WriteReturns rws arr res) = + WriteReturns + (substituteNames subst rws) (substituteNames subst arr) + (substituteNames subst res) + substituteNames subst (ConcatReturns o w per_thread_elems v) = + ConcatReturns + (substituteNames subst o) + (substituteNames subst w) + (substituteNames subst per_thread_elems) + (substituteNames subst v) + substituteNames subst (TileReturns dims v) = + TileReturns (substituteNames subst dims) (substituteNames subst v) + +instance Attributes lore => Rename (KernelBody lore) where + rename (KernelBody attr stms res) = do + attr' <- rename attr + renamingStms stms $ \stms' -> + KernelBody attr' stms' <$> rename res + +instance Rename KernelResult where + rename = substituteRename + +aliasAnalyseKernelBody :: (Attributes lore, + CanBeAliased (Op lore)) => + KernelBody lore + -> KernelBody (Aliases lore) +aliasAnalyseKernelBody (KernelBody attr stms res) = + let Body attr' stms' _ = Alias.analyseBody $ Body attr stms [] + in KernelBody attr' stms' res + +removeKernelBodyAliases :: CanBeAliased (Op lore) => + KernelBody (Aliases lore) -> KernelBody lore +removeKernelBodyAliases (KernelBody (_, attr) stms res) = + KernelBody attr (fmap removeStmAliases stms) res + +addKernelBodyRanges :: (Attributes lore, CanBeRanged (Op lore)) => + KernelBody lore -> Range.RangeM (KernelBody (Ranges lore)) +addKernelBodyRanges (KernelBody attr stms res) = + Range.analyseStms stms $ \stms' -> do + let attr' = (mkBodyRanges stms $ map kernelResultSubExp res, attr) + return $ KernelBody attr' stms' res + +removeKernelBodyRanges :: CanBeRanged (Op lore) => + KernelBody (Ranges lore) -> KernelBody lore +removeKernelBodyRanges (KernelBody (_, attr) stms res) = + KernelBody attr (fmap removeStmRanges stms) res + +removeKernelBodyWisdom :: CanBeWise (Op lore) => + KernelBody (Wise lore) -> KernelBody lore +removeKernelBodyWisdom (KernelBody attr stms res) = + let Body attr' stms' _ = removeBodyWisdom $ Body attr stms [] + in KernelBody attr' stms' res + +consumedInKernelBody :: Aliased lore => + KernelBody lore -> Names +consumedInKernelBody (KernelBody attr stms res) = + consumedInBody (Body attr stms []) <> mconcat (map consumedByReturn res) + where consumedByReturn (WriteReturns _ a _) = oneName a + consumedByReturn _ = mempty + +checkKernelBody :: TC.Checkable lore => + [Type] -> KernelBody (Aliases lore) -> TC.TypeM lore () +checkKernelBody ts (KernelBody (_, attr) stms kres) = do + TC.checkBodyLore attr + TC.checkStms stms $ do + unless (length ts == length kres) $ + TC.bad $ TC.TypeError $ "Kernel return type is " ++ prettyTuple ts ++ + ", but body returns " ++ show (length kres) ++ " values." + zipWithM_ checkKernelResult kres ts + + where checkKernelResult (Returns what) t = + TC.require [t] what + checkKernelResult (WriteReturns rws arr res) t = do + mapM_ (TC.require [Prim int32]) rws + arr_t <- lookupType arr + forM_ res $ \(is, e) -> do + mapM_ (TC.require [Prim int32]) is + TC.require [t] e + unless (arr_t == t `arrayOfShape` Shape rws) $ + TC.bad $ TC.TypeError $ "WriteReturns returning " ++ + pretty e ++ " of type " ++ pretty t ++ ", shape=" ++ pretty rws ++ + ", but destination array has type " ++ pretty arr_t + TC.consume =<< TC.lookupAliases arr + checkKernelResult (ConcatReturns o w per_thread_elems v) t = do + case o of + SplitContiguous -> return () + SplitStrided stride -> TC.require [Prim int32] stride + TC.require [Prim int32] w + TC.require [Prim int32] per_thread_elems + vt <- lookupType v + unless (vt == t `arrayOfRow` arraySize 0 vt) $ + TC.bad $ TC.TypeError $ "Invalid type for ConcatReturns " ++ pretty v + checkKernelResult (TileReturns dims v) t = do + forM_ dims $ \(dim, tile) -> do + TC.require [Prim int32] dim + TC.require [Prim int32] tile + vt <- lookupType v + unless (vt == t `arrayOfShape` Shape (map snd dims)) $ + TC.bad $ TC.TypeError $ "Invalid type for TileReturns " ++ pretty v + +kernelBodyMetrics :: OpMetrics (Op lore) => KernelBody lore -> MetricsM () +kernelBodyMetrics = mapM_ bindingMetrics . kernelBodyStms + +instance PrettyLore lore => Pretty (KernelBody lore) where + ppr (KernelBody _ stms res) = + PP.stack (map ppr (stmsToList stms)) + text "return" <+> PP.braces (PP.commasep $ map ppr res) + +instance Pretty KernelResult where + ppr (Returns what) = + text "thread returns" <+> ppr what + ppr (WriteReturns rws arr res) = + ppr arr <+> text "with" <+> PP.apply (map ppRes res) + where ppRes (is, e) = + PP.brackets (PP.commasep $ zipWith f is rws) <+> text "<-" <+> ppr e + f i rw = ppr i <+> text "<" <+> ppr rw + ppr (ConcatReturns o w per_thread_elems v) = + text "concat" <> suff <> + parens (commasep [ppr w, ppr per_thread_elems]) <+> ppr v + where suff = case o of SplitContiguous -> mempty + SplitStrided stride -> text "Strided" <> parens (ppr stride) + ppr (TileReturns dims v) = + text "tile" <> + parens (commasep $ map onDim dims) <+> ppr v + where onDim (dim, tile) = ppr dim <+> text "/" <+> ppr tile + +--- Segmented operations + +-- | Do we need group-virtualisation when generating code for the +-- segmented operation? In most cases, we do, but for some simple +-- kernels, we compute the full number of groups in advance, and then +-- virtualisation is an unnecessary (but generally very small) +-- overhead. This only really matters for fairly trivial but very +-- wide @map@ kernels where each thread performs constant-time work on +-- scalars. +data SegVirt = SegVirt | SegNoVirt + deriving (Eq, Ord, Show) + +-- | At which level the *body* of a 'SegOp' executes. +data SegLevel = SegThread { segNumGroups :: Count NumGroups SubExp + , segGroupSize :: Count GroupSize SubExp + , segVirt :: SegVirt } + | SegGroup { segNumGroups :: Count NumGroups SubExp + , segGroupSize :: Count GroupSize SubExp + , segVirt :: SegVirt } + | SegThreadScalar { segNumGroups :: Count NumGroups SubExp + , segGroupSize :: Count GroupSize SubExp + , segVirt :: SegVirt } + -- ^ Like 'SegThread', but with the invariant that the + -- results produced are only used within the same + -- physical thread later on, and can thus be kept in + -- registers. May only occur within an enclosing + -- 'SegGroup' construct. + deriving (Eq, Ord, Show) + +-- | Index space of a 'SegOp'. +data SegSpace = SegSpace { segFlat :: VName + -- ^ Flat physical index corresponding to the + -- dimensions (at code generation used for a + -- thread ID or similar). + , unSegSpace :: [(VName, SubExp)] + } + deriving (Eq, Ord, Show) + + +segSpaceDims :: SegSpace -> [SubExp] +segSpaceDims (SegSpace _ space) = map snd space + +scopeOfSegSpace :: SegSpace -> Scope lore +scopeOfSegSpace (SegSpace phys space) = + M.fromList $ zip (phys : map fst space) $ repeat $ IndexInfo Int32 + +checkSegSpace :: TC.Checkable lore => SegSpace -> TC.TypeM lore () +checkSegSpace (SegSpace _ dims) = + mapM_ (TC.require [Prim int32] . snd) dims + +data SegOp lore = SegMap SegLevel SegSpace [Type] (KernelBody lore) + | SegRed SegLevel SegSpace [SegRedOp lore] [Type] (KernelBody lore) + -- ^ The KernelSpace must always have at least two dimensions, + -- implying that the result of a SegRed is always an array. + | SegScan SegLevel SegSpace (Lambda lore) [SubExp] [Type] (KernelBody lore) + | SegGenRed SegLevel SegSpace [GenReduceOp lore] [Type] (KernelBody lore) + deriving (Eq, Ord, Show) + +segLevel :: SegOp lore -> SegLevel +segLevel (SegMap lvl _ _ _) = lvl +segLevel (SegRed lvl _ _ _ _) = lvl +segLevel (SegScan lvl _ _ _ _ _) = lvl +segLevel (SegGenRed lvl _ _ _ _) = lvl + +segSpace :: SegOp lore -> SegSpace +segSpace (SegMap _ lvl _ _) = lvl +segSpace (SegRed _ lvl _ _ _) = lvl +segSpace (SegScan _ lvl _ _ _ _) = lvl +segSpace (SegGenRed _ lvl _ _ _) = lvl + +segResultShape :: SegSpace -> Type -> KernelResult -> Type +segResultShape _ t (WriteReturns rws _ _) = + t `arrayOfShape` Shape rws +segResultShape space t (Returns _) = + foldr (flip arrayOfRow) t $ segSpaceDims space +segResultShape _ t (ConcatReturns _ w _ _) = + t `arrayOfRow` w +segResultShape _ t (TileReturns dims _) = + t `arrayOfShape` Shape (map fst dims) + +segOpType :: SegOp lore -> [Type] +segOpType (SegMap _ space ts kbody) = + zipWith (segResultShape space) ts $ kernelBodyResult kbody +segOpType (SegRed _ space reds ts kbody) = + red_ts ++ + zipWith (segResultShape space) map_ts + (drop (length red_ts) $ kernelBodyResult kbody) + where map_ts = drop (length red_ts) ts + segment_dims = init $ segSpaceDims space + red_ts = do + op <- reds + let shape = Shape segment_dims <> segRedShape op + map (`arrayOfShape` shape) (lambdaReturnType $ segRedLambda op) +segOpType (SegScan _ space _ nes ts kbody) = + map (`arrayOfShape` Shape dims) scan_ts ++ + zipWith (segResultShape space) map_ts + (drop (length scan_ts) $ kernelBodyResult kbody) + where dims = segSpaceDims space + (scan_ts, map_ts) = splitAt (length nes) ts +segOpType (SegGenRed _ space ops _ _) = do + op <- ops + let shape = Shape (segment_dims <> [genReduceWidth op]) <> genReduceShape op + map (`arrayOfShape` shape) (lambdaReturnType $ genReduceOp op) + where dims = segSpaceDims space + segment_dims = init dims + +instance TypedOp (SegOp lore) where + opType = pure . staticShapes . segOpType + +instance (Attributes lore, Aliased lore) => AliasedOp (SegOp lore) where + opAliases = map (const mempty) . segOpType + + consumedInOp (SegMap _ _ _ kbody) = + consumedInKernelBody kbody + consumedInOp (SegRed _ _ _ _ kbody) = + consumedInKernelBody kbody + consumedInOp (SegScan _ _ _ _ _ kbody) = + consumedInKernelBody kbody + consumedInOp (SegGenRed _ _ ops _ kbody) = + namesFromList (concatMap genReduceDest ops) <> consumedInKernelBody kbody + +checkSegLevel :: Maybe SegLevel -> SegLevel -> TC.TypeM lore () +checkSegLevel Nothing SegThreadScalar{} = + TC.bad $ TC.TypeError "SegThreadScalar at top level." +checkSegLevel Nothing _ = + return () +checkSegLevel (Just SegThread{}) _ = + TC.bad $ TC.TypeError "SegOps cannot occur when already at thread level." +checkSegLevel (Just x) y + | x == y = TC.bad $ TC.TypeError $ "Already at at level " ++ pretty x + | segNumGroups x /= segNumGroups y || segGroupSize x /= segGroupSize y = + TC.bad $ TC.TypeError "Physical layout for SegLevel does not match parent SegLevel." + | otherwise = + return () + +checkSegBasics :: TC.Checkable lore => + Maybe SegLevel -> SegLevel -> SegSpace -> [Type] -> TC.TypeM lore () +checkSegBasics cur_lvl lvl space ts = do + checkSegLevel cur_lvl lvl + checkSegSpace space + mapM_ TC.checkType ts + +typeCheckSegOp :: TC.Checkable lore => + Maybe SegLevel -> SegOp (Aliases lore) -> TC.TypeM lore () +typeCheckSegOp cur_lvl (SegMap lvl space ts kbody) = + checkScanRed cur_lvl lvl space [] ts kbody + +typeCheckSegOp cur_lvl (SegRed lvl space reds ts body) = + checkScanRed cur_lvl lvl space reds' ts body + where reds' = zip3 + (map segRedLambda reds) + (map segRedNeutral reds) + (map segRedShape reds) + +typeCheckSegOp cur_lvl (SegScan lvl space scan_op nes ts body) = + checkScanRed cur_lvl lvl space [(scan_op, nes, mempty)] ts body + +typeCheckSegOp cur_lvl (SegGenRed lvl space ops ts kbody) = do + checkSegBasics cur_lvl lvl space ts + + TC.binding (scopeOfSegSpace space) $ do + nes_ts <- forM ops $ \(GenReduceOp dest_w dests nes shape op) -> do + TC.require [Prim int32] dest_w + nes' <- mapM TC.checkArg nes + mapM_ (TC.require [Prim int32]) $ shapeDims shape + + -- Operator type must match the type of neutral elements. + let stripVecDims = stripArray $ shapeRank shape + TC.checkLambda op $ map (TC.noArgAliases . first stripVecDims) $ nes' ++ nes' + let nes_t = map TC.argType nes' + unless (nes_t == lambdaReturnType op) $ + TC.bad $ TC.TypeError $ "SegGenRed operator has return type " ++ + prettyTuple (lambdaReturnType op) ++ " but neutral element has type " ++ + prettyTuple nes_t + + -- Arrays must have proper type. + let dest_shape = Shape (segment_dims <> [dest_w]) <> shape + forM_ (zip nes_t dests) $ \(t, dest) -> do + TC.requireI [t `arrayOfShape` dest_shape] dest + TC.consume =<< TC.lookupAliases dest + + return $ map (`arrayOfShape` shape) nes_t + + checkKernelBody ts kbody + + -- Return type of bucket function must be an index for each + -- operation followed by the values to write. + let bucket_ret_t = replicate (length ops) (Prim int32) ++ concat nes_ts + unless (bucket_ret_t == ts) $ + TC.bad $ TC.TypeError $ "SegGenRed body has return type " ++ + prettyTuple ts ++ " but should have type " ++ + prettyTuple bucket_ret_t + + where segment_dims = init $ segSpaceDims space + +checkScanRed :: TC.Checkable lore => + Maybe SegLevel -> SegLevel + -> SegSpace + -> [(Lambda (Aliases lore), [SubExp], Shape)] + -> [Type] + -> KernelBody (Aliases lore) + -> TC.TypeM lore () +checkScanRed cur_lvl lvl space ops ts kbody = do + checkSegBasics cur_lvl lvl space ts + + TC.binding (scopeOfSegSpace space) $ do + ne_ts <- forM ops $ \(lam, nes, shape) -> do + mapM_ (TC.require [Prim int32]) $ shapeDims shape + nes' <- mapM TC.checkArg nes + + -- Operator type must match the type of neutral elements. + let stripVecDims = stripArray $ shapeRank shape + TC.checkLambda lam $ map (TC.noArgAliases . first stripVecDims) $ nes' ++ nes' + let nes_t = map TC.argType nes' + + unless (lambdaReturnType lam == nes_t) $ + TC.bad $ TC.TypeError "wrong type for operator or neutral elements." + + return $ map (`arrayOfShape` shape) nes_t + + let expecting = concat ne_ts + got = take (length expecting) ts + unless (expecting == got) $ + TC.bad $ TC.TypeError $ + "Wrong return for body (does not match neutral elements; expected " ++ + pretty expecting ++ "; found " ++ + pretty got ++ ")" + + checkKernelBody ts kbody + +-- | Like 'Mapper', but just for 'SegOp's. +data SegOpMapper flore tlore m = SegOpMapper { + mapOnSegOpSubExp :: SubExp -> m SubExp + , mapOnSegOpLambda :: Lambda flore -> m (Lambda tlore) + , mapOnSegOpBody :: KernelBody flore -> m (KernelBody tlore) + , mapOnSegOpVName :: VName -> m VName + } + +-- | A mapper that simply returns the 'SegOp' verbatim. +identitySegOpMapper :: Monad m => SegOpMapper lore lore m +identitySegOpMapper = SegOpMapper { mapOnSegOpSubExp = return + , mapOnSegOpLambda = return + , mapOnSegOpBody = return + , mapOnSegOpVName = return + } + +mapOnSegSpace :: Monad f => + SegOpMapper flore tlore f -> SegSpace -> f SegSpace +mapOnSegSpace tv (SegSpace phys dims) = + SegSpace phys <$> traverse (traverse $ mapOnSegOpSubExp tv) dims + +mapSegOpM :: (Applicative m, Monad m) => + SegOpMapper flore tlore m -> SegOp flore -> m (SegOp tlore) +mapSegOpM tv (SegMap lvl space ts body) = + SegMap + <$> mapOnSegLevel tv lvl + <*> mapOnSegSpace tv space + <*> mapM (mapOnSegOpType tv) ts + <*> mapOnSegOpBody tv body +mapSegOpM tv (SegRed lvl space reds ts lam) = + SegRed + <$> mapOnSegLevel tv lvl + <*> mapOnSegSpace tv space + <*> mapM onSegOp reds + <*> mapM (mapOnType $ mapOnSegOpSubExp tv) ts + <*> mapOnSegOpBody tv lam + where onSegOp (SegRedOp comm red_op nes shape) = + SegRedOp comm + <$> mapOnSegOpLambda tv red_op + <*> mapM (mapOnSegOpSubExp tv) nes + <*> (Shape <$> mapM (mapOnSegOpSubExp tv) (shapeDims shape)) +mapSegOpM tv (SegScan lvl space scan_op nes ts body) = + SegScan + <$> mapOnSegLevel tv lvl + <*> mapOnSegSpace tv space + <*> mapOnSegOpLambda tv scan_op + <*> mapM (mapOnSegOpSubExp tv) nes + <*> mapM (mapOnType $ mapOnSegOpSubExp tv) ts + <*> mapOnSegOpBody tv body +mapSegOpM tv (SegGenRed lvl space ops ts body) = + SegGenRed + <$> mapOnSegLevel tv lvl + <*> mapOnSegSpace tv space + <*> mapM onGenRedOp ops + <*> mapM (mapOnType $ mapOnSegOpSubExp tv) ts + <*> mapOnSegOpBody tv body + where onGenRedOp (GenReduceOp w arrs nes shape op) = + GenReduceOp <$> mapOnSegOpSubExp tv w + <*> mapM (mapOnSegOpVName tv) arrs + <*> mapM (mapOnSegOpSubExp tv) nes + <*> (Shape <$> mapM (mapOnSegOpSubExp tv) (shapeDims shape)) + <*> mapOnSegOpLambda tv op + +mapOnSegLevel :: Monad m => + SegOpMapper flore tlore m -> SegLevel -> m SegLevel +mapOnSegLevel tv (SegThread num_groups group_size virt) = + SegThread + <$> traverse (mapOnSegOpSubExp tv) num_groups + <*> traverse (mapOnSegOpSubExp tv) group_size + <*> pure virt +mapOnSegLevel tv (SegGroup num_groups group_size virt) = + SegGroup + <$> traverse (mapOnSegOpSubExp tv) num_groups + <*> traverse (mapOnSegOpSubExp tv) group_size + <*> pure virt +mapOnSegLevel tv (SegThreadScalar num_groups group_size virt) = + SegThreadScalar + <$> traverse (mapOnSegOpSubExp tv) num_groups + <*> traverse (mapOnSegOpSubExp tv) group_size + <*> pure virt + +mapOnSegOpType :: Monad m => + SegOpMapper flore tlore m -> Type -> m Type +mapOnSegOpType _tv (Prim pt) = pure $ Prim pt +mapOnSegOpType tv (Array pt shape u) = Array pt <$> f shape <*> pure u + where f (Shape dims) = Shape <$> mapM (mapOnSegOpSubExp tv) dims +mapOnSegOpType _tv (Mem s) = pure $ Mem s + +-- | Like 'Walker', but just for 'SegOp's. +data SegOpWalker lore m = SegOpWalker { + walkOnSegOpSubExp :: SubExp -> m () + , walkOnSegOpLambda :: Lambda lore -> m () + , walkOnSegOpBody :: KernelBody lore -> m () + , walkOnSegOpVName :: VName -> m () + } + +-- | A no-op traversal. +identitySegOpWalker :: Monad m => SegOpWalker lore m +identitySegOpWalker = SegOpWalker { + walkOnSegOpSubExp = const $ return () + , walkOnSegOpLambda = const $ return () + , walkOnSegOpBody = const $ return () + , walkOnSegOpVName = const $ return () + } + +walkSegOpMapper :: forall lore m. Monad m => + SegOpWalker lore m -> SegOpMapper lore lore m +walkSegOpMapper f = SegOpMapper { + mapOnSegOpSubExp = wrap walkOnSegOpSubExp + , mapOnSegOpLambda = wrap walkOnSegOpLambda + , mapOnSegOpBody = wrap walkOnSegOpBody + , mapOnSegOpVName = wrap walkOnSegOpVName + } + where wrap :: (SegOpWalker lore m -> a -> m ()) -> a -> m a + wrap op k = op f k >> return k + +-- | As 'mapSegOpM', but ignoring the results. +walkSegOpM :: Monad m => SegOpWalker lore m -> SegOp lore -> m () +walkSegOpM f = void . mapSegOpM m + where m = walkSegOpMapper f + +instance Attributes lore => Substitute (SegOp lore) where + substituteNames subst = runIdentity . mapSegOpM substitute + where substitute = + SegOpMapper { mapOnSegOpSubExp = return . substituteNames subst + , mapOnSegOpLambda = return . substituteNames subst + , mapOnSegOpBody = return . substituteNames subst + , mapOnSegOpVName = return . substituteNames subst + } + +instance Attributes lore => Rename (SegOp lore) where + rename = mapSegOpM renamer + where renamer = SegOpMapper rename rename rename rename + +instance (Attributes lore, FreeIn (LParamAttr lore)) => + FreeIn (SegOp lore) where + freeIn' e = flip execState mempty $ mapSegOpM free e + where walk f x = modify (<>f x) >> return x + free = SegOpMapper { mapOnSegOpSubExp = walk freeIn' + , mapOnSegOpLambda = walk freeIn' + , mapOnSegOpBody = walk freeIn' + , mapOnSegOpVName = walk freeIn' + } + +instance OpMetrics (Op lore) => OpMetrics (SegOp lore) where + opMetrics (SegMap _ _ _ body) = + inside "SegMap" $ kernelBodyMetrics body + opMetrics (SegRed _ _ reds _ body) = + inside "SegRed" $ do mapM_ (lambdaMetrics . segRedLambda) reds + kernelBodyMetrics body + opMetrics (SegScan _ _ scan_op _ _ body) = + inside "SegScan" $ lambdaMetrics scan_op >> kernelBodyMetrics body + opMetrics (SegGenRed _ _ ops _ body) = + inside "SegGenRed" $ do mapM_ (lambdaMetrics . genReduceOp) ops + kernelBodyMetrics body + +instance Pretty SegSpace where + ppr (SegSpace phys dims) = parens (commasep $ do (i,d) <- dims + return $ ppr i <+> "<" <+> ppr d) <+> + parens (text "~" <> ppr phys) + +instance PP.Pretty SegLevel where + ppr SegThread{} = "thread" + ppr SegThreadScalar{} = "scalar" + ppr SegGroup{} = "group" + +ppSegLevel :: SegLevel -> PP.Doc +ppSegLevel lvl = + PP.parens $ + text "#groups=" <> ppr (segNumGroups lvl) <> PP.semi <+> + text "groupsize=" <> ppr (segGroupSize lvl) <> + case segVirt lvl of + SegNoVirt -> mempty + SegVirt -> PP.semi <+> text "virtualise" + +instance PrettyLore lore => PP.Pretty (SegOp lore) where + ppr (SegMap lvl space ts body) = + text "segmap_" <> ppr lvl + ppSegLevel lvl + PP.align (ppr space) <+> + PP.colon <+> ppTuple' ts <+> PP.nestedBlock "{" "}" (ppr body) + + ppr (SegRed lvl space reds ts body) = + text "segred_" <> ppr lvl + ppSegLevel lvl + PP.parens (PP.braces (mconcat $ intersperse (PP.comma <> PP.line) $ map ppOp reds)) + PP.align (ppr space) <+> PP.colon <+> ppTuple' ts <+> + PP.nestedBlock "{" "}" (ppr body) + where ppOp (SegRedOp comm lam nes shape) = + PP.braces (PP.commasep $ map ppr nes) <> PP.comma + ppr shape <> PP.comma + comm' <> ppr lam + where comm' = case comm of Commutative -> text "commutative " + Noncommutative -> mempty + + ppr (SegScan lvl space scan_op nes ts body) = + text "segscan_" <> ppr lvl + ppSegLevel lvl + PP.parens (ppr scan_op <> PP.comma + PP.braces (PP.commasep $ map ppr nes)) + PP.align (ppr space) <+> PP.colon <+> ppTuple' ts <+> + PP.nestedBlock "{" "}" (ppr body) + + ppr (SegGenRed lvl space ops ts body) = + text "seggenred_" <> ppr lvl + ppSegLevel lvl + PP.parens (PP.braces (mconcat $ intersperse (PP.comma <> PP.line) $ map ppOp ops)) + PP.align (ppr space) <+> PP.colon <+> ppTuple' ts <+> + PP.nestedBlock "{" "}" (ppr body) + where ppOp (GenReduceOp w dests nes shape op) = + ppr w <> PP.comma + PP.braces (PP.commasep $ map ppr dests) <> PP.comma + PP.braces (PP.commasep $ map ppr nes) <> PP.comma + ppr shape <> PP.comma + ppr op + +instance Attributes inner => RangedOp (SegOp inner) where + opRanges op = replicate (length $ segOpType op) unknownRange + +instance (Attributes lore, CanBeRanged (Op lore)) => CanBeRanged (SegOp lore) where + type OpWithRanges (SegOp lore) = SegOp (Ranges lore) + + removeOpRanges = runIdentity . mapSegOpM remove + where remove = SegOpMapper return (return . removeLambdaRanges) + (return . removeKernelBodyRanges) return + addOpRanges = Range.runRangeM . mapSegOpM add + where add = SegOpMapper return Range.analyseLambda + addKernelBodyRanges return + +instance (Attributes lore, + Attributes (Aliases lore), + CanBeAliased (Op lore)) => CanBeAliased (SegOp lore) where + type OpWithAliases (SegOp lore) = SegOp (Aliases lore) + + addOpAliases = runIdentity . mapSegOpM alias + where alias = SegOpMapper return (return . Alias.analyseLambda) + (return . aliasAnalyseKernelBody) return + + removeOpAliases = runIdentity . mapSegOpM remove + where remove = SegOpMapper return (return . removeLambdaAliases) + (return . removeKernelBodyAliases) return + +instance (CanBeWise (Op lore), Attributes lore) => CanBeWise (SegOp lore) where + type OpWithWisdom (SegOp lore) = SegOp (Wise lore) + + removeOpWisdom = runIdentity . mapSegOpM remove + where remove = SegOpMapper return + (return . removeLambdaWisdom) + (return . removeKernelBodyWisdom) + return + +instance Attributes lore => ST.IndexOp (SegOp lore) where + indexOp vtable k (SegMap _ space _ kbody) is = do + Returns se <- maybeNth k $ kernelBodyResult kbody + let (gtids, _) = unzip $ unSegSpace space + guard $ length gtids == length is + let prim_table = M.fromList $ zip gtids $ zip is $ repeat mempty + prim_table' = foldl expandPrimExpTable prim_table $ kernelBodyStms kbody + case se of + Var v -> M.lookup v prim_table' + _ -> Nothing + where expandPrimExpTable table stm + | [v] <- patternNames $ stmPattern stm, + Just (pe,cs) <- + runWriterT $ primExpFromExp (asPrimExp table) $ stmExp stm = + M.insert v (pe, stmCerts stm <> cs) table + | otherwise = + table + + asPrimExp table v + | Just (e,cs) <- M.lookup v table = tell cs >> return e + | Just (Prim pt) <- ST.lookupType v vtable = + return $ LeafExp v pt + | otherwise = lift Nothing + indexOp _ _ _ _ = Nothing + +instance Attributes lore => IsOp (SegOp lore) where + cheapOp _ = False + safeOp _ = True + +--- Host operations + +-- | A host-level operation; parameterised by what else it can do. +data HostOp lore op + = SplitSpace SplitOrdering SubExp SubExp SubExp + -- ^ @SplitSpace o w i elems_per_thread@. + -- + -- Computes how to divide array elements to + -- threads in a kernel. Returns the number of + -- elements in the chunk that the current thread + -- should take. + -- + -- @w@ is the length of the outer dimension in + -- the array. @i@ is the current thread + -- index. Each thread takes at most + -- @elems_per_thread@ elements. + -- + -- If the order @o@ is 'SplitContiguous', thread with index @i@ + -- should receive elements + -- @i*elems_per_tread, i*elems_per_thread + 1, + -- ..., i*elems_per_thread + (elems_per_thread-1)@. + -- + -- If the order @o@ is @'SplitStrided' stride@, + -- the thread will receive elements @i, + -- i+stride, i+2*stride, ..., + -- i+(elems_per_thread-1)*stride@. + | GetSize Name SizeClass + -- ^ Produce some runtime-configurable size. + | GetSizeMax SizeClass + -- ^ The maximum size of some class. + | CmpSizeLe Name SizeClass SubExp + -- ^ Compare size (likely a threshold) with some Int32 value. + | SegOp (SegOp lore) + -- ^ A segmented operation. + | OtherOp op + deriving (Eq, Ord, Show) + +instance (Attributes lore, Substitute op) => Substitute (HostOp lore op) where + substituteNames substs (SegOp op) = + SegOp $ substituteNames substs op + substituteNames substs (OtherOp op) = + OtherOp $ substituteNames substs op + substituteNames subst (SplitSpace o w i elems_per_thread) = + SplitSpace + (substituteNames subst o) + (substituteNames subst w) + (substituteNames subst i) + (substituteNames subst elems_per_thread) + substituteNames substs (CmpSizeLe name sclass x) = + CmpSizeLe name sclass $ substituteNames substs x + substituteNames _ x = x + +instance (Attributes lore, Rename op) => Rename (HostOp lore op) where + rename (SplitSpace o w i elems_per_thread) = + SplitSpace + <$> rename o + <*> rename w + <*> rename i + <*> rename elems_per_thread + rename (SegOp op) = SegOp <$> rename op + rename (OtherOp op) = OtherOp <$> rename op + rename (CmpSizeLe name sclass x) = CmpSizeLe name sclass <$> rename x + rename x = pure x + +instance (Attributes lore, IsOp op) => IsOp (HostOp lore op) where + safeOp (SegOp op) = safeOp op + safeOp (OtherOp op) = safeOp op + safeOp _ = True + cheapOp (SegOp op) = cheapOp op + cheapOp (OtherOp op) = cheapOp op + cheapOp _ = True + +instance TypedOp op => TypedOp (HostOp lore op) where + opType SplitSpace{} = pure [Prim int32] + opType GetSize{} = pure [Prim int32] + opType GetSizeMax{} = pure [Prim int32] + opType CmpSizeLe{} = pure [Prim Bool] + opType (SegOp op) = opType op + opType (OtherOp op) = opType op + +instance (Aliased lore, AliasedOp op, Attributes lore) => AliasedOp (HostOp lore op) where + opAliases (SegOp op) = opAliases op + opAliases (OtherOp op) = opAliases op + opAliases _ = [mempty] + + consumedInOp (SegOp op) = consumedInOp op + consumedInOp (OtherOp op) = consumedInOp op + consumedInOp _ = mempty + +instance (Attributes lore, RangedOp op) => RangedOp (HostOp lore op) where + opRanges (SplitSpace _ _ _ elems_per_thread) = + [(Just (ScalarBound 0), + Just (ScalarBound (SE.subExpToScalExp elems_per_thread int32)))] + opRanges (SegOp op) = opRanges op + opRanges (OtherOp op) = opRanges op + opRanges _ = [unknownRange] + +instance (Attributes lore, FreeIn op) => FreeIn (HostOp lore op) where + freeIn' (SplitSpace o w i elems_per_thread) = + freeIn' o <> freeIn' [w, i, elems_per_thread] + freeIn' (SegOp op) = freeIn' op + freeIn' (OtherOp op) = freeIn' op + freeIn' (CmpSizeLe _ _ x) = freeIn' x + freeIn' _ = mempty + +instance (CanBeAliased (Op lore), CanBeAliased op, Attributes lore) => CanBeAliased (HostOp lore op) where + type OpWithAliases (HostOp lore op) = HostOp (Aliases lore) (OpWithAliases op) + + addOpAliases (SplitSpace o w i elems_per_thread) = + SplitSpace o w i elems_per_thread + addOpAliases (SegOp op) = SegOp $ addOpAliases op + addOpAliases (OtherOp op) = OtherOp $ addOpAliases op + addOpAliases (GetSize name sclass) = GetSize name sclass + addOpAliases (GetSizeMax sclass) = GetSizeMax sclass + addOpAliases (CmpSizeLe name sclass x) = CmpSizeLe name sclass x + + removeOpAliases (SplitSpace o w i elems_per_thread) = + SplitSpace o w i elems_per_thread + removeOpAliases (SegOp op) = SegOp $ removeOpAliases op + removeOpAliases (OtherOp op) = OtherOp $ removeOpAliases op + removeOpAliases (GetSize name sclass) = GetSize name sclass + removeOpAliases (GetSizeMax sclass) = GetSizeMax sclass + removeOpAliases (CmpSizeLe name sclass x) = CmpSizeLe name sclass x + +instance (CanBeRanged (Op lore), CanBeRanged op, Attributes lore) => CanBeRanged (HostOp lore op) where + type OpWithRanges (HostOp lore op) = HostOp (Ranges lore) (OpWithRanges op) + + addOpRanges (SplitSpace o w i elems_per_thread) = + SplitSpace o w i elems_per_thread + addOpRanges (SegOp op) = SegOp $ addOpRanges op + addOpRanges (OtherOp op) = OtherOp $ addOpRanges op + addOpRanges (GetSize name sclass) = GetSize name sclass + addOpRanges (GetSizeMax sclass) = GetSizeMax sclass + addOpRanges (CmpSizeLe name sclass x) = CmpSizeLe name sclass x + + removeOpRanges (SplitSpace o w i elems_per_thread) = + SplitSpace o w i elems_per_thread + removeOpRanges (SegOp op) = SegOp $ removeOpRanges op + removeOpRanges (OtherOp op) = OtherOp $ removeOpRanges op + removeOpRanges (GetSize name sclass) = GetSize name sclass + removeOpRanges (GetSizeMax sclass) = GetSizeMax sclass + removeOpRanges (CmpSizeLe name sclass x) = CmpSizeLe name sclass x + +instance (CanBeWise (Op lore), CanBeWise op, Attributes lore) => CanBeWise (HostOp lore op) where + type OpWithWisdom (HostOp lore op) = HostOp (Wise lore) (OpWithWisdom op) + + removeOpWisdom (SplitSpace o w i elems_per_thread) = + SplitSpace o w i elems_per_thread + removeOpWisdom (SegOp op) = SegOp $ removeOpWisdom op + removeOpWisdom (GetSize name sclass) = GetSize name sclass + removeOpWisdom (GetSizeMax sclass) = GetSizeMax sclass + removeOpWisdom (CmpSizeLe name sclass x) = CmpSizeLe name sclass x + removeOpWisdom (OtherOp op) = OtherOp $ removeOpWisdom op + +instance (Attributes lore, ST.IndexOp op) => ST.IndexOp (HostOp lore op) where + indexOp vtable k (SegOp op) is = ST.indexOp vtable k op is + indexOp vtable k (OtherOp op) is = ST.indexOp vtable k op is + indexOp _ _ _ _ = Nothing + +instance (PrettyLore lore, PP.Pretty op) => PP.Pretty (HostOp lore op) where + ppr (SplitSpace o w i elems_per_thread) = + text "splitSpace" <> suff <> + parens (commasep [ppr w, ppr i, ppr elems_per_thread]) + where suff = case o of SplitContiguous -> mempty + SplitStrided stride -> text "Strided" <> parens (ppr stride) + + ppr (GetSize name size_class) = + text "get_size" <> parens (commasep [ppr name, ppr size_class]) + + ppr (GetSizeMax size_class) = + text "get_size_max" <> parens (ppr size_class) + + ppr (CmpSizeLe name size_class x) = + text "get_size" <> parens (commasep [ppr name, ppr size_class]) <+> + text "<=" <+> ppr x + + ppr (SegOp op) = ppr op + + ppr (OtherOp op) = ppr op + +instance (OpMetrics (Op lore), OpMetrics op) => OpMetrics (HostOp lore op) where + opMetrics SplitSpace{} = seen "SplitSpace" + opMetrics GetSize{} = seen "GetSize" + opMetrics GetSizeMax{} = seen "GetSizeMax" + opMetrics CmpSizeLe{} = seen "CmpSizeLe" + opMetrics (SegOp op) = opMetrics op + opMetrics (OtherOp op) = opMetrics op + +typeCheckHostOp :: TC.Checkable lore => + (SegLevel -> OpWithAliases (Op lore) -> TC.TypeM lore ()) + -> Maybe SegLevel + -> (op -> TC.TypeM lore ()) + -> HostOp (Aliases lore) op + -> TC.TypeM lore () +typeCheckHostOp _ _ _ (SplitSpace o w i elems_per_thread) = do + case o of + SplitContiguous -> return () + SplitStrided stride -> TC.require [Prim int32] stride + mapM_ (TC.require [Prim int32]) [w, i, elems_per_thread] +typeCheckHostOp _ _ _ GetSize{} = return () +typeCheckHostOp _ _ _ GetSizeMax{} = return () +typeCheckHostOp _ _ _ (CmpSizeLe _ _ x) = TC.require [Prim int32] x +typeCheckHostOp checker lvl _ (SegOp op) = + TC.checkOpWith (checker $ segLevel op) $ + typeCheckSegOp lvl op +typeCheckHostOp _ _ f (OtherOp op) = f op diff --git a/src/Futhark/Representation/Kernels/KernelExp.hs b/src/Futhark/Representation/Kernels/KernelExp.hs --- a/src/Futhark/Representation/Kernels/KernelExp.hs +++ b/src/Futhark/Representation/Kernels/KernelExp.hs @@ -22,7 +22,6 @@ import Control.Monad import Data.Monoid ((<>)) import Data.Maybe -import qualified Data.Set as S import qualified Data.Map.Strict as M import qualified Futhark.Analysis.Alias as Alias @@ -176,28 +175,28 @@ opType (Barrier ses) = staticShapes <$> traverse subExpType ses instance FreeIn SplitOrdering where - freeIn SplitContiguous = mempty - freeIn (SplitStrided stride) = freeIn stride + freeIn' SplitContiguous = mempty + freeIn' (SplitStrided stride) = freeIn' stride instance Attributes lore => FreeIn (KernelExp lore) where - freeIn (SplitSpace o w i elems_per_thread) = - freeIn o <> freeIn [w, i, elems_per_thread] - freeIn (Combine (CombineSpace scatter cspace) ts active body) = - freeIn scatter <> freeIn (map snd cspace) <> freeIn ts <> freeIn active <> freeIn body - freeIn (GroupReduce w lam input) = - freeIn w <> freeIn lam <> freeIn input - freeIn (GroupScan w lam input) = - freeIn w <> freeIn lam <> freeIn input - freeIn (GroupStream w maxchunk lam accs arrs) = - freeIn w <> freeIn maxchunk <> freeIn lam <> freeIn accs <> freeIn arrs - freeIn (GroupGenReduce w dests op bucket values locks) = - freeIn w <> freeIn dests <> freeIn op <> freeIn bucket <> freeIn values <> freeIn locks - freeIn (Barrier ses) = freeIn ses + freeIn' (SplitSpace o w i elems_per_thread) = + freeIn' o <> freeIn' [w, i, elems_per_thread] + freeIn' (Combine (CombineSpace scatter cspace) ts active body) = + freeIn' scatter <> freeIn' (map snd cspace) <> freeIn' ts <> freeIn' active <> freeIn' body + freeIn' (GroupReduce w lam input) = + freeIn' w <> freeIn' lam <> freeIn' input + freeIn' (GroupScan w lam input) = + freeIn' w <> freeIn' lam <> freeIn' input + freeIn' (GroupStream w maxchunk lam accs arrs) = + freeIn' w <> freeIn' maxchunk <> freeIn' lam <> freeIn' accs <> freeIn' arrs + freeIn' (GroupGenReduce w dests op bucket values locks) = + freeIn' w <> freeIn' dests <> freeIn' op <> freeIn' bucket <> freeIn' values <> freeIn' locks + freeIn' (Barrier ses) = freeIn' ses instance Attributes lore => FreeIn (GroupStreamLambda lore) where - freeIn (GroupStreamLambda chunk_size chunk_offset acc_params arr_params body) = - freeIn body `S.difference` bound_here - where bound_here = S.fromList $ + freeIn' (GroupStreamLambda chunk_size chunk_offset acc_params arr_params body) = + fvBind bound_here $ freeIn' body + where bound_here = namesFromList $ chunk_offset : chunk_size : map paramName (acc_params ++ arr_params) @@ -219,25 +218,25 @@ opAliases (GroupStream _ _ lam _ _) = map (const mempty) $ groupStreamAccParams lam opAliases (GroupGenReduce _ dests _ _ _ _) = - map S.singleton dests + map oneName dests opAliases (Barrier ses) = map subExpAliases ses consumedInOp (GroupReduce _ _ input) = - S.fromList $ map snd input + namesFromList $ map snd input consumedInOp (GroupScan _ _ input) = - S.fromList $ map snd input + namesFromList $ map snd input consumedInOp (GroupStream _ _ lam accs arrs) = -- GroupStream always consumes array-typed accumulators. This -- guarantees that we can use their storage for the result of the -- lambda. - S.map consumedArray $ - S.fromList (map paramName acc_params) <> consumedInBody body + namesFromList $ map consumedArray $ + map paramName acc_params <> namesToList (consumedInBody body) where GroupStreamLambda _ _ acc_params arr_params body = lam consumedArray v = fromMaybe v $ subExpVar =<< lookup v params_to_arrs params_to_arrs = zip (map paramName $ acc_params ++ arr_params) $ accs ++ map Var arrs consumedInOp (GroupGenReduce _ dests _ _ _ _) = - S.fromList dests + namesFromList dests consumedInOp SplitSpace{} = mempty consumedInOp Barrier{} = mempty diff --git a/src/Futhark/Representation/Kernels/Simplify.hs b/src/Futhark/Representation/Kernels/Simplify.hs --- a/src/Futhark/Representation/Kernels/Simplify.hs +++ b/src/Futhark/Representation/Kernels/Simplify.hs @@ -10,17 +10,14 @@ -- * Building blocks , simplifyKernelOp - , simplifyKernelExp ) where import Control.Monad -import Data.Either import Data.Foldable import Data.List import Data.Maybe import qualified Data.Map.Strict as M -import qualified Data.Set as S import Futhark.Representation.Kernels import qualified Futhark.Optimise.Simplify.Engine as Engine @@ -29,119 +26,69 @@ import Futhark.MonadFreshNames import Futhark.Tools import Futhark.Pass +import Futhark.Representation.SOACS.Simplify (simplifySOAC) import qualified Futhark.Optimise.Simplify as Simplify import Futhark.Optimise.Simplify.Rule import qualified Futhark.Analysis.SymbolTable as ST import qualified Futhark.Analysis.UsageTable as UT -import Futhark.Analysis.Rephrase (castStm) import Futhark.Util (chunks) +import qualified Futhark.Transform.FirstOrderTransform as FOT simpleKernels :: Simplify.SimpleOps Kernels -simpleKernels = Simplify.bindableSimpleOps (simplifyKernelOp simpleInKernel inKernelEnv) - -simpleInKernel :: KernelSpace -> Simplify.SimpleOps InKernel -simpleInKernel = Simplify.bindableSimpleOps . simplifyKernelExp +simpleKernels = Simplify.bindableSimpleOps $ simplifyKernelOp simplifySOAC simplifyKernels :: Prog Kernels -> PassM (Prog Kernels) simplifyKernels = Simplify.simplifyProg simpleKernels kernelRules Simplify.noExtraHoistBlockers -simplifyLambda :: (HasScope InKernel m, MonadFreshNames m) => - KernelSpace -> Lambda InKernel -> [Maybe VName] -> m (Lambda InKernel) -simplifyLambda kspace = - Simplify.simplifyLambda (simpleInKernel kspace) - inKernelRules Engine.noExtraHoistBlockers +simplifyLambda :: (HasScope Kernels m, MonadFreshNames m) => + Lambda Kernels -> [Maybe VName] -> m (Lambda Kernels) +simplifyLambda = + Simplify.simplifyLambda simpleKernels kernelRules Engine.noExtraHoistBlockers simplifyKernelOp :: (Engine.SimplifiableLore lore, - Engine.SimplifiableLore outerlore, - BodyAttr outerlore ~ (), BodyAttr lore ~ (), - ExpAttr lore ~ ExpAttr outerlore, - SameScope lore outerlore, - RetType lore ~ RetType outerlore, - BranchType lore ~ BranchType outerlore) => - (KernelSpace -> Engine.SimpleOps lore) -> Engine.Env lore - -> HostOp outerlore (Kernel lore) - -> Engine.SimpleM outerlore (HostOp (Wise outerlore) (Kernel (Wise lore)), - Stms (Wise outerlore)) - -simplifyKernelOp mk_ops env (HostOp (Kernel desc space ts kbody)) = do - space' <- Engine.simplify space - ts' <- mapM Engine.simplify ts - outer_vtable <- Engine.askVtable - ((kbody_stms, kbody_res), kbody_hoisted) <- - Engine.subSimpleM (mk_ops space) env outer_vtable $ do - par_blocker <- Engine.asksEngineEnv $ Engine.blockHoistPar . Engine.envHoistBlockers - Engine.localVtable (<>scope_vtable) $ - Engine.blockIf (Engine.hasFree bound_here - `Engine.orIf` Engine.isOp - `Engine.orIf` par_blocker - `Engine.orIf` Engine.isConsumed) $ - simplifyKernelBodyM kbody - kbody_hoisted' <- mapM processHoistedStm kbody_hoisted - return (HostOp $ Kernel desc space' ts' $ mkWiseKernelBody () kbody_stms kbody_res, - kbody_hoisted') - where scope = scopeOfKernelSpace space - scope_vtable = ST.fromScope scope - bound_here = S.fromList $ M.keys scope - -simplifyKernelOp mk_ops env (HostOp (SegMap space ts body)) = do - space' <- Engine.simplify space - ts' <- mapM Engine.simplify ts + BodyAttr lore ~ ()) => + Simplify.SimplifyOp lore op + -> HostOp lore op + -> Engine.SimpleM lore (HostOp (Wise lore) (OpWithWisdom op), Stms (Wise lore)) - (body', body_hoisted) <- hoistFromBody space' (mk_ops space') env body +simplifyKernelOp f (OtherOp op) = do + (op', stms) <- f op + return (OtherOp op', stms) - return (HostOp $ SegMap space' ts' body', +simplifyKernelOp _ (SegOp (SegMap lvl space ts kbody)) = do + (lvl', space', ts') <- Engine.simplify (lvl, space, ts) + (kbody', body_hoisted) <- simplifyKernelBody space kbody + return (SegOp $ SegMap lvl' space' ts' kbody', body_hoisted) -simplifyKernelOp mk_ops env (HostOp (SegRed space reds ts body)) = do - ts' <- mapM Engine.simplify ts - space' <- Engine.simplify space - outer_vtable <- Engine.askVtable - +simplifyKernelOp _ (SegOp (SegRed lvl space reds ts kbody)) = do + (lvl', space', ts') <- Engine.simplify (lvl, space, ts) (reds', reds_hoisted) <- fmap unzip $ forM reds $ \(SegRedOp comm lam nes shape) -> do - (lam', hoisted) <- Engine.subSimpleM (mk_ops space) env outer_vtable $ - Engine.localVtable (<>scope_vtable) $ - Engine.simplifyLambda lam $ - replicate (length nes * 2) Nothing + (lam', hoisted) <- + Engine.localVtable (<>scope_vtable) $ + Engine.simplifyLambda lam $ replicate (length nes * 2) Nothing shape' <- Engine.simplify shape nes' <- mapM Engine.simplify nes return (SegRedOp comm lam' nes' shape', hoisted) - red_op_hoisted' <- mapM processHoistedStm $ mconcat reds_hoisted - (body', body_hoisted) <- hoistFromBody space' (mk_ops space') env body - - return (HostOp $ SegRed space' reds' ts' body', - red_op_hoisted' <> body_hoisted) - - where scope_vtable = ST.fromScope scope - scope = scopeOfKernelSpace space - -simplifyKernelOp mk_ops env (HostOp (SegScan space scan_op nes ts body)) = do - outer_vtable <- Engine.askVtable - space' <- Engine.simplify space - - (scan_op', scan_op_hoisted) <- - Engine.subSimpleM (mk_ops space) env outer_vtable $ - Engine.localVtable (<>scope_vtable) $ - Engine.simplifyLambda scan_op $ replicate (length nes * 2) Nothing - scan_op_hoisted' <- mapM processHoistedStm scan_op_hoisted - - nes' <- mapM Engine.simplify nes - ts' <- mapM Engine.simplify ts - - (body', body_hoisted) <- hoistFromBody space' (mk_ops space') env body + (kbody', body_hoisted) <- simplifyKernelBody space kbody - return (HostOp $ SegScan space' scan_op' nes' ts' body', - scan_op_hoisted' <> body_hoisted) + return (SegOp $ SegRed lvl' space' reds' ts' kbody', + mconcat reds_hoisted <> body_hoisted) + where scope = scopeOfSegSpace space + scope_vtable = ST.fromScope scope - where scope_vtable = ST.fromScope scope - scope = scopeOfKernelSpace space +simplifyKernelOp _ (SegOp (SegScan lvl space scan_op nes ts kbody)) = do + lvl' <- Engine.simplify lvl + (space', scan_op', nes', ts', kbody', hoisted) <- + simplifyRedOrScan space scan_op nes ts kbody -simplifyKernelOp mk_ops env (HostOp (SegGenRed space ops ts body)) = do - outer_vtable <- Engine.askVtable + return (SegOp $ SegScan lvl' space' scan_op' nes' ts' kbody', + hoisted) - space' <- Engine.simplify space - ts' <- mapM Engine.simplify ts +simplifyKernelOp _ (SegOp (SegGenRed lvl space ops ts kbody)) = do + (lvl', space', ts') <- Engine.simplify (lvl, space, ts) (ops', ops_hoisted) <- fmap unzip $ forM ops $ \(GenReduceOp w arrs nes dims lam) -> do @@ -150,76 +97,80 @@ nes' <- Engine.simplify nes dims' <- Engine.simplify dims (lam', op_hoisted) <- - Engine.subSimpleM (mk_ops space) env outer_vtable $ Engine.localVtable (<>scope_vtable) $ Engine.simplifyLambda lam $ replicate (length nes * 2) Nothing return (GenReduceOp w' arrs' nes' dims' lam', op_hoisted) - red_op_hoisted' <- mapM processHoistedStm $ mconcat ops_hoisted - - (body', body_hoisted) <- hoistFromBody space' (mk_ops space') env body + (kbody', body_hoisted) <- simplifyKernelBody space kbody - return (HostOp $ SegGenRed space' ops' ts' body', - red_op_hoisted' <> body_hoisted) + return (SegOp $ SegGenRed lvl' space' ops' ts' kbody', + mconcat ops_hoisted <> body_hoisted) - where scope_vtable = ST.fromScope scope - scope = scopeOfKernelSpace space + where scope = scopeOfSegSpace space + scope_vtable = ST.fromScope scope -simplifyKernelOp _ _ (GetSize key size_class) = +simplifyKernelOp _ (SplitSpace o w i elems_per_thread) = + (,) <$> (SplitSpace <$> Engine.simplify o <*> Engine.simplify w + <*> Engine.simplify i <*> Engine.simplify elems_per_thread) + <*> pure mempty +simplifyKernelOp _ (GetSize key size_class) = return (GetSize key size_class, mempty) -simplifyKernelOp _ _ (GetSizeMax size_class) = +simplifyKernelOp _ (GetSizeMax size_class) = return (GetSizeMax size_class, mempty) -simplifyKernelOp _ _ (CmpSizeLe key size_class x) = do +simplifyKernelOp _ (CmpSizeLe key size_class x) = do x' <- Engine.simplify x return (CmpSizeLe key size_class x', mempty) -hoistFromBody :: (Engine.SimplifiableLore lore, - SameScope lore outerlore, - BodyAttr outerlore ~ (), BodyAttr lore ~ (), - ExpAttr lore ~ ExpAttr outerlore, - RetType lore ~ RetType outerlore, - BranchType lore ~ BranchType outerlore) => - KernelSpace -> Simplify.SimpleOps lore -> Engine.Env lore - -> KernelBody lore - -> Engine.SimpleM outerlore (KernelBody (Wise lore), Stms (Wise outerlore)) -hoistFromBody kspace ops env kbody = do - outer_vtable <- Engine.askVtable +simplifyRedOrScan :: (Engine.SimplifiableLore lore, BodyAttr lore ~ ()) => + SegSpace + -> Lambda lore -> [SubExp] -> [Type] + -> KernelBody lore + -> Simplify.SimpleM lore + (SegSpace, Lambda (Wise lore), [SubExp], [Type], KernelBody (Wise lore), + Stms (Wise lore)) +simplifyRedOrScan space scan_op nes ts kbody = do + space' <- Engine.simplify space + nes' <- mapM Engine.simplify nes + ts' <- mapM Engine.simplify ts - ((body_stms, body_res), body_hoisted) <- - Engine.subSimpleM ops env outer_vtable $ do - par_blocker <- Engine.asksEngineEnv $ Engine.blockHoistPar . Engine.envHoistBlockers - Engine.localVtable (<>scope_vtable) $ - Engine.blockIf (Engine.hasFree bound_here - `Engine.orIf` Engine.isOp - `Engine.orIf` par_blocker - `Engine.orIf` Engine.isConsumed) $ - simplifyKernelBodyM kbody + (scan_op', scan_op_hoisted) <- + Engine.localVtable (<>scope_vtable) $ + Engine.simplifyLambda scan_op $ replicate (length nes * 2) Nothing - body_hoisted' <- mapM processHoistedStm body_hoisted + (kbody', body_hoisted) <- simplifyKernelBody space kbody - return (mkWiseKernelBody () body_stms body_res, - body_hoisted') + return (space', scan_op', nes', ts', kbody', + scan_op_hoisted <> body_hoisted) - where scope_vtable = ST.fromScope scope - scope = scopeOfKernelSpace kspace - bound_here = S.fromList $ M.keys scope + where scope = scopeOfSegSpace space + scope_vtable = ST.fromScope scope -processHoistedStm :: (Monad m, - PrettyLore from, - ExpAttr from ~ ExpAttr to, - BodyAttr from ~ BodyAttr to, - RetType from ~ RetType to, - BranchType from ~ BranchType to, - LetAttr from ~ LetAttr to, - FParamAttr from ~ FParamAttr to, - LParamAttr from ~ LParamAttr to) => - Stm from -> m (Stm to) -processHoistedStm bnd - | Just bnd' <- castStm bnd = return bnd' - | otherwise = fail $ "Cannot hoist binding: " ++ pretty bnd +simplifyKernelBody :: (Engine.SimplifiableLore lore, BodyAttr lore ~ ()) => + SegSpace -> KernelBody lore + -> Engine.SimpleM lore (KernelBody (Wise lore), Stms (Wise lore)) +simplifyKernelBody space (KernelBody _ stms res) = do + par_blocker <- Engine.asksEngineEnv $ Engine.blockHoistPar . Engine.envHoistBlockers + ((body_stms, body_res), hoisted) <- + Engine.localVtable (<>scope_vtable) $ + Engine.localVtable (\vtable -> vtable { ST.simplifyMemory = True }) $ + Engine.blockIf (Engine.hasFree bound_here + `Engine.orIf` Engine.isOp + `Engine.orIf` par_blocker + `Engine.orIf` Engine.isConsumed) $ + Engine.simplifyStms stms $ do + res' <- Engine.localVtable (ST.hideCertified $ namesFromList $ M.keys $ scopeOf stms) $ + mapM Engine.simplify res + return ((res', UT.usages $ freeIn res'), mempty) + + return (mkWiseKernelBody () body_stms body_res, + hoisted) + + where scope_vtable = ST.fromScope $ scopeOfSegSpace space + bound_here = namesFromList $ M.keys $ scopeOfSegSpace space + mkWiseKernelBody :: (Attributes lore, CanBeWise (Op lore)) => BodyAttr lore -> Stms (Wise lore) -> [KernelResult] -> KernelBody (Wise lore) mkWiseKernelBody attr bnds res = @@ -227,220 +178,64 @@ in KernelBody attr' bnds res where res_vs = map kernelResultSubExp res -inKernelEnv :: Engine.Env InKernel -inKernelEnv = Engine.emptyEnv inKernelRules Simplify.noExtraHoistBlockers - instance Engine.Simplifiable SplitOrdering where simplify SplitContiguous = return SplitContiguous simplify (SplitStrided stride) = SplitStrided <$> Engine.simplify stride -instance Engine.Simplifiable CombineSpace where - simplify (CombineSpace scatter cspace) = - CombineSpace <$> mapM Engine.simplify scatter - <*> mapM (traverse Engine.simplify) cspace - -simplifyKernelExp :: Engine.SimplifiableLore lore => - KernelSpace -> KernelExp lore - -> Engine.SimpleM lore (KernelExp (Wise lore), Stms (Wise lore)) - -simplifyKernelExp _ (Barrier se) = - (,) <$> (Barrier <$> Engine.simplify se) <*> pure mempty - -simplifyKernelExp _ (SplitSpace o w i elems_per_thread) = - (,) <$> (SplitSpace <$> Engine.simplify o <*> Engine.simplify w - <*> Engine.simplify i <*> Engine.simplify elems_per_thread) - <*> pure mempty - -simplifyKernelExp kspace (Combine cspace ts active body) = do - ((body_stms', body_res'), hoisted) <- - wrapbody $ Engine.blockIf (Engine.hasFree bound_here `Engine.orIf` - maybeBlockUnsafe) $ - localScope (scopeOfCombineSpace cspace) $ - Engine.simplifyBody (map (const Observe) ts) body - body' <- Engine.constructBody body_stms' body_res' - (,) <$> (Combine <$> Engine.simplify cspace - <*> mapM Engine.simplify ts - <*> mapM Engine.simplify active - <*> pure body') <*> pure hoisted - where bound_here = S.fromList $ M.keys $ scopeOfCombineSpace cspace - - protectCombineHoisted checkIfActive m = do - (x, stms) <- m - runBinder $ do - if any (not . safeExp . stmExp) stms - then do is_active <- checkIfActive - mapM_ (Engine.protectIf (not . safeExp) is_active) stms - else addStms stms - return x - - (maybeBlockUnsafe, wrapbody) - | [d] <- map snd $ cspaceDims cspace, - d == spaceGroupSize kspace = - (Engine.isFalse True, - protectCombineHoisted $ - letSubExp "active" =<< - foldBinOp LogAnd (constant True) =<< - mapM (uncurry check) active) - | otherwise = - (Engine.isNotSafe, id) - - check v se = - letSubExp "is_active" $ BasicOp $ CmpOp (CmpSlt Int32) (Var v) se - -simplifyKernelExp _ (GroupReduce w lam input) = do - arrs' <- mapM Engine.simplify arrs - nes' <- mapM Engine.simplify nes - w' <- Engine.simplify w - (lam', hoisted) <- Engine.simplifyLambdaSeq lam (map (const Nothing) arrs') - return (GroupReduce w' lam' $ zip nes' arrs', hoisted) - where (nes,arrs) = unzip input - -simplifyKernelExp _ (GroupScan w lam input) = do - w' <- Engine.simplify w - nes' <- mapM Engine.simplify nes - arrs' <- mapM Engine.simplify arrs - (lam', hoisted) <- Engine.simplifyLambdaSeq lam (map (const Nothing) arrs') - return (GroupScan w' lam' $ zip nes' arrs', hoisted) - where (nes,arrs) = unzip input - -simplifyKernelExp _ (GroupGenReduce w dests op bucket vs locks) = do - w' <- Engine.simplify w - dests' <- mapM Engine.simplify dests - (op', hoisted) <- Engine.simplifyLambdaSeq op (map (const Nothing) vs) - bucket' <- Engine.simplify bucket - vs' <- mapM Engine.simplify vs - locks' <- Engine.simplify locks - return (GroupGenReduce w' dests' op' bucket' vs' locks', hoisted) - -simplifyKernelExp _ (GroupStream w maxchunk lam accs arrs) = do - w' <- Engine.simplify w - maxchunk' <- Engine.simplify maxchunk - accs' <- mapM Engine.simplify accs - arrs' <- mapM Engine.simplify arrs - (lam', hoisted) <- simplifyGroupStreamLambda lam w' maxchunk' arrs' - return (GroupStream w' maxchunk' lam' accs' arrs', hoisted) - -simplifyKernelBodyM :: Engine.SimplifiableLore lore => - KernelBody lore - -> Engine.SimpleM lore (Engine.SimplifiedBody lore [KernelResult]) -simplifyKernelBodyM (KernelBody _ stms res) = - Engine.simplifyStms stms $ do res' <- mapM Engine.simplify res - return ((res', UT.usages $ freeIn res'), mempty) - -simplifyGroupStreamLambda :: Engine.SimplifiableLore lore => - GroupStreamLambda lore - -> SubExp -> SubExp -> [VName] - -> Engine.SimpleM lore (GroupStreamLambda (Wise lore), Stms (Wise lore)) -simplifyGroupStreamLambda lam w max_chunk arrs = do - let GroupStreamLambda block_size block_offset acc_params arr_params body = lam - bound_here = S.fromList $ block_size : block_offset : - map paramName (acc_params ++ arr_params) - ((body_stms', body_res'), hoisted) <- - Engine.enterLoop $ - Engine.bindLoopVar block_size Int32 max_chunk $ - Engine.bindLoopVar block_offset Int32 w $ - Engine.bindLParams acc_params $ - Engine.bindChunkLParams block_offset (zip arr_params arrs) $ - Engine.blockIf (Engine.hasFree bound_here `Engine.orIf` Engine.isConsumed) $ - Engine.simplifyBody (replicate (length (bodyResult body)) Observe) body - acc_params' <- mapM (Engine.simplifyParam Engine.simplify) acc_params - arr_params' <- mapM (Engine.simplifyParam Engine.simplify) arr_params - body' <- Engine.constructBody body_stms' body_res' - return (GroupStreamLambda block_size block_offset acc_params' arr_params' body', hoisted) - -instance Engine.Simplifiable KernelSpace where - simplify (KernelSpace gtid ltid gid num_threads num_groups group_size virt_groups structure) = - KernelSpace gtid ltid gid - <$> Engine.simplify num_threads - <*> Engine.simplify num_groups - <*> Engine.simplify group_size - <*> Engine.simplify virt_groups - <*> Engine.simplify structure +instance Engine.Simplifiable SegLevel where + simplify (SegThread num_groups group_size virt) = + SegThread <$> traverse Engine.simplify num_groups <*> + traverse Engine.simplify group_size <*> pure virt + simplify (SegGroup num_groups group_size virt) = + SegGroup <$> traverse Engine.simplify num_groups <*> + traverse Engine.simplify group_size <*> pure virt + simplify (SegThreadScalar num_groups group_size virt) = + SegThreadScalar <$> traverse Engine.simplify num_groups <*> + traverse Engine.simplify group_size <*> pure virt -instance Engine.Simplifiable SpaceStructure where - simplify (FlatThreadSpace dims) = - FlatThreadSpace <$> (zip gtids <$> mapM Engine.simplify gdims) - where (gtids, gdims) = unzip dims - simplify (NestedThreadSpace dims) = - NestedThreadSpace - <$> (zip4 gtids - <$> mapM Engine.simplify gdims - <*> pure ltids - <*> mapM Engine.simplify ldims) - where (gtids, gdims, ltids, ldims) = unzip4 dims +instance Engine.Simplifiable SegSpace where + simplify (SegSpace phys dims) = + SegSpace phys <$> mapM (traverse Engine.simplify) dims instance Engine.Simplifiable KernelResult where - simplify (GroupsReturn what) = - GroupsReturn <$> Engine.simplify what - simplify (ThreadsReturn what) = - ThreadsReturn <$> Engine.simplify what - simplify (WriteReturn ws a res) = - WriteReturn <$> Engine.simplify ws <*> Engine.simplify a <*> Engine.simplify res - simplify (ConcatReturns o w pte moffset what) = + simplify (Returns what) = + Returns <$> Engine.simplify what + simplify (WriteReturns ws a res) = + WriteReturns <$> Engine.simplify ws <*> Engine.simplify a <*> Engine.simplify res + simplify (ConcatReturns o w pte what) = ConcatReturns <$> Engine.simplify o <*> Engine.simplify w <*> Engine.simplify pte - <*> Engine.simplify moffset <*> Engine.simplify what + simplify (TileReturns dims what) = + TileReturns <$> Engine.simplify dims <*> Engine.simplify what instance BinderOps (Wise Kernels) where mkExpAttrB = bindableMkExpAttrB mkBodyB = bindableMkBodyB mkLetNamesB = bindableMkLetNamesB -instance BinderOps (Wise InKernel) where - mkExpAttrB = bindableMkExpAttrB - mkBodyB = bindableMkBodyB - mkLetNamesB = bindableMkLetNamesB - kernelRules :: RuleBook (Wise Kernels) kernelRules = standardRules <> ruleBook [ RuleOp removeInvariantKernelResults - , RuleOp mergeSegRedOps] + , RuleOp mergeSegRedOps + , RuleOp redomapIotaToLoop ] [ RuleOp distributeKernelResults , RuleBasicOp removeUnnecessaryCopy] -fuseStreamIota :: TopDownRuleOp (Wise InKernel) -fuseStreamIota vtable pat _ (GroupStream w max_chunk lam accs arrs) - | ([(iota_cs, iota_param, iota_start, iota_stride, iota_t)], params_and_arrs) <- - partitionEithers $ zipWith (isIota vtable) (groupStreamArrParams lam) arrs = do - - let (arr_params', arrs') = unzip params_and_arrs - chunk_size = groupStreamChunkSize lam - offset = groupStreamChunkOffset lam - - body' <- insertStmsM $ inScopeOf lam $ certifying iota_cs $ do - -- Convert index to appropriate type. - offset' <- asIntS iota_t $ Var offset - offset'' <- letSubExp "offset_by_stride" $ - BasicOp $ BinOp (Mul iota_t) offset' iota_stride - start <- letSubExp "iota_start" $ - BasicOp $ BinOp (Add iota_t) offset'' iota_start - letBindNames_ [paramName iota_param] $ - BasicOp $ Iota (Var chunk_size) start iota_stride iota_t - return $ groupStreamLambdaBody lam - let lam' = lam { groupStreamArrParams = arr_params', - groupStreamLambdaBody = body' - } - letBind_ pat $ Op $ GroupStream w max_chunk lam' accs arrs' -fuseStreamIota _ _ _ _ = cannotSimplify - -isIota :: ST.SymbolTable lore -> a -> VName - -> Either (Certificates, a, SubExp, SubExp, IntType) (a, VName) -isIota vtable chunk arr - | Just (BasicOp (Iota _ x s it), cs) <- ST.lookupExp arr vtable = - Left (cs, chunk, x, s, it) - | otherwise = - Right (chunk, arr) - -- If a kernel produces something invariant to the kernel, turn it -- into a replicate. removeInvariantKernelResults :: TopDownRuleOp (Wise Kernels) removeInvariantKernelResults vtable (Pattern [] kpes) attr - (HostOp (Kernel desc space ts (KernelBody _ kstms kres))) = do + (SegOp (SegMap lvl space ts (KernelBody _ kstms kres))) = Simplify $ do + + case lvl of + SegThreadScalar{} -> cannotSimplify + _ -> return () + (ts', kpes', kres') <- unzip3 <$> filterM checkForInvarianceResult (zip3 ts kpes kres) @@ -448,55 +243,52 @@ when (kres == kres') cannotSimplify - addStm $ Let (Pattern [] kpes') attr $ Op $ HostOp $ Kernel desc space ts' $ + addStm $ Let (Pattern [] kpes') attr $ Op $ SegOp $ SegMap lvl space ts' $ mkWiseKernelBody () kstms kres' where isInvariant Constant{} = True isInvariant (Var v) = isJust $ ST.lookup v vtable - space_dims = map snd $ spaceDimensions space - - checkForInvarianceResult (_, pe, ThreadsReturn se) + checkForInvarianceResult (_, pe, Returns se) | isInvariant se = do - let rep a d = BasicOp . Replicate (Shape [d]) <$> letSubExp "rep" a - letBindNames_ [patElemName pe] =<< - foldM rep (BasicOp (SubExp se)) (reverse space_dims) + letBindNames_ [patElemName pe] $ + BasicOp $ Replicate (Shape $ segSpaceDims space) se return False checkForInvarianceResult _ = return True -removeInvariantKernelResults _ _ _ _ = cannotSimplify +removeInvariantKernelResults _ _ _ _ = Skip -- Some kernel results can be moved outside the kernel, which can -- simplify further analysis. distributeKernelResults :: BottomUpRuleOp (Wise Kernels) distributeKernelResults (vtable, used) - (Pattern [] kpes) attr (HostOp (Kernel desc kspace kts (KernelBody _ kstms kres))) = do + (Pattern [] kpes) attr (SegOp (SegMap lvl space kts (KernelBody _ kstms kres))) = Simplify $ do -- Iterate through the bindings. For each, we check whether it is -- in kres and can be moved outside. If so, we remove it from kres -- and kpes and make it a binding outside. - (kpes', kts', kres', kstms_rev) <- localScope (scopeOfKernelSpace kspace) $ + (kpes', kts', kres', kstms_rev) <- localScope (scopeOfSegSpace space) $ foldM distribute (kpes, kts, kres, []) kstms when (kpes' == kpes) cannotSimplify - addStm $ Let (Pattern [] kpes') attr $ Op $ HostOp $ - Kernel desc kspace kts' $ mkWiseKernelBody () (stmsFromList $ reverse kstms_rev) kres' + addStm $ Let (Pattern [] kpes') attr $ Op $ SegOp $ + SegMap lvl space kts' $ mkWiseKernelBody () (stmsFromList $ reverse kstms_rev) kres' where free_in_kstms = fold $ fmap freeIn kstms distribute (kpes', kts', kres', kstms_rev) bnd | Let (Pattern [] [pe]) _ (BasicOp (Index arr slice)) <- bnd, - kspace_slice <- map (DimFix . Var . fst) $ spaceDimensions kspace, - kspace_slice `isPrefixOf` slice, - remaining_slice <- drop (length kspace_slice) slice, - all (isJust . flip ST.lookup vtable) $ S.toList $ + space_slice <- map (DimFix . Var . fst) $ unSegSpace space, + space_slice `isPrefixOf` slice, + remaining_slice <- drop (length space_slice) slice, + all (isJust . flip ST.lookup vtable) $ namesToList $ freeIn arr <> freeIn remaining_slice, Just (kpe, kpes'', kts'', kres'') <- isResult kpes' kts' kres' pe = do - let outer_slice = map (\(_, d) -> DimSlice - (constant (0::Int32)) - d - (constant (1::Int32))) $ - spaceDimensions kspace + let outer_slice = map (\d -> DimSlice + (constant (0::Int32)) + d + (constant (1::Int32))) $ + segSpaceDims space index kpe' = letBind_ (Pattern [] [kpe']) $ BasicOp $ Index arr $ outer_slice <> remaining_slice if patElemName kpe `UT.isConsumed` used @@ -505,7 +297,7 @@ letBind_ (Pattern [] [kpe]) $ BasicOp $ Copy precopy else index kpe return (kpes'', kts'', kres'', - if patElemName pe `S.member` free_in_kstms + if patElemName pe `nameIn` free_in_kstms then bnd : kstms_rev else kstms_rev) @@ -518,32 +310,32 @@ | (kpes'', kts'', kres'') <- unzip3 kpes_and_kres -> Just (kpe, kpes'', kts'', kres'') _ -> Nothing - where matches (_, _, kre) = kre == ThreadsReturn (Var $ patElemName pe) -distributeKernelResults _ _ _ _ = cannotSimplify + where matches (_, _, kre) = kre == Returns (Var $ patElemName pe) +distributeKernelResults _ _ _ _ = Skip -- If a SegRed contains two reduction operations that have the same -- vector shape, merge them together. This saves on communication -- overhead, but can in principle lead to more local memory usage. mergeSegRedOps :: TopDownRuleOp (Wise Kernels) -mergeSegRedOps _ (Pattern [] pes) _ (HostOp (SegRed space ops ts kbody)) +mergeSegRedOps _ (Pattern [] pes) _ (SegOp (SegRed lvl space ops ts kbody)) | length ops > 1, op_groupings <- groupBy sameShape $ zip ops $ chunks (map (length . segRedNeutral) ops) $ zip3 red_pes red_ts red_res, - any ((>1) . length) op_groupings = do + any ((>1) . length) op_groupings = Simplify $ do let (ops', aux) = unzip $ mapMaybe combineOps op_groupings (red_pes', red_ts', red_res') = unzip3 $ concat aux pes' = red_pes' ++ map_pes ts' = red_ts' ++ map_ts kbody' = kbody { kernelBodyResult = red_res' ++ map_res } - letBind_ (Pattern [] pes') $ Op $ HostOp $ SegRed space ops' ts' kbody' + letBind_ (Pattern [] pes') $ Op $ SegOp $ SegRed lvl space ops' ts' kbody' where (red_pes, map_pes) = splitAt (segRedResults ops) pes (red_ts, map_ts) = splitAt (segRedResults ops) ts (red_res, map_res) = splitAt (segRedResults ops) $ kernelBodyResult kbody sameShape (op1, _) (op2, _) = segRedShape op1 == segRedShape op2 - combineOps :: [(SegRedOp (Wise InKernel), [a])] - -> Maybe (SegRedOp (Wise InKernel), [a]) + combineOps :: [(SegRedOp (Wise Kernels), [a])] + -> Maybe (SegRedOp (Wise Kernels), [a]) combineOps [] = Nothing combineOps (x:xs) = Just $ foldl' combine x xs @@ -567,46 +359,16 @@ , segRedShape = segRedShape op1 -- Same as shape of op2 due to the grouping. }, op1_aux ++ op2_aux) -mergeSegRedOps _ _ _ _ = cannotSimplify - -simplifyKnownIterationStream :: TopDownRuleOp (Wise InKernel) --- Remove GroupStreams over single-element arrays. Not much to stream --- here, and no information to exploit. -simplifyKnownIterationStream _ pat _ (GroupStream (Constant v) _ lam accs arrs) - | oneIsh v = do - let GroupStreamLambda chunk_size chunk_offset acc_params arr_params body = lam - - letBindNames_ [chunk_size] $ BasicOp $ SubExp $ constant (1::Int32) - - letBindNames_ [chunk_offset] $ BasicOp $ SubExp $ constant (0::Int32) - - forM_ (zip acc_params accs) $ \(p,a) -> - letBindNames_ [paramName p] $ BasicOp $ SubExp a - - forM_ (zip arr_params arrs) $ \(p,a) -> - letBindNames_ [paramName p] $ BasicOp $ Index a $ - fullSlice (paramType p) - [DimSlice (Var chunk_offset) (Var chunk_size) (constant (1::Int32))] - - res <- bodyBind body - forM_ (zip (patternElements pat) res) $ \(pe,r) -> - letBindNames_ [patElemName pe] $ BasicOp $ SubExp r -simplifyKnownIterationStream _ _ _ _ = cannotSimplify - -removeUnusedStreamInputs :: TopDownRuleOp (Wise InKernel) -removeUnusedStreamInputs _ pat _ (GroupStream w maxchunk lam accs arrs) - | (used,unused) <- partition (isUsed . paramName . fst) $ zip arr_params arrs, - not $ null unused = do - let (arr_params', arrs') = unzip used - lam' = GroupStreamLambda chunk_size chunk_offset acc_params arr_params' body - letBind_ pat $ Op $ GroupStream w maxchunk lam' accs arrs' - where GroupStreamLambda chunk_size chunk_offset acc_params arr_params body = lam - - isUsed = (`S.member` freeIn body) -removeUnusedStreamInputs _ _ _ _ = cannotSimplify +mergeSegRedOps _ _ _ _ = Skip -inKernelRules :: RuleBook (Wise InKernel) -inKernelRules = standardRules <> - ruleBook [RuleOp fuseStreamIota, - RuleOp simplifyKnownIterationStream, - RuleOp removeUnusedStreamInputs] [] +-- We turn reductions over (solely) iotas into do-loops, because there +-- is no useful structure here anyway. This is mostly a hack to work +-- around the fact that loop tiling would otherwise pointlessly tile +-- them. +redomapIotaToLoop :: TopDownRuleOp (Wise Kernels) +redomapIotaToLoop vtable pat aux (OtherOp soac@(Screma _ form [arr])) + | Just _ <- isRedomapSOAC form, + Just (Iota{}, _) <- ST.lookupBasicOp arr vtable = + Simplify $ certifying (stmAuxCerts aux) $ FOT.transformSOAC pat soac +redomapIotaToLoop _ _ _ _ = + Skip diff --git a/src/Futhark/Representation/Kernels/Sizes.hs b/src/Futhark/Representation/Kernels/Sizes.hs --- a/src/Futhark/Representation/Kernels/Sizes.hs +++ b/src/Futhark/Representation/Kernels/Sizes.hs @@ -1,10 +1,18 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} module Futhark.Representation.Kernels.Sizes - ( SizeClass (..), KernelPath ) + ( SizeClass (..) + , KernelPath + , Count(..) + , NumGroups, GroupSize, NumThreads + ) where +import Data.Traversable + import Futhark.Util.Pretty import Language.Futhark.Core (Name) -import Futhark.Representation.AST.Pretty () +import Futhark.Util.IntegralExp (IntegralExp) +import Futhark.Representation.AST.Attributes.Names (FreeIn) -- | An indication of which comparisons have been performed to get to -- this point, as well as the result of each comparison. @@ -29,3 +37,25 @@ ppr SizeNumGroups = text "num_groups" ppr SizeTile = text "tile_size" ppr SizeLocalMemory = text "local_memory" + +-- | A wrapper supporting a phantom type for indicating what we are counting. +newtype Count u e = Count { unCount :: e } + deriving (Eq, Ord, Show, Num, IntegralExp, FreeIn, Pretty) + +instance Functor (Count u) where + fmap = fmapDefault + +instance Foldable (Count u) where + foldMap = foldMapDefault + +instance Traversable (Count u) where + traverse f (Count x) = Count <$> f x + +-- | Phantom type for the number of groups of some kernel. +data NumGroups + +-- | Phantom type for the group size of some kernel. +data GroupSize + +-- | Phantom type for number of threads. +data NumThreads diff --git a/src/Futhark/Representation/Ranges.hs b/src/Futhark/Representation/Ranges.hs --- a/src/Futhark/Representation/Ranges.hs +++ b/src/Futhark/Representation/Ranges.hs @@ -34,7 +34,6 @@ import Control.Monad.Identity import Control.Monad.Reader -import qualified Data.Set as S import Data.Monoid ((<>)) import Data.Foldable @@ -167,18 +166,15 @@ mkBodyRanges :: Stms lore -> Result -> [Range] mkBodyRanges bnds = map $ removeUnknownBounds . rangeOf where boundInBnds = - fold $ fmap (S.fromList . patternNames . stmPattern) bnds + fold $ fmap (namesFromList . patternNames . stmPattern) bnds removeUnknownBounds (lower,upper) = (removeUnknownBound lower, removeUnknownBound upper) removeUnknownBound (Just bound) - | freeIn bound `intersects` boundInBnds = Nothing - | otherwise = Just bound + | freeIn bound `namesIntersect` boundInBnds = Nothing + | otherwise = Just bound removeUnknownBound Nothing = Nothing - -intersects :: Ord a => S.Set a -> S.Set a -> Bool -intersects a b = not $ S.null $ a `S.intersection` b mkRangedLetStm :: (Attributes lore, CanBeRanged (Op lore)) => Pattern lore diff --git a/src/Futhark/Representation/SOACS.hs b/src/Futhark/Representation/SOACS.hs --- a/src/Futhark/Representation/SOACS.hs +++ b/src/Futhark/Representation/SOACS.hs @@ -5,14 +5,12 @@ ( -- * The Lore definition SOACS -- * Syntax types - , Prog , Body , Stm , Pattern , BasicOp , Exp , Lambda - , FunDef , FParam , LParam , RetType @@ -27,17 +25,14 @@ , AST.BodyT(Body) , AST.PatternT(Pattern) , AST.PatElemT(PatElem) - , AST.ProgT(Prog) , AST.ExpT(BasicOp) - , AST.FunDefT(FunDef) - , AST.ParamT(Param) ) where import qualified Futhark.Representation.AST.Syntax as AST import Futhark.Representation.AST.Syntax - hiding (Prog, BasicOp, Exp, Body, Stm, - Pattern, Lambda, FunDef, FParam, LParam, RetType, PatElem) + hiding (BasicOp, Exp, Body, Stm, + Pattern, Lambda, FParam, LParam, RetType, PatElem) import Futhark.Representation.SOACS.SOAC import Futhark.Representation.AST.Attributes import Futhark.Representation.AST.Traversals @@ -59,14 +54,12 @@ instance Attributes SOACS where expTypesFromPattern = return . expExtTypesFromPattern -type Prog = AST.Prog SOACS type BasicOp = AST.BasicOp SOACS type Exp = AST.Exp SOACS type Body = AST.Body SOACS type Stm = AST.Stm SOACS type Pattern = AST.Pattern SOACS type Lambda = AST.Lambda SOACS -type FunDef = AST.FunDefT SOACS type FParam = AST.FParam SOACS type LParam = AST.LParam SOACS type RetType = AST.RetType SOACS diff --git a/src/Futhark/Representation/SOACS/SOAC.hs b/src/Futhark/Representation/SOACS/SOAC.hs --- a/src/Futhark/Representation/SOACS/SOAC.hs +++ b/src/Futhark/Representation/SOACS/SOAC.hs @@ -47,10 +47,10 @@ ) where +import Control.Monad.State.Strict import Control.Monad.Writer import Control.Monad.Identity import qualified Data.Map.Strict as M -import qualified Data.Set as S import Data.Maybe import Data.List @@ -318,11 +318,11 @@ mapSOACM tv (CmpThreshold what s) = CmpThreshold <$> mapOnSOACSubExp tv what <*> pure s instance Attributes lore => FreeIn (SOAC lore) where - freeIn = execWriter . mapSOACM free - where walk f x = tell (f x) >> return x - free = SOACMapper { mapOnSOACSubExp = walk freeIn - , mapOnSOACLambda = walk freeIn - , mapOnSOACVName = walk freeIn + freeIn' = flip execState mempty . mapSOACM free + where walk f x = modify (<>f x) >> return x + free = SOACMapper { mapOnSOACSubExp = walk freeIn' + , mapOnSOACLambda = walk freeIn' + , mapOnSOACVName = walk freeIn' } instance Attributes lore => Substitute (SOAC lore) where @@ -368,24 +368,24 @@ -- Only map functions can consume anything. The operands to scan -- and reduce functions are always considered "fresh". consumedInOp (Screma _ (ScremaForm _ _ map_lam) arrs) = - S.map consumedArray $ consumedByLambda map_lam + mapNames consumedArray $ consumedByLambda map_lam where consumedArray v = fromMaybe v $ lookup v params_to_arrs params_to_arrs = zip (map paramName $ lambdaParams map_lam) arrs consumedInOp (Stream _ form lam arrs) = - S.fromList $ subExpVars $ + namesFromList $ subExpVars $ case form of Sequential accs -> - map (consumedArray accs) $ S.toList $ consumedByLambda lam + map (consumedArray accs) $ namesToList $ consumedByLambda lam Parallel _ _ _ accs -> - map (consumedArray accs) $ S.toList $ consumedByLambda lam + map (consumedArray accs) $ namesToList $ consumedByLambda lam where consumedArray accs v = fromMaybe (Var v) $ lookup v $ paramsToInput accs -- Drop the chunk parameter, which cannot alias anything. paramsToInput accs = zip (map paramName $ drop 1 $ lambdaParams lam) (accs++map Var arrs) consumedInOp (Scatter _ _ _ as) = - S.fromList $ map (\(_, _, a) -> a) as + namesFromList $ map (\(_, _, a) -> a) as consumedInOp (GenReduce _ ops _ _) = - S.fromList $ concatMap genReduceDest ops + namesFromList $ concatMap genReduceDest ops consumedInOp CmpThreshold{} = mempty mapGenReduceOp :: (Lambda flore -> Lambda tlore) diff --git a/src/Futhark/Representation/SOACS/Simplify.hs b/src/Futhark/Representation/SOACS/Simplify.hs --- a/src/Futhark/Representation/SOACS/Simplify.hs +++ b/src/Futhark/Representation/SOACS/Simplify.hs @@ -7,9 +7,11 @@ module Futhark.Representation.SOACS.Simplify ( simplifySOACS , simplifyLambda + , simplifyFun , simplifyStms , simpleSOACS + , simplifySOAC , soacRules ) @@ -47,7 +49,7 @@ simpleSOACS :: Simplify.SimpleOps SOACS simpleSOACS = Simplify.bindableSimpleOps simplifySOAC -simplifySOACS :: Prog -> PassM Prog +simplifySOACS :: Prog SOACS -> PassM (Prog SOACS) simplifySOACS = Simplify.simplifyProg simpleSOACS soacRules blockers where blockers = Engine.noExtraHoistBlockers { Engine.getArraySizes = getShapeNames } @@ -58,8 +60,12 @@ getShapeNames bnd = let tps1 = map patElemType $ patternElements $ stmPattern bnd tps2 = map (snd . patElemAttr) $ patternElements $ stmPattern bnd - in S.fromList $ subExpVars $ concatMap arrayDims (tps1 ++ tps2) + in namesFromList $ subExpVars $ concatMap arrayDims (tps1 ++ tps2) +simplifyFun :: MonadFreshNames m => FunDef SOACS -> m (FunDef SOACS) +simplifyFun = + Simplify.simplifyFun simpleSOACS soacRules Engine.noExtraHoistBlockers + simplifyLambda :: (HasScope SOACS m, MonadFreshNames m) => Lambda -> [Maybe VName] -> m Lambda simplifyLambda = @@ -70,7 +76,8 @@ simplifyStms = Simplify.simplifyStms simpleSOACS soacRules Engine.noExtraHoistBlockers -simplifySOAC :: Simplify.SimplifyOp SOACS +simplifySOAC :: Simplify.SimplifiableLore lore => + Simplify.SimplifyOp lore (SOAC lore) simplifySOAC (CmpThreshold what s) = do what' <- Engine.simplify what return (CmpThreshold what' s, mempty) @@ -181,7 +188,7 @@ hoistCertificates vtable pat aux soac | (soac', hoisted) <- runState (mapSOACM mapper soac) mempty, hoisted /= mempty = - certifying (hoisted <> stmAuxCerts aux) $ letBind_ pat $ Op soac' + Simplify $ certifying (hoisted <> stmAuxCerts aux) $ letBind_ pat $ Op soac' where mapper = identitySOACMapper { mapOnSOACLambda = onLambda } onLambda lam = do stms' <- mapM onStm $ bodyStms $ lambdaBody lam @@ -196,7 +203,7 @@ return $ Let se_pat se_aux' $ BasicOp $ SubExp se onStm stm = return stm hoistCertificates _ _ _ _ = - cannotSimplify + Skip liftIdentityMapping :: BottomUpRuleOp (Wise SOACS) liftIdentityMapping (_, usages) pat _ (Screma w form arrs) @@ -206,7 +213,7 @@ rettype = lambdaReturnType fun ses = bodyResult $ lambdaBody fun - freeOrConst (Var v) = v `S.member` free + freeOrConst (Var v) = v `nameIn` free freeOrConst Constant{} = True checkInvariance (outId, Var v, _) (invariant, mapresult, rettype') @@ -229,21 +236,21 @@ case foldr checkInvariance ([], [], []) $ zip3 (patternElements pat) ses rettype of - ([], _, _) -> cannotSimplify - (invariant, mapresult, rettype') -> do + ([], _, _) -> Skip + (invariant, mapresult, rettype') -> Simplify $ do let (pat', ses') = unzip mapresult fun' = fun { lambdaBody = (lambdaBody fun) { bodyResult = ses' } , lambdaReturnType = rettype' } mapM_ (uncurry letBind) invariant letBindNames_ (map patElemName pat') $ Op $ Screma w (mapSOAC fun') arrs -liftIdentityMapping _ _ _ _ = cannotSimplify +liftIdentityMapping _ _ _ _ = Skip liftIdentityStreaming :: BottomUpRuleOp (Wise SOACS) liftIdentityStreaming _ (Pattern [] pes) _ (Stream w form lam arrs) | (variant_map, invariant_map) <- partitionEithers $ map isInvariantRes $ zip3 map_ts map_pes map_res, - not $ null invariant_map = do + not $ null invariant_map = Simplify $ do forM_ invariant_map $ \(pe, arr) -> letBind_ (Pattern [] [pe]) $ BasicOp $ Copy arr @@ -266,25 +273,25 @@ Right (pe, arr) isInvariantRes x = Left x -liftIdentityStreaming _ _ _ _ = cannotSimplify +liftIdentityStreaming _ _ _ _ = Skip -- | Remove all arguments to the map that are simply replicates. -- These can be turned into free variables instead. removeReplicateMapping :: TopDownRuleOp (Wise SOACS) removeReplicateMapping vtable pat _ (Screma w form arrs) | Just fun <- isMapSOAC form, - Just (bnds, fun', arrs') <- removeReplicateInput vtable fun arrs = do + Just (bnds, fun', arrs') <- removeReplicateInput vtable fun arrs = Simplify $ do forM_ bnds $ \(vs,cs,e) -> certifying cs $ letBindNames vs e letBind_ pat $ Op $ Screma w (mapSOAC fun') arrs' -removeReplicateMapping _ _ _ _ = cannotSimplify +removeReplicateMapping _ _ _ _ = Skip -- | Like 'removeReplicateMapping', but for 'Scatter'. removeReplicateWrite :: TopDownRuleOp (Wise SOACS) removeReplicateWrite vtable pat _ (Scatter len lam ivs as) - | Just (bnds, lam', ivs') <- removeReplicateInput vtable lam ivs = do + | Just (bnds, lam', ivs') <- removeReplicateInput vtable lam ivs = Simplify $ do forM_ bnds $ \(vs,cs,e) -> certifying cs $ letBindNames vs e letBind_ pat $ Op $ Scatter len lam' ivs' as -removeReplicateWrite _ _ _ _ = cannotSimplify +removeReplicateWrite _ _ _ _ = Skip removeReplicateInput :: Aliased lore => ST.SymbolTable lore @@ -307,7 +314,7 @@ isReplicateAndNotConsumed p v | Just (BasicOp (Replicate (Shape (_:ds)) e), v_cs) <- ST.lookupExp v vtable, - not $ paramName p `S.member` consumedByLambda fun = + not $ paramName p `nameIn` consumedByLambda fun = Right ([paramName p], v_cs, case ds of @@ -320,14 +327,14 @@ removeUnusedSOACInput :: TopDownRuleOp (Wise SOACS) removeUnusedSOACInput _ pat _ (Screma w (ScremaForm scan reduce map_lam) arrs) | (used,unused) <- partition usedInput params_and_arrs, - not (null unused) = do + not (null unused) = Simplify $ do let (used_params, used_arrs) = unzip used map_lam' = map_lam { lambdaParams = used_params } letBind_ pat $ Op $ Screma w (ScremaForm scan reduce map_lam') used_arrs where params_and_arrs = zip (lambdaParams map_lam) arrs used_in_body = freeIn $ lambdaBody map_lam - usedInput (param, _) = paramName param `S.member` used_in_body -removeUnusedSOACInput _ _ _ _ = cannotSimplify + usedInput (param, _) = paramName param `nameIn` used_in_body +removeUnusedSOACInput _ _ _ _ = Skip removeDeadMapping :: BottomUpRuleOp (Wise SOACS) removeDeadMapping (_, used) pat _ (Screma w form arrs) @@ -340,9 +347,9 @@ , lambdaReturnType = ts' } in if pat /= Pattern [] pat' - then letBind_ (Pattern [] pat') $ Op $ Screma w (mapSOAC fun') arrs - else cannotSimplify -removeDeadMapping _ _ _ _ = cannotSimplify + then Simplify $ letBind_ (Pattern [] pat') $ Op $ Screma w (mapSOAC fun') arrs + else Skip +removeDeadMapping _ _ _ _ = Skip removeDuplicateMapOutput :: BottomUpRuleOp (Wise SOACS) removeDuplicateMapOutput (_, used) pat _ (Screma w form arrs) @@ -353,8 +360,8 @@ ses_ts_pes = zip3 ses ts pes (ses_ts_pes', copies) = foldl checkForDuplicates (mempty,mempty) ses_ts_pes - in if null copies then cannotSimplify - else do + in if null copies then Skip + else Simplify $ do let (ses', ts', pes') = unzip3 ses_ts_pes' pat' = Pattern [] pes' fun' = fun { lambdaBody = (lambdaBody fun) { bodyResult = ses' } @@ -371,7 +378,7 @@ -- array pe'. (ses_ts_pes', (pe', pe) : copies) | otherwise = (ses_ts_pes' ++ [(se,t,pe)], copies) -removeDuplicateMapOutput _ _ _ _ = cannotSimplify +removeDuplicateMapOutput _ _ _ _ = Skip -- Mapping some operations becomes an extension of that operation. mapOpToOp :: BottomUpRuleOp (Wise SOACS) @@ -380,7 +387,7 @@ | Just (map_pe, cs, w, BasicOp (Reshape newshape reshape_arr), [p], [arr]) <- isMapWithOp pat e, paramName p == reshape_arr, - not $ UT.isConsumed (patElemName map_pe) used = do + not $ UT.isConsumed (patElemName map_pe) used = Simplify $ do let redim | isJust $ shapeCoercion newshape = DimCoercion w | otherwise = DimNew w certifying (stmAuxCerts aux1 <> cs) $ letBind_ pat $ @@ -390,7 +397,7 @@ BasicOp (Concat d arr arrs dw), ps, outer_arr : outer_arrs) <- isMapWithOp pat e, (arr:arrs) == map paramName ps = - certifying (stmAuxCerts aux1 <> cs) $ letBind_ pat $ + Simplify $ certifying (stmAuxCerts aux1 <> cs) $ letBind_ pat $ BasicOp $ Concat (d+1) outer_arr outer_arrs dw | Just (map_pe, cs, _, @@ -398,22 +405,22 @@ isMapWithOp pat e, paramName p == rearrange_arr, not $ UT.isConsumed (patElemName map_pe) used = - certifying (stmAuxCerts aux1 <> cs) $ letBind_ pat $ + Simplify $ certifying (stmAuxCerts aux1 <> cs) $ letBind_ pat $ BasicOp $ Rearrange (0 : map (1+) perm) arr | Just (map_pe, cs, _, BasicOp (Rotate rots rotate_arr), [p], [arr]) <- isMapWithOp pat e, paramName p == rotate_arr, not $ UT.isConsumed (patElemName map_pe) used = - certifying (stmAuxCerts aux1 <> cs) $ letBind_ pat $ + Simplify $ certifying (stmAuxCerts aux1 <> cs) $ letBind_ pat $ BasicOp $ Rotate (intConst Int32 0 : rots) arr -mapOpToOp _ _ _ _ = cannotSimplify +mapOpToOp _ _ _ _ = Skip isMapWithOp :: PatternT attr -> SOAC (Wise SOACS) -> Maybe (PatElemT attr, Certificates, SubExp, - AST.Exp (Wise SOACS), [ParamT Type], [VName]) + AST.Exp (Wise SOACS), [Param Type], [VName]) isMapWithOp pat e | Pattern [] [map_pe] <- pat, Screma w form arrs <- e, @@ -442,14 +449,14 @@ zip red_pes redlam_params, let necessary = findNecessaryForReturned (`elem` used_after) (zip redlam_params $ redlam_res <> redlam_res) redlam_deps, - let alive_mask = map ((`S.member` necessary) . paramName) redlam_params, + let alive_mask = map ((`nameIn` necessary) . paramName) redlam_params, - not $ all (==True) alive_mask = do + not $ all (==True) alive_mask = Simplify $ do let fixDeadToNeutral lives ne = if lives then Nothing else Just ne dead_fix = zipWith fixDeadToNeutral alive_mask nes (used_red_pes, _, used_nes) = - unzip3 $ filter (\(_,x,_) -> paramName x `S.member` necessary) $ + unzip3 $ filter (\(_,x,_) -> paramName x `nameIn` necessary) $ zip3 red_pes redlam_params nes let maplam' = removeLambdaResults (take (length nes) alive_mask) maplam @@ -458,7 +465,7 @@ certifying cs $ letBind_ (Pattern [] $ used_red_pes ++ map_pes) $ Op $ Screma w (redomapSOAC [Reduce comm redlam' used_nes] maplam') arrs -removeDeadReduction _ _ _ _ = cannotSimplify +removeDeadReduction _ _ _ _ = Skip -- | If we are writing to an array that is never used, get rid of it. removeDeadWrite :: BottomUpRuleOp (Wise SOACS) @@ -473,16 +480,16 @@ , lambdaReturnType = i_ts' ++ v_ts' } in if pat /= Pattern [] pat' - then letBind_ (Pattern [] pat') $ Op $ Scatter w fun' arrs dests' - else cannotSimplify -removeDeadWrite _ _ _ _ = cannotSimplify + then Simplify $ letBind_ (Pattern [] pat') $ Op $ Scatter w fun' arrs dests' + else Skip +removeDeadWrite _ _ _ _ = Skip -- handles now concatenation of more than two arrays fuseConcatScatter :: TopDownRuleOp (Wise SOACS) fuseConcatScatter vtable pat _ (Scatter _ fun arrs dests) | Just (ws@(w':_), xss, css) <- unzip3 <$> mapM isConcat arrs, xivs <- transpose xss, - all (w'==) ws = do + all (w'==) ws = Simplify $ do let r = length xivs fun2s <- mapM (\_ -> renameLambda fun) [1 .. r-1] let fun_n = length $ lambdaReturnType fun @@ -515,25 +522,25 @@ return (a, b, cs <> cs') _ -> Nothing -fuseConcatScatter _ _ _ _ = cannotSimplify +fuseConcatScatter _ _ _ _ = Skip simplifyClosedFormReduce :: TopDownRuleOp (Wise SOACS) simplifyClosedFormReduce _ pat _ (Screma (Constant w) form _) | Just nes <- concatMap redNeutral . fst <$> isRedomapSOAC form, zeroIsh w = - forM_ (zip (patternNames pat) nes) $ \(v, ne) -> + Simplify $ forM_ (zip (patternNames pat) nes) $ \(v, ne) -> letBindNames_ [v] $ BasicOp $ SubExp ne simplifyClosedFormReduce vtable pat _ (Screma _ form arrs) | Just [Reduce _ red_fun nes] <- isReduceSOAC form = - foldClosedForm (`ST.lookupExp` vtable) pat red_fun nes arrs -simplifyClosedFormReduce _ _ _ _ = cannotSimplify + Simplify $ foldClosedForm (`ST.lookupExp` vtable) pat red_fun nes arrs +simplifyClosedFormReduce _ _ _ _ = Skip -- For now we just remove singleton SOACs. simplifyKnownIterationSOAC :: TopDownRuleOp (Wise SOACS) simplifyKnownIterationSOAC _ pat _ (Screma (Constant k) (ScremaForm (scan_lam, scan_nes) reds map_lam) arrs) - | oneIsh k = do + | oneIsh k = Simplify $ do zipWithM_ bindMapParam (lambdaParams map_lam) arrs (to_scan, to_red, map_res) <- splitAt3 (length scan_nes) (length red_nes) <$> bodyBind (lambdaBody map_lam) @@ -556,21 +563,24 @@ BasicOp $ ArrayLit [se] $ rowType $ patElemType pe bindResult pe se = letBindNames_ [patElemName pe] $ BasicOp $ SubExp se -simplifyKnownIterationSOAC _ _ _ _ = cannotSimplify +simplifyKnownIterationSOAC _ _ _ _ = Skip data ArrayOp = ArrayIndexing Certificates VName (Slice SubExp) | ArrayRearrange Certificates VName [Int] + | ArrayRotate Certificates VName [SubExp] | ArrayVar Certificates VName -- ^ Never constructed. deriving (Eq, Ord, Show) arrayOpArr :: ArrayOp -> VName arrayOpArr (ArrayIndexing _ arr _) = arr arrayOpArr (ArrayRearrange _ arr _) = arr +arrayOpArr (ArrayRotate _ arr _) = arr arrayOpArr (ArrayVar _ arr) = arr arrayOpCerts :: ArrayOp -> Certificates arrayOpCerts (ArrayIndexing cs _ _) = cs arrayOpCerts (ArrayRearrange cs _ _) = cs +arrayOpCerts (ArrayRotate cs _ _) = cs arrayOpCerts (ArrayVar cs _) = cs isArrayOp :: Certificates -> AST.Exp (Wise SOACS) -> Maybe ArrayOp @@ -578,12 +588,15 @@ Just $ ArrayIndexing cs arr slice isArrayOp cs (BasicOp (Rearrange perm arr)) = Just $ ArrayRearrange cs arr perm +isArrayOp cs (BasicOp (Rotate rots arr)) = + Just $ ArrayRotate cs arr rots isArrayOp _ _ = Nothing fromArrayOp :: ArrayOp -> (Certificates, AST.Exp (Wise SOACS)) fromArrayOp (ArrayIndexing cs arr slice) = (cs, BasicOp $ Index arr slice) fromArrayOp (ArrayRearrange cs arr perm) = (cs, BasicOp $ Rearrange perm arr) +fromArrayOp (ArrayRotate cs arr rots) = (cs, BasicOp $ Rotate rots arr) fromArrayOp (ArrayVar cs arr) = (cs, BasicOp $ SubExp $ Var arr) arrayOps :: AST.Body (Wise SOACS) -> S.Set ArrayOp @@ -591,12 +604,12 @@ where onStm (Let _ aux e) = case isArrayOp (stmAuxCerts aux) e of Just op -> S.singleton op - Nothing -> execWriter $ walkExpM walker e + Nothing -> execState (walkExpM walker e) mempty onOp = execWriter . mapSOACM identitySOACMapper { mapOnSOACLambda = onLambda } onLambda lam = do tell $ arrayOps $ lambdaBody lam return lam - walker = identityWalker { walkOnBody = tell . arrayOps - , walkOnOp = tell . onOp } + walker = identityWalker { walkOnBody = modify . (<>) . arrayOps + , walkOnOp = modify . (<>) . onOp } replaceArrayOps :: M.Map ArrayOp ArrayOp -> AST.Body (Wise SOACS) -> AST.Body (Wise SOACS) @@ -635,7 +648,7 @@ | Just (p, _) <- find isIota (zip (lambdaParams map_lam) arrs), indexings <- filter (indexesWith (paramName p)) $ S.toList $ arrayOps $ lambdaBody map_lam, - not $ null indexings = do + not $ null indexings = Simplify $ do -- For each indexing with iota, add the corresponding array to -- the Screma, and construct a new lambda parameter. (more_arrs, more_params, replacements) <- @@ -671,7 +684,7 @@ mapOverArr _ = return Nothing -simplifyMapIota _ _ _ _ = cannotSimplify +simplifyMapIota _ _ _ _ = Skip -- If a Screma's map function contains a transformation -- (e.g. transpose) on a parameter, create a new parameter @@ -680,7 +693,7 @@ moveTransformToInput :: TopDownRuleOp (Wise SOACS) moveTransformToInput vtable pat _ (Screma w (ScremaForm scan reduce map_lam) arrs) | ops <- filter arrayIsMapParam $ S.toList $ arrayOps $ lambdaBody map_lam, - not $ null ops = do + not $ null ops = Simplify $ do (more_arrs, more_params, replacements) <- unzip3 . catMaybes <$> mapM mapOverArr ops @@ -700,12 +713,15 @@ -- everything else must be map-invariant. arrayIsMapParam (ArrayIndexing cs arr slice) = arr `elem` map_param_names && - all (`ST.elem` vtable) (S.toList $ freeIn cs <> freeIn slice) && + all (`ST.elem` vtable) (namesToList $ freeIn cs <> freeIn slice) && not (null slice) && not (null $ sliceDims slice) arrayIsMapParam (ArrayRearrange cs arr perm) = arr `elem` map_param_names && - all (`ST.elem` vtable) (S.toList $ freeIn cs) && + all (`ST.elem` vtable) (namesToList $ freeIn cs) && not (null perm) + arrayIsMapParam (ArrayRotate cs arr rots) = + arr `elem` map_param_names && + all (`ST.elem` vtable) (namesToList $ freeIn cs <> freeIn rots) arrayIsMapParam ArrayVar{} = False @@ -720,6 +736,8 @@ BasicOp $ Index arr $ whole_dim : slice ArrayRearrange _ _ perm -> BasicOp $ Rearrange (0 : map (+1) perm) arr + ArrayRotate _ _ rots -> + BasicOp $ Rotate (intConst Int32 0 : rots) arr ArrayVar{} -> BasicOp $ SubExp $ Var arr arr_transformed_t <- lookupType arr_transformed @@ -731,4 +749,4 @@ mapOverArr _ = return Nothing moveTransformToInput _ _ _ _ = - cannotSimplify + Skip diff --git a/src/Futhark/Test.hs b/src/Futhark/Test.hs --- a/src/Futhark/Test.hs +++ b/src/Futhark/Test.hs @@ -178,14 +178,14 @@ lexeme :: Parser a -> Parser a lexeme p = p <* space --- | Like 'lexeme', but does not consume trailing linebreaks. +-- Like 'lexeme', but does not consume trailing linebreaks. lexeme' :: Parser a -> Parser a lexeme' p = p <* many (oneOf (" \t" :: String)) lexstr :: T.Text -> Parser () lexstr = void . try . lexeme . string --- | Like 'lexstr', but does not consume trailing linebreaks. +-- Like 'lexstr', but does not consume trailing linebreaks. lexstr' :: T.Text -> Parser () lexstr' = void . try . lexeme' . string @@ -451,7 +451,7 @@ -- parameter is used for error messages. valuesFromByteString :: String -> BS.ByteString -> Either String [Value] valuesFromByteString srcname = - maybe (Left $ "Cannot parse values from " ++ srcname) Right . readValues + maybe (Left $ "Cannot parse values from '" ++ srcname ++ "'") Right . readValues -- | Get the actual core Futhark values corresponding to a 'Values' -- specification. The 'FilePath' is the directory which file paths @@ -461,9 +461,12 @@ return vs getValues dir v = do s <- getValuesBS dir v - case valuesFromByteString "file" s of + case valuesFromByteString file s of Left e -> fail $ show e Right vs -> return vs + where file = case v of Values{} -> "" + InFile f -> f + GenValues{} -> "" -- | Extract a pretty representation of some 'Values'. In the IO -- monad because this might involve reading from a file. There is no diff --git a/src/Futhark/Test/Values.hs b/src/Futhark/Test/Values.hs --- a/src/Futhark/Test/Values.hs +++ b/src/Futhark/Test/Values.hs @@ -7,6 +7,7 @@ -- for your test programs. module Futhark.Test.Values ( Value(..) + , Vector -- * Reading Values , readValues @@ -48,6 +49,8 @@ import Futhark.Util (maybeHead) type STVector s = UMVec.STVector s + +-- | An Unboxed vector. type Vector = UVec.Vector -- | An efficiently represented Futhark value. Use 'pretty' to get a diff --git a/src/Futhark/Transform/FirstOrderTransform.hs b/src/Futhark/Transform/FirstOrderTransform.hs --- a/src/Futhark/Transform/FirstOrderTransform.hs +++ b/src/Futhark/Transform/FirstOrderTransform.hs @@ -14,31 +14,25 @@ , transformLambda , transformSOAC , transformBody - - -- * Utility - , doLoopMapAccumL - , doLoopMapAccumL' ) where import Control.Monad.Except import Control.Monad.State import qualified Data.Map.Strict as M -import qualified Data.Set as S import Data.List (zip4) import qualified Futhark.Representation.AST as AST import Futhark.Representation.SOACS import Futhark.MonadFreshNames import Futhark.Tools -import Futhark.Representation.Aliases (Aliases, removeLambdaAliases) import Futhark.Representation.AST.Attributes.Aliases import Futhark.Util (chunks, splitAt3) transformFunDef :: (MonadFreshNames m, Bindable tolore, BinderOps tolore, LetAttr SOACS ~ LetAttr tolore, CanBeAliased (Op tolore)) => - FunDef -> m (AST.FunDef tolore) + FunDef SOACS -> m (AST.FunDef tolore) transformFunDef (FunDef entry fname rettype params body) = do (body',_) <- modifyNameSource $ runState $ runBinderT m mempty return $ FunDef entry fname rettype params body' @@ -49,11 +43,10 @@ type Transformer m = (MonadBinder m, Bindable (Lore m), BinderOps (Lore m), LocalScope (Lore m) m, - LetAttr SOACS ~ LetAttr (Lore m), LParamAttr SOACS ~ LParamAttr (Lore m), CanBeAliased (Op (Lore m))) -transformBody :: Transformer m => +transformBody :: (Transformer m, LetAttr (Lore m) ~ LetAttr SOACS) => Body -> m (AST.Body (Lore m)) transformBody (Body () bnds res) = insertStmsM $ do mapM_ transformStmRecursively bnds @@ -61,8 +54,8 @@ -- | First transform any nested 'Body' or 'Lambda' elements, then -- apply 'transformSOAC' if the expression is a SOAC. -transformStmRecursively :: Transformer m => - Stm -> m () +transformStmRecursively :: (Transformer m, LetAttr (Lore m) ~ LetAttr SOACS) => + Stm -> m () transformStmRecursively (Let pat aux (Op soac)) = certifying (stmAuxCerts aux) $ @@ -159,8 +152,9 @@ -- We need to discard the final scan accumulators, as they are not -- bound in the original pattern. - pat' <- discardPattern (map paramType scanacc_params) pat - letBind_ pat' $ DoLoop [] merge loopform loop_body + names <- (++patternNames pat) + <$> replicateM (length scanacc_params) (newVName "discard") + letBindNames_ names $ DoLoop [] merge loopform loop_body transformSOAC pat (Stream w form lam arrs) = sequentialStreamWholeArray pat w nes lam arrs @@ -258,7 +252,7 @@ Bindable lore, BinderOps lore, LocalScope somelore m, SameScope somelore lore, - LetAttr SOACS ~ LetAttr lore, + LetAttr lore ~ LetAttr SOACS, CanBeAliased (Op lore)) => Lambda -> m (AST.Lambda lore) transformLambda (Lambda params body rettype) = do @@ -267,36 +261,6 @@ transformBody body return $ Lambda params body' rettype -newFold :: Transformer m => - String -> [(SubExp,Type)] -> [VName] - -> m ([Ident], [SubExp], [Ident]) -newFold what accexps_and_types arrexps = do - initacc <- mapM copyIfArray acc_exps - acc <- mapM (newIdent "acc") acc_types - arrts <- mapM lookupType arrexps - inarrs <- mapM (newIdent $ what ++ "_inarr") arrts - return (acc, initacc, inarrs) - where (acc_exps, acc_types) = unzip accexps_and_types - -copyIfArray :: Transformer m => - SubExp -> m SubExp -copyIfArray (Constant v) = return $ Constant v -copyIfArray (Var v) = Var <$> copyIfArrayName v - -copyIfArrayName :: Transformer m => - VName -> m VName -copyIfArrayName v = do - t <- lookupType v - case t of - Array {} -> letExp (baseString v ++ "_first_order_copy") $ BasicOp $ Copy v - _ -> return v - -index :: (HasScope lore m, Monad m) => - [VName] -> SubExp -> m [AST.Exp lore] -index arrs i = forM arrs $ \arr -> do - arr_t <- lookupType arr - return $ BasicOp $ Index arr $ fullSlice arr_t [DimFix i] - resultArray :: Transformer m => [Type] -> m [VName] resultArray = mapM oneArray where oneArray t = letExp "result" $ BasicOp $ Scratch (elemType t) (arrayDims t) @@ -331,64 +295,3 @@ loopMerge' :: [(Ident,Uniqueness)] -> [SubExp] -> [(Param DeclType, SubExp)] loopMerge' vars vals = [ (Param pname $ toDecl ptype u, val) | ((Ident pname ptype, u),val) <- zip vars vals ] - -discardPattern :: (MonadFreshNames m, LetAttr (Lore m) ~ LetAttr SOACS) => - [Type] -> AST.Pattern (Lore m) -> m (AST.Pattern (Lore m)) -discardPattern discard pat = do - discard_pat <- basicPattern [] <$> mapM (newIdent "discard") discard - return $ discard_pat <> pat - --- | Turn a Haskell-style mapAccumL into a sequential do-loop. This --- is the guts of transforming a 'Redomap'. -doLoopMapAccumL :: (LocalScope (Lore m) m, MonadBinder m, - Bindable (Lore m), BinderOps (Lore m), - LetAttr (Lore m) ~ Type, - CanBeAliased (Op (Lore m))) => - SubExp - -> AST.Lambda (Aliases (Lore m)) - -> [SubExp] - -> [VName] - -> [VName] - -> m (AST.Exp (Lore m)) -doLoopMapAccumL width innerfun accexps arrexps mapout_arrs = do - (merge, i, loopbody) <- - doLoopMapAccumL' width innerfun accexps arrexps mapout_arrs - return $ DoLoop [] merge (ForLoop i Int32 width []) loopbody - -doLoopMapAccumL' :: (LocalScope (Lore m) m, MonadBinder m, - Bindable (Lore m), BinderOps (Lore m), - LetAttr (Lore m) ~ Type, - CanBeAliased (Op (Lore m))) => - SubExp - -> AST.Lambda (Aliases (Lore m)) - -> [SubExp] - -> [VName] - -> [VName] - -> m ([(AST.FParam (Lore m), SubExp)], VName, AST.Body (Lore m)) -doLoopMapAccumL' width innerfun accexps arrexps mapout_arrs = do - i <- newVName "i" - -- for the MAP part - let acc_num = length accexps - let res_tps = lambdaReturnType innerfun - let map_arr_tps = drop acc_num res_tps - let res_ts = [ arrayOf t (Shape [width]) NoUniqueness - | t <- map_arr_tps ] - let accts = map paramType $ fst $ splitAt acc_num $ lambdaParams innerfun - outarrs <- mapM (newIdent "mapaccum_outarr") res_ts - -- for the REDUCE part - (acc, initacc, inarrs) <- newFold "mapaccum" (zip accexps accts) arrexps - let consumed = consumedInBody $ lambdaBody innerfun - withUniqueness p | identName p `S.member` consumed = (p, Unique) - | p `elem` outarrs = (p, Unique) - | otherwise = (p, Nonunique) - merge = loopMerge' (map withUniqueness $ inarrs++acc++outarrs) - (map Var arrexps++initacc++map Var mapout_arrs) - loopbody <- runBodyBinder $ localScope (scopeOfFParams $ map fst merge) $ do - accxis<- bindLambda (removeLambdaAliases innerfun) . - (map (BasicOp . SubExp . Var . identName) acc ++) =<< - index (map identName inarrs) (Var i) - let (acc', xis) = splitAt acc_num accxis - dests <- letwith (map identName outarrs) (pexp (Var i)) $ - map (BasicOp . SubExp) xis - return $ resultBody (map (Var . identName) inarrs ++ acc' ++ map Var dests) - return (merge, i, loopbody) diff --git a/src/Futhark/Transform/Rename.hs b/src/Futhark/Transform/Rename.hs --- a/src/Futhark/Transform/Rename.hs +++ b/src/Futhark/Transform/Rename.hs @@ -2,6 +2,7 @@ {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} -- | This module provides facilities for transforming Futhark programs such -- that names are unique, via the 'renameProg' function. -- Additionally, the module also supports adding integral \"tags\" to @@ -36,19 +37,19 @@ import Control.Monad.State import Control.Monad.Reader import qualified Data.Map.Strict as M -import qualified Data.Set as S import Data.Maybe import Futhark.Representation.AST.Syntax import Futhark.Representation.AST.Traversals +import Futhark.Representation.AST.Attributes.Names import Futhark.Representation.AST.Attributes.Patterns -import Futhark.FreshNames -import Futhark.MonadFreshNames (MonadFreshNames(..), modifyNameSource) +import Futhark.FreshNames hiding (newName) +import Futhark.MonadFreshNames (MonadFreshNames(..), modifyNameSource, newName) import Futhark.Transform.Substitute runRenamer :: RenameM a -> VNameSource -> (a, VNameSource) -runRenamer m src = runReader (runStateT m src) env - where env = RenameEnv M.empty newName +runRenamer (RenameM m) src = runReader (runStateT m src) env + where env = RenameEnv M.empty -- | Rename variables such that each is unique. The semantics of the -- program are unaffected, under the assumption that the program was @@ -107,18 +108,17 @@ renamePattern = modifyNameSource . runRenamer . rename' where rename' pat = bind (patternNames pat) $ rename pat -data RenameEnv = RenameEnv { - envNameMap :: M.Map VName VName - , envNameFn :: VNameSource -> VName -> (VName, VNameSource) - } +newtype RenameEnv = RenameEnv { envNameMap :: M.Map VName VName } -- | The monad in which renaming is performed. -type RenameM = StateT VNameSource (Reader RenameEnv) +newtype RenameM a = RenameM (StateT VNameSource (Reader RenameEnv) a) + deriving (Functor, Applicative, Monad, + MonadFreshNames, MonadReader RenameEnv) -- | Produce a map of the substitutions that should be performed by -- the renamer. renamerSubstitutions :: RenameM Substitutions -renamerSubstitutions = lift $ asks envNameMap +renamerSubstitutions = asks envNameMap -- | Perform a renaming using the 'Substitute' instance. This only -- works if the argument does not itself perform any name binding, but @@ -128,13 +128,6 @@ substs <- renamerSubstitutions return $ substituteNames substs x --- | Return a fresh, unique name. The @VName@ is prepended to the --- name. -new :: VName -> RenameM VName -new k = do (k', src') <- asks envNameFn <*> get <*> pure k - put src' - return k' - -- | Members of class 'Rename' can be uniquely renamed. class Rename a where -- | Rename the given value such that it does not contain shadowing, @@ -177,7 +170,7 @@ bind :: [VName] -> RenameM a -> RenameM a bind vars body = do - vars' <- mapM new vars + vars' <- mapM newName vars -- This works because map union prefers elements from left -- operand. local (bind' vars') body @@ -206,7 +199,7 @@ rename (Var v) = Var <$> rename v rename (Constant v) = return $ Constant v -instance Rename attr => Rename (ParamT attr) where +instance Rename attr => Rename (Param attr) where rename (Param name attr) = Param <$> rename name <*> rename attr instance Rename attr => Rename (PatternT attr) where @@ -292,7 +285,7 @@ return $ Lambda params' body' ret' instance Rename Names where - rename = fmap S.fromList . mapM rename . S.toList + rename = fmap namesFromList . mapM rename . namesToList instance Rename Rank where rename = return diff --git a/src/Futhark/Transform/Substitute.hs b/src/Futhark/Transform/Substitute.hs --- a/src/Futhark/Transform/Substitute.hs +++ b/src/Futhark/Transform/Substitute.hs @@ -14,12 +14,12 @@ import Control.Monad.Identity import qualified Data.Map.Strict as M -import qualified Data.Set as S import Futhark.Representation.AST.Syntax import Futhark.Representation.AST.Traversals import Futhark.Representation.AST.Attributes.Scope import Futhark.Analysis.PrimExp +import Futhark.Representation.AST.Attributes.Names -- | The substitutions to be made are given by a mapping from names to -- names. @@ -80,7 +80,7 @@ substituteNames substs (StmAux cs attr) = StmAux (substituteNames substs cs) (substituteNames substs attr) -instance Substitute attr => Substitute (ParamT attr) where +instance Substitute attr => Substitute (Param attr) where substituteNames substs (Param name attr) = Param (substituteNames substs name) @@ -135,7 +135,7 @@ substituteNames _ (Ext x) = Ext x instance Substitute Names where - substituteNames = S.map . substituteNames + substituteNames = mapNames . substituteNames instance Substitute shape => Substitute (TypeBase shape u) where substituteNames _ (Prim et) = Prim et @@ -175,6 +175,9 @@ LParamInfo $ substituteNames subst attr substituteNames _ (IndexInfo it) = IndexInfo it + +instance Substitute FV where + substituteNames subst = fvNames . substituteNames subst . freeIn -- | Lores in which all annotations support name -- substitution. diff --git a/src/Futhark/TypeCheck.hs b/src/Futhark/TypeCheck.hs --- a/src/Futhark/TypeCheck.hs +++ b/src/Futhark/TypeCheck.hs @@ -1,5 +1,6 @@ {-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts, FlexibleInstances, MultiParamTypeClasses, TypeFamilies, ScopedTypeVariables #-} {-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE TupleSections #-} -- | The type checker checks whether the program is type-consistent. module Futhark.TypeCheck ( -- * Interface @@ -18,7 +19,7 @@ , lookupAliases , Occurences , collectOccurences - , subCheck + , checkOpWith -- * Checkers , require @@ -172,39 +173,39 @@ deriving (Eq, Show) observation :: Names -> Occurence -observation = flip Occurence S.empty +observation = flip Occurence mempty consumption :: Names -> Occurence -consumption = Occurence S.empty +consumption = Occurence mempty nullOccurence :: Occurence -> Bool -nullOccurence occ = S.null (observed occ) && S.null (consumed occ) +nullOccurence occ = observed occ == mempty && consumed occ == mempty type Occurences = [Occurence] allConsumed :: Occurences -> Names -allConsumed = S.unions . map consumed +allConsumed = mconcat . map consumed seqOccurences :: Occurences -> Occurences -> Occurences seqOccurences occurs1 occurs2 = filter (not . nullOccurence) (map filt occurs1) ++ occurs2 where filt occ = - occ { observed = observed occ `S.difference` postcons } + occ { observed = observed occ `namesSubtract` postcons } postcons = allConsumed occurs2 altOccurences :: Occurences -> Occurences -> Occurences altOccurences occurs1 occurs2 = filter (not . nullOccurence) (map filt occurs1) ++ occurs2 where filt occ = - occ { consumed = consumed occ `S.difference` postcons - , observed = observed occ `S.difference` postcons } + occ { consumed = consumed occ `namesSubtract` postcons + , observed = observed occ `namesSubtract` postcons } postcons = allConsumed occurs2 unOccur :: Names -> Occurences -> Occurences unOccur to_be_removed = filter (not . nullOccurence) . map unOccur' where unOccur' occ = - occ { observed = observed occ `S.difference` to_be_removed - , consumed = consumed occ `S.difference` to_be_removed + occ { observed = observed occ `namesSubtract` to_be_removed + , consumed = consumed occ `namesSubtract` to_be_removed } -- | The 'Consumption' data structure is used to keep track of which @@ -217,7 +218,7 @@ ConsumptionError e <> _ = ConsumptionError e _ <> ConsumptionError e = ConsumptionError e Consumption o1 <> Consumption o2 - | v:_ <- S.toList $ consumed_in_o1 `S.intersection` used_in_o2 = + | v:_ <- namesToList $ consumed_in_o1 `namesIntersection` used_in_o2 = ConsumptionError $ "Variable " <> pretty v <> " referenced after being consumed." | otherwise = Consumption $ o1 `seqOccurences` o2 @@ -235,6 +236,7 @@ data Env lore = Env { envVtable :: M.Map VName (VarBinding lore) , envFtable :: M.Map Name (FunBinding lore) + , envCheckOp :: OpWithAliases (Op lore) -> TypeM lore () , envContext :: [String] } @@ -282,10 +284,10 @@ -- | Mark a name as bound. If the name has been bound previously in -- the program, report a type error. bound :: VName -> TypeM lore () -bound name = do already_seen <- gets $ S.member name +bound name = do already_seen <- gets $ nameIn name when already_seen $ bad $ TypeError $ "Name " ++ pretty name ++ " bound twice" - modify $ S.insert name + modify (<>oneName name) occur :: Occurences -> TypeM lore () occur = tell . Consumption . filter (not . nullOccurence) @@ -297,14 +299,14 @@ observe name = do attr <- lookupVar name unless (primType $ typeOf attr) $ - occur [observation $ S.insert name $ aliases attr] + occur [observation $ oneName name <> aliases attr] -- | Proclaim that we have written to the given variables. consume :: Checkable lore => Names -> TypeM lore () consume als = do scope <- askScope let isArray = maybe False ((>0) . arrayRank . typeOf) . (`M.lookup` scope) - occur [consumption $ S.filter isArray als] + occur [consumption $ namesFromList $ filter isArray $ namesToList als] collectOccurences :: TypeM lore a -> TypeM lore (a, Occurences) collectOccurences m = pass $ do @@ -312,6 +314,10 @@ o <- checkConsumption c return ((x, o), const mempty) +checkOpWith :: (OpWithAliases (Op lore) -> TypeM lore ()) + -> TypeM lore a -> TypeM lore a +checkOpWith checker = local $ \env -> env { envCheckOp = checker } + checkConsumption :: Consumption -> TypeM lore Occurences checkConsumption (ConsumptionError e) = bad $ TypeError e checkConsumption (Consumption os) = return os @@ -335,7 +341,7 @@ tell . Consumption =<< mapM inspect os return x where inspect o = do - new_consumed <- mconcat <$> mapM wasConsumed (S.toList $ consumed o) + new_consumed <- mconcat <$> mapM wasConsumed (namesToList $ consumed o) return o { consumed = new_consumed } wasConsumed v | Just als <- lookup v consumable = return als @@ -349,8 +355,8 @@ -- | Given the immediate aliases, compute the full transitive alias -- set (including the immediate aliases). expandAliases :: Names -> Env lore -> Names -expandAliases names env = names `S.union` aliasesOfAliases - where aliasesOfAliases = mconcat . map look . S.toList $ names +expandAliases names env = names <> aliasesOfAliases + where aliasesOfAliases = mconcat . map look . namesToList $ names look k = case M.lookup k $ envVtable env of Just (LetInfo (als, _)) -> unNames als _ -> mempty @@ -362,7 +368,6 @@ binding bnds = check . local (`bindVars` bnds) where bindVars = M.foldlWithKey' bindVar boundnames = M.keys bnds - boundnameset = S.fromList boundnames bindVar env name (LetInfo (Names' als, attr)) = let als' | primType (typeOf attr) = mempty @@ -378,7 +383,7 @@ check m = do mapM_ bound $ M.keys bnds (a, os) <- collectOccurences m - tell $ Consumption $ unOccur boundnameset os + tell $ Consumption $ unOccur (namesFromList boundnames) os return a lookupVar :: VName -> TypeM lore (NameInfo (Aliases lore)) @@ -393,7 +398,7 @@ info <- lookupVar name return $ if primType $ typeOf info then mempty - else S.insert name $ aliases info + else oneName name <> aliases info aliases :: NameInfo (Aliases lore) -> Names aliases (LetInfo (als, _)) = unNames als @@ -456,6 +461,7 @@ let typeenv = Env { envVtable = M.empty , envFtable = mempty , envContext = [] + , envCheckOp = checkOp } let onFunction ftable fun = fmap fst $ runTypeM typeenv $ @@ -531,8 +537,10 @@ check scope <- askScope let isArray = maybe False ((>0) . arrayRank . typeOf) . (`M.lookup` scope) - context ("When checking the body aliases: " ++ pretty (bodyAliases body)) $ - checkReturnAlias $ map (S.filter isArray) $ bodyAliases body + context ("When checking the body aliases: " ++ + pretty (map namesToList $ bodyAliases body)) $ + checkReturnAlias $ map (namesFromList . filter isArray . namesToList) $ + bodyAliases body where param_names = map fst params checkNoDuplicateParams = foldM_ expand [] param_names @@ -546,47 +554,26 @@ -- | Check that unique return values do not alias a -- non-consumed parameter. checkReturnAlias = - foldM_ checkReturnAlias' S.empty . returnAliasing rettype + foldM_ checkReturnAlias' mempty . returnAliasing rettype checkReturnAlias' seen (Unique, names) - | any (`S.member` S.map snd seen) $ S.toList names = + | any (`S.member` S.map fst seen) $ namesToList names = bad $ UniqueReturnAliased fname | otherwise = do consume names - return $ seen `S.union` tag Unique names + return $ seen <> tag Unique names checkReturnAlias' seen (Nonunique, names) - | any (`S.member` seen) $ S.toList $ tag Unique names = + | any (`S.member` seen) $ tag Unique names = bad $ UniqueReturnAliased fname - | otherwise = return $ seen `S.union` tag Nonunique names + | otherwise = return $ seen <> tag Nonunique names - tag u = S.map $ \name -> (u, name) + tag u = S.fromList . map (,u) . namesToList returnAliasing expected got = reverse $ zip (reverse (map uniqueness expected) ++ repeat Nonunique) $ reverse got -subCheck :: forall lore newlore a. - (Checkable newlore, - RetType lore ~ RetType newlore, - LetAttr lore ~ LetAttr newlore, - FParamAttr lore ~ FParamAttr newlore, - LParamAttr lore ~ LParamAttr newlore) => - TypeM newlore a -> - TypeM lore a -subCheck m = do - typeenv <- asks newEnv - case runTypeM typeenv m of - Left err -> bad $ TypeError $ show err - Right (x, cons) -> tell cons >> return x - where newEnv :: Env lore -> Env newlore - newEnv (Env vtable ftable ctx) = - Env (M.map coerceVar vtable) ftable ctx - coerceVar (LetInfo x) = LetInfo x - coerceVar (FParamInfo x) = FParamInfo x - coerceVar (LParamInfo x) = LParamInfo x - coerceVar (IndexInfo it) = IndexInfo it - checkSubExp :: Checkable lore => SubExp -> TypeM lore Type checkSubExp (Constant val) = return $ Prim $ primValueType val @@ -694,7 +681,7 @@ bad $ SlicingError (arrayRank src_t) (length idxes) se_aliases <- subExpAliasesM se - when (src `S.member` se_aliases) $ + when (src `nameIn` se_aliases) $ bad $ TypeError "The target of an Update must not alias the value to be written." mapM_ checkDimIndex idxes @@ -855,7 +842,8 @@ bad $ ReturnTypeError (nameFromString "") (staticShapes rettype') (staticShapes bodyt) -checkExp (Op op) = checkOp op +checkExp (Op op) = do checker <- asks envCheckOp + checker op checkSOACArrayArgs :: Checkable lore => SubExp -> [VName] -> TypeM lore [Arg] @@ -1061,6 +1049,7 @@ class Attributes lore => CheckableOp lore where checkOp :: OpWithAliases (Op lore) -> TypeM lore () + -- ^ Used at top level; can be locally changed with 'checkOpWith'. -- | The class of lores that can be type-checked. class (Attributes lore, CanBeAliased (Op lore), CheckableOp lore) => Checkable lore where diff --git a/src/Futhark/Util.hs b/src/Futhark/Util.hs --- a/src/Futhark/Util.hs +++ b/src/Futhark/Util.hs @@ -31,6 +31,8 @@ toPOSIX, trim, pmapIO, + UserString, + EncodedString, zEncodeString ) where diff --git a/src/Futhark/Util/Log.hs b/src/Futhark/Util/Log.hs --- a/src/Futhark/Util/Log.hs +++ b/src/Futhark/Util/Log.hs @@ -1,6 +1,5 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TypeSynonymInstances #-} {-# LANGUAGE OverloadedStrings #-} -- | Opaque type for an operations log that provides fast O(1) -- appends. diff --git a/src/Language/Futhark.hs b/src/Language/Futhark.hs --- a/src/Language/Futhark.hs +++ b/src/Language/Futhark.hs @@ -8,7 +8,7 @@ , ModExp, ModParam, SigExp, ModBind, SigBind , ValBind, Dec, Spec, Prog , TypeBind, TypeDecl - , StructTypeArg, ArrayElemType + , StructTypeArg, ScalarType , TypeParam, Case ) where @@ -68,8 +68,8 @@ -- | A type-checked type parameter. type TypeParam = TypeParamBase VName --- | A known array element type with no shape annotations. -type ArrayElemType = ArrayElemTypeBase () +-- | A known scalar type with no shape annotations. +type ScalarType = ScalarTypeBase () -- | A type-checked case (of a match expression). type Case = CaseBase Info VName diff --git a/src/Language/Futhark/Attributes.hs b/src/Language/Futhark/Attributes.hs --- a/src/Language/Futhark/Attributes.hs +++ b/src/Language/Futhark/Attributes.hs @@ -29,7 +29,6 @@ , patternIdents , patternType , patternStructType - , patternPatternType , patternParam , patternOrderZero , patternDimNames @@ -61,12 +60,12 @@ , removeShapeAnnotations , vacuousShapeAnnotations , anyDimShapeAnnotations - , recordArrayElemToType , tupleRecord , isTupleRecord , areTupleFields , tupleFieldNames , sortFields + , sortConstrs , isTypeParam , combineTypeShapes , unscopeType @@ -77,7 +76,6 @@ , NoInfo(..) , UncheckedType , UncheckedTypeExp - , UncheckedArrayElemType , UncheckedIdent , UncheckedTypeDecl , UncheckedDimIndex @@ -93,7 +91,7 @@ ) where -import Control.Monad.Writer +import Control.Monad.Writer hiding (Sum) import Data.Char import Data.Foldable import qualified Data.Map.Strict as M @@ -127,30 +125,24 @@ -- removed. nestedDims :: TypeBase (DimDecl VName) as -> [DimDecl VName] nestedDims t = - case t of Array _ _ a ds -> nub $ arrayNestedDims a <> shapeDims ds - Record fs -> nub $ fold $ fmap nestedDims fs - Prim{} -> mempty - TypeVar _ _ _ targs -> concatMap typeArgDims targs - Arrow _ v t1 t2 -> filter (notV v) $ nestedDims t1 <> nestedDims t2 - Enum{} -> [] - where arrayNestedDims ArrayPrimElem{} = - mempty - arrayNestedDims (ArrayPolyElem _ targs) = - concatMap typeArgDims targs - arrayNestedDims (ArrayRecordElem ts) = - fold (fmap recordArrayElemNestedDims ts) - arrayNestedDims ArrayEnumElem{} = mempty - - recordArrayElemNestedDims (RecordArrayArrayElem a ds) = - arrayNestedDims a <> shapeDims ds - recordArrayElemNestedDims (RecordArrayElem et) = - arrayNestedDims et + case t of Array _ _ a ds -> + nub $ nestedDims (Scalar a) <> shapeDims ds + Scalar (Record fs) -> + nub $ fold $ fmap nestedDims fs + Scalar Prim{} -> + mempty + Scalar (Sum cs) -> + nub $ fold $ (fmap . concatMap) nestedDims cs + Scalar (Arrow _ v t1 t2) -> + filter (notV v) $ nestedDims t1 <> nestedDims t2 + Scalar (TypeVar _ _ _ targs) -> + concatMap typeArgDims targs - typeArgDims (TypeArgDim d _) = [d] + where typeArgDims (TypeArgDim d _) = [d] typeArgDims (TypeArgType at _) = nestedDims at - notV Nothing = const True - notV (Just v) = (/=NamedDim (qualName v)) + notV Unnamed = const True + notV (Named v) = (/=NamedDim (qualName v)) -- | Change the shape of a type to be just the 'Rank'. removeShapeAnnotations :: TypeBase (DimDecl vn) as -> TypeBase () as @@ -173,7 +165,8 @@ -- | Return the uniqueness of a type. uniqueness :: TypeBase shape as -> Uniqueness uniqueness (Array _ u _ _) = u -uniqueness (TypeVar _ u _ _) = u +uniqueness (Scalar (TypeVar _ u _ _)) = u +uniqueness (Scalar (Sum ts)) = mconcat $ map (mconcat . map uniqueness) $ M.elems ts uniqueness _ = Nonunique -- | @unique t@ is 'True' if the type of the argument is unique. @@ -188,13 +181,13 @@ -- | @diet t@ returns a description of how a function parameter of -- type @t@ might consume its argument. diet :: TypeBase shape as -> Diet -diet (Record ets) = RecordDiet $ fmap diet ets -diet (Prim _) = Observe -diet TypeVar{} = Observe -diet (Arrow _ _ t1 t2) = FuncDiet (diet t1) (diet t2) -diet (Array _ Unique _ _) = Consume -diet (Array _ Nonunique _ _) = Observe -diet (Enum _) = Observe +diet (Scalar (Record ets)) = RecordDiet $ fmap diet ets +diet (Scalar (Prim _)) = Observe +diet (Scalar TypeVar{}) = Observe +diet (Scalar (Arrow _ _ t1 t2)) = FuncDiet (diet t1) (diet t2) +diet (Array _ Unique _ _) = Consume +diet (Array _ Nonunique _ _) = Observe +diet (Scalar Sum{}) = Observe -- | Convert any type to one that has rank information, no alias -- information, and no embedded names. @@ -216,27 +209,11 @@ -- @n@ array dimensions from @t@. Returns @Nothing@ if @t@ has less -- than @n@ dimensions. peelArray :: Int -> TypeBase dim as -> Maybe (TypeBase dim as) -peelArray 0 t = Just t -peelArray n (Array _ _ (ArrayPrimElem et) shape) - | shapeRank shape == n = - Just $ Prim et -peelArray n (Array als u (ArrayPolyElem et targs) shape) - | shapeRank shape == n = - Just $ TypeVar als u et targs -peelArray n (Array als u (ArrayRecordElem ts) shape) - | shapeRank shape == n = - Just $ Record $ fmap asType ts - where asType (RecordArrayElem (ArrayPrimElem bt)) = Prim bt - asType (RecordArrayElem (ArrayPolyElem bt targs)) = TypeVar als u bt targs - asType (RecordArrayElem (ArrayRecordElem ts')) = Record $ fmap asType ts' - asType (RecordArrayElem (ArrayEnumElem cs)) = Enum cs - asType (RecordArrayArrayElem et e_shape) = Array als u et e_shape -peelArray n (Array _ _ (ArrayEnumElem cs) shape) +peelArray n (Array als u t shape) | shapeRank shape == n = - Just $ Enum cs -peelArray n (Array als u et shape) = do - shape' <- stripDims n shape - return $ Array als u et shape' + Just $ Scalar t `addAliases` const als + | otherwise = + Array als u t <$> stripDims n shape peelArray _ _ = Nothing -- | @arrayOf t s u@ constructs an array type. The convenience @@ -244,80 +221,43 @@ -- itself be an array. If @t@ is an @n@-dimensional array, and @s@ is -- a list of length @n@, the resulting type is of an @n+m@ dimensions. -- The uniqueness of the new array will be @u@, no matter the --- uniqueness of @t@. The function returns 'Nothing' in case an --- attempt is made to create an array of functions. +-- uniqueness of @t@. arrayOf :: Monoid as => TypeBase dim as -> ShapeDecl dim -> Uniqueness - -> Maybe (TypeBase dim as) -arrayOf t = arrayOfWithAliases t mempty + -> TypeBase dim as +arrayOf t = arrayOfWithAliases (t `setUniqueness` Nonunique) mempty arrayOfWithAliases :: Monoid as => TypeBase dim as -> as -> ShapeDecl dim -> Uniqueness - -> Maybe (TypeBase dim as) + -> TypeBase dim as arrayOfWithAliases (Array as1 _ et shape1) as2 shape2 u = - Just $ Array (as1<>as2) u et (shape2 <> shape1) -arrayOfWithAliases (Prim et) as shape u = - Just $ Array as u (ArrayPrimElem et) shape -arrayOfWithAliases (TypeVar _ _ x targs) as shape u = - Just $ Array as u (ArrayPolyElem x targs) shape -arrayOfWithAliases (Record ts) as shape u = do - ts' <- traverse typeToRecordArrayElem ts - return $ Array as u (ArrayRecordElem ts') shape -arrayOfWithAliases Arrow{} _ _ _ = Nothing -arrayOfWithAliases (Enum cs) as shape u = - Just $ Array as u (ArrayEnumElem cs) shape - -typeToRecordArrayElem :: Monoid as => - TypeBase dim as -> Maybe (RecordArrayElemTypeBase dim) -typeToRecordArrayElem (Prim bt) = - Just $ RecordArrayElem $ ArrayPrimElem bt -typeToRecordArrayElem (TypeVar _ _ bt targs) = - Just $ RecordArrayElem $ ArrayPolyElem bt targs -typeToRecordArrayElem (Record ts') = - RecordArrayElem . ArrayRecordElem <$> - traverse typeToRecordArrayElem ts' -typeToRecordArrayElem (Array _ _ et shape) = - Just $ RecordArrayArrayElem et shape -typeToRecordArrayElem Arrow{} = Nothing -typeToRecordArrayElem (Enum cs) = - Just $ RecordArrayElem $ ArrayEnumElem cs - -recordArrayElemToType :: Monoid as => - RecordArrayElemTypeBase dim - -> TypeBase dim as -recordArrayElemToType (RecordArrayElem et) = arrayElemToType et -recordArrayElemToType (RecordArrayArrayElem et shape) = Array mempty Nonunique et shape - -arrayElemToType :: Monoid as => ArrayElemTypeBase dim -> TypeBase dim as -arrayElemToType (ArrayPolyElem bt targs) = - TypeVar mempty Nonunique bt targs -arrayElemToType (ArrayRecordElem ts) = - Record $ fmap recordArrayElemToType ts -arrayElemToType (ArrayPrimElem bt) = Prim bt -arrayElemToType (ArrayEnumElem cs) = Enum cs + Array (as1<>as2) u et (shape2 <> shape1) +arrayOfWithAliases (Scalar t) as shape u = + Array as u (bimap id (const ()) t) shape -- | @stripArray n t@ removes the @n@ outermost layers of the array. -- Essentially, it is the type of indexing an array of type @t@ with -- @n@ indexes. -stripArray :: Monoid as => Int -> TypeBase dim as -> TypeBase dim as +stripArray :: Int -> TypeBase dim as -> TypeBase dim as stripArray n (Array als u et shape) | Just shape' <- stripDims n shape = - Array als u et shape' - | otherwise = arrayElemToType et `setUniqueness` u `addAliases` (<>als) + Array als u et shape' + | otherwise = + Scalar et `setUniqueness` u `setAliases` als stripArray _ t = t -- | Create a record type corresponding to a tuple with the given -- element types. tupleRecord :: [TypeBase dim as] -> TypeBase dim as -tupleRecord = Record . M.fromList . zip tupleFieldNames +tupleRecord = Scalar . Record . M.fromList . zip tupleFieldNames isTupleRecord :: TypeBase dim as -> Maybe [TypeBase dim as] -isTupleRecord (Record fs) = areTupleFields fs +isTupleRecord (Scalar (Record fs)) = areTupleFields fs isTupleRecord _ = Nothing areTupleFields :: M.Map Name a -> Maybe [a] @@ -340,6 +280,8 @@ fieldish s = case reads $ nameToString s of [(x, "")] -> Left (x::Int) _ -> Right s +sortConstrs :: M.Map Name a -> [(Name, a)] +sortConstrs cs = sortOn fst $ M.toList cs isTypeParam :: TypeParamBase vn -> Bool isTypeParam TypeParamType{} = True @@ -351,43 +293,27 @@ -- information (e.g., shape restrictions) from the user given annotation. combineTypeShapes :: (Monoid as, ArrayDim dim) => TypeBase dim as -> TypeBase dim as -> TypeBase dim as -combineTypeShapes (Record ts1) (Record ts2) +combineTypeShapes (Scalar (Record ts1)) (Scalar (Record ts2)) | M.keys ts1 == M.keys ts2 = - Record $ M.map (uncurry combineTypeShapes) (M.intersectionWith (,) ts1 ts2) + Scalar $ Record $ M.map (uncurry combineTypeShapes) (M.intersectionWith (,) ts1 ts2) combineTypeShapes (Array als1 u1 et1 shape1) (Array als2 _u2 et2 shape2) | Just new_shape <- unifyShapes shape1 shape2 = - Array (als1<>als2) u1 (combineElemTypeInfo et1 et2) new_shape + arrayOfWithAliases (combineTypeShapes (Scalar et1) (Scalar et2) + `setAliases` mempty) + (als1<>als2) new_shape u1 combineTypeShapes _ new_tp = new_tp -combineElemTypeInfo :: ArrayDim dim => - ArrayElemTypeBase dim - -> ArrayElemTypeBase dim -> ArrayElemTypeBase dim -combineElemTypeInfo (ArrayRecordElem et1) (ArrayRecordElem et2) = - ArrayRecordElem $ M.map (uncurry combineRecordArrayTypeInfo) - (M.intersectionWith (,) et1 et2) -combineElemTypeInfo _ new_tp = new_tp - -combineRecordArrayTypeInfo :: ArrayDim dim => - RecordArrayElemTypeBase dim - -> RecordArrayElemTypeBase dim - -> RecordArrayElemTypeBase dim -combineRecordArrayTypeInfo (RecordArrayElem et1) (RecordArrayElem et2) = - RecordArrayElem $ combineElemTypeInfo et1 et2 -combineRecordArrayTypeInfo (RecordArrayArrayElem et1 shape1) - (RecordArrayArrayElem et2 shape2) - | Just new_shape <- unifyShapes shape1 shape2 = - RecordArrayArrayElem (combineElemTypeInfo et1 et2) new_shape -combineRecordArrayTypeInfo _ new_tp = new_tp - --- | Set the uniqueness attribute of a type. If the type is a tuple, --- the uniqueness of its components will be modified. +-- | Set the uniqueness attribute of a type. If the type is a record +-- or sum type, the uniqueness of its components will be modified. setUniqueness :: TypeBase dim as -> Uniqueness -> TypeBase dim as setUniqueness (Array als _ et shape) u = Array als u et shape -setUniqueness (TypeVar als _ t targs) u = - TypeVar als u t targs -setUniqueness (Record ets) u = - Record $ fmap (`setUniqueness` u) ets +setUniqueness (Scalar (TypeVar als _ t targs)) u = + Scalar $ TypeVar als u t targs +setUniqueness (Scalar (Record ets)) u = + Scalar $ Record $ fmap (`setUniqueness` u) ets +setUniqueness (Scalar (Sum ets)) u = + Scalar $ Sum $ fmap (map (`setUniqueness` u)) ets setUniqueness t _ = t -- | @t \`setAliases\` als@ returns @t@, but with @als@ substituted for @@ -419,7 +345,7 @@ primValueType BoolValue{} = Bool valueType :: Value -> TypeBase () () -valueType (PrimValue bv) = Prim $ primValueType bv +valueType (PrimValue bv) = Scalar $ Prim $ primValueType bv valueType (ArrayValue _ t) = t -- | Construct a 'ShapeDecl' with the given number of zero-information @@ -440,7 +366,7 @@ -- | The type of an Futhark term. The aliasing will refer to itself, if -- the term is a non-tuple-typed variable. typeOf :: ExpBase Info VName -> PatternType -typeOf (Literal val _) = Prim $ primValueType val +typeOf (Literal val _) = Scalar $ Prim $ primValueType val typeOf (IntLit _ (Info t) _) = t typeOf (FloatLit _ (Info t) _) = t typeOf (Parens e _) = typeOf e @@ -448,7 +374,7 @@ typeOf (TupLit es _) = tupleRecord $ map typeOf es typeOf (RecordLit fs _) = -- Reverse, because M.unions is biased to the left. - Record $ M.unions $ reverse $ map record fs + Scalar $ Record $ M.unions $ reverse $ map record fs where record (RecordFieldExplicit name e _) = M.singleton name $ typeOf e record (RecordFieldImplicit name (Info t) _) = M.singleton (baseName name) $ t @@ -463,7 +389,7 @@ typeOf (Apply _ _ _ (Info t) _) = t typeOf (Negate e _) = typeOf e typeOf (LetPat _ _ _ (Info t) _) = t -typeOf (LetFun _ _ body _) = typeOf body +typeOf (LetFun name _ body _) = unscopeType (S.singleton name) $ typeOf body typeOf (LetWith _ _ _ _ _ (Info t) _) = t typeOf (Index _ _ (Info t) _) = t typeOf (Update e _ _ _) = typeOf e `setAliases` mempty @@ -472,9 +398,9 @@ typeOf (Assert _ e _ _) = typeOf e typeOf (DoLoop pat _ _ _ _) = patternType pat typeOf (Lambda params _ _ (Info (als, t)) _) = - unscopeType bound_here $ - foldr (uncurry (Arrow ()) . patternParam) t params `setAliases` als + unscopeType bound_here $ foldr (arrow . patternParam) t params `setAliases` als where bound_here = S.map identName (mconcat $ map patternIdents params) + arrow (px, tx) y = Scalar $ Arrow () px tx y typeOf (OpSection _ (Info t) _) = t typeOf (OpSectionLeft _ _ _ (_, Info pt2) (Info ret) _) = @@ -483,38 +409,34 @@ foldFunType [fromStruct pt1] ret typeOf (ProjectSection _ (Info t) _) = t typeOf (IndexSection _ (Info t) _) = t -typeOf (VConstr0 _ (Info t) _) = t -typeOf (Match _ _ (Info t) _) = t +typeOf (Constr _ _ (Info t) _) = t +typeOf (Match _ cs (Info t) _) = + unscopeType (foldMap unscopeSet cs) t + where unscopeSet (CasePat p _ _) = S.map identName $ patternIdents p foldFunType :: Monoid as => [TypeBase dim as] -> TypeBase dim as -> TypeBase dim as -foldFunType ps ret = foldr (Arrow mempty Nothing) ret ps +foldFunType ps ret = foldr arrow ret ps + where arrow t1 t2 = Scalar $ Arrow mempty Unnamed t1 t2 -- | Extract the parameter types and return type from a type. -- If the type is not an arrow type, the list of parameter types is empty. unfoldFunType :: TypeBase dim as -> ([TypeBase dim as], TypeBase dim as) -unfoldFunType (Arrow _ _ t1 t2) = let (ps, r) = unfoldFunType t2 - in (t1 : ps, r) +unfoldFunType (Scalar (Arrow _ _ t1 t2)) = + let (ps, r) = unfoldFunType t2 + in (t1 : ps, r) unfoldFunType t = ([], t) -- | The type names mentioned in a type. typeVars :: Monoid as => TypeBase dim as -> S.Set VName typeVars t = case t of - Prim{} -> mempty - TypeVar _ _ tn targs -> - mconcat $ typeVarFree tn : map typeArgFree targs - Arrow _ _ t1 t2 -> typeVars t1 <> typeVars t2 - Record fields -> foldMap typeVars fields - Array _ _ ArrayPrimElem{} _ -> mempty - Array _ _ (ArrayPolyElem tn targs) _ -> + Scalar Prim{} -> mempty + Scalar (TypeVar _ _ tn targs) -> mconcat $ typeVarFree tn : map typeArgFree targs - Array _ _ (ArrayRecordElem fields) _ -> - foldMap (typeVars . f) fields - -- This local function is to avoid an ambiguous type. - where f :: RecordArrayElemTypeBase dim -> TypeBase dim () - f = recordArrayElemToType - Array _ _ ArrayEnumElem{} _ -> mempty - Enum{} -> mempty + Scalar (Arrow _ _ t1 t2) -> typeVars t1 <> typeVars t2 + Scalar (Record fields) -> foldMap typeVars fields + Scalar (Sum cs) -> mconcat $ (foldMap . fmap) typeVars cs + Array _ _ rt _ -> typeVars $ Scalar rt where typeVarFree = S.singleton . typeLeaf typeArgFree (TypeArgType ta _) = typeVars ta typeArgFree TypeArgDim{} = mempty @@ -523,12 +445,12 @@ -- a function type, does not contain a function type as a subcomponent, and may -- not be instantiated with a function type. orderZero :: TypeBase dim as -> Bool -orderZero (Prim _) = True -orderZero Array{} = True -orderZero (Record fs) = all orderZero $ M.elems fs -orderZero TypeVar{} = True -orderZero Arrow{} = False -orderZero Enum{} = True +orderZero Array{} = True +orderZero (Scalar (Prim _)) = True +orderZero (Scalar (Record fs)) = all orderZero $ M.elems fs +orderZero (Scalar TypeVar{}) = True +orderZero (Scalar Arrow{}) = False +orderZero (Scalar (Sum cs)) = all (all orderZero) cs -- | Extract all the shape names that occur in a given pattern. patternDimNames :: PatternBase Info VName -> S.Set VName @@ -540,6 +462,7 @@ patternDimNames (PatternAscription p (TypeDecl _ (Info t)) _) = patternDimNames p <> typeDimNames t patternDimNames (PatternLit _ (Info tp) _) = typeDimNames tp +patternDimNames (PatternConstr _ _ ps _) = foldMap patternDimNames ps -- | Extract all the shape names that occur in a given type. typeDimNames :: TypeBase (DimDecl VName) als -> S.Set VName @@ -559,6 +482,7 @@ Wildcard (Info t) _ -> orderZero t PatternAscription p _ _ -> patternOrderZero p PatternLit _ (Info t) _ -> orderZero t + PatternConstr _ _ ps _ -> all patternOrderZero ps -- | The set of identifiers bound in a pattern. patternIdents :: (Functor f, Ord vn) => PatternBase f vn -> S.Set (IdentBase f vn) @@ -569,42 +493,34 @@ patternIdents Wildcard{} = mempty patternIdents (PatternAscription p _ _) = patternIdents p patternIdents PatternLit{} = mempty +patternIdents (PatternConstr _ _ ps _ ) = mconcat $ map patternIdents ps -- | The type of values bound by the pattern. patternType :: PatternBase Info VName -> PatternType -patternType (Wildcard (Info t) _) = t -patternType (PatternParens p _) = patternType p -patternType (Id _ (Info t) _) = t -patternType (TuplePattern pats _) = tupleRecord $ map patternType pats -patternType (RecordPattern fs _) = Record $ patternType <$> M.fromList fs -patternType (PatternAscription p _ _) = patternType p -patternType (PatternLit _ (Info t) _) = t - --- | The type of a pattern, including shape annotations. -patternPatternType :: PatternBase Info VName -> PatternType -patternPatternType (Wildcard (Info t) _) = t -patternPatternType (PatternParens p _) = patternPatternType p -patternPatternType (Id _ (Info t) _) = t -patternPatternType (TuplePattern pats _) = tupleRecord $ map patternPatternType pats -patternPatternType (RecordPattern fs _) = Record $ patternPatternType <$> M.fromList fs -patternPatternType (PatternAscription p _ _) = patternPatternType p -patternPatternType (PatternLit _ (Info t) _) = t +patternType (Wildcard (Info t) _) = t +patternType (PatternParens p _) = patternType p +patternType (Id _ (Info t) _) = t +patternType (TuplePattern pats _) = tupleRecord $ map patternType pats +patternType (RecordPattern fs _) = Scalar $ Record $ patternType <$> M.fromList fs +patternType (PatternAscription p _ _) = patternType p +patternType (PatternLit _ (Info t) _) = t +patternType (PatternConstr _ (Info t) _ _) = t -- | The type matched by the pattern, including shape declarations if present. patternStructType :: PatternBase Info VName -> StructType -patternStructType = toStruct . patternPatternType +patternStructType = toStruct . patternType -- | When viewed as a function parameter, does this pattern correspond -- to a named parameter of some type? -patternParam :: PatternBase Info VName -> (Maybe VName, StructType) +patternParam :: PatternBase Info VName -> (PName, StructType) patternParam (PatternParens p _) = patternParam p patternParam (PatternAscription (Id v _ _) td _) = - (Just v, unInfo $ expandedType td) + (Named v, unInfo $ expandedType td) patternParam (Id v (Info t) _) = - (Just v, toStruct t) + (Named v, toStruct t) patternParam p = - (Nothing, patternStructType p) + (Unnamed, patternStructType p) -- | Names of primitive types to types. This is only valid if no -- shadowing is going on, but useful for tools. @@ -633,12 +549,6 @@ map primFun (M.toList Primitive.primFuns) ++ - [ ("~", IntrinsicOverloadedFun - (map Signed [minBound..maxBound] ++ - map Unsigned [minBound..maxBound]) - [Nothing] Nothing) - , ("!", IntrinsicMonoFun [Bool] Bool)] ++ - [("opaque", IntrinsicOpaque)] ++ map unOpFun Primitive.allUnOps ++ @@ -658,101 +568,119 @@ map FloatType [minBound..maxBound] ++ [Bool]) ++ + -- This overrides the ! from Primitive. + [ ("!", IntrinsicOverloadedFun + (map Signed [minBound..maxBound] ++ + map Unsigned [minBound..maxBound] ++ + [Bool]) + [Nothing] Nothing) ] ++ + -- The reason for the loop formulation is to ensure that we -- get a missing case warning if we forget a case. mapMaybe mkIntrinsicBinOp [minBound..maxBound] ++ [("flatten", IntrinsicPolyFun [tp_a] - [Array () Nonunique (ArrayPolyElem tv_a' []) (rank 2)] $ - Array () Nonunique (ArrayPolyElem tv_a' []) (rank 1)), + [Array () Nonunique t_a (rank 2)] $ + Array () Nonunique t_a (rank 1)), ("unflatten", IntrinsicPolyFun [tp_a] - [Prim $ Signed Int32, - Prim $ Signed Int32, - Array () Nonunique (ArrayPolyElem tv_a' []) (rank 1)] $ - Array () Nonunique (ArrayPolyElem tv_a' []) (rank 2)), + [Scalar $ Prim $ Signed Int32, + Scalar $ Prim $ Signed Int32, + Array () Nonunique t_a (rank 1)] $ + Array () Nonunique t_a (rank 2)), ("concat", IntrinsicPolyFun [tp_a] [arr_a, arr_a] uarr_a), ("rotate", IntrinsicPolyFun [tp_a] - [Prim $ Signed Int32, arr_a] arr_a), + [Scalar $ Prim $ Signed Int32, arr_a] arr_a), ("transpose", IntrinsicPolyFun [tp_a] [arr_2d_a] arr_2d_a), ("cmp_threshold", IntrinsicPolyFun [] - [Prim $ Signed Int32, - Array () Nonunique (ArrayPrimElem (Signed Int32)) (rank 1)] $ - Prim Bool), + [Scalar $ Prim $ Signed Int32, + Array () Nonunique (Prim $ Signed Int32) (rank 1)] $ + Scalar $ Prim Bool), ("scatter", IntrinsicPolyFun [tp_a] - [Array () Unique (ArrayPolyElem tv_a' []) (rank 1), - Array () Nonunique (ArrayPrimElem (Signed Int32)) (rank 1), - Array () Nonunique (ArrayPolyElem tv_a' []) (rank 1)] $ - Array () Unique (ArrayPolyElem tv_a' []) (rank 1)), + [Array () Unique t_a (rank 1), + Array () Nonunique (Prim $ Signed Int32) (rank 1), + Array () Nonunique t_a (rank 1)] $ + Array () Unique t_a (rank 1)), ("zip", IntrinsicPolyFun [tp_a, tp_b] [arr_a, arr_b] arr_a_b), ("unzip", IntrinsicPolyFun [tp_a, tp_b] [arr_a_b] t_arr_a_arr_b), ("gen_reduce", IntrinsicPolyFun [tp_a] [uarr_a, - t_a `arr` (t_a `arr` t_a), - t_a, - Array () Nonunique (ArrayPrimElem (Signed Int32)) (rank 1), + Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), + Scalar t_a, + Array () Nonunique (Prim $ Signed Int32) (rank 1), arr_a] uarr_a), - ("map", IntrinsicPolyFun [tp_a, tp_b] [t_a `arr` t_b, arr_a] uarr_b), + ("map", IntrinsicPolyFun [tp_a, tp_b] [Scalar t_a `arr` Scalar t_b, arr_a] uarr_b), ("reduce", IntrinsicPolyFun [tp_a] - [t_a `arr` (t_a `arr` t_a), t_a, arr_a] t_a), + [Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), Scalar t_a, arr_a] $ + Scalar t_a), ("reduce_comm", IntrinsicPolyFun [tp_a] - [t_a `arr` (t_a `arr` t_a), t_a, arr_a] t_a), + [Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), Scalar t_a, arr_a] $ + Scalar t_a), ("scan", IntrinsicPolyFun [tp_a] - [t_a `arr` (t_a `arr` t_a), t_a, arr_a] uarr_a), + [Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), Scalar t_a, arr_a] uarr_a), ("partition", IntrinsicPolyFun [tp_a] - [Prim (Signed Int32), t_a `arr` Prim (Signed Int32), arr_a] $ - tupleRecord [uarr_a, Array () Unique (ArrayPrimElem (Signed Int32)) (rank 1)]), + [Scalar (Prim $ Signed Int32), + Scalar t_a `arr` Scalar (Prim $ Signed Int32), arr_a] $ + tupleRecord [uarr_a, Array () Unique (Prim $ Signed Int32) (rank 1)]), - ("stream_map", - IntrinsicPolyFun [tp_a, tp_b] [Prim (Signed Int32) `arr` (arr_a `arr` arr_b), arr_a] uarr_b), + ("map_stream", + IntrinsicPolyFun [tp_a, tp_b] + [Scalar (Prim $ Signed Int32) `arr` (arr_a `arr` arr_b), arr_a] + uarr_b), - ("stream_map_per", - IntrinsicPolyFun [tp_a, tp_b] [Prim (Signed Int32) `arr` (arr_a `arr` arr_b), arr_a] uarr_b), + ("map_stream_per", + IntrinsicPolyFun [tp_a, tp_b] + [Scalar (Prim $ Signed Int32) `arr` (arr_a `arr` arr_b), arr_a] + uarr_b), - ("stream_red", - IntrinsicPolyFun [tp_a, tp_b] [t_b `arr` (t_b `arr` t_b), Prim (Signed Int32) `arr` (arr_a `arr` t_b), arr_a] t_b), + ("reduce_stream", + IntrinsicPolyFun [tp_a, tp_b] + [Scalar t_b `arr` (Scalar t_b `arr` Scalar t_b), + Scalar (Prim $ Signed Int32) `arr` (arr_a `arr` Scalar t_b), + arr_a] $ + Scalar t_b), - ("stream_red_per", - IntrinsicPolyFun [tp_a, tp_b] [t_b `arr` (t_b `arr` t_b), Prim (Signed Int32) `arr` (arr_a `arr` t_b), arr_a] t_b), + ("reduce_stream_per", + IntrinsicPolyFun [tp_a, tp_b] + [Scalar t_b `arr` (Scalar t_b `arr` Scalar t_b), + Scalar (Prim $ Signed Int32) `arr` (arr_a `arr` Scalar t_b), + arr_a] $ + Scalar t_b), - ("trace", IntrinsicPolyFun [tp_a] [t_a] t_a), - ("break", IntrinsicPolyFun [tp_a] [t_a] t_a)] + ("trace", IntrinsicPolyFun [tp_a] [Scalar t_a] $ Scalar t_a), + ("break", IntrinsicPolyFun [tp_a] [Scalar t_a] $ Scalar t_a)] where tv_a = VName (nameFromString "a") 0 - tv_a' = typeName tv_a - t_a = TypeVar () Nonunique tv_a' [] - arr_a = Array () Nonunique (ArrayPolyElem tv_a' []) (rank 1) - arr_2d_a = Array () Nonunique (ArrayPolyElem tv_a' []) (rank 2) - uarr_a = Array () Unique (ArrayPolyElem tv_a' []) (rank 1) + t_a = TypeVar () Nonunique (typeName tv_a) [] + arr_a = Array () Nonunique t_a (rank 1) + arr_2d_a = Array () Nonunique t_a (rank 2) + uarr_a = Array () Unique t_a (rank 1) tp_a = TypeParamType Unlifted tv_a noLoc tv_b = VName (nameFromString "b") 1 - tv_b' = typeName tv_b - t_b = TypeVar () Nonunique tv_b' [] - arr_b = Array () Nonunique (ArrayPolyElem tv_b' []) (rank 1) - uarr_b = Array () Unique (ArrayPolyElem tv_b' []) (rank 1) + t_b = TypeVar () Nonunique (typeName tv_b) [] + arr_b = Array () Nonunique t_b (rank 1) + uarr_b = Array () Unique t_b (rank 1) tp_b = TypeParamType Unlifted tv_b noLoc arr_a_b = Array () Nonunique - (ArrayRecordElem (M.fromList $ zip tupleFieldNames - [RecordArrayElem $ ArrayPolyElem tv_a' [], - RecordArrayElem $ ArrayPolyElem tv_b' []])) + (Record (M.fromList $ zip tupleFieldNames [Scalar t_a, Scalar t_b])) (rank 1) - t_arr_a_arr_b = Record $ M.fromList $ zip tupleFieldNames [arr_a, arr_b] + t_arr_a_arr_b = Scalar $ Record $ M.fromList $ zip tupleFieldNames [arr_a, arr_b] - arr = Arrow mempty Nothing + arr x y = Scalar $ Arrow mempty Unnamed x y namify i (k,v) = (VName (nameFromString k) i, v) @@ -928,9 +856,6 @@ type UncheckedType = TypeBase (ShapeDecl Name) () type UncheckedTypeExp = TypeExp Name - --- | An array element type with no aliasing information. -type UncheckedArrayElemType = ArrayElemTypeBase (ShapeDecl Name) -- | A type declaration with no expanded type. type UncheckedTypeDecl = TypeDeclBase NoInfo Name diff --git a/src/Language/Futhark/Interpreter.hs b/src/Language/Futhark/Interpreter.hs --- a/src/Language/Futhark/Interpreter.hs +++ b/src/Language/Futhark/Interpreter.hs @@ -2,7 +2,7 @@ {-# LANGUAGE GeneralizedNewtypeDeriving #-} module Language.Futhark.Interpreter ( Ctx(..) - , Env(..) + , Env , InterpreterError , initialCtx , interpretExp @@ -27,7 +27,8 @@ import Data.List hiding (break) import Data.Maybe import qualified Data.Map as M -import Data.Monoid +import qualified Data.List.NonEmpty as NE +import Data.Monoid hiding (Sum) import Data.Loc import Language.Futhark hiding (Value) @@ -85,13 +86,13 @@ | ValueArray !(Array Int Value) | ValueRecord (M.Map Name Value) | ValueFun (Value -> EvalM Value) - | ValueEnum Name + | ValueSum Name [Value] instance Eq Value where ValuePrim x == ValuePrim y = x == y ValueArray x == ValueArray y = x == y ValueRecord x == ValueRecord y = x == y - ValueEnum x == ValueEnum y = x == y + (ValueSum n1 vs1) == (ValueSum n2 vs2) = n1 == n2 && vs1 == vs2 _ == _ = False prettyRecord :: Pretty a => M.Map Name a -> Doc @@ -114,7 +115,7 @@ ppr (ValueRecord m) = prettyRecord m ppr ValueFun{} = text "#" - ppr (ValueEnum n) = text "#" <> ppr n + ppr (ValueSum n vs) = text "#" <> ppr n <+> sep (map ppr vs) -- | Create an array value; failing if that would result in an -- irregular array. @@ -179,7 +180,7 @@ asBool :: Value -> Bool asBool (ValuePrim (BoolValue x)) = x -asBool v = error $ "Unexpectedly not an integer: " ++ pretty v +asBool v = error $ "Unexpectedly not a boolean: " ++ pretty v lookupInEnv :: (Env -> M.Map VName x) -> QualName VName -> Env -> Maybe x @@ -326,7 +327,9 @@ if v == v' then pure env else mzero - +patternMatch env (PatternConstr n _ ps _) (ValueSum n' vs) + | n == n' = + foldM (\env' (p,v) -> patternMatch env' p v) env $ zip ps vs patternMatch _ _ _ = mzero -- | For matching size annotations (the actual type will have been @@ -337,7 +340,7 @@ -> Value -> Either String Env -matchValueToType env t@(TypeVar _ _ tn []) val +matchValueToType env t@(Scalar (TypeVar _ _ tn [])) val | Just shape <- M.lookup (typeLeaf tn) $ envShapes env, shape /= valueShape val = Left $ "Value passed for type parameter `" <> prettyName (typeLeaf tn) <> @@ -357,7 +360,7 @@ | otherwise -> continue $ valEnv (M.singleton (qualLeaf v) - (Just $ T.BoundV [] $ Prim $ Signed Int32, + (Just $ T.BoundV [] $ Scalar $ Prim $ Signed Int32, ValuePrim $ SignedValue $ Int32Value arr_n)) <> env AnyDim -> continue env @@ -388,7 +391,7 @@ zeroDim AnyDim = True zeroDim (ConstDim x) = x == 0 -matchValueToType env (Record fs) (ValueRecord arr) = +matchValueToType env (Scalar (Record fs)) (ValueRecord arr) = foldM (\env' (t, v) -> matchValueToType env' t v) env $ M.intersectionWith (,) fs arr @@ -396,7 +399,7 @@ bindToZero :: [VName] -> Env bindToZero = valEnv . M.fromList . map f - where f v = (v, (Just $ T.BoundV [] $ Prim $ Signed Int32, + where f v = (v, (Just $ T.BoundV [] $ Scalar $ Prim $ Signed Int32, ValuePrim $ SignedValue $ Int32Value 0)) data Indexing = IndexingFix Int32 @@ -504,22 +507,21 @@ -- | Expand type based on information that was not available at -- type-checking time (the structure of abstract types). evalType :: Env -> StructType -> StructType -evalType _ (Prim pt) = Prim pt -evalType env (Record fs) = Record $ fmap (evalType env) fs -evalType env (Arrow () p t1 t2) = - Arrow () p (evalType env t1) (evalType env t2) +evalType _ (Scalar (Prim pt)) = Scalar $ Prim pt +evalType env (Scalar (Record fs)) = Scalar $ Record $ fmap (evalType env) fs +evalType env (Scalar (Arrow () p t1 t2)) = + Scalar $ Arrow () p (evalType env t1) (evalType env t2) evalType env t@(Array _ u _ shape) = let et = stripArray (shapeRank shape) t et' = evalType env et shape' = fmap evalDim shape - in fromMaybe (error "Cannot construct array after substitution") $ - arrayOf et' shape' u + in arrayOf et' shape' u where evalDim (NamedDim qn) | Just (TermValue _ (ValuePrim (SignedValue (Int32Value x)))) <- lookupVar qn env = ConstDim $ fromIntegral x evalDim d = d -evalType env t@(TypeVar () _ tn args) = +evalType env t@(Scalar (TypeVar () _ tn args)) = case lookupType (qualNameFromTypeName tn) env of Just (T.TypeAbbr _ ps t') -> let (substs, types) = mconcat $ zipWith matchPtoA ps args @@ -537,7 +539,7 @@ let t'' = evalType env t' in (mempty, M.singleton p $ T.TypeAbbr l [] t'') matchPtoA _ _ = mempty -evalType _ (Enum cs) = Enum cs +evalType env (Scalar (Sum cs)) = Scalar $ Sum $ (fmap . fmap) (evalType env) cs evalFunction :: Env -> [TypeParam] -> [Pattern] -> Exp -> (Aliasing, StructType) -> SrcLoc -> EvalM Value @@ -554,7 +556,7 @@ let unbound_dims = bindToZero $ map typeParamName $ filter isDimParam tparams v <- eval (env <> unbound_dims) body case (t, v) of - (Arrow _ _ _ rt, ValueFun f) -> + (Scalar (Arrow _ _ _ rt), ValueFun f) -> return $ ValueFun $ \arg -> do r <- f arg match (evalType env rt) r _ -> match t v @@ -610,9 +612,9 @@ where toInt = case stripArray 1 t of - Prim (Signed t') -> + Scalar (Prim (Signed t')) -> ValuePrim . SignedValue . intValue t' - Prim (Unsigned t') -> + Scalar (Prim (Unsigned t')) -> ValuePrim . UnsignedValue . intValue t' _ -> error $ "Nonsensical range type: " ++ show t @@ -633,22 +635,23 @@ eval env (LetFun f (tparams, pats, _, Info ret, fbody) body loc) = do v <- evalFunction env tparams pats fbody (mempty, ret) loc - let ftype = T.BoundV [] $ foldr (uncurry (Arrow ()) . patternParam) ret pats + let arrow (xp, xt) yt = Scalar $ Arrow () xp xt yt + ftype = T.BoundV [] $ foldr (arrow . patternParam) ret pats eval (valEnv (M.singleton f (Just ftype, v)) <> env) body eval _ (IntLit v (Info t) _) = case t of - Prim (Signed it) -> + Scalar (Prim (Signed it)) -> return $ ValuePrim $ SignedValue $ intValue it v - Prim (Unsigned it) -> + Scalar (Prim (Unsigned it)) -> return $ ValuePrim $ UnsignedValue $ intValue it v - Prim (FloatType ft) -> + Scalar (Prim (FloatType ft)) -> return $ ValuePrim $ FloatValue $ floatValue ft v _ -> error $ "eval: nonsensical type for integer literal: " ++ pretty t eval _ (FloatLit v (Info t) _) = case t of - Prim (FloatType ft) -> + Scalar (Prim (FloatType ft)) -> return $ ValuePrim $ FloatValue $ floatValue ft v _ -> error $ "eval: nonsensical type for float literal: " ++ pretty t @@ -766,7 +769,7 @@ | otherwise = do env' <- withLoopParams v forLoop iv bound (inc i) =<< - eval (valEnv (M.singleton iv (Just $ T.BoundV [] $ Prim $ Signed Int32, + eval (valEnv (M.singleton iv (Just $ T.BoundV [] $ Scalar $ Prim $ Signed Int32, ValuePrim (SignedValue i))) <> env') body whileLoop cond v = do @@ -794,11 +797,13 @@ unless cond $ bad loc env s eval env e -eval _ (VConstr0 c _ _) = return $ ValueEnum c +eval env (Constr c es _ _) = do + vs <- mapM (eval env) es + return $ ValueSum c vs eval env (Match e cs _ _) = do v <- eval env e - match v cs + match v $ NE.toList cs where match _ [] = fail "Pattern match failure." match v (c:cs') = do @@ -880,7 +885,8 @@ evalDec env (ValDec (ValBind _ v _ (Info t) tps ps def _ loc)) = do let t' = evalType env t - ftype = T.BoundV [] $ foldr (uncurry (Arrow ()) . patternParam) t' ps + arrow (xp, xt) yt = Scalar $ Arrow () xp xt yt + ftype = T.BoundV [] $ foldr (arrow . patternParam) t' ps val <- evalFunction env tps ps def (mempty, t') loc return $ valEnv (M.singleton v (Just ftype, val)) <> env @@ -973,6 +979,8 @@ getB (BoolValue x) = Just $ P.BoolValue x getB _ = Nothing + putB (P.BoolValue x) = Just $ BoolValue x + putB _ = Nothing fun1 f = TermValue Nothing $ ValueFun $ \x -> f x @@ -1017,15 +1025,15 @@ x' <- valf x retf =<< op x' - def "~" = Just $ unopDef [ (getS, putS, P.doUnOp $ P.Complement Int8) + def "!" = Just $ unopDef [ (getS, putS, P.doUnOp $ P.Complement Int8) , (getS, putS, P.doUnOp $ P.Complement Int16) , (getS, putS, P.doUnOp $ P.Complement Int32) , (getS, putS, P.doUnOp $ P.Complement Int64) , (getU, putU, P.doUnOp $ P.Complement Int8) , (getU, putU, P.doUnOp $ P.Complement Int16) , (getU, putU, P.doUnOp $ P.Complement Int32) - , (getU, putU, P.doUnOp $ P.Complement Int64)] - def "!" = Just $ fun1 $ return . ValuePrim . BoolValue . not . asBool + , (getU, putU, P.doUnOp $ P.Complement Int64) + , (getB, putB, P.doUnOp P.Not) ] def "+" = arithOp P.Add P.FAdd def "-" = arithOp P.Sub P.FSub @@ -1091,6 +1099,12 @@ _ -> error $ "Cannot unsign: " ++ pretty x where bool = Just . BoolValue + def s | "map_stream" `isPrefixOf` s = + Just $ fun2t stream + + def s | "reduce_stream" `isPrefixOf` s = + Just $ fun3t $ \_ f arg -> stream f arg + def "map" = Just $ fun2t $ \f xs -> toArray =<< mapM (apply noLoc mempty f) (fromArray xs) @@ -1103,12 +1117,6 @@ return (x':out, x') toArray . reverse . fst =<< foldM next ([], ne) (fromArray xs) - def s | "stream_map" `isPrefixOf` s = - Just $ fun2t stream - - def s | "stream_red" `isPrefixOf` s = - Just $ fun3t $ \_ f arg -> stream f arg - def "scatter" = Just $ fun3t $ \arr is vs -> case arr of ValueArray arr' -> @@ -1196,7 +1204,7 @@ tdef s = do t <- nameFromString s `M.lookup` namesToPrimTypes - return $ T.TypeAbbr Unlifted [] $ Prim t + return $ T.TypeAbbr Unlifted [] $ Scalar $ Prim t stream f arg@(ValueArray xs) = let n = ValuePrim $ SignedValue $ Int32Value $ arrayLength xs diff --git a/src/Language/Futhark/Parser/Lexer.x b/src/Language/Futhark/Parser/Lexer.x --- a/src/Language/Futhark/Parser/Lexer.x +++ b/src/Language/Futhark/Parser/Lexer.x @@ -47,7 +47,7 @@ @identifier = [a-zA-Z] [a-zA-Z0-9_']* | "_" [a-zA-Z0-9] [a-zA-Z0-9_']* @qualidentifier = (@identifier ".")+ @identifier -@unop = ("!"|"~") +@unop = "!" @qualunop = (@identifier ".")+ @unop @opchar = ("+"|"-"|"*"|"/"|"%"|"="|"!"|">"|"<"|"|"|"&"|"^"|".") diff --git a/src/Language/Futhark/Parser/Parser.y b/src/Language/Futhark/Parser/Parser.y --- a/src/Language/Futhark/Parser/Parser.y +++ b/src/Language/Futhark/Parser/Parser.y @@ -29,6 +29,7 @@ import Data.Char (ord) import Data.Maybe (fromMaybe, fromJust) import Data.Loc hiding (L) -- Lexer has replacements. +import qualified Data.List.NonEmpty as NE import qualified Data.Map.Strict as M import Data.Monoid @@ -162,9 +163,9 @@ doc { L _ (DOC _) } %left bottom -%left ifprec letprec unsafe caseprec typeprec enumprec -%left ',' case -%left ':' +%left ifprec letprec unsafe caseprec typeprec enumprec sumprec +%left ',' case id constructor '(' '{' +%right ':' %right '...' '..<' '..>' '..' %left '`' %right '->' @@ -424,7 +425,7 @@ { TEArray $4 (fst $2) (srcspan $1 $>) } | '[' ']' TypeExpTerm %prec indexprec { TEArray $3 AnyDim (srcspan $1 $>) } - | TypeExpApply { $1 } + | TypeExpApply %prec sumprec { $1 } -- Errors | '[' DimDecl ']' %prec bottom @@ -433,6 +434,22 @@ "Did you mean []" ++ pretty (fst $2) ++ "?"] } +SumType :: { UncheckedTypeExp } +SumType : SumClauses %prec sumprec { let (cs, loc) = $1 + in TESum cs loc } + +SumClauses :: { ([(Name, [UncheckedTypeExp])], SrcLoc) } + : SumClauses '|' SumClause %prec sumprec { let (cs, loc1) = $1; + (c, ts, loc2) = $3 + in (cs++[(c, ts)], srcspan loc1 loc2) } + | SumClause %prec sumprec { let (n, ts, loc) = $1 + in ([(n, ts)], loc) } + +SumClause :: { (Name, [UncheckedTypeExp], SrcLoc) } + : SumClause TypeExpAtom { let (n, ts, loc) = $1 + in (n, ts ++ [$2], srcspan loc $>)} + | Constr { (fst $1, [], snd $1) } + TypeExpApply :: { UncheckedTypeExp } : TypeExpApply TypeArg { TEApply $1 $2 (srcspan $1 $>) } @@ -452,16 +469,10 @@ | '{' '}' { TERecord [] (srcspan $1 $>) } | '{' FieldTypes1 '}' { TERecord $2 (srcspan $1 $>) } | QualName { TEVar (fst $1) (snd $1) } - | Enum { TEEnum (fst $1) (snd $1)} - -Enum :: { ([Name], SrcLoc) } - : VConstr0 %prec enumprec { ([fst $1], snd $1) } - | VConstr0 '|' Enum - { let names = fst $1 : fst $3; loc = srcspan (snd $1) (snd $3) - in (names, loc) } + | SumType { $1 } -VConstr0 :: { (Name, SrcLoc) } - : constructor { let L _ (CONSTRUCTOR c) = $1 in (c, srclocOf $1) } +Constr :: { (Name, SrcLoc) } + : constructor { let L _ (CONSTRUCTOR c) = $1 in (c, srclocOf $1) } TypeArg :: { TypeArgExp Name } : '[' DimDecl ']' { TypeArgExpDim (fst $2) (srcspan $1 $>) } @@ -577,22 +588,27 @@ | Exp2 with FieldAccesses_ '=' Exp2 { RecordUpdate $1 (map fst $3) $5 NoInfo (srcspan $1 $>) } - | '\\' FunParams1 maybeAscription(TypeExpTerm) '->' Exp + | '\\' FunParams1 maybeAscription(TypeExpTerm) '->' Exp %prec letprec { Lambda (fst $2 : snd $2) $5 $3 NoInfo (srcspan $1 $>) } - | Apply { $1 } + | Apply_ { $1 } -Apply :: { UncheckedExp } - : Apply Atom %prec juxtprec - { Apply $1 $2 NoInfo NoInfo (srcspan $1 $>) } - | UnOp Atom %prec juxtprec - { Apply (Var (fst $1) NoInfo (snd $1)) $2 NoInfo NoInfo (srcspan (snd $1) $>) } - | Atom %prec juxtprec - { $1 } +Apply_ :: { UncheckedExp } + : ApplyList { case $1 of + ((Constr n [] _ loc1):_) -> Constr n (tail $1) NoInfo (srcspan loc1 (last $1)) + _ -> foldl1 (\f x -> Apply f x NoInfo NoInfo (srcspan f x)) $1 } +ApplyList :: { [UncheckedExp] } + : ApplyList Atom %prec juxtprec + { $1 ++ [$2] } + | UnOp Atom %prec juxtprec + { [Var (fst $1) NoInfo (snd $1), $2] } + | Atom %prec juxtprec + { [$1] } + Atom :: { UncheckedExp } Atom : PrimLit { Literal (fst $1) (snd $1) } - | VConstr0 { VConstr0 (fst $1) NoInfo (snd $1) } + | Constr { Constr (fst $1) [] NoInfo (snd $1) } | intlit { let L loc (INTLIT x) = $1 in IntLit x NoInfo loc } | floatlit { let L loc (FLOATLIT x) = $1 in FloatLit x NoInfo loc } | stringlit { let L loc (STRINGLIT s) = $1 in @@ -708,12 +724,12 @@ | LetExp %prec letprec { $1 } MatchExp :: { UncheckedExp } - : match Exp Cases { let loc = srcspan $1 $> + : match Exp Cases { let loc = srcspan $1 (NE.toList $>) in Match $2 $> NoInfo loc } -Cases :: { [CaseBase NoInfo Name] } - : Case %prec caseprec { [$1] } - | Case Cases { $1 : $2 } +Cases :: { NE.NonEmpty (CaseBase NoInfo Name) } + : Case %prec caseprec { $1 NE.:| [] } + | Case Cases { NE.cons $1 $2 } Case :: { CaseBase NoInfo Name } : case CPattern '->' Exp { let loc = srcspan $1 $> @@ -722,6 +738,9 @@ CPattern :: { PatternBase NoInfo Name } : CInnerPattern ':' TypeExpDecl { PatternAscription $1 $3 (srcspan $1 $>) } | CInnerPattern { $1 } + | Constr ConstrFields { let (n, loc) = $1; + loc' = srcspan loc $> + in PatternConstr n NoInfo $2 loc'} CPatterns1 :: { [PatternBase NoInfo Name] } : CPattern { [$1] } @@ -737,7 +756,13 @@ | '(' CPattern ',' CPatterns1 ')' { TuplePattern ($2:$4) (srcspan $1 $>) } | '{' CFieldPatterns '}' { RecordPattern $2 (srcspan $1 $>) } | CaseLiteral { PatternLit (fst $1) NoInfo (snd $1) } + | Constr { let (n, loc) = $1 + in PatternConstr n NoInfo [] loc } +ConstrFields :: { [PatternBase NoInfo Name] } + : CInnerPattern { [$1] } + | ConstrFields CInnerPattern { $1 ++ [$2] } + CFieldPattern :: { (Name, PatternBase NoInfo Name) } : FieldId '=' CPattern { (fst $1, $3) } @@ -760,7 +785,6 @@ | floatlit { let L loc (FLOATLIT x) = $1 in (FloatLit x NoInfo loc, loc) } | stringlit { let L loc (STRINGLIT s) = $1 in (ArrayLit (map (flip Literal loc . UnsignedValue . Int8Value . fromIntegral) $ encode s) NoInfo loc, loc) } - | VConstr0 { (VConstr0 (fst $1) NoInfo (snd $1), snd $1) } LoopForm :: { LoopFormBase NoInfo Name } LoopForm : for VarId '<' Exp @@ -870,7 +894,7 @@ StringValue :: { Value } StringValue : stringlit { let L pos (STRINGLIT s) = $1 in - ArrayValue (arrayFromList $ map (PrimValue . UnsignedValue . Int8Value . fromIntegral) $ encode s) $ Prim $ Signed Int32 } + ArrayValue (arrayFromList $ map (PrimValue . UnsignedValue . Int8Value . fromIntegral) $ encode s) $ Scalar $ Prim $ Signed Int32 } BoolValue :: { Value } BoolValue : true { PrimValue $ BoolValue True } @@ -913,7 +937,7 @@ Right v -> return $ ArrayValue (arrayFromList $ $2:$4) $ valueType v } | id '(' PrimType ')' - {% ($1 `mustBe` "empty") >> return (ArrayValue (listArray (0,-1) []) (Prim $3)) } + {% ($1 `mustBe` "empty") >> return (ArrayValue (listArray (0,-1) []) (Scalar (Prim $3))) } | id '(' RowType ')' {% ($1 `mustBe` "empty") >> return (ArrayValue (listArray (0,-1) []) $3) } @@ -922,8 +946,8 @@ {% emptyArrayError $1 } RowType :: { TypeBase () () } -RowType : '[' ']' RowType { fromJust $ arrayOf $3 (rank 1) Nonunique } - | '[' ']' PrimType { fromJust $ arrayOf (Prim $3) (rank 1) Nonunique } +RowType : '[' ']' RowType { arrayOf $3 (rank 1) Nonunique } + | '[' ']' PrimType { arrayOf (Scalar (Prim $3)) (rank 1) Nonunique } Values :: { [Value] } Values : Value ',' Values { $1 : $3 } diff --git a/src/Language/Futhark/Pretty.hs b/src/Language/Futhark/Pretty.hs --- a/src/Language/Futhark/Pretty.hs +++ b/src/Language/Futhark/Pretty.hs @@ -10,7 +10,7 @@ , leadingOperator , IsName(..) , prettyName - , Annot + , Annot(..) ) where @@ -19,8 +19,9 @@ import Data.Functor import qualified Data.Map.Strict as M import Data.List +import qualified Data.List.NonEmpty as NE import Data.Maybe -import Data.Monoid +import Data.Monoid hiding (Sum) import Data.Ord import Data.Word @@ -64,6 +65,7 @@ -- the prettyprinter to either print the original AST, or the computed -- attribute. class Annot f where + -- | Extract value, if any. unAnnot :: f a -> Maybe a instance Annot NoInfo where @@ -108,44 +110,34 @@ instance Pretty (ShapeDecl ()) where ppr (ShapeDecl ds) = mconcat $ replicate (length ds) $ text "[]" -instance Pretty (ShapeDecl dim) => Pretty (RecordArrayElemTypeBase dim) where - ppr (RecordArrayElem et) = ppr et - ppr (RecordArrayArrayElem et shape) = - ppr shape <> ppr et - -instance Pretty (ShapeDecl dim) => Pretty (ArrayElemTypeBase dim) where - ppr (ArrayPrimElem pt) = ppr pt - ppr (ArrayPolyElem v args) = - ppr (qualNameFromTypeName v) <+> spread (map ppr args) - ppr (ArrayRecordElem fs) - | Just ts <- areTupleFields fs = - parens (commasep $ map ppr ts) - | otherwise = - braces (commasep $ map ppField $ M.toList fs) - where ppField (name, t) = text (nameToString name) <> colon <+> ppr t - ppr (ArrayEnumElem cs) = - cat $ punctuate (text " | ") $ map ((text "#" <>) . ppr) cs - -instance Pretty (ShapeDecl dim) => Pretty (TypeBase dim as) where +instance Pretty (ShapeDecl dim) => Pretty (ScalarTypeBase dim as) where ppr = pprPrec 0 pprPrec _ (Prim et) = ppr et pprPrec _ (TypeVar _ u et targs) = ppr u <> ppr (qualNameFromTypeName et) <+> spread (map ppr targs) - pprPrec _ (Array _ u at shape) = ppr u <> ppr shape <> ppr at pprPrec _ (Record fs) | Just ts <- areTupleFields fs = parens $ commasep $ map ppr ts | otherwise = - braces $ commasep $ map ppField $ M.toList fs + oneLine (braces $ commasep fs') + <|> braces (mconcat $ punctuate (text "," <> line) fs') where ppField (name, t) = text (nameToString name) <> colon <+> ppr t - pprPrec p (Arrow _ (Just v) t1 t2) = + fs' = map ppField $ M.toList fs + pprPrec p (Arrow _ (Named v) t1 t2) = parensIf (p > 0) $ parens (pprName v <> colon <+> ppr t1) <+> text "->" <+> ppr t2 - pprPrec p (Arrow _ Nothing t1 t2) = + pprPrec p (Arrow _ Unnamed t1 t2) = parensIf (p > 0) $ pprPrec 1 t1 <+> text "->" <+> ppr t2 - pprPrec _ (Enum cs) = - cat $ punctuate (text " | ") $ map ((text "#" <>) . ppr) cs + pprPrec _ (Sum cs) = + oneLine (mconcat $ punctuate (text " | ") cs') + <|> align (mconcat $ punctuate (text " |" <> line) cs') + where ppConstr (name, fs) = sep $ (text "#" <> ppr name) : map ppr fs + cs' = map ppConstr $ M.toList cs +instance Pretty (ShapeDecl dim) => Pretty (TypeBase dim as) where + ppr (Array _ u at shape) = ppr u <> ppr shape <> ppr at + ppr (Scalar t) = ppr t + instance Pretty (ShapeDecl dim) => Pretty (TypeArg dim) where ppr (TypeArgDim d _) = ppr $ ShapeDecl [d] ppr (TypeArgType t _) = ppr t @@ -161,8 +153,9 @@ ppr (TEArrow (Just v) t1 t2 _) = parens v' <+> text "->" <+> ppr t2 where v' = pprName v <> colon <+> ppr t1 ppr (TEArrow Nothing t1 t2 _) = ppr t1 <+> text "->" <+> ppr t2 - ppr (TEEnum cs _) = - cat $ punctuate (text " | ") $ map ((text "#" <>) . ppr) cs + ppr (TESum cs _) = + align $ cat $ punctuate (text " |" <> softline) $ map ppConstr cs + where ppConstr (name, fs) = text "#" <> ppr name <+> sep (map ppr fs) instance (Eq vn, IsName vn) => Pretty (TypeArgExp vn) where ppr (TypeArgExpDim d _) = ppr $ ShapeDecl [d] @@ -246,6 +239,7 @@ LetPat{} -> True LetWith{} -> True If{} -> True + Match{} -> True ArrayLit{} -> False _ -> hasArrayLit e pprPrec _ (LetFun fname (tparams, params, retdecl, rettype, e) body _) = @@ -263,16 +257,16 @@ | otherwise = text "let" <+> ppr dest <+> equals <+> ppr src <+> text "with" <+> brackets (commasep (map ppr idxs)) <+> - text "<-" <+> align (ppr ve) + text "=" <+> align (ppr ve) letBody body pprPrec _ (Update src idxs ve _) = ppr src <+> text "with" <+> brackets (commasep (map ppr idxs)) <+> - text "<-" <+> align (ppr ve) + text "=" <+> align (ppr ve) pprPrec _ (RecordUpdate src fs ve _ _) = ppr src <+> text "with" <+> mconcat (intersperse (text ".") (map ppr fs)) <+> - text "<-" <+> align (ppr ve) + text "=" <+> align (ppr ve) pprPrec _ (Index e idxs _ _) = pprPrec 9 e <> brackets (commasep (map ppr idxs)) pprPrec _ (Unsafe e _) = text "unsafe" <+> pprPrec (-1) e @@ -296,15 +290,15 @@ text "loop" <+> ppr pat <+> equals <+> ppr initexp <+> ppr form <+> text "do" indent 2 (ppr loopbody) - pprPrec _ (VConstr0 n _ _) = text "#" <> ppr n - pprPrec _ (Match e cs _ _) = text "match" <+> ppr e ppr cs + pprPrec _ (Constr n cs _ _) = text "#" <> ppr n <+> sep (map ppr cs) + pprPrec _ (Match e cs _ _) = text "match" <+> ppr e (stack . map ppr) (NE.toList cs) instance (Eq vn, IsName vn, Annot f) => Pretty (FieldBase f vn) where ppr (RecordFieldExplicit name e _) = ppr name <> equals <> ppr e ppr (RecordFieldImplicit name _ _) = pprName name instance (Eq vn, IsName vn, Annot f) => Pretty (CaseBase f vn) where - ppr (CasePat p e _) = ppr p <+> text "->" <+> ppr e + ppr (CasePat p e _) = text "case" <+> ppr p <+> text "->" indent 2 (ppr e) instance (Eq vn, IsName vn, Annot f) => Pretty (LoopFormBase f vn) where ppr (For i ubound) = @@ -327,6 +321,7 @@ Just t' -> parens $ text "_" <> colon <+> ppr t' Nothing -> text "_" ppr (PatternLit e _ _) = ppr e + ppr (PatternConstr n _ ps _) = text "#" <> ppr n <+> sep (map ppr ps) ppAscription :: Pretty t => Maybe t -> Doc ppAscription Nothing = mempty diff --git a/src/Language/Futhark/Syntax.hs b/src/Language/Futhark/Syntax.hs --- a/src/Language/Futhark/Syntax.hs +++ b/src/Language/Futhark/Syntax.hs @@ -29,8 +29,8 @@ , TypeArg(..) , TypeExp(..) , TypeArgExp(..) - , RecordArrayElemTypeBase(..) - , ArrayElemTypeBase(..) + , PName(..) + , ScalarTypeBase(..) , PatternType , StructType , Diet(..) @@ -74,6 +74,7 @@ , DecBase(..) -- * Miscellaneous + , Showable , NoInfo(..) , Info(..) , Alias(..) @@ -91,10 +92,11 @@ import Data.Foldable import Data.Loc import qualified Data.Map.Strict as M -import Data.Monoid +import Data.Monoid hiding (Sum) import Data.Ord import qualified Data.Set as S import Data.Traversable +import qualified Data.List.NonEmpty as NE import Data.List import Prelude @@ -115,7 +117,6 @@ Show (f StructType), Show (f (Aliasing, StructType)), Show (f (M.Map VName VName)), - Show (f [RecordArrayElemTypeBase ()]), Show (f Uniqueness)) => Showable f vn where -- | No information functor. Usually used for placeholder type- or @@ -286,80 +287,60 @@ qualNameFromTypeName :: TypeName -> QualName VName qualNameFromTypeName (TypeName qs x) = QualName qs x --- | Types that can be elements of tuple-arrays. -data RecordArrayElemTypeBase dim = - RecordArrayElem (ArrayElemTypeBase dim) - | RecordArrayArrayElem (ArrayElemTypeBase dim) (ShapeDecl dim) - deriving (Eq, Show) - -instance Traversable RecordArrayElemTypeBase where - traverse f (RecordArrayElem t) = RecordArrayElem <$> traverse f t - traverse f (RecordArrayArrayElem a shape) = - RecordArrayArrayElem <$> traverse f a <*> traverse f shape - -instance Functor RecordArrayElemTypeBase where - fmap = fmapDefault - -instance Foldable RecordArrayElemTypeBase where - foldMap = foldMapDefault - -data ArrayElemTypeBase dim = - ArrayPrimElem PrimType - | ArrayPolyElem TypeName [TypeArg dim] - | ArrayRecordElem (M.Map Name (RecordArrayElemTypeBase dim)) - | ArrayEnumElem [Name] - deriving (Eq, Show) - -instance Traversable ArrayElemTypeBase where - traverse _ (ArrayPrimElem t) = - pure $ ArrayPrimElem t - traverse f (ArrayPolyElem t args) = - ArrayPolyElem t <$> traverse (traverse f) args - traverse f (ArrayRecordElem fs) = - ArrayRecordElem <$> traverse (traverse f) fs - traverse _ (ArrayEnumElem cs) = - pure $ ArrayEnumElem cs - -instance Functor ArrayElemTypeBase where - fmap = fmapDefault +-- | The name (if any) of a function parameter. The 'Eq' and 'Ord' +-- instances always compare values of this type equal. +data PName = Named VName | Unnamed + deriving (Show) -instance Foldable ArrayElemTypeBase where - foldMap = foldMapDefault +instance Eq PName where + _ == _ = True --- | An expanded Futhark type is either an array, a prim type, a --- tuple, or a type variable. When comparing types for equality with --- '==', aliases are ignored, but dimensions much match. Function --- parameter names are ignored. -data TypeBase dim as = Prim PrimType - | Enum [Name] - | Array as Uniqueness (ArrayElemTypeBase dim) (ShapeDecl dim) - | Record (M.Map Name (TypeBase dim as)) - | TypeVar as Uniqueness TypeName [TypeArg dim] - | Arrow as (Maybe VName) (TypeBase dim as) (TypeBase dim as) - -- ^ The aliasing corresponds to the lexical - -- closure of the function. - deriving (Show) +instance Ord PName where + _ <= _ = True -instance (Eq dim, Eq as) => Eq (TypeBase dim as) where - Prim x1 == Prim y1 = x1 == y1 - Array x1 y1 z1 v1 == Array x2 y2 z2 v2 = x1 == x2 && y1 == y2 && z1 == z2 && v1 == v2 - Record x1 == Record x2 = x1 == x2 - TypeVar _ u1 x1 y1 == TypeVar _ u2 x2 y2 = u1 == u2 && x1 == x2 && y1 == y2 - Arrow _ _ x1 y1 == Arrow _ _ x2 y2 = x1 == x2 && y1 == y2 - Enum ns1 == Enum ns2 = sort ns1 == sort ns2 - _ == _ = False +-- | Types that can be elements of arrays. This representation does +-- allow arrays of records of functions, which is nonsensical, but it +-- convolutes the code too much if we try to statically rule it out. +data ScalarTypeBase dim as + = Prim PrimType + | TypeVar as Uniqueness TypeName [TypeArg dim] + | Record (M.Map Name (TypeBase dim as)) + | Sum (M.Map Name [TypeBase dim as]) + | Arrow as PName (TypeBase dim as) (TypeBase dim as) + -- ^ The aliasing corresponds to the lexical + -- closure of the function. + deriving (Eq, Ord, Show) -instance Bitraversable TypeBase where +instance Bitraversable ScalarTypeBase where bitraverse _ _ (Prim t) = pure $ Prim t - bitraverse f g (Array a u t shape) = - Array <$> g a <*> pure u <*> traverse f t <*> traverse f shape bitraverse f g (Record fs) = Record <$> traverse (bitraverse f g) fs bitraverse f g (TypeVar als u t args) = TypeVar <$> g als <*> pure u <*> pure t <*> traverse (traverse f) args bitraverse f g (Arrow als v t1 t2) = Arrow <$> g als <*> pure v <*> bitraverse f g t1 <*> bitraverse f g t2 - bitraverse _ _ (Enum n) = pure $ Enum n + bitraverse f g (Sum cs) = Sum <$> (traverse . traverse) (bitraverse f g) cs +instance Bifunctor ScalarTypeBase where + bimap = bimapDefault + +instance Bifoldable ScalarTypeBase where + bifoldMap = bifoldMapDefault + +-- | An expanded Futhark type is either an array, or something that +-- can be an element of an array. When comparing types for equality, +-- function parameter names are ignored. This representation permits +-- some malformed types (arrays of functions), but importantly rules +-- out arrays-of-arrays. +data TypeBase dim as + = Scalar (ScalarTypeBase dim as) + | Array as Uniqueness (ScalarTypeBase dim ()) (ShapeDecl dim) + deriving (Eq, Ord, Show) + +instance Bitraversable TypeBase where + bitraverse f g (Scalar t) = Scalar <$> bitraverse f g t + bitraverse f g (Array a u t shape) = + Array <$> g a <*> pure u <*> bitraverse f pure t <*> traverse f shape + instance Bifunctor TypeBase where bimap = bimapDefault @@ -368,7 +349,7 @@ data TypeArg dim = TypeArgDim dim SrcLoc | TypeArgType (TypeBase dim ()) SrcLoc - deriving (Eq, Show) + deriving (Eq, Ord, Show) instance Traversable TypeArg where traverse f (TypeArgDim v loc) = TypeArgDim <$> f v <*> pure loc @@ -409,7 +390,7 @@ | TEUnique (TypeExp vn) SrcLoc | TEApply (TypeExp vn) (TypeArgExp vn) SrcLoc | TEArrow (Maybe vn) (TypeExp vn) (TypeExp vn) SrcLoc - | TEEnum [Name] SrcLoc + | TESum [(Name, [TypeExp vn])] SrcLoc deriving (Show) deriving instance Eq (TypeExp Name) deriving instance Eq (TypeExp VName) @@ -422,7 +403,7 @@ locOf (TEUnique _ loc) = locOf loc locOf (TEApply _ _ loc) = locOf loc locOf (TEArrow _ _ _ loc) = locOf loc - locOf (TEEnum _ loc) = locOf loc + locOf (TESum _ loc) = locOf loc data TypeArgExp vn = TypeArgExpDim (DimDecl vn) SrcLoc | TypeArgExpType (TypeExp vn) @@ -677,10 +658,10 @@ -- and return the value of the second expression if it -- does. - | VConstr0 Name (f PatternType) SrcLoc - -- ^ An enum element, e.g., @#foo@. + | Constr Name [ExpBase f vn] (f PatternType) SrcLoc + -- ^ An n-ary value constructor. - | Match (ExpBase f vn) [CaseBase f vn] (f PatternType) SrcLoc + | Match (ExpBase f vn) (NE.NonEmpty (CaseBase f vn)) (f PatternType) SrcLoc -- ^ A match expression. deriving instance Showable f vn => Show (ExpBase f vn) @@ -717,7 +698,7 @@ locOf (DoLoop _ _ _ _ pos) = locOf pos locOf (Unsafe _ loc) = locOf loc locOf (Assert _ _ _ loc) = locOf loc - locOf (VConstr0 _ _ loc) = locOf loc + locOf (Constr _ _ _ loc) = locOf loc locOf (Match _ _ _ loc) = locOf loc -- | An entry in a record literal. @@ -753,6 +734,7 @@ | Wildcard (f PatternType) SrcLoc -- Nothing, i.e. underscore. | PatternAscription (PatternBase f vn) (TypeDeclBase f vn) SrcLoc | PatternLit (ExpBase f vn) (f PatternType) SrcLoc + | PatternConstr Name (f PatternType) [PatternBase f vn] SrcLoc deriving instance Showable f vn => Show (PatternBase f vn) instance Located (PatternBase f vn) where @@ -763,6 +745,7 @@ locOf (Wildcard _ loc) = locOf loc locOf (PatternAscription _ _ loc) = locOf loc locOf (PatternLit _ _ loc) = locOf loc + locOf (PatternConstr _ _ _ loc) = locOf loc -- | Documentation strings, including source location. data DocComment = DocComment String SrcLoc diff --git a/src/Language/Futhark/Traversals.hs b/src/Language/Futhark/Traversals.hs --- a/src/Language/Futhark/Traversals.hs +++ b/src/Language/Futhark/Traversals.hs @@ -25,6 +25,7 @@ ) where import qualified Data.Set as S +import qualified Data.List.NonEmpty as NE import Language.Futhark.Syntax @@ -96,10 +97,10 @@ mapOnExp tv e) <*> mapOnExp tv body <*> pure loc astMap tv (LetWith dest src idxexps vexp body t loc) = - pure LetWith <*> - astMap tv dest <*> astMap tv src <*> - mapM (astMap tv) idxexps <*> mapOnExp tv vexp <*> - mapOnExp tv body <*> traverse (mapOnPatternType tv) t <*> pure loc + LetWith <$> + astMap tv dest <*> astMap tv src <*> + mapM (astMap tv) idxexps <*> mapOnExp tv vexp <*> + mapOnExp tv body <*> traverse (mapOnPatternType tv) t <*> pure loc astMap tv (Update src slice v loc) = Update <$> mapOnExp tv src <*> mapM (astMap tv) slice <*> mapOnExp tv v <*> pure loc @@ -109,11 +110,11 @@ astMap tv (Project field e t loc) = Project field <$> mapOnExp tv e <*> traverse (mapOnPatternType tv) t <*> pure loc astMap tv (Index arr idxexps t loc) = - pure Index <*> - astMap tv arr <*> - mapM (astMap tv) idxexps <*> - traverse (mapOnPatternType tv) t <*> - pure loc + Index <$> + astMap tv arr <*> + mapM (astMap tv) idxexps <*> + traverse (mapOnPatternType tv) t <*> + pure loc astMap tv (Unsafe e loc) = Unsafe <$> mapOnExp tv e <*> pure loc astMap tv (Assert e1 e2 desc loc) = @@ -146,8 +147,8 @@ DoLoop <$> astMap tv mergepat <*> mapOnExp tv mergeexp <*> astMap tv form <*> mapOnExp tv loopbody <*> pure loc - astMap tv (VConstr0 name t loc) = - VConstr0 name <$> traverse (mapOnPatternType tv) t <*> pure loc + astMap tv (Constr name es ts loc) = + Constr name <$> traverse (mapOnExp tv) es <*> traverse (mapOnPatternType tv) ts <*> pure loc astMap tv (Match e cases t loc) = Match <$> mapOnExp tv e <*> astMap tv cases <*> traverse (mapOnPatternType tv) t <*> pure loc @@ -169,7 +170,8 @@ TEApply <$> astMap tv t1 <*> astMap tv t2 <*> pure loc astMap tv (TEArrow v t1 t2 loc) = TEArrow v <$> astMap tv t1 <*> astMap tv t2 <*> pure loc - astMap _ te@TEEnum{} = pure te + astMap tv (TESum cs loc) = + TESum <$> traverse (traverse $ astMap tv) cs <*> pure loc instance ASTMappable (TypeArgExp VName) where astMap tv (TypeArgExpDim dim loc) = @@ -204,37 +206,22 @@ (TypeName -> f TypeName) -> (dim1 -> f dim2) -> (als1 -> f als2) -> t dim1 als1 -> f (t dim2 als2) -traverseType :: Applicative f => - TypeTraverser f TypeBase dim1 als1 dims als2 -traverseType _ _ _ (Prim t) = pure $ Prim t -traverseType f g h (Array als u et shape) = - Array <$> h als <*> pure u <*> traverseArrayElemType f g et <*> traverse g shape -traverseType f g h (Record fs) = Record <$> traverse (traverseType f g h) fs -traverseType f g h (TypeVar als u t args) = +traverseScalarType :: Applicative f => + TypeTraverser f ScalarTypeBase dim1 als1 dims als2 +traverseScalarType _ _ _ (Prim t) = pure $ Prim t +traverseScalarType f g h (Record fs) = Record <$> traverse (traverseType f g h) fs +traverseScalarType f g h (TypeVar als u t args) = TypeVar <$> h als <*> pure u <*> f t <*> traverse (traverseTypeArg f g) args -traverseType f g h (Arrow als v t1 t2) = +traverseScalarType f g h (Arrow als v t1 t2) = Arrow <$> h als <*> pure v <*> traverseType f g h t1 <*> traverseType f g h t2 -traverseType _ _ _ (Enum cs) = pure $ Enum cs - -traverseArrayElemType :: Applicative f => - (TypeName -> f TypeName) -> (dim1 -> f dim2) - -> ArrayElemTypeBase dim1 -> f (ArrayElemTypeBase dim2) -traverseArrayElemType _ _ (ArrayPrimElem t) = - pure $ ArrayPrimElem t -traverseArrayElemType f g (ArrayPolyElem t args) = - ArrayPolyElem <$> f t <*> traverse (traverseTypeArg f g) args -traverseArrayElemType f g (ArrayRecordElem fs) = - ArrayRecordElem <$> traverse (traverseRecordArrayElemType f g) fs -traverseArrayElemType _ _ (ArrayEnumElem cs) = - pure $ ArrayEnumElem cs +traverseScalarType f g h (Sum cs) = Sum <$> (traverse . traverse) (traverseType f g h) cs -traverseRecordArrayElemType :: Applicative f => - (TypeName -> f TypeName) -> (dim1 -> f dim2) - -> RecordArrayElemTypeBase dim1 -> f (RecordArrayElemTypeBase dim2) -traverseRecordArrayElemType f g (RecordArrayElem et) = - RecordArrayElem <$> traverseArrayElemType f g et -traverseRecordArrayElemType f g (RecordArrayArrayElem et shape) = - RecordArrayArrayElem <$> traverseArrayElemType f g et <*> traverse g shape +traverseType :: Applicative f => + TypeTraverser f TypeBase dim1 als1 dims als2 +traverseType f g h (Array als u et shape) = + Array <$> h als <*> pure u <*> traverseScalarType f g pure et <*> traverse g shape +traverseType f g h (Scalar t) = + Scalar <$> traverseScalarType f g h t traverseTypeArg :: Applicative f => (TypeName -> f TypeName) -> (dim1 -> f dim2) @@ -273,6 +260,8 @@ Wildcard <$> (Info <$> mapOnPatternType tv t) <*> pure loc astMap tv (PatternLit e (Info t) loc) = PatternLit <$> astMap tv e <*> (Info <$> mapOnPatternType tv t) <*> pure loc + astMap tv (PatternConstr n (Info t) ps loc) = + PatternConstr n <$> (Info <$> mapOnPatternType tv t) <*> mapM (astMap tv) ps <*> pure loc instance ASTMappable (FieldBase Info VName) where astMap tv (RecordFieldExplicit name e loc) = @@ -289,6 +278,9 @@ astMap tv = traverse $ astMap tv instance ASTMappable a => ASTMappable [a] where + astMap tv = traverse $ astMap tv + +instance ASTMappable a => ASTMappable (NE.NonEmpty a) where astMap tv = traverse $ astMap tv instance (ASTMappable a, ASTMappable b) => ASTMappable (a,b) where diff --git a/src/Language/Futhark/TypeChecker.hs b/src/Language/Futhark/TypeChecker.hs --- a/src/Language/Futhark/TypeChecker.hs +++ b/src/Language/Futhark/TypeChecker.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE FlexibleContexts, TupleSections #-} +{-# LANGUAGE FlexibleContexts #-} -- | The type checker checks whether the program is type-consistent -- and adds type annotations and various other elaborations. The -- program does not need to have any particular properties for the @@ -16,14 +16,13 @@ where import Control.Monad.Except -import Control.Monad.Writer +import Control.Monad.Writer hiding (Sum) import Data.List import Data.Loc import Data.Maybe import Data.Either import Data.Ord import qualified Data.Map.Strict as M -import qualified Data.Set as S import Prelude hiding (abs, mod) @@ -31,8 +30,8 @@ import Language.Futhark.Semantic import Futhark.FreshNames hiding (newName) import Language.Futhark.TypeChecker.Monad +import Language.Futhark.TypeChecker.Modules import Language.Futhark.TypeChecker.Terms -import Language.Futhark.TypeChecker.Unify (doUnification) import Language.Futhark.TypeChecker.Types --- The main checker @@ -115,7 +114,7 @@ intrinsicsModule = Env mempty initialTypeTable mempty mempty intrinsicsNameMap addIntrinsicT (name, IntrinsicType t) = - Just (name, TypeAbbr Unlifted [] $ Prim t) + Just (name, TypeAbbr Unlifted [] $ Scalar $ Prim t) addIntrinsicT _ = Nothing @@ -163,10 +162,11 @@ typeParamEnv (TypeParamDim v _) = mempty { envVtable = - M.singleton v $ BoundV [] (Prim (Signed Int32)) } + M.singleton v $ BoundV [] (Scalar $ Prim $ Signed Int32) } typeParamEnv (TypeParamType l v _) = mempty { envTypeTable = - M.singleton v $ TypeAbbr l [] $ TypeVar () Nonunique (typeName v) [] } + M.singleton v $ TypeAbbr l [] $ + Scalar $ TypeVar () Nonunique (typeName v) [] } -- In this function, after the recursion, we add the Env of the -- current Spec *after* the one that is returned from the recursive @@ -213,7 +213,8 @@ M.singleton (Type, name) $ qualName name' , envTypeTable = M.singleton name' $ TypeAbbr l ps' $ - TypeVar () Nonunique (typeName name') $ map typeParamToArg ps' + Scalar $ TypeVar () Nonunique (typeName name') $ + map typeParamToArg ps' } (abstypes, env, specs') <- localEnv tenv $ checkSpecs specs return (M.insert (qualName name') l abstypes, @@ -378,25 +379,6 @@ fsig_subst <- badOnLeft $ matchMTys body_mty fsig_mty loc return (Just (fsig_e', Info fsig_subst), body_e', fsig_mty) -applyFunctor :: SrcLoc - -> FunSig - -> MTy - -> TypeM (MTy, - M.Map VName VName, - M.Map VName VName) -applyFunctor applyloc (FunSig p_abs p_mod body_mty) a_mty = do - p_subst <- badOnLeft $ matchMTys a_mty (MTy p_abs p_mod) applyloc - - -- Apply type abbreviations from a_mty to body_mty. - let a_abbrs = mtyTypeAbbrs a_mty - isSub v = case M.lookup v a_abbrs of - Just abbr -> Just $ TypeSub abbr - _ -> Just $ DimSub $ NamedDim $ qualName v - type_subst = M.mapMaybe isSub p_subst - body_mty' = substituteTypesInMTy type_subst body_mty - (body_mty'', body_subst) <- newNamesForMTy body_mty' - return (body_mty'', p_subst, body_subst) - checkModBind :: ModBindBase NoInfo Name -> TypeM (TySet, Env, ModBindBase Info VName) checkModBind (ModBind name [] maybe_fsig_e e doc loc) = do (maybe_fsig_e', e', mty) <- checkModBody (fst <$> maybe_fsig_e) e loc @@ -496,21 +478,22 @@ _ -> return () + let arrow (xp, xt) yt = Scalar $ Arrow () xp xt yt return (mempty { envVtable = M.singleton fname' $ - BoundV tparams' $ foldr (uncurry (Arrow ()) . patternParam) rettype params' + BoundV tparams' $ foldr (arrow . patternParam) rettype params' , envNameMap = M.singleton (Term, fname) $ qualName fname' }, ValBind entry' fname' maybe_tdecl' (Info rettype) tparams' params' body' doc loc) nastyType :: Monoid als => TypeBase dim als -> Bool -nastyType Prim{} = False +nastyType (Scalar Prim{}) = False nastyType t@Array{} = nastyType $ stripArray 1 t nastyType _ = True nastyReturnType :: Monoid als => Maybe (TypeExp VName) -> TypeBase dim als -> Bool -nastyReturnType _ (Arrow _ _ t1 t2) = +nastyReturnType _ (Scalar (Arrow _ _ t1 t2)) = nastyType t1 || nastyReturnType Nothing t2 nastyReturnType (Just te) _ | niceTypeExp te = False @@ -583,402 +566,3 @@ checkDecs [] = return (mempty, mempty, []) - ---- Signature matching - --- Return new renamed/abstracted env, as well as a mapping from --- names in the signature to names in the new env. This is used for --- functor application. The first env is the module env, and the --- second the env it must match. -matchMTys :: MTy -> MTy -> SrcLoc - -> Either TypeError (M.Map VName VName) -matchMTys = matchMTys' mempty - where - matchMTys' :: TypeSubs -> MTy -> MTy -> SrcLoc - -> Either TypeError (M.Map VName VName) - - matchMTys' _ (MTy _ ModFun{}) (MTy _ ModEnv{}) loc = - Left $ TypeError loc "Cannot match parametric module with non-paramatric module type." - - matchMTys' _ (MTy _ ModEnv{}) (MTy _ ModFun{}) loc = - Left $ TypeError loc "Cannot match non-parametric module with paramatric module type." - - matchMTys' old_abs_subst_to_type (MTy mod_abs mod) (MTy sig_abs sig) loc = do - -- Check that abstract types in 'sig' have an implementation in - -- 'mod'. This also gives us a substitution that we use to check - -- the types of values. - abs_substs <- resolveAbsTypes mod_abs mod sig_abs loc - - let abs_subst_to_type = old_abs_subst_to_type <> - M.map (TypeSub . snd) abs_substs - abs_name_substs = M.map (qualLeaf . fst) abs_substs - substs <- matchMods abs_subst_to_type mod sig loc - return (substs <> abs_name_substs) - - matchMods :: TypeSubs -> Mod -> Mod -> SrcLoc - -> Either TypeError (M.Map VName VName) - matchMods _ ModEnv{} ModFun{} loc = - Left $ TypeError loc "Cannot match non-parametric module with paramatric module type." - matchMods _ ModFun{} ModEnv{} loc = - Left $ TypeError loc "Cannot match parametric module with non-paramatric module type." - - matchMods abs_subst_to_type (ModEnv mod) (ModEnv sig) loc = - matchEnvs abs_subst_to_type mod sig loc - - matchMods old_abs_subst_to_type - (ModFun (FunSig mod_abs mod_pmod mod_mod)) - (ModFun (FunSig sig_abs sig_pmod sig_mod)) - loc = do - abs_substs <- resolveAbsTypes mod_abs mod_pmod sig_abs loc - let abs_subst_to_type = old_abs_subst_to_type <> - M.map (TypeSub . snd) abs_substs - abs_name_substs = M.map (qualLeaf . fst) abs_substs - pmod_substs <- matchMods abs_subst_to_type mod_pmod sig_pmod loc - mod_substs <- matchMTys' abs_subst_to_type mod_mod sig_mod loc - return (pmod_substs <> mod_substs <> abs_name_substs) - - matchEnvs :: TypeSubs - -> Env -> Env -> SrcLoc - -> Either TypeError (M.Map VName VName) - matchEnvs abs_subst_to_type env sig loc = do - -- XXX: we only want to create substitutions for visible names. - -- This must be wrong in some cases. Probably we need to - -- rethink how we do shadowing for module types. - let visible = S.fromList $ map qualLeaf $ M.elems $ envNameMap sig - isVisible name = name `S.member` visible - - -- Check that all values are defined correctly, substituting the - -- abstract types first. - val_substs <- fmap M.fromList $ forM (M.toList $ envVtable sig) $ \(name, spec_bv) -> do - let spec_bv' = substituteTypesInBoundV abs_subst_to_type spec_bv - case findBinding envVtable Term (baseName name) env of - Just (name', bv) -> matchVal loc name spec_bv' name' bv - _ -> missingVal loc (baseName name) - - -- Check that all type abbreviations are correctly defined. - abbr_name_substs <- fmap M.fromList $ - forM (filter (isVisible . fst) $ M.toList $ - envTypeTable sig) $ \(name, TypeAbbr _ spec_ps spec_t) -> - case findBinding envTypeTable Type (baseName name) env of - Just (name', TypeAbbr _ ps t) -> - matchTypeAbbr loc abs_subst_to_type val_substs name spec_ps spec_t name' ps t - Nothing -> missingType loc $ baseName name - - -- Check for correct modules. - mod_substs <- fmap M.unions $ forM (M.toList $ envModTable sig) $ \(name, modspec) -> - case findBinding envModTable Term (baseName name) env of - Just (name', mod) -> - M.insert name name' <$> matchMods abs_subst_to_type mod modspec loc - Nothing -> - missingMod loc $ baseName name - - return $ val_substs <> mod_substs <> abbr_name_substs - - matchTypeAbbr :: SrcLoc -> TypeSubs -> M.Map VName VName - -> VName -> [TypeParam] -> StructType - -> VName -> [TypeParam] -> StructType - -> Either TypeError (VName, VName) - matchTypeAbbr loc abs_subst_to_type val_substs spec_name spec_ps spec_t name ps t = do - -- We have to create substitutions for the type parameters, too. - unless (length spec_ps == length ps) nomatch - param_substs <- mconcat <$> zipWithM matchTypeParam spec_ps ps - let val_substs' = M.map (DimSub . NamedDim . qualName) val_substs - spec_t' = substituteTypes (val_substs'<>param_substs<>abs_subst_to_type) spec_t - if spec_t' == t - then return (spec_name, name) - else nomatch - where nomatch = mismatchedType loc (M.keys abs_subst_to_type) - (baseName spec_name) (spec_ps, spec_t) (ps, t) - - matchTypeParam (TypeParamDim x _) (TypeParamDim y _) = - pure $ M.singleton x $ DimSub $ NamedDim $ qualName y - matchTypeParam (TypeParamType Unlifted x _) (TypeParamType Unlifted y _) = - pure $ M.singleton x $ TypeSub $ TypeAbbr Unlifted [] $ - TypeVar () Nonunique (typeName y) [] - matchTypeParam (TypeParamType _ x _) (TypeParamType Lifted y _) = - pure $ M.singleton x $ TypeSub $ TypeAbbr Lifted [] $ - TypeVar () Nonunique (typeName y) [] - matchTypeParam _ _ = - nomatch - - matchVal :: SrcLoc - -> VName -> BoundV - -> VName -> BoundV - -> Either TypeError (VName, VName) - matchVal loc spec_name spec_t name t - | matchFunBinding loc spec_t t = return (spec_name, name) - matchVal loc spec_name spec_v _ v = - Left $ TypeError loc $ unlines $ - ["Module type specifies"] ++ - map (" "++) (lines $ ppValBind spec_name spec_v) ++ - ["but module provides"] ++ - map (" "++) (lines $ppValBind spec_name v) - - matchFunBinding :: SrcLoc -> BoundV -> BoundV -> Bool - matchFunBinding loc (BoundV _ orig_spec_t) (BoundV tps orig_t) = - -- Would be nice if we could propagate the actual error here. - case doUnification loc tps - (toStructural orig_spec_t) (toStructural orig_t) of - Left _ -> False - Right t -> t `subtypeOf` toStructural orig_spec_t - - missingType loc name = - Left $ TypeError loc $ - "Module does not define a type named " ++ pretty name ++ "." - - missingVal loc name = - Left $ TypeError loc $ - "Module does not define a value named " ++ pretty name ++ "." - - missingMod loc name = - Left $ TypeError loc $ - "Module does not define a module named " ++ pretty name ++ "." - - mismatchedType loc abs name spec_t env_t = - Left $ TypeError loc $ - unlines ["Module defines", - indent $ ppTypeAbbr abs name env_t, - "but module type requires", - indent $ ppTypeAbbr abs name spec_t] - - indent = intercalate "\n" . map (" "++) . lines - - resolveAbsTypes :: TySet -> Mod -> TySet -> SrcLoc - -> Either TypeError (M.Map VName (QualName VName, TypeBinding)) - resolveAbsTypes mod_abs mod sig_abs loc = do - let abs_mapping = M.fromList $ zip - (map (fmap baseName . fst) $ M.toList mod_abs) (M.toList mod_abs) - fmap M.fromList $ forM (M.toList sig_abs) $ \(name, name_l) -> - case findTypeDef (fmap baseName name) mod of - Just (name', TypeAbbr mod_l ps t) - | Unlifted <- name_l, - not (orderZero t) || mod_l == Lifted -> - mismatchedLiftedness loc (map qualLeaf $ M.keys mod_abs) name (ps, t) - | Just (abs_name, _) <- M.lookup (fmap baseName name) abs_mapping -> - return (qualLeaf name, (abs_name, TypeAbbr name_l ps t)) - | otherwise -> - return (qualLeaf name, (name', TypeAbbr name_l ps t)) - _ -> - missingType loc $ fmap baseName name - - mismatchedLiftedness loc abs name mod_t = - Left $ TypeError loc $ - unlines ["Module defines", - indent $ ppTypeAbbr abs name mod_t, - "but module type requires this type to be non-functional."] - - ppValBind v (BoundV tps t) = - unwords $ ["val", prettyName v] ++ map pretty tps ++ [":", pretty t] - - ppTypeAbbr abs name (ps, t) = - "type " ++ unwords (pretty name : map pretty ps) ++ t' - where t' = case t of - TypeVar () _ tn args - | typeLeaf tn `elem` abs, - map typeParamToArg ps == args -> "" - _ -> " = " ++ pretty t - -findBinding :: (Env -> M.Map VName v) - -> Namespace -> Name - -> Env - -> Maybe (VName, v) -findBinding table namespace name the_env = do - QualName _ name' <- M.lookup (namespace, name) $ envNameMap the_env - (name',) <$> M.lookup name' (table the_env) - -findTypeDef :: QualName Name -> Mod -> Maybe (QualName VName, TypeBinding) -findTypeDef _ ModFun{} = Nothing -findTypeDef (QualName [] name) (ModEnv the_env) = do - (name', tb) <- findBinding envTypeTable Type name the_env - return (qualName name', tb) -findTypeDef (QualName (q:qs) name) (ModEnv the_env) = do - (q', q_mod) <- findBinding envModTable Term q the_env - (QualName qs' name', tb) <- findTypeDef (QualName qs name) q_mod - return (QualName (q':qs') name', tb) - -typeParamToArg :: TypeParam -> StructTypeArg -typeParamToArg (TypeParamDim v ploc) = - TypeArgDim (NamedDim $ qualName v) ploc -typeParamToArg (TypeParamType _ v ploc) = - TypeArgType (TypeVar () Nonunique (typeName v) []) ploc - -substituteTypesInMod :: TypeSubs -> Mod -> Mod -substituteTypesInMod substs (ModEnv e) = - ModEnv $ substituteTypesInEnv substs e -substituteTypesInMod substs (ModFun (FunSig abs mod mty)) = - ModFun $ FunSig abs (substituteTypesInMod substs mod) (substituteTypesInMTy substs mty) - -substituteTypesInMTy :: TypeSubs -> MTy -> MTy -substituteTypesInMTy substs (MTy abs mod) = MTy abs $ substituteTypesInMod substs mod - -substituteTypesInEnv :: TypeSubs -> Env -> Env -substituteTypesInEnv substs env = - env { envVtable = M.map (substituteTypesInBoundV substs) $ envVtable env - , envTypeTable = M.mapWithKey subT $ envTypeTable env - , envModTable = M.map (substituteTypesInMod substs) $ envModTable env - } - where subT name _ - | Just (TypeSub (TypeAbbr l ps t)) <- M.lookup name substs = TypeAbbr l ps t - subT _ (TypeAbbr l ps t) = TypeAbbr l ps $ substituteTypes substs t - -substituteTypesInBoundV :: TypeSubs -> BoundV -> BoundV -substituteTypesInBoundV substs (BoundV tps t) = - BoundV tps (substituteTypes substs t) - -allNamesInMTy :: MTy -> S.Set VName -allNamesInMTy (MTy abs mod) = - S.fromList (map qualLeaf $ M.keys abs) <> allNamesInMod mod - -allNamesInMod :: Mod -> S.Set VName -allNamesInMod (ModEnv env) = allNamesInEnv env -allNamesInMod ModFun{} = mempty - --- All names defined anywhere in the env. -allNamesInEnv :: Env -> S.Set VName -allNamesInEnv (Env vtable ttable stable modtable _names) = - S.fromList (M.keys vtable ++ M.keys ttable ++ - M.keys stable ++ M.keys modtable) <> - mconcat (map allNamesInMTy (M.elems stable) ++ - map allNamesInMod (M.elems modtable) ++ - map allNamesInType (M.elems ttable)) - where allNamesInType (TypeAbbr _ ps _) = S.fromList $ map typeParamName ps - -newNamesForMTy :: MTy -> TypeM (MTy, M.Map VName VName) -newNamesForMTy orig_mty = do - -- Create unique renames for the module type. - pairs <- forM (S.toList $ allNamesInMTy orig_mty) $ \v -> do - v' <- newName v - return (v, v') - let substs = M.fromList pairs - rev_substs = M.fromList $ map (uncurry $ flip (,)) pairs - - return (substituteInMTy substs orig_mty, rev_substs) - - where - substituteInMTy :: M.Map VName VName -> MTy -> MTy - substituteInMTy substs (MTy mty_abs mty_mod) = - MTy (M.mapKeys (fmap substitute) mty_abs) (substituteInMod mty_mod) - where - substituteInEnv (Env vtable ttable _stable modtable names) = - let vtable' = substituteInMap substituteInBinding vtable - ttable' = substituteInMap substituteInTypeBinding ttable - mtable' = substituteInMap substituteInMod modtable - in Env { envVtable = vtable' - , envTypeTable = ttable' - , envSigTable = mempty - , envModTable = mtable' - , envNameMap = M.map (fmap substitute) names - } - - substitute v = - fromMaybe v $ M.lookup v substs - - substituteInMap f m = - let (ks, vs) = unzip $ M.toList m - in M.fromList $ - zip (map (\k -> fromMaybe k $ M.lookup k substs) ks) - (map f vs) - - substituteInBinding (BoundV ps t) = - BoundV (map substituteInTypeParam ps) (substituteInType t) - - substituteInMod (ModEnv env) = - ModEnv $ substituteInEnv env - substituteInMod (ModFun funsig) = - ModFun $ substituteInFunSig funsig - - substituteInFunSig (FunSig abs mod mty) = - FunSig (M.mapKeys (fmap substitute) abs) - (substituteInMod mod) (substituteInMTy substs mty) - - substituteInTypeBinding (TypeAbbr l ps t) = - TypeAbbr l (map substituteInTypeParam ps) $ substituteInType t - - substituteInTypeParam (TypeParamDim p loc) = - TypeParamDim (substitute p) loc - substituteInTypeParam (TypeParamType l p loc) = - TypeParamType l (substitute p) loc - - substituteInType :: StructType -> StructType - substituteInType (TypeVar () u (TypeName qs v) targs) = - TypeVar () u (TypeName (map substitute qs) $ substitute v) $ map substituteInTypeArg targs - substituteInType (Prim t) = - Prim t - substituteInType (Record ts) = - Record $ fmap substituteInType ts - substituteInType (Enum cs) = - Enum cs - substituteInType (Array () u (ArrayPrimElem t) shape) = - Array () u (ArrayPrimElem t) (substituteInShape shape) - substituteInType (Array () u (ArrayPolyElem (TypeName qs v) targs) shape) = - Array () u (ArrayPolyElem - (TypeName (map substitute qs) $ substitute v) - (map substituteInTypeArg targs)) - (substituteInShape shape) - substituteInType (Array () u (ArrayRecordElem ts) shape) = - let ts' = fmap (substituteInType . recordArrayElemToType) ts - in case arrayOf (Record ts') (substituteInShape shape) u of - Just t' -> t' - _ -> error "substituteInType: Cannot create array after substitution." - substituteInType (Array () u (ArrayEnumElem cs) shape) = - Array () u (ArrayEnumElem cs) (substituteInShape shape) - substituteInType (Arrow als v t1 t2) = - Arrow als v (substituteInType t1) (substituteInType t2) - - substituteInShape (ShapeDecl ds) = - ShapeDecl $ map substituteInDim ds - substituteInDim (NamedDim (QualName qs v)) = - NamedDim $ QualName (map substitute qs) $ substitute v - substituteInDim d = d - - substituteInTypeArg (TypeArgDim (NamedDim (QualName qs v)) loc) = - TypeArgDim (NamedDim $ QualName (map substitute qs) $ substitute v) loc - substituteInTypeArg (TypeArgDim (ConstDim x) loc) = - TypeArgDim (ConstDim x) loc - substituteInTypeArg (TypeArgDim AnyDim loc) = - TypeArgDim AnyDim loc - substituteInTypeArg (TypeArgType t loc) = - TypeArgType (substituteInType t) loc - -mtyTypeAbbrs :: MTy -> M.Map VName TypeBinding -mtyTypeAbbrs (MTy _ mod) = modTypeAbbrs mod - -modTypeAbbrs :: Mod -> M.Map VName TypeBinding -modTypeAbbrs (ModEnv env) = - envTypeAbbrs env -modTypeAbbrs (ModFun (FunSig _ mod mty)) = - modTypeAbbrs mod <> mtyTypeAbbrs mty - -envTypeAbbrs :: Env -> M.Map VName TypeBinding -envTypeAbbrs env = - envTypeTable env <> - (mconcat . map modTypeAbbrs . M.elems . envModTable) env - --- | Refine the given type name in the given env. -refineEnv :: SrcLoc -> TySet -> Env -> QualName Name -> [TypeParam] -> StructType - -> TypeM (QualName VName, TySet, Env) -refineEnv loc tset env tname ps t - | Just (tname', TypeAbbr l cur_ps (TypeVar () _ (TypeName qs v) _)) <- - findTypeDef tname (ModEnv env), - QualName (qualQuals tname') v `M.member` tset = - if paramsMatch cur_ps ps then - return (tname', - QualName qs v `M.delete` tset, - substituteTypesInEnv - (M.fromList [(qualLeaf tname', - TypeSub $ TypeAbbr l cur_ps t), - (v, TypeSub $ TypeAbbr l ps t)]) - env) - else throwError $ TypeError loc $ "Cannot refine a type having " <> - tpMsg ps <> " with a type having " <> tpMsg cur_ps <> "." - | otherwise = - throwError $ TypeError loc $ - pretty tname ++ " is not an abstract type in the module type." - where tpMsg [] = "no type parameters" - tpMsg xs = "type parameters " <> unwords (map pretty xs) - -paramsMatch :: [TypeParam] -> [TypeParam] -> Bool -paramsMatch ps1 ps2 = length ps1 == length ps2 && all match (zip ps1 ps2) - where match (TypeParamType l1 _ _, TypeParamType l2 _ _) = l1 <= l2 - match (TypeParamDim _ _, TypeParamDim _ _) = True - match _ = False diff --git a/src/Language/Futhark/TypeChecker/Modules.hs b/src/Language/Futhark/TypeChecker/Modules.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/TypeChecker/Modules.hs @@ -0,0 +1,438 @@ +{-# LANGUAGE TupleSections #-} +module Language.Futhark.TypeChecker.Modules + ( matchMTys + , newNamesForMTy + , refineEnv + , applyFunctor + ) where + +import Control.Monad.Except +import Control.Monad.Writer hiding (Sum) +import Data.List +import Data.Loc +import Data.Maybe +import Data.Either +import Data.Ord +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import Prelude hiding (abs, mod) + +import Language.Futhark +import Language.Futhark.Semantic +import Language.Futhark.TypeChecker.Monad +import Language.Futhark.TypeChecker.Unify (doUnification) +import Language.Futhark.TypeChecker.Types +import Futhark.Util.Pretty (Pretty) + + +substituteTypesInMod :: TypeSubs -> Mod -> Mod +substituteTypesInMod substs (ModEnv e) = + ModEnv $ substituteTypesInEnv substs e +substituteTypesInMod substs (ModFun (FunSig abs mod mty)) = + ModFun $ FunSig abs (substituteTypesInMod substs mod) (substituteTypesInMTy substs mty) + +substituteTypesInMTy :: TypeSubs -> MTy -> MTy +substituteTypesInMTy substs (MTy abs mod) = MTy abs $ substituteTypesInMod substs mod + +substituteTypesInEnv :: TypeSubs -> Env -> Env +substituteTypesInEnv substs env = + env { envVtable = M.map (substituteTypesInBoundV substs) $ envVtable env + , envTypeTable = M.mapWithKey subT $ envTypeTable env + , envModTable = M.map (substituteTypesInMod substs) $ envModTable env + } + where subT name _ + | Just (TypeSub (TypeAbbr l ps t)) <- M.lookup name substs = TypeAbbr l ps t + subT _ (TypeAbbr l ps t) = TypeAbbr l ps $ substituteTypes substs t + +substituteTypesInBoundV :: TypeSubs -> BoundV -> BoundV +substituteTypesInBoundV substs (BoundV tps t) = + BoundV tps (substituteTypes substs t) + +allNamesInMTy :: MTy -> S.Set VName +allNamesInMTy (MTy abs mod) = + S.fromList (map qualLeaf $ M.keys abs) <> allNamesInMod mod + +allNamesInMod :: Mod -> S.Set VName +allNamesInMod (ModEnv env) = allNamesInEnv env +allNamesInMod ModFun{} = mempty + +-- All names defined anywhere in the env. +allNamesInEnv :: Env -> S.Set VName +allNamesInEnv (Env vtable ttable stable modtable _names) = + S.fromList (M.keys vtable ++ M.keys ttable ++ + M.keys stable ++ M.keys modtable) <> + mconcat (map allNamesInMTy (M.elems stable) ++ + map allNamesInMod (M.elems modtable) ++ + map allNamesInType (M.elems ttable)) + where allNamesInType (TypeAbbr _ ps _) = S.fromList $ map typeParamName ps + +newNamesForMTy :: MTy -> TypeM (MTy, M.Map VName VName) +newNamesForMTy orig_mty = do + -- Create unique renames for the module type. + pairs <- forM (S.toList $ allNamesInMTy orig_mty) $ \v -> do + v' <- newName v + return (v, v') + let substs = M.fromList pairs + rev_substs = M.fromList $ map (uncurry $ flip (,)) pairs + + return (substituteInMTy substs orig_mty, rev_substs) + + where + substituteInMTy :: M.Map VName VName -> MTy -> MTy + substituteInMTy substs (MTy mty_abs mty_mod) = + MTy (M.mapKeys (fmap substitute) mty_abs) (substituteInMod mty_mod) + where + substituteInEnv (Env vtable ttable _stable modtable names) = + let vtable' = substituteInMap substituteInBinding vtable + ttable' = substituteInMap substituteInTypeBinding ttable + mtable' = substituteInMap substituteInMod modtable + in Env { envVtable = vtable' + , envTypeTable = ttable' + , envSigTable = mempty + , envModTable = mtable' + , envNameMap = M.map (fmap substitute) names + } + + substitute v = + fromMaybe v $ M.lookup v substs + + substituteInMap f m = + let (ks, vs) = unzip $ M.toList m + in M.fromList $ + zip (map (\k -> fromMaybe k $ M.lookup k substs) ks) + (map f vs) + + substituteInBinding (BoundV ps t) = + BoundV (map substituteInTypeParam ps) (substituteInType t) + + substituteInMod (ModEnv env) = + ModEnv $ substituteInEnv env + substituteInMod (ModFun funsig) = + ModFun $ substituteInFunSig funsig + + substituteInFunSig (FunSig abs mod mty) = + FunSig (M.mapKeys (fmap substitute) abs) + (substituteInMod mod) (substituteInMTy substs mty) + + substituteInTypeBinding (TypeAbbr l ps t) = + TypeAbbr l (map substituteInTypeParam ps) $ substituteInType t + + substituteInTypeParam (TypeParamDim p loc) = + TypeParamDim (substitute p) loc + substituteInTypeParam (TypeParamType l p loc) = + TypeParamType l (substitute p) loc + + substituteInType :: StructType -> StructType + substituteInType (Scalar (TypeVar () u (TypeName qs v) targs)) = + Scalar $ TypeVar () u (TypeName (map substitute qs) $ substitute v) $ map substituteInTypeArg targs + substituteInType (Scalar (Prim t)) = + Scalar $ Prim t + substituteInType (Scalar (Record ts)) = + Scalar $ Record $ fmap substituteInType ts + substituteInType (Scalar (Sum ts)) = + Scalar $ Sum $ (fmap . fmap) substituteInType ts + substituteInType (Array () u t shape) = + arrayOf (substituteInType $ Scalar t) (substituteInShape shape) u + substituteInType (Scalar (Arrow als v t1 t2)) = + Scalar $ Arrow als v (substituteInType t1) (substituteInType t2) + + substituteInShape (ShapeDecl ds) = + ShapeDecl $ map substituteInDim ds + substituteInDim (NamedDim (QualName qs v)) = + NamedDim $ QualName (map substitute qs) $ substitute v + substituteInDim d = d + + substituteInTypeArg (TypeArgDim (NamedDim (QualName qs v)) loc) = + TypeArgDim (NamedDim $ QualName (map substitute qs) $ substitute v) loc + substituteInTypeArg (TypeArgDim (ConstDim x) loc) = + TypeArgDim (ConstDim x) loc + substituteInTypeArg (TypeArgDim AnyDim loc) = + TypeArgDim AnyDim loc + substituteInTypeArg (TypeArgType t loc) = + TypeArgType (substituteInType t) loc + +mtyTypeAbbrs :: MTy -> M.Map VName TypeBinding +mtyTypeAbbrs (MTy _ mod) = modTypeAbbrs mod + +modTypeAbbrs :: Mod -> M.Map VName TypeBinding +modTypeAbbrs (ModEnv env) = + envTypeAbbrs env +modTypeAbbrs (ModFun (FunSig _ mod mty)) = + modTypeAbbrs mod <> mtyTypeAbbrs mty + +envTypeAbbrs :: Env -> M.Map VName TypeBinding +envTypeAbbrs env = + envTypeTable env <> + (mconcat . map modTypeAbbrs . M.elems . envModTable) env + +-- | Refine the given type name in the given env. +refineEnv :: SrcLoc -> TySet -> Env -> QualName Name -> [TypeParam] -> StructType + -> TypeM (QualName VName, TySet, Env) +refineEnv loc tset env tname ps t + | Just (tname', TypeAbbr l cur_ps (Scalar (TypeVar () _ (TypeName qs v) _))) <- + findTypeDef tname (ModEnv env), + QualName (qualQuals tname') v `M.member` tset = + if paramsMatch cur_ps ps then + return (tname', + QualName qs v `M.delete` tset, + substituteTypesInEnv + (M.fromList [(qualLeaf tname', + TypeSub $ TypeAbbr l cur_ps t), + (v, TypeSub $ TypeAbbr l ps t)]) + env) + else throwError $ TypeError loc $ "Cannot refine a type having " <> + tpMsg ps <> " with a type having " <> tpMsg cur_ps <> "." + | otherwise = + throwError $ TypeError loc $ + pretty tname ++ " is not an abstract type in the module type." + where tpMsg [] = "no type parameters" + tpMsg xs = "type parameters " <> unwords (map pretty xs) + +paramsMatch :: [TypeParam] -> [TypeParam] -> Bool +paramsMatch ps1 ps2 = length ps1 == length ps2 && all match (zip ps1 ps2) + where match (TypeParamType l1 _ _, TypeParamType l2 _ _) = l1 <= l2 + match (TypeParamDim _ _, TypeParamDim _ _) = True + match _ = False + + +findBinding :: (Env -> M.Map VName v) + -> Namespace -> Name + -> Env + -> Maybe (VName, v) +findBinding table namespace name the_env = do + QualName _ name' <- M.lookup (namespace, name) $ envNameMap the_env + (name',) <$> M.lookup name' (table the_env) + +findTypeDef :: QualName Name -> Mod -> Maybe (QualName VName, TypeBinding) +findTypeDef _ ModFun{} = Nothing +findTypeDef (QualName [] name) (ModEnv the_env) = do + (name', tb) <- findBinding envTypeTable Type name the_env + return (qualName name', tb) +findTypeDef (QualName (q:qs) name) (ModEnv the_env) = do + (q', q_mod) <- findBinding envModTable Term q the_env + (QualName qs' name', tb) <- findTypeDef (QualName qs name) q_mod + return (QualName (q':qs') name', tb) + +resolveAbsTypes :: TySet -> Mod -> TySet -> SrcLoc + -> Either TypeError (M.Map VName (QualName VName, TypeBinding)) +resolveAbsTypes mod_abs mod sig_abs loc = do + let abs_mapping = M.fromList $ zip + (map (fmap baseName . fst) $ M.toList mod_abs) (M.toList mod_abs) + fmap M.fromList $ forM (M.toList sig_abs) $ \(name, name_l) -> + case findTypeDef (fmap baseName name) mod of + Just (name', TypeAbbr mod_l ps t) + | Unlifted <- name_l, + not (orderZero t) || mod_l == Lifted -> + mismatchedLiftedness (map qualLeaf $ M.keys mod_abs) name (ps, t) + | Just (abs_name, _) <- M.lookup (fmap baseName name) abs_mapping -> + return (qualLeaf name, (abs_name, TypeAbbr name_l ps t)) + | otherwise -> + return (qualLeaf name, (name', TypeAbbr name_l ps t)) + _ -> + missingType loc $ fmap baseName name + where mismatchedLiftedness abs name mod_t = + Left $ TypeError loc $ + unlines ["Module defines", + indent $ ppTypeAbbr abs name mod_t, + "but module type requires this type to be non-functional."] + +missingType :: Pretty a => SrcLoc -> a -> Either TypeError b +missingType loc name = + Left $ TypeError loc $ + "Module does not define a type named " ++ pretty name ++ "." + +missingVal :: Pretty a => SrcLoc -> a -> Either TypeError b +missingVal loc name = + Left $ TypeError loc $ + "Module does not define a value named " ++ pretty name ++ "." + +missingMod :: Pretty a => SrcLoc -> a -> Either TypeError b +missingMod loc name = + Left $ TypeError loc $ + "Module does not define a module named " ++ pretty name ++ "." + +mismatchedType :: Pretty a => + SrcLoc + -> [VName] + -> a + -> ([TypeParam], StructType) + -> ([TypeParam], StructType) + -> Either TypeError b +mismatchedType loc abs name spec_t env_t = + Left $ TypeError loc $ + unlines ["Module defines", + indent $ ppTypeAbbr abs name env_t, + "but module type requires", + indent $ ppTypeAbbr abs name spec_t] + +indent :: String -> String +indent = intercalate "\n" . map (" "++) . lines + +ppTypeAbbr :: Pretty a => [VName] -> a -> ([TypeParam], StructType) -> String +ppTypeAbbr abs name (ps, t) = + "type " ++ unwords (pretty name : map pretty ps) ++ t' + where t' = case t of + Scalar (TypeVar () _ tn args) + | typeLeaf tn `elem` abs, + map typeParamToArg ps == args -> "" + _ -> " = " ++ pretty t + + +-- Return new renamed/abstracted env, as well as a mapping from +-- names in the signature to names in the new env. This is used for +-- functor application. The first env is the module env, and the +-- second the env it must match. +matchMTys :: MTy -> MTy -> SrcLoc + -> Either TypeError (M.Map VName VName) +matchMTys = matchMTys' mempty + where + matchMTys' :: TypeSubs -> MTy -> MTy -> SrcLoc + -> Either TypeError (M.Map VName VName) + + matchMTys' _ (MTy _ ModFun{}) (MTy _ ModEnv{}) loc = + Left $ TypeError loc "Cannot match parametric module with non-paramatric module type." + + matchMTys' _ (MTy _ ModEnv{}) (MTy _ ModFun{}) loc = + Left $ TypeError loc "Cannot match non-parametric module with paramatric module type." + + matchMTys' old_abs_subst_to_type (MTy mod_abs mod) (MTy sig_abs sig) loc = do + -- Check that abstract types in 'sig' have an implementation in + -- 'mod'. This also gives us a substitution that we use to check + -- the types of values. + abs_substs <- resolveAbsTypes mod_abs mod sig_abs loc + + let abs_subst_to_type = old_abs_subst_to_type <> + M.map (TypeSub . snd) abs_substs + abs_name_substs = M.map (qualLeaf . fst) abs_substs + substs <- matchMods abs_subst_to_type mod sig loc + return (substs <> abs_name_substs) + + matchMods :: TypeSubs -> Mod -> Mod -> SrcLoc + -> Either TypeError (M.Map VName VName) + matchMods _ ModEnv{} ModFun{} loc = + Left $ TypeError loc "Cannot match non-parametric module with paramatric module type." + matchMods _ ModFun{} ModEnv{} loc = + Left $ TypeError loc "Cannot match parametric module with non-paramatric module type." + + matchMods abs_subst_to_type (ModEnv mod) (ModEnv sig) loc = + matchEnvs abs_subst_to_type mod sig loc + + matchMods old_abs_subst_to_type + (ModFun (FunSig mod_abs mod_pmod mod_mod)) + (ModFun (FunSig sig_abs sig_pmod sig_mod)) + loc = do + abs_substs <- resolveAbsTypes mod_abs mod_pmod sig_abs loc + let abs_subst_to_type = old_abs_subst_to_type <> + M.map (TypeSub . snd) abs_substs + abs_name_substs = M.map (qualLeaf . fst) abs_substs + pmod_substs <- matchMods abs_subst_to_type mod_pmod sig_pmod loc + mod_substs <- matchMTys' abs_subst_to_type mod_mod sig_mod loc + return (pmod_substs <> mod_substs <> abs_name_substs) + + matchEnvs :: TypeSubs + -> Env -> Env -> SrcLoc + -> Either TypeError (M.Map VName VName) + matchEnvs abs_subst_to_type env sig loc = do + -- XXX: we only want to create substitutions for visible names. + -- This must be wrong in some cases. Probably we need to + -- rethink how we do shadowing for module types. + let visible = S.fromList $ map qualLeaf $ M.elems $ envNameMap sig + isVisible name = name `S.member` visible + + -- Check that all values are defined correctly, substituting the + -- abstract types first. + val_substs <- fmap M.fromList $ forM (M.toList $ envVtable sig) $ \(name, spec_bv) -> do + let spec_bv' = substituteTypesInBoundV abs_subst_to_type spec_bv + case findBinding envVtable Term (baseName name) env of + Just (name', bv) -> matchVal loc name spec_bv' name' bv + _ -> missingVal loc (baseName name) + + -- Check that all type abbreviations are correctly defined. + abbr_name_substs <- fmap M.fromList $ + forM (filter (isVisible . fst) $ M.toList $ + envTypeTable sig) $ \(name, TypeAbbr _ spec_ps spec_t) -> + case findBinding envTypeTable Type (baseName name) env of + Just (name', TypeAbbr _ ps t) -> + matchTypeAbbr loc abs_subst_to_type val_substs name spec_ps spec_t name' ps t + Nothing -> missingType loc $ baseName name + + -- Check for correct modules. + mod_substs <- fmap M.unions $ forM (M.toList $ envModTable sig) $ \(name, modspec) -> + case findBinding envModTable Term (baseName name) env of + Just (name', mod) -> + M.insert name name' <$> matchMods abs_subst_to_type mod modspec loc + Nothing -> + missingMod loc $ baseName name + + return $ val_substs <> mod_substs <> abbr_name_substs + + matchTypeAbbr :: SrcLoc -> TypeSubs -> M.Map VName VName + -> VName -> [TypeParam] -> StructType + -> VName -> [TypeParam] -> StructType + -> Either TypeError (VName, VName) + matchTypeAbbr loc abs_subst_to_type val_substs spec_name spec_ps spec_t name ps t = do + -- We have to create substitutions for the type parameters, too. + unless (length spec_ps == length ps) nomatch + param_substs <- mconcat <$> zipWithM matchTypeParam spec_ps ps + let val_substs' = M.map (DimSub . NamedDim . qualName) val_substs + spec_t' = substituteTypes (val_substs'<>param_substs<>abs_subst_to_type) spec_t + if spec_t' == t + then return (spec_name, name) + else nomatch + where nomatch = mismatchedType loc (M.keys abs_subst_to_type) + (baseName spec_name) (spec_ps, spec_t) (ps, t) + + matchTypeParam (TypeParamDim x _) (TypeParamDim y _) = + pure $ M.singleton x $ DimSub $ NamedDim $ qualName y + matchTypeParam (TypeParamType Unlifted x _) (TypeParamType Unlifted y _) = + pure $ M.singleton x $ TypeSub $ TypeAbbr Unlifted [] $ + Scalar $ TypeVar () Nonunique (typeName y) [] + matchTypeParam (TypeParamType _ x _) (TypeParamType Lifted y _) = + pure $ M.singleton x $ TypeSub $ TypeAbbr Lifted [] $ + Scalar $ TypeVar () Nonunique (typeName y) [] + matchTypeParam _ _ = + nomatch + + matchVal :: SrcLoc + -> VName -> BoundV + -> VName -> BoundV + -> Either TypeError (VName, VName) + matchVal loc spec_name spec_t name t + | matchFunBinding loc spec_t t = return (spec_name, name) + matchVal loc spec_name spec_v _ v = + Left $ TypeError loc $ unlines $ + ["Module type specifies"] ++ + map (" "++) (lines $ ppValBind spec_name spec_v) ++ + ["but module provides"] ++ + map (" "++) (lines $ppValBind spec_name v) + + matchFunBinding :: SrcLoc -> BoundV -> BoundV -> Bool + matchFunBinding loc (BoundV _ orig_spec_t) (BoundV tps orig_t) = + -- Would be nice if we could propagate the actual error here. + case doUnification loc tps + (toStructural orig_spec_t) (toStructural orig_t) of + Left _ -> False + Right t -> t `subtypeOf` toStructural orig_spec_t + + ppValBind v (BoundV tps t) = + unwords $ ["val", prettyName v] ++ map pretty tps ++ [":", pretty t] + +applyFunctor :: SrcLoc + -> FunSig + -> MTy + -> TypeM (MTy, + M.Map VName VName, + M.Map VName VName) +applyFunctor applyloc (FunSig p_abs p_mod body_mty) a_mty = do + p_subst <- badOnLeft $ matchMTys a_mty (MTy p_abs p_mod) applyloc + + -- Apply type abbreviations from a_mty to body_mty. + let a_abbrs = mtyTypeAbbrs a_mty + isSub v = case M.lookup v a_abbrs of + Just abbr -> Just $ TypeSub abbr + _ -> Just $ DimSub $ NamedDim $ qualName v + type_subst = M.mapMaybe isSub p_subst + body_mty' = substituteTypesInMTy type_subst body_mty + (body_mty'', body_subst) <- newNamesForMTy body_mty' + return (body_mty'', p_subst, body_subst) diff --git a/src/Language/Futhark/TypeChecker/Monad.hs b/src/Language/Futhark/TypeChecker/Monad.hs --- a/src/Language/Futhark/TypeChecker/Monad.hs +++ b/src/Language/Futhark/TypeChecker/Monad.hs @@ -11,7 +11,6 @@ , checkQualNameWithEnv , bindSpaced , qualifyTypeVars - , getType , TypeError(..) , unexpectedType @@ -111,13 +110,14 @@ Namespace -> QualName Name -> SrcLoc -> m a unknownVariableError space name loc = throwError $ TypeError loc $ - "Unknown " ++ ppSpace space ++ " " ++ pretty name + "Unknown " ++ ppSpace space ++ " " ++ quote (pretty name) underscoreUse :: MonadTypeChecker m => SrcLoc -> QualName Name -> m a underscoreUse loc name = throwError $ TypeError loc $ - "Use of " ++ pretty name ++ ": variables prefixed with underscore must not be accessed." + "Use of " ++ quote (pretty name) ++ + ": variables prefixed with underscore may not be accessed." instance Show TypeError where show (TypeError pos msg) = @@ -185,16 +185,16 @@ getBreadCrumbs :: m [BreadCrumb] getBreadCrumbs = return [] -typeError :: (MonadError TypeError m, MonadBreadCrumbs m) => - SrcLoc -> String -> m a +typeError :: (Located loc, MonadError TypeError m, MonadBreadCrumbs m) => + loc -> String -> m a typeError loc s = do bc <- getBreadCrumbs let bc' | null bc = "" | otherwise = "\n" ++ unlines (map show bc) - throwError $ TypeError loc $ s ++ bc' + throwError $ TypeError (srclocOf loc) $ s ++ bc' class MonadError TypeError m => MonadTypeChecker m where - warn :: SrcLoc -> String -> m () + warn :: Located loc => loc -> String -> m () newName :: VName -> m VName newID :: Name -> m VName @@ -214,10 +214,9 @@ checkNamedDim loc v = do (v', t) <- lookupVar loc v case t of - Prim (Signed Int32) -> return v' - _ -> throwError $ TypeError loc $ - "Dimension declaration " ++ pretty v ++ - " should be of type `i32`." + Scalar (Prim (Signed Int32)) -> return v' + _ -> throwError $ TypeError loc $ + "Dimension declaration " ++ pretty v ++ " should be of type `i32`." checkName :: MonadTypeChecker m => Namespace -> Name -> SrcLoc -> m VName checkName space name loc = qualLeaf <$> checkQualName space (qualName name) loc @@ -229,7 +228,7 @@ bindNameMap mapping body instance MonadTypeChecker TypeM where - warn loc problem = tell $ singleWarning loc problem + warn loc problem = tell $ singleWarning (srclocOf loc) problem newName s = do src <- get let (s', src') = Futhark.FreshNames.newName src s @@ -293,9 +292,9 @@ -- | Extract from a type either a function type comprising a list of -- parameter types and a return type, or a first-order type. getType :: TypeBase dim as - -> Either ([(Maybe VName, TypeBase dim as)], TypeBase dim as) + -> Either ([(PName, TypeBase dim as)], TypeBase dim as) (TypeBase dim as) -getType (Arrow _ v t1 t2) = +getType (Scalar (Arrow _ v t1 t2)) = case getType t2 of Left (ps, r) -> Left ((v, t1) : ps, r) Right _ -> Left ([(v, t1)], t2) @@ -346,7 +345,7 @@ reachable [] name env = isJust $ find matches $ M.elems (envTypeTable env) - where matches (TypeAbbr _ [] (TypeVar _ _ (TypeName x_qs name') [])) = + where matches (TypeAbbr _ _ (Scalar (TypeVar _ _ (TypeName x_qs name') _))) = null x_qs && name == name' matches _ = False @@ -403,6 +402,6 @@ where type_names = S.fromList $ map (nameFromString . pretty) anyPrimType binop_names = S.fromList $ map (nameFromString . pretty) [minBound..(maxBound::BinOp)] - unop_names = S.fromList $ map nameFromString ["~", "!"] + unop_names = S.fromList $ map nameFromString ["!"] fun_names = S.fromList $ map nameFromString ["shape"] atTopLevel _ = False diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs --- a/src/Language/Futhark/TypeChecker/Terms.hs +++ b/src/Language/Futhark/TypeChecker/Terms.hs @@ -1,5 +1,6 @@ {-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts #-} {-# LANGUAGE FlexibleInstances, DeriveFunctor #-} +{-# Language TupleSections #-} -- | Facilities for type-checking Futhark terms. Checking a term -- requires a little more context to track uniqueness and such. -- @@ -16,10 +17,11 @@ import Control.Monad.Except import Control.Monad.State -import Control.Monad.RWS +import Control.Monad.RWS hiding (Sum) import qualified Control.Monad.Fail as Fail import Data.Char (isAlpha) import Data.List +import qualified Data.List.NonEmpty as NE import Data.Loc import Data.Maybe import qualified Data.Map.Strict as M @@ -32,10 +34,10 @@ import Language.Futhark.Traversals import Language.Futhark.TypeChecker.Monad hiding (BoundV, checkQualNameWithEnv) import Language.Futhark.TypeChecker.Types hiding (checkTypeDecl) -import Language.Futhark.TypeChecker.Unify +import Language.Futhark.TypeChecker.Unify hiding (Usage) import qualified Language.Futhark.TypeChecker.Types as Types import qualified Language.Futhark.TypeChecker.Monad as TypeM -import Futhark.Util.Pretty hiding (space, bool) +import Futhark.Util.Pretty hiding (space, bool, group) --- Uniqueness @@ -204,7 +206,7 @@ MonadError TypeError) instance Fail.MonadFail TermTypeM where - fail = typeError noLoc . ("unknown failure (likely a bug): "++) + fail = typeError (noLoc :: SrcLoc) . ("unknown failure (likely a bug): "++) instance MonadUnify TermTypeM where getConstraints = gets fst @@ -213,8 +215,8 @@ newTypeVar loc desc = do i <- incCounter v <- newID $ mkTypeVarName desc i - modifyConstraints $ M.insert v $ NoConstraint Nothing loc - return $ TypeVar mempty Nonunique (typeName v) [] + modifyConstraints $ M.insert v $ NoConstraint Nothing $ mkUsage' loc + return $ Scalar $ TypeVar mempty Nonunique (typeName v) [] instance MonadBreadCrumbs TermTypeM where breadCrumb bc = local $ \env -> @@ -238,7 +240,8 @@ initialTermScope = TermScope initialVtable mempty topLevelNameMap mempty where initialVtable = M.fromList $ mapMaybe addIntrinsicF $ M.toList intrinsics - funF ts t = foldr (Arrow mempty Nothing . Prim) (Prim t) ts + funF ts t = foldr (arrow . Scalar . Prim) (Scalar $ Prim t) ts + arrow x y = Scalar $ Arrow mempty Unnamed x y addIntrinsicF (name, IntrinsicMonoFun ts t) = Just (name, BoundV Global [] $ funF ts t) @@ -247,7 +250,7 @@ addIntrinsicF (name, IntrinsicPolyFun tvs pts rt) = Just (name, BoundV Global tvs $ fromStruct $ vacuousShapeAnnotations $ - Arrow mempty Nothing pts' rt) + Scalar $ Arrow mempty Unnamed pts' rt) where pts' = case pts of [pt] -> pt _ -> tupleRecord pts addIntrinsicF (name, IntrinsicEquality) = @@ -292,6 +295,7 @@ lookupVar loc qn = do outer_env <- liftTypeM askRootEnv (scope, qn'@(QualName qs name)) <- checkQualNameWithEnv Term qn loc + let usage = mkUsage loc $ "use of " ++ quote (pretty qn) t <- case M.lookup name $ scopeVtable scope of Nothing -> throwError $ TypeError loc $ @@ -308,31 +312,33 @@ Just OpaqueF -> do argtype <- newTypeVar loc "t" - return $ Arrow mempty Nothing argtype argtype + return $ Scalar $ Arrow mempty Unnamed argtype argtype Just EqualityF -> do argtype <- newTypeVar loc "t" - equalityType loc argtype - return $ Arrow mempty Nothing argtype $ - Arrow mempty Nothing argtype $ Prim Bool + equalityType usage argtype + return $ + Scalar $ Arrow mempty Unnamed argtype $ + Scalar $ Arrow mempty Unnamed argtype $ Scalar $ Prim Bool Just (OverloadedF ts pts rt) -> do argtype <- newTypeVar loc "t" - mustBeOneOf ts loc argtype + mustBeOneOf ts usage argtype let (pts', rt') = instOverloaded argtype pts rt - return $ fromStruct $ vacuousShapeAnnotations $ - foldr (Arrow mempty Nothing) rt' pts' + arrow xt yt = Scalar $ Arrow mempty Unnamed xt yt + return $ fromStruct $ vacuousShapeAnnotations $ foldr arrow rt' pts' observe $ Ident name (Info t) loc return (qn', t) where instOverloaded argtype pts rt = - (map (maybe (toStruct argtype) Prim) pts, - maybe (toStruct argtype) Prim rt) + (map (maybe (toStruct argtype) (Scalar . Prim)) pts, + maybe (toStruct argtype) (Scalar . Prim) rt) checkNamedDim loc v = do (v', t) <- lookupVar loc v - unify loc (toStructural t) (Prim $ Signed Int32) + unify (mkUsage loc "use as array size") (toStructural t) $ + Scalar $ Prim $ Signed Int32 return v' checkQualNameWithEnv :: Namespace -> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName) @@ -374,7 +380,8 @@ (tdecl', _) <- Types.checkTypeDecl [] tdecl mapM_ observeDim $ nestedDims $ unInfo $ expandedType tdecl' return tdecl' - where observeDim (NamedDim v) = observe $ Ident (qualLeaf v) (Info $ Prim $ Signed Int32) noLoc + where observeDim (NamedDim v) = + observe $ Ident (qualLeaf v) (Info $ Scalar $ Prim $ Signed Int32) noLoc observeDim _ = return () -- | Instantiate a type scheme with fresh type variables for its type @@ -396,25 +403,26 @@ instantiateTypeParam loc tparam = do i <- incCounter v <- newID $ mkTypeVarName (takeWhile isAlpha (baseString (typeParamName tparam))) i - modifyConstraints $ M.insert v $ NoConstraint (Just l) loc - return (v, Subst $ TypeVar mempty Nonunique (typeName v) []) + modifyConstraints $ M.insert v $ NoConstraint (Just l) $ mkUsage' loc + return (v, Subst $ Scalar $ TypeVar mempty Nonunique (typeName v) []) where l = case tparam of TypeParamType x _ _ -> x _ -> Lifted newArrayType :: SrcLoc -> String -> Int -> TermTypeM (TypeBase () (), TypeBase () ()) newArrayType loc desc r = do v <- newID $ nameFromString desc - modifyConstraints $ M.insert v $ NoConstraint Nothing loc - return (Array () Nonunique - (ArrayPolyElem (typeName v) []) (ShapeDecl $ replicate r ()), - TypeVar () Nonunique (typeName v) []) + modifyConstraints $ M.insert v $ NoConstraint Nothing $ mkUsage' loc + let rowt = TypeVar () Nonunique (typeName v) [] + return (Array () Nonunique rowt (ShapeDecl $ replicate r ()), + Scalar rowt) --- Errors useAfterConsume :: MonadTypeChecker m => Name -> SrcLoc -> SrcLoc -> m a useAfterConsume name rloc wloc = throwError $ TypeError rloc $ - "Variable " ++ pretty name ++ " previously consumed at " ++ locStr wloc ++ ". (Possibly through aliasing)" + "Variable " ++ quote (pretty name) ++ " previously consumed at " ++ + locStr wloc ++ ". (Possibly through aliasing)" consumeAfterConsume :: MonadTypeChecker m => Name -> SrcLoc -> SrcLoc -> m a consumeAfterConsume name loc1 loc2 = @@ -447,35 +455,21 @@ unifyExpTypes e1 e2 = do e1_t <- expType e1 e2_t <- expType e2 - unify (srclocOf e2) (toStructural e1_t) (toStructural e2_t) + unify (mkUsage (srclocOf e2) "requiring equality of types") (toStructural e1_t) (toStructural e2_t) return $ unifyTypeAliases e1_t e2_t -- | Assumes that the two types have already been unified. unifyTypeAliases :: PatternType -> PatternType -> PatternType unifyTypeAliases t1 t2 = case (t1, t2) of - (Array als1 u1 et1 shape1, Array als2 u2 et2 _) -> - Array (als1<>als2) (min u1 u2) (unifyArrayElems et1 et2) shape1 - (Record f1, Record f2) -> - Record $ M.intersectionWith unifyTypeAliases f1 f2 - (TypeVar als1 u v targs1, TypeVar als2 _ _ targs2) -> - TypeVar (als1 <> als2) u v $ zipWith unifyTypeArg targs1 targs2 + (Array als1 u1 et1 shape1, Array als2 u2 _ _) -> + Array (als1<>als2) (min u1 u2) et1 shape1 + (Scalar (Record f1), Scalar (Record f2)) -> + Scalar $ Record $ M.intersectionWith unifyTypeAliases f1 f2 + (Scalar (TypeVar als1 u v targs1), Scalar (TypeVar als2 _ _ targs2)) -> + Scalar $ TypeVar (als1 <> als2) u v $ zipWith unifyTypeArg targs1 targs2 _ -> t1 - where unifyArrayElems (ArrayPrimElem pt1) (ArrayPrimElem _) = - ArrayPrimElem pt1 - unifyArrayElems (ArrayPolyElem v targs1) (ArrayPolyElem _ _targs2) = - ArrayPolyElem v targs1 - unifyArrayElems (ArrayRecordElem fields1) (ArrayRecordElem fields2) = - ArrayRecordElem $ M.intersectionWith unifyRecordArray fields1 fields2 - unifyArrayElems x _ = x - - unifyRecordArray (RecordArrayElem at1) (RecordArrayElem at2) = - RecordArrayElem $ unifyArrayElems at1 at2 - unifyRecordArray (RecordArrayArrayElem at1 shape1) (RecordArrayArrayElem at2 _) = - RecordArrayArrayElem (unifyArrayElems at1 at2) shape1 - unifyRecordArray x _ = x - - unifyTypeArg (TypeArgType t1' loc) (TypeArgType _ _) = + where unifyTypeArg (TypeArgType t1' loc) (TypeArgType _ _) = TypeArgType t1' loc unifyTypeArg a _ = a @@ -510,13 +504,19 @@ TuplePattern <$> zipWithM checkPattern' ps (map Ascribed ts) <*> pure loc checkPattern' p@(TuplePattern ps loc) (Ascribed t) = do ps_t <- replicateM (length ps) (newTypeVar loc "t") - unify loc (tupleRecord ps_t) $ toStructural t + unify (mkUsage loc "matching a tuple pattern") (tupleRecord ps_t) $ toStructural t t' <- normaliseType t checkPattern' p $ Ascribed t' checkPattern' (TuplePattern ps loc) NoneInferred = TuplePattern <$> mapM (`checkPattern'` NoneInferred) ps <*> pure loc -checkPattern' (RecordPattern p_fs loc) (Ascribed (Record t_fs)) +checkPattern' (RecordPattern p_fs _) _ + | Just (f, fp) <- find (("_" `isPrefixOf`) . nameToString . fst) p_fs = + typeError fp $ unlines [ "Underscore-prefixed fields are not allowed." + , "Did you mean " ++ + quote (drop 1 (nameToString f) ++ "=_") ++ "?"] + +checkPattern' (RecordPattern p_fs loc) (Ascribed (Scalar (Record t_fs))) | sort (map fst p_fs) == sort (M.keys t_fs) = RecordPattern . M.toList <$> check <*> pure loc where check = traverse (uncurry checkPattern') $ M.intersectionWith (,) @@ -527,7 +527,7 @@ when (sort (M.keys fields') /= sort (map fst fields)) $ typeError loc $ "Duplicate fields in record pattern " ++ pretty p - unify loc (Record fields') $ toStructural t + unify (mkUsage loc "matching a record pattern") (Scalar (Record fields')) $ toStructural t t' <- normaliseType t checkPattern' p $ Ascribed t' checkPattern' (RecordPattern fs loc) NoneInferred = @@ -539,7 +539,7 @@ let st' = fromStruct st case maybe_outer_t of Ascribed outer_t -> do - unify loc (toStructural st) (toStructural outer_t) + unify (mkUsage loc "explicit type ascription") (toStructural st) (toStructural outer_t) -- We also have to make sure that uniqueness and shapes match. -- This is done explicitly, because they are ignored by @@ -562,7 +562,7 @@ checkPattern' (PatternLit e NoInfo loc) (Ascribed t) = do e' <- checkExp e t' <- expType e' - unify loc (toStructural t') (toStructural t) + unify (mkUsage loc "matching against literal") (toStructural t') (toStructural t) return $ PatternLit e' (Info t') loc checkPattern' (PatternLit e NoInfo loc) NoneInferred = do @@ -570,6 +570,27 @@ t' <- expType e' return $ PatternLit e' (Info t') loc +checkPattern' (PatternConstr n NoInfo ps loc) (Ascribed (Scalar (Sum cs))) + | Just ts <- M.lookup n cs = do + ps' <- zipWithM checkPattern' ps $ map Ascribed ts + return $ PatternConstr n (Info (Scalar (Sum cs))) ps' loc + +checkPattern' (PatternConstr n NoInfo ps loc) (Ascribed t) = do + t' <- newTypeVar loc "t" + ps' <- mapM (`checkPattern'` NoneInferred) ps + mustHaveConstr usage n t' (toStructural . patternType <$> ps') + unify usage t' (toStructural t) + t'' <- normaliseType t + return $ PatternConstr n (Info t'') ps' loc + where usage = mkUsage loc "matching against constructor" + +checkPattern' (PatternConstr n NoInfo ps loc) NoneInferred = do + ps' <- mapM (`checkPattern'` NoneInferred) ps + t <- newTypeVar loc "t" + mustHaveConstr usage n t (toStructural . patternType <$> ps') + return $ PatternConstr n (Info t) ps' loc + where usage = mkUsage loc "matching against constructor" + bindPatternNames :: PatternBase NoInfo Name -> TermTypeM a -> TermTypeM a bindPatternNames = bindSpaced . map asTerm . S.toList . patternIdents where asTerm v = (Term, identName v) @@ -590,11 +611,11 @@ bindVar scope (Ident name (Info tp) _) = let inedges = boundAliases $ aliases tp update (BoundV l tparams in_t) - -- If 'name' is record-typed, don't alias the components - -- to 'name', because records have no identity beyond - -- their components. - | Record _ <- tp = BoundV l tparams in_t - | otherwise = BoundV l tparams (in_t `addAliases` S.insert (AliasBound name)) + -- If 'name' is record or sum-typed, don't alias the + -- components to 'name', because these no identity + -- beyond their components. + | Array{} <- tp = BoundV l tparams (in_t `addAliases` S.insert (AliasBound name)) + | otherwise = BoundV l tparams in_t update b = b tp' = tp `addAliases` S.insert (AliasBound name) @@ -644,14 +665,14 @@ bindingTypeParams tparams = binding (mapMaybe typeParamIdent tparams) . bindingTypes (mapMaybe typeParamType tparams) where typeParamType (TypeParamType l v loc) = - Just (v, (TypeAbbr l [] (TypeVar () Nonunique (typeName v) []), + Just (v, (TypeAbbr l [] (Scalar (TypeVar () Nonunique (typeName v) [])), ParamType l loc)) typeParamType TypeParamDim{} = Nothing typeParamIdent :: TypeParam -> Maybe Ident typeParamIdent (TypeParamDim v loc) = - Just $ Ident v (Info (Prim (Signed Int32))) loc + Just $ Ident v (Info $ Scalar $ Prim $ Signed Int32) loc typeParamIdent _ = Nothing bindingIdent :: IdentBase NoInfo Name -> PatternType -> (Ident -> TermTypeM a) @@ -705,6 +726,7 @@ patternUses (RecordPattern fs _) = foldMap (patternUses . snd) fs patternUses (PatternAscription p (TypeDecl declte _) _) = patternUses p <> typeExpUses declte +patternUses (PatternConstr _ _ ps _) = foldMap patternUses ps patternDims :: Pattern -> [Ident] patternDims (PatternParens p _) = patternDims p @@ -720,13 +742,13 @@ -- | @require ts e@ causes a 'TypeError' if @expType e@ is not one of -- the types in @ts@. Otherwise, simply returns @e@. -require :: [PrimType] -> Exp -> TermTypeM Exp -require ts e = do mustBeOneOf ts (srclocOf e) . toStructural =<< expType e - return e +require :: String -> [PrimType] -> Exp -> TermTypeM Exp +require why ts e = do mustBeOneOf ts (mkUsage (srclocOf e) why) . toStructural =<< expType e + return e -unifies :: TypeBase () () -> Exp -> TermTypeM Exp -unifies t e = do - unify (srclocOf e) t =<< toStructural <$> expType e +unifies :: String -> TypeBase () () -> Exp -> TermTypeM Exp +unifies why t e = do + unify (mkUsage (srclocOf e) why) t =<< toStructural <$> expType e return e -- The closure of a lambda or local function are those variables that @@ -748,12 +770,12 @@ checkExp (IntLit val NoInfo loc) = do t <- newTypeVar loc "t" - mustBeOneOf anyNumberType loc t + mustBeOneOf anyNumberType (mkUsage loc "integer literal") t return $ IntLit val (Info $ vacuousShapeAnnotations $ fromStruct t) loc checkExp (FloatLit val NoInfo loc) = do t <- newTypeVar loc "t" - mustBeOneOf anyFloatType loc t + mustBeOneOf anyFloatType (mkUsage loc "float literal") t return $ FloatLit val (Info $ vacuousShapeAnnotations $ fromStruct t) loc checkExp (TupLit es loc) = @@ -794,13 +816,13 @@ e:es -> do e' <- checkExp e et <- expType e' - es' <- mapM (unifies (toStructural et) <=< checkExp) es + es' <- mapM (unifies "type of first array element" (toStructural et) <=< checkExp) es et' <- normaliseType et t <- arrayOfM loc et' (ShapeDecl [AnyDim]) Unique return $ ArrayLit (e':es') (Info t) loc checkExp (Range start maybe_step end NoInfo loc) = do - start' <- require anyIntType =<< checkExp start + start' <- require "use in range expression" anyIntType =<< checkExp start start_t <- toStructural <$> expType start' maybe_step' <- case maybe_step of Nothing -> return Nothing @@ -810,12 +832,15 @@ (Literal x _, Literal y _) -> when (x == y) warning (Var x_name _ _, Var y_name _ _) -> when (x_name == y_name) warning _ -> return () - Just <$> (unifies start_t =<< checkExp step) + Just <$> (unifies "use in range expression" start_t =<< checkExp step) end' <- case end of - DownToExclusive e -> DownToExclusive <$> (unifies start_t =<< checkExp e) - UpToExclusive e -> UpToExclusive <$> (unifies start_t =<< checkExp e) - ToInclusive e -> ToInclusive <$> (unifies start_t =<< checkExp e) + DownToExclusive e -> DownToExclusive <$> + (unifies "use in range expression" start_t =<< checkExp e) + UpToExclusive e -> UpToExclusive <$> + (unifies "use in range expression" start_t =<< checkExp e) + ToInclusive e -> ToInclusive <$> + (unifies "use in range expression" start_t =<< checkExp e) t <- arrayOfM loc start_t (rank 1) Unique @@ -827,7 +852,7 @@ e' <- checkExp e t <- expType e' let decl_t = unInfo $ expandedType decl' - unify loc (toStructural decl_t) (toStructural t) + unify (mkUsage loc "explicit type ascription") (toStructural decl_t) (toStructural t) -- We also have to make sure that uniqueness matches. This is done -- explicitly, because uniqueness is ignored by unification. @@ -854,7 +879,7 @@ checkExp (Project k e NoInfo loc) = do e' <- checkExp e t <- expType e' - kt <- mustHaveField loc k t + kt <- mustHaveField (mkUsage loc $ "projection of field " ++ quote (pretty k)) k t return $ Project k e' (Info kt) loc checkExp (If e1 e2 e3 _ loc) = @@ -862,11 +887,11 @@ ((e2', e3'), dflow) <- tapOccurences $ checkExp e2 `alternative` checkExp e3 brancht <- unifyExpTypes e2' e3' let t' = addAliases brancht (`S.difference` S.map AliasBound (allConsumed dflow)) - zeroOrderType loc "returned from branch" t' + zeroOrderType (mkUsage loc "returning value of this type from 'if' expression") "returned from branch" t' return $ If e1' e2' e3' (Info t') loc where checkCond = do e1' <- checkExp e1 - unify (srclocOf e1') (Prim Bool) . toStructural =<< expType e1' + unify (mkUsage (srclocOf e1') "use as 'if' condition") (Scalar $ Prim Bool) . toStructural =<< expType e1' return e1' checkExp (Parens e loc) = @@ -910,11 +935,12 @@ checkField e k = do t <- expType e - kt <- mustHaveField loc k t + let usage = mkUsage loc $ "projection of field " ++ quote (pretty k) + kt <- mustHaveField usage k t return $ Project k e (Info kt) loc checkExp (Negate arg loc) = do - arg' <- require anyNumberType =<< checkExp arg + arg' <- require "numeric negation" anyNumberType =<< checkExp arg return $ Negate arg' loc checkExp (Apply e1 e2 NoInfo NoInfo loc) = do @@ -932,7 +958,7 @@ case anyConsumption e_occs of Just c -> let msg = "of value computed with consumption at " ++ locStr (location c) - in zeroOrderType loc msg t + in zeroOrderType (mkUsage loc "consumption in right-hand side of 'let'-binding") msg t _ -> return () bindingPattern pat (Ascribed $ anyDimShapeAnnotations t) $ \pat' -> do body' <- checkExp body @@ -945,7 +971,8 @@ closure' <- lexicalClosure params' closure - let ftype = foldr (uncurry (Arrow ()) . patternParam) rettype params' + let arrow (xp, xt) yt = Scalar $ Arrow () xp xt yt + ftype = foldr (arrow . patternParam) rettype params' entry = BoundV Local tparams' $ ftype `setAliases` closure' bindF scope = scope { scopeVtable = M.insert name' entry $ scopeVtable scope , scopeNameMap = M.insert (Term, name) (qualName name') $ @@ -959,7 +986,7 @@ let elemt = stripArray (length $ filter isFix idxes) t sequentially (checkIdent src) $ \src' _ -> do let src'' = Var (qualName $ identName src') (identType src') (srclocOf src) - void $ unifies t src'' + void $ unifies "type of target array" t src'' unless (unique $ unInfo $ identType src') $ typeError loc $ "Source " ++ quote (pretty (identName src)) ++ @@ -974,7 +1001,7 @@ _ -> return () idxes' <- mapM checkDimIndex idxes - sequentially (unifies elemt =<< checkExp ve) $ \ve' _ -> do + sequentially (unifies "type of target array" elemt =<< checkExp ve) $ \ve' _ -> do ve_t <- expType ve' when (AliasBound (identName src') `S.member` aliases ve_t) $ badLetWithValue loc @@ -989,8 +1016,8 @@ checkExp (Update src idxes ve loc) = do (t, _) <- newArrayType (srclocOf src) "src" $ length idxes let elemt = stripArray (length $ filter isFix idxes) t - sequentially (checkExp ve >>= unifies elemt) $ \ve' _ -> - sequentially (checkExp src >>= unifies t) $ \src' _ -> do + sequentially (checkExp ve >>= unifies "type of target array" elemt) $ \ve' _ -> + sequentially (checkExp src >>= unifies "type of target array" t) $ \src' _ -> do idxes' <- mapM checkDimIndex idxes @@ -1012,13 +1039,14 @@ src' <- checkExp src ve' <- checkExp ve a <- expType src' - r <- foldM (flip $ mustHaveField loc) a fields - unify loc (toStructural r) . toStructural =<< expType ve' - return $ RecordUpdate src' fields ve' (Info $ fromStruct a) loc + let usage = mkUsage loc "record update" + r <- foldM (flip $ mustHaveField usage) a fields + unify usage (toStructural r) . toStructural =<< expType ve' + return $ RecordUpdate src' fields ve' (Info a) loc checkExp (Index e idxes NoInfo loc) = do (t, _) <- newArrayType (srclocOf e) "e" $ length idxes - e' <- unifies t =<< checkExp e + e' <- unifies "being indexed at" t =<< checkExp e idxes' <- mapM checkDimIndex idxes t' <- anyDimShapeAnnotations . stripArray (length $ filter isFix idxes) <$> @@ -1031,7 +1059,7 @@ Unsafe <$> checkExp e <*> pure loc checkExp (Assert e1 e2 NoInfo loc) = do - e1' <- require [Bool] =<< checkExp e1 + e1' <- require "being asserted" [Bool] =<< checkExp e1 e2' <- checkExp e2 return $ Assert e1' e2' (Info (pretty e1)) loc @@ -1065,7 +1093,7 @@ (e', e_arg) <- checkArg e (t1, rt) <- checkApply loc ftype e_arg case rt of - Arrow _ _ t2 rettype -> + Scalar (Arrow _ _ t2 rettype) -> return $ OpSectionLeft op' (Info ftype) e' (Info $ toStruct t1, Info $ toStruct t2) (Info rettype) loc _ -> typeError loc $ @@ -1075,9 +1103,9 @@ (op', ftype) <- lookupVar loc op (e', e_arg) <- checkArg e case ftype of - Arrow as1 m1 t1 (Arrow as2 m2 t2 ret) -> do - (t2', Arrow _ _ t1' rettype) <- - checkApply loc (Arrow as2 m2 t2 (Arrow as1 m1 t1 ret)) e_arg + Scalar (Arrow as1 m1 t1 (Scalar (Arrow as2 m2 t2 ret))) -> do + (t2', Scalar (Arrow _ _ t1' rettype)) <- + checkApply loc (Scalar $ Arrow as2 m2 t2 $ Scalar $ Arrow as1 m1 t1 ret) e_arg return $ OpSectionRight op' (Info ftype) e' (Info $ toStruct t1', Info $ toStruct t2') (Info rettype) loc _ -> typeError loc $ @@ -1085,22 +1113,24 @@ checkExp (ProjectSection fields NoInfo loc) = do a <- newTypeVar loc "a" - b <- foldM (flip $ mustHaveField loc) a fields - return $ ProjectSection fields (Info $ Arrow mempty Nothing a b) loc + let usage = mkUsage loc "projection at" + b <- foldM (flip $ mustHaveField usage) a fields + return $ ProjectSection fields (Info $ Scalar $ Arrow mempty Unnamed a b) loc checkExp (IndexSection idxes NoInfo loc) = do (t, _) <- newArrayType loc "e" (length idxes) idxes' <- mapM checkDimIndex idxes let t' = stripArray (length $ filter isFix idxes) t return $ IndexSection idxes' (Info $ vacuousShapeAnnotations $ fromStruct $ - Arrow mempty Nothing t t') loc + Scalar $ Arrow mempty Unnamed t t') loc where isFix DimFix{} = True isFix _ = False checkExp (DoLoop mergepat mergeexp form loopbody loc) = sequentially (checkExp mergeexp) $ \mergeexp' _ -> do - zeroOrderType (srclocOf mergeexp) "used as loop variable" (typeOf mergeexp') + zeroOrderType (mkUsage (srclocOf mergeexp) "use as loop variable") + "used as loop variable" (typeOf mergeexp') merge_t <- do merge_t <- expType mergeexp' @@ -1116,7 +1146,7 @@ ((mergepat', form', loopbody'), bodyflow) <- case form of For i uboundexp -> do - uboundexp' <- require anySignedType =<< checkExp uboundexp + uboundexp' <- require "being the bound in a 'for' loop" anySignedType =<< checkExp uboundexp bound_t <- expType uboundexp' bindingIdent i bound_t $ \i' -> noUnique $ bindingPattern mergepat merge_t $ @@ -1128,7 +1158,7 @@ ForIn xpat e -> do (arr_t, _) <- newArrayType (srclocOf e) "e" 1 - e' <- unifies arr_t =<< checkExp e + e' <- unifies "being iterated in a 'for-in' loop" arr_t =<< checkExp e t <- expType e' case t of _ | Just t' <- peelArray 1 t -> @@ -1146,7 +1176,8 @@ While cond -> noUnique $ bindingPattern mergepat merge_t $ \mergepat' -> onlySelfAliasing $ tapOccurences $ - sequentially (unifies (Prim Bool) =<< checkExp cond) $ \cond' _ -> do + sequentially (checkExp cond >>= + unifies "being the condition of a 'while' loop" (Scalar $ Prim Bool)) $ \cond' _ -> do loopbody' <- checkExp loopbody return (mergepat', While cond', @@ -1154,7 +1185,8 @@ mergepat'' <- do loop_t <- expType loopbody' - convergePattern mergepat' (allConsumed bodyflow) loop_t (srclocOf loopbody') + convergePattern mergepat' (allConsumed bodyflow) loop_t $ + mkUsage (srclocOf loopbody') "being (part of) the result of the loop body" let consumeMerge (Id _ (Info pt) ploc) mt | unique pt = consume ploc $ aliases mt @@ -1182,8 +1214,8 @@ let t' = t `setUniqueness` Unique `setAliases` mempty in Id name (Info t') iloc | otherwise = - let t' = case t of Record{} -> t - _ -> t `setUniqueness` Nonunique + let t' = case t of Scalar Record{} -> t + _ -> t `setUniqueness` Nonunique in Id name (Info t') iloc uniquePat (TuplePattern pats ploc) = TuplePattern (map uniquePat pats) ploc @@ -1192,6 +1224,8 @@ uniquePat (PatternAscription p t ploc) = PatternAscription p t ploc uniquePat p@PatternLit{} = p + uniquePat (PatternConstr n t ps ploc) = + PatternConstr n t (map uniquePat ps) ploc -- Make the pattern unique where needed. pat' = uniquePat pat @@ -1201,7 +1235,7 @@ body_t' <- normaliseType body_t pat_t <- normaliseType $ patternType pat' unless (body_t' `subtypeOf` pat_t) $ - unexpectedType body_loc + unexpectedType (srclocOf body_loc) (toStructural body_t') [toStructural pat_t] @@ -1231,7 +1265,7 @@ checkMergeReturn p t checkMergeReturn (PatternAscription p _ _) t = checkMergeReturn p t - checkMergeReturn (RecordPattern pfs _) (Record tfs) = + checkMergeReturn (RecordPattern pfs _) (Scalar (Record tfs)) = sequence_ $ M.elems $ M.intersectionWith checkMergeReturn (M.fromList pfs) tfs checkMergeReturn (TuplePattern pats _) t | Just ts <- isTupleRecord t = zipWithM_ checkMergeReturn pats ts @@ -1239,39 +1273,41 @@ return () (pat_cons, _) <- execStateT (checkMergeReturn pat' body_t') (mempty, mempty) let body_cons' = body_cons <> S.map aliasVar pat_cons - if body_cons' == body_cons && patternPatternType pat' == patternPatternType pat + if body_cons' == body_cons && patternType pat' == patternType pat then return pat' else convergePattern pat' body_cons' body_t' body_loc -checkExp (VConstr0 name NoInfo loc) = do +checkExp (Constr name es NoInfo loc) = do t <- newTypeVar loc "t" - mustHaveConstr loc name t - return $ VConstr0 name (Info t) loc - -checkExp (Match _ [] NoInfo loc) = - typeError loc "Match expressions must have at least one case." + es' <- mapM checkExp es + ets <- mapM expType es' + mustHaveConstr (mkUsage loc "use of constructor") name t (toStructural <$> ets) + -- A sum value aliases *anything* that went into its construction. + let als = mconcat (map aliases ets) + return $ Constr name es' (Info $ t `addAliases` (<>als)) loc -checkExp (Match e (c:cs) NoInfo loc) = +checkExp (Match e cs NoInfo loc) = sequentially (checkExp e) $ \e' _ -> do mt <- expType e' - (cs', t) <- checkCases mt c cs - zeroOrderType loc "returned from pattern match" t + (cs', t) <- checkCases mt cs + zeroOrderType (mkUsage loc "being returned 'match'") "returned from pattern match" t return $ Match e' cs' (Info t) loc checkCases :: PatternType - -> CaseBase NoInfo Name - -> [CaseBase NoInfo Name] - -> TermTypeM ([CaseBase Info VName], PatternType) -checkCases mt c [] = do - (c', t) <- checkCase mt c - return ([c'], t) -checkCases mt c (c2:cs) = do - (((c', c_t), (cs', cs_t)), dflow) <- - tapOccurences $ checkCase mt c `alternative` checkCases mt c2 cs - unify (srclocOf c) (toStructural c_t) (toStructural cs_t) - let t = unifyTypeAliases c_t cs_t `addAliases` - (`S.difference` S.map AliasBound (allConsumed dflow)) - return (c':cs', t) + -> NE.NonEmpty (CaseBase NoInfo Name) + -> TermTypeM (NE.NonEmpty (CaseBase Info VName), PatternType) +checkCases mt rest_cs = + case NE.uncons rest_cs of + (c, Nothing) -> do + (c', t) <- checkCase mt c + return (c' NE.:| [], t) + (c, Just cs) -> do + (((c', c_t), (cs', cs_t)), dflow) <- + tapOccurences $ checkCase mt c `alternative` checkCases mt cs + unify (mkUsage (srclocOf c) "pattern match") (toStructural c_t) (toStructural cs_t) + let t = unifyTypeAliases c_t cs_t `addAliases` + (`S.difference` S.map AliasBound (allConsumed dflow)) + return (NE.cons c' cs', t) checkCase :: PatternType -> CaseBase NoInfo Name -> TermTypeM (CaseBase Info VName, PatternType) @@ -1285,7 +1321,7 @@ -- unmatched pattern warnings by the type checker. data Unmatched p = UnmatchedNum p [ExpBase Info VName] | UnmatchedBool p - | UnmatchedEnum p + | UnmatchedConstr p | Unmatched p deriving (Functor, Show) @@ -1293,7 +1329,7 @@ ppr um = case um of (UnmatchedNum p nums) -> ppr' p <+> text "where p is not one of" <+> ppr nums (UnmatchedBool p) -> ppr' p - (UnmatchedEnum p) -> ppr' p + (UnmatchedConstr p) -> ppr' p (Unmatched p) -> ppr' p where ppr' (PatternAscription p t _) = ppr p <> text ":" <+> ppr t @@ -1304,6 +1340,7 @@ where ppField (name, t) = text (nameToString name) <> equals <> ppr' t ppr' Wildcard{} = text "_" ppr' (PatternLit e _ _) = ppr e + ppr' (PatternConstr n _ ps _) = text "#" <> ppr n <+> sep (map ppr' ps) unpackPat :: Pattern -> [Maybe Pattern] unpackPat Wildcard{} = [Nothing] @@ -1313,26 +1350,31 @@ unpackPat (RecordPattern fs _) = Just . snd <$> sortFields (M.fromList fs) unpackPat (PatternAscription p _ _) = unpackPat p unpackPat p@PatternLit{} = [Just p] +unpackPat p@PatternConstr{} = [Just p] wildPattern :: Pattern -> Int -> Unmatched Pattern -> Unmatched Pattern -wildPattern (TuplePattern ps loc) pos um = f <$> um - where f p = TuplePattern (take (pos - 1) ps' ++ [p] ++ drop pos ps') loc +wildPattern (TuplePattern ps loc) pos um = wildTuple <$> um + where wildTuple p = TuplePattern (take (pos - 1) ps' ++ [p] ++ drop pos ps') loc ps' = map wildOut ps - wildOut p = Wildcard (Info (patternPatternType p)) (srclocOf p) + wildOut p = Wildcard (Info (patternType p)) (srclocOf p) wildPattern (RecordPattern fs loc) pos um = wildRecord <$> um - where wildRecord p = - RecordPattern (take (pos - 1) fs' ++ [(fst (fs!!(pos - 1)), p)] ++ drop pos fs') loc - fs' = map wildOut fs - wildOut (f,p) = (f, Wildcard (Info (patternPatternType p)) (srclocOf p)) + where wildRecord p = + RecordPattern (take (pos - 1) fs' ++ [(fst (fs!!(pos - 1)), p)] ++ drop pos fs') loc + fs' = map wildOut fs + wildOut (f,p) = (f, Wildcard (Info (patternType p)) (srclocOf p)) wildPattern (PatternAscription p _ _) pos um = wildPattern p pos um wildPattern (PatternParens p _) pos um = wildPattern p pos um +wildPattern (PatternConstr n t ps loc) pos um = wildConstr <$> um + where wildConstr p = PatternConstr n t (take (pos - 1) ps' ++ [p] ++ drop pos ps') loc + ps' = map wildOut ps + wildOut p = Wildcard (Info (patternType p)) (srclocOf p) wildPattern _ _ um = um checkUnmatched :: (MonadBreadCrumbs m, MonadTypeChecker m) => Exp -> m () checkUnmatched e = void $ checkUnmatched' e >> astMap tv e where checkUnmatched' (Match _ cs _ loc) = - let ps = map (\(CasePat p _ _) -> p) cs - in case unmatched id ps of + let ps = fmap (\(CasePat p _ _) -> p) cs + in case unmatched id $ NE.toList ps of [] -> return () ps' -> typeError loc $ "Unmatched cases in match expression: \n" ++ unlines (map ((" " ++) . pretty) ps') @@ -1345,37 +1387,88 @@ , mapOnPatternType = pure } +-- | A data type for constructor patterns. This is used to make the +-- code for detecting unmatched constructors cleaner, by separating +-- the constructor-pattern cases from other cases. +data ConstrPat = ConstrPat { constrName :: Name + , constrType :: PatternType + , constrPayload :: [Pattern] + , constrSrcLoc :: SrcLoc + } + +-- Be aware of these fishy equality instances! + +instance Eq ConstrPat where + ConstrPat c1 _ _ _ == ConstrPat c2 _ _ _ = c1 == c2 + +instance Ord ConstrPat where + ConstrPat c1 _ _ _ `compare` ConstrPat c2 _ _ _ = c1 `compare` c2 + unmatched :: (Unmatched Pattern -> Unmatched Pattern) -> [Pattern] -> [Unmatched Pattern] -unmatched hole (p:ps) - | sameStructure labeledCols = do +unmatched hole orig_ps + | p:_ <- orig_ps, + sameStructure labeledCols = do (i, cols) <- labeledCols - let hole' p' = hole $ wildPattern p i p' + let hole' = if isConstr p then hole else hole . wildPattern p i case sequence cols of - Nothing -> [] + Nothing -> [] Just cs | all isPatternLit cs -> map hole' $ localUnmatched cs | otherwise -> unmatched hole' cs + | otherwise = [] - where labeledCols = zip [1..] $ transpose $ map unpackPat (p:ps) + where labeledCols = zip [1..] $ transpose $ map unpackPat orig_ps localUnmatched :: [Pattern] -> [Unmatched Pattern] localUnmatched [] = [] localUnmatched ps'@(p':_) = case patternType p' of - Enum cs'' -> - let matched = nub $ mapMaybe (pExp >=> constr) ps' - in map (UnmatchedEnum . buildEnum (Enum cs'')) $ cs'' \\ matched - Prim t - | not (any idOrWild ps') -> + Scalar (Sum cs'') -> + -- We now know that we are matching a sum type, and thus + -- that all patterns ps' are constructors (checked by + -- 'all isPatternLit' before this function is called). + let constrs = M.keys cs'' + matched = mapMaybe constr ps' + unmatched' = map (UnmatchedConstr . buildConstr cs'') $ + constrs \\ map constrName matched + in case unmatched' of + [] -> + let constrGroups = group (sort matched) + removedConstrs = mapMaybe stripConstrs constrGroups + transposed = (fmap . fmap) transpose removedConstrs + findUnmatched (pc, trans) = do + col <- trans + case col of + [] -> [] + ((i, _):_) -> unmatched (wilder i pc) (map snd col) + wilder i pc s = (`PatternParens` noLoc) <$> wildPattern pc i s + in concatMap findUnmatched transposed + _ -> unmatched' + Scalar (Prim t) | not (any idOrWild ps') -> + -- We now know that we are matching a sum type, and thus + -- that all patterns ps' are literals (checked by 'all + -- isPatternLit' before this function is called). case t of Bool -> let matched = nub $ mapMaybe (pExp >=> bool) $ filter isPatternLit ps' - in map (UnmatchedBool . buildBool (Prim t)) $ [True, False] \\ matched + in map (UnmatchedBool . buildBool (Scalar (Prim t))) $ [True, False] \\ matched _ -> let matched = mapMaybe pExp $ filter isPatternLit ps' - in [UnmatchedNum (buildId (Info (Prim t)) "p") matched] + in [UnmatchedNum (buildId (Info $ Scalar $ Prim t) "p") matched] _ -> [] + isConstr PatternConstr{} = True + isConstr (PatternParens p _) = isConstr p + isConstr _ = False + + + stripConstrs :: [ConstrPat] -> Maybe (Pattern, [[(Int, Pattern)]]) + stripConstrs (pc@ConstrPat{} : cs') = Just (unConstr pc, stripConstr pc : map stripConstr cs') + stripConstrs [] = Nothing + + stripConstr :: ConstrPat -> [(Int, Pattern)] + stripConstr (ConstrPat _ _ ps' _) = zip [1..] ps' + sameStructure [] = True sameStructure (x:xs) = all (\y -> length y == length x' ) xs' where (x':xs') = map snd (x:xs) @@ -1383,13 +1476,18 @@ pExp (PatternLit e' _ _) = Just e' pExp _ = Nothing - constr (VConstr0 c _ _) = Just c - constr (Ascript e' _ _ _) = constr e' + constr (PatternConstr c (Info t) ps loc) = Just $ ConstrPat c t ps loc + constr (PatternParens p _) = constr p + constr (PatternAscription p' _ _) = constr p' constr _ = Nothing + unConstr p = + PatternConstr (constrName p) (Info $ constrType p) (constrPayload p) (constrSrcLoc p) + isPatternLit PatternLit{} = True isPatternLit (PatternAscription p' _ _) = isPatternLit p' isPatternLit (PatternParens p' _) = isPatternLit p' + isPatternLit PatternConstr{} = True isPatternLit _ = False idOrWild Id{} = True @@ -1401,8 +1499,13 @@ bool (Literal (BoolValue b) _ ) = Just b bool _ = Nothing - buildEnum t c = - PatternLit (VConstr0 c (Info t) noLoc) (Info t) noLoc + buildConstr m c = + let t = Scalar $ Sum m + cs = m M.! c + wildCS = map (\ct -> Wildcard (Info ct) noLoc) cs + in if null wildCS + then PatternConstr c (Info t) [] noLoc + else PatternParens (PatternConstr c (Info t) wildCS noLoc) noLoc buildBool t b = PatternLit (Literal (BoolValue b) noLoc) (Info (vacuousShapeAnnotations t)) noLoc buildId t n = @@ -1410,8 +1513,6 @@ -- exists exclusively for printing warnings. Id (VName (nameFromString n) (-1)) t noLoc -unmatched _ _ = [] - checkIdent :: IdentBase NoInfo Name -> TermTypeM Ident checkIdent (Ident name _ loc) = do (QualName _ name', vt) <- lookupVar loc (qualName name) @@ -1419,12 +1520,11 @@ checkDimIndex :: DimIndexBase NoInfo Name -> TermTypeM DimIndex checkDimIndex (DimFix i) = - DimFix <$> (unifies (Prim $ Signed Int32) =<< checkExp i) + DimFix <$> (unifies "use as index" (Scalar $ Prim $ Signed Int32) =<< checkExp i) checkDimIndex (DimSlice i j s) = - DimSlice - <$> maybe (return Nothing) (fmap Just . unifies (Prim $ Signed Int32) <=< checkExp) i - <*> maybe (return Nothing) (fmap Just . unifies (Prim $ Signed Int32) <=< checkExp) j - <*> maybe (return Nothing) (fmap Just . unifies (Prim $ Signed Int32) <=< checkExp) s + DimSlice <$> check i <*> check j <*> check s + where check = maybe (return Nothing) $ + fmap Just . unifies "use as index" (Scalar $ Prim $ Signed Int32) <=< checkExp sequentially :: TermTypeM a -> (a -> Occurences -> TermTypeM b) -> TermTypeM b sequentially m1 m2 = do @@ -1446,8 +1546,8 @@ checkApply :: SrcLoc -> PatternType -> Arg -> TermTypeM (PatternType, PatternType) -checkApply loc (Arrow as _ tp1 tp2) (argtype, dflow, argloc) = do - unify argloc (toStructural tp1) (toStructural argtype) +checkApply loc (Scalar (Arrow as _ tp1 tp2)) (argtype, dflow, argloc) = do + unify (mkUsage argloc "use as function argument") (toStructural tp1) (toStructural argtype) -- Perform substitutions of instantiated variables in the types. tp1' <- normaliseType tp1 @@ -1462,16 +1562,17 @@ case anyConsumption dflow of Just c -> let msg = "of value computed with consumption at " ++ locStr (location c) - in zeroOrderType argloc msg tp1 + in zeroOrderType (mkUsage argloc "potential consumption in expression") msg tp1 _ -> return () occur $ dflow `seqOccurences` occurs let tp2'' = anyDimShapeAnnotations $ returnType tp2' (diet tp1') argtype' return (tp1', tp2'') -checkApply loc tfun@TypeVar{} arg = do +checkApply loc tfun@(Scalar TypeVar{}) arg = do tv <- newTypeVar loc "b" - unify loc (toStructural tfun) $ Arrow mempty Nothing (toStructural (argType arg)) tv + unify (mkUsage loc "use as function") (toStructural tfun) $ + Scalar $ Arrow mempty Unnamed (toStructural (argType arg)) tv constraints <- getConstraints checkApply loc (applySubst (`lookupSubst` constraints) tfun) arg @@ -1492,51 +1593,53 @@ returnType (Array als Nonunique et shape) d arg = Array (als<>arg_als) Unique et shape -- Intentional! where arg_als = aliases $ maskAliases arg d -returnType (Record fs) d arg = - Record $ fmap (\et -> returnType et d arg) fs -returnType (Prim t) _ _ = Prim t -returnType (TypeVar _ Unique t targs) _ _ = - TypeVar mempty Unique t targs -returnType (TypeVar als Nonunique t targs) d arg = - TypeVar (als<>arg_als) Unique t targs -- Intentional! +returnType (Scalar (Record fs)) d arg = + Scalar $ Record $ fmap (\et -> returnType et d arg) fs +returnType (Scalar (Prim t)) _ _ = + Scalar $ Prim t +returnType (Scalar (TypeVar _ Unique t targs)) _ _ = + Scalar $ TypeVar mempty Unique t targs +returnType (Scalar (TypeVar als Nonunique t targs)) d arg = + Scalar $ TypeVar (als<>arg_als) Unique t targs -- Intentional! where arg_als = aliases $ maskAliases arg d -returnType (Arrow _ v t1 t2) d arg = - Arrow als v (t1 `setAliases` mempty) (t2 `setAliases` als) +returnType (Scalar (Arrow _ v t1 t2)) d arg = + Scalar $ Arrow als v (t1 `setAliases` mempty) (t2 `setAliases` als) where als = aliases $ maskAliases arg d -returnType (Enum cs) _ _ = Enum cs +returnType (Scalar (Sum cs)) d arg = + Scalar $ Sum $ (fmap . fmap) (\et -> returnType et d arg) cs -- | @t `maskAliases` d@ removes aliases (sets them to 'mempty') from --- the parts of @t@ that are denoted as 'Consumed' by the 'Diet' @d@. +-- the parts of @t@ that are denoted as consumed by the 'Diet' @d@. maskAliases :: Monoid as => TypeBase shape as -> Diet -> TypeBase shape as maskAliases t Consume = t `setAliases` mempty maskAliases t Observe = t -maskAliases (Record ets) (RecordDiet ds) = - Record $ M.intersectionWith maskAliases ets ds +maskAliases (Scalar (Record ets)) (RecordDiet ds) = + Scalar $ Record $ M.intersectionWith maskAliases ets ds maskAliases t FuncDiet{} = t maskAliases _ _ = error "Invalid arguments passed to maskAliases." consumeArg :: SrcLoc -> PatternType -> Diet -> TermTypeM [Occurence] -consumeArg loc (Record ets) (RecordDiet ds) = +consumeArg loc (Scalar (Record ets)) (RecordDiet ds) = concat . M.elems <$> traverse (uncurry $ consumeArg loc) (M.intersectionWith (,) ets ds) consumeArg loc (Array _ Nonunique _ _) Consume = typeError loc "Consuming parameter passed non-unique argument." -consumeArg loc (Arrow _ _ t1 _) (FuncDiet d _) +consumeArg loc (Scalar (Arrow _ _ t1 _)) (FuncDiet d _) | not $ contravariantArg t1 d = typeError loc "Non-consuming higher-order parameter passed consuming argument." where contravariantArg (Array _ Unique _ _) Observe = False - contravariantArg (TypeVar _ Unique _ _) Observe = + contravariantArg (Scalar (TypeVar _ Unique _ _)) Observe = False - contravariantArg (Record ets) (RecordDiet ds) = + contravariantArg (Scalar (Record ets)) (RecordDiet ds) = and (M.intersectionWith contravariantArg ets ds) - contravariantArg (Arrow _ _ tp tr) (FuncDiet dp dr) = + contravariantArg (Scalar (Arrow _ _ tp tr)) (FuncDiet dp dr) = contravariantArg tp dp && contravariantArg tr dr contravariantArg _ _ = True -consumeArg loc (Arrow _ _ _ t2) (FuncDiet _ pd) = +consumeArg loc (Scalar (Arrow _ _ _ t2)) (FuncDiet _ pd) = consumeArg loc t2 pd consumeArg loc at Consume = return [consumption (aliases at) loc] consumeArg loc at _ = return [observation (aliases at) loc] @@ -1580,36 +1683,38 @@ -- only make very conservative fixing. fixOverloadedTypes :: TermTypeM () fixOverloadedTypes = getConstraints >>= mapM_ fixOverloaded . M.toList - where fixOverloaded (v, Overloaded ots loc) + where fixOverloaded (v, Overloaded ots usage) | Signed Int32 `elem` ots = do - unify loc (TypeVar () Nonunique (typeName v) []) $ Prim $ Signed Int32 - warn loc "Defaulting ambiguous type to `i32`." + unify usage (Scalar (TypeVar () Nonunique (typeName v) [])) $ + Scalar $ Prim $ Signed Int32 + warn usage "Defaulting ambiguous type to `i32`." | FloatType Float64 `elem` ots = do - unify loc (TypeVar () Nonunique (typeName v) []) $ Prim $ FloatType Float64 - warn loc "Defaulting ambiguous type to `f64`." + unify usage (Scalar (TypeVar () Nonunique (typeName v) [])) $ + Scalar $ Prim $ FloatType Float64 + warn usage "Defaulting ambiguous type to `f64`." | otherwise = - typeError loc $ + typeError usage $ unlines ["Type is ambiguous (could be one of " ++ intercalate ", " (map pretty ots) ++ ").", "Add a type annotation to disambiguate the type."] - fixOverloaded (_, NoConstraint _ loc) = - typeError loc $ unlines ["Type of expression is ambiguous.", - "Add a type annotation to disambiguate the type."] + fixOverloaded (_, NoConstraint _ usage) = + typeError usage $ unlines ["Type of expression is ambiguous.", + "Add a type annotation to disambiguate the type."] - fixOverloaded (_, Equality loc) = - typeError loc $ unlines ["Type is ambiguous (must be equality type).", - "Add a type annotation to disambiguate the type."] + fixOverloaded (_, Equality usage) = + typeError usage $ unlines ["Type is ambiguous (must be equality type).", + "Add a type annotation to disambiguate the type."] - fixOverloaded (_, HasFields fs loc) = - typeError loc $ unlines ["Type is ambiguous (must be record with fields {" ++ fs' ++ "}).", - "Add a type annotation to disambiguate the type."] - where fs' = intercalate ", " $ map field $ M.toList fs - field (l, t) = pretty l ++ ": " ++ pretty t + fixOverloaded (_, HasFields fs usage) = + typeError usage $ unlines ["Type is ambiguous. Must be record with fields:", + unlines $ map field $ M.toList fs, + "Add a type annotation to disambiguate the type."] + where field (l, t) = pretty $ indent 2 $ ppr l <> colon <+> align (ppr t) - fixOverloaded (_, HasConstrs cs loc) = - typeError loc $ unlines [ "Type is ambiguous (must be an enum with constructors: " ++ cs' ++ ")." - ,"Add a type annotation to disambiguate the type."] - where cs' = intercalate " | " $ map (\c -> '#' : pretty c) cs + fixOverloaded (_, HasConstrs cs usage) = + typeError usage $ unlines [ "Type is ambiguous (must be a sum type with constructors: " ++ + pretty (Sum cs) ++ ")." + , "Add a type annotation to disambiguate the type."] fixOverloaded _ = return () @@ -1682,9 +1787,9 @@ Nothing -> return () - tag u = S.map $ \name -> (u, name) + tag u = S.map (u,) - returnAliasing (Record ets1) (Record ets2) = + returnAliasing (Scalar (Record ets1)) (Scalar (Record ets2)) = concat $ M.elems $ M.intersectionWith returnAliasing ets1 ets2 returnAliasing expected got = [(uniqueness expected, S.map aliasVar $ aliases got)] @@ -1694,7 +1799,9 @@ onDubiousNames $ S.filter patternNameButNotParamName $ mconcat $ map typeDimNames $ rettype : map patternStructType params - where param_names = S.fromList $ mapMaybe (fst . patternParam) params + where named (Named v) = Just v + named Unnamed = Nothing + param_names = S.fromList $ mapMaybe (named . fst . patternParam) params all_pattern_names = S.map identName $ mconcat $ map patternIdents params patternNameButNotParamName v = v `S.member` all_pattern_names && not (v `S.member` param_names) onDubiousNames dubious @@ -1728,8 +1835,8 @@ inferReturnUniqueness params t = let forbidden = aliasesMultipleTimes t uniques = uniqueParamNames params - delve (Record fs) = - Record $ M.map delve fs + delve (Scalar (Record fs)) = + Scalar $ Record $ M.map delve fs delve t' | all (`S.member` uniques) (boundArrayAliases t'), not $ any ((`S.member` forbidden) . aliasVar) (aliases t') = @@ -1741,7 +1848,7 @@ -- An alias inhibits uniqueness if it is used in disjoint values. aliasesMultipleTimes :: PatternType -> Names aliasesMultipleTimes = S.fromList . map fst . filter ((>1) . snd) . M.toList . delve - where delve (Record fs) = + where delve (Scalar (Record fs)) = foldl' (M.unionWith (+)) mempty $ map delve $ M.elems fs delve t = M.fromList $ zip (map aliasVar $ S.toList (aliases t)) $ repeat (1::Int) @@ -1754,11 +1861,12 @@ boundArrayAliases :: PatternType -> S.Set VName boundArrayAliases (Array als _ _ _) = boundAliases als -boundArrayAliases Prim{} = mempty -boundArrayAliases Enum{} = mempty -boundArrayAliases (Record fs) = foldMap boundArrayAliases fs -boundArrayAliases (TypeVar als _ _ _) = boundAliases als -boundArrayAliases Arrow{} = mempty +boundArrayAliases (Scalar Prim{}) = mempty +boundArrayAliases (Scalar (Record fs)) = foldMap boundArrayAliases fs +boundArrayAliases (Scalar (TypeVar als _ _ _)) = boundAliases als +boundArrayAliases (Scalar Arrow{}) = mempty +boundArrayAliases (Scalar (Sum fs)) = + mconcat $ concatMap (map boundArrayAliases) $ M.elems fs -- | The set of in-scope variables that are being aliased. boundAliases :: Aliasing -> S.Set VName @@ -1769,8 +1877,8 @@ nothingMustBeUnique :: SrcLoc -> TypeBase () () -> TermTypeM () nothingMustBeUnique loc = check where check (Array _ Unique _ _) = bad - check (TypeVar _ Unique _ _) = bad - check (Record fs) = mapM_ check fs + check (Scalar (TypeVar _ Unique _ _)) = bad + check (Scalar (Record fs)) = mapM_ check fs check _ = return () bad = typeError loc "A top-level constant cannot have a unique type." @@ -1818,7 +1926,7 @@ case maybe_rettype of Just rettype -> do let rettype_structural = toStructural rettype - void $ unifies rettype_structural body' + void $ unifies "return type annotation" rettype_structural body' -- We also have to make sure that uniqueness matches. This is done -- explicitly, because uniqueness is ignored by unification. rettype' <- normaliseType rettype @@ -1834,33 +1942,27 @@ -- | Find at all type variables in the given type that are covered by -- the constraints, and produce type parameters that close over them. --- Produce an error if the given list of type parameters is non-empty, --- yet does not cover all type variables in the type. +-- +-- The passed-in list of type parameters is always prepended to the +-- produced list of type parameters. closeOverTypes :: Constraints -> [TypeParam] -> StructType -> TermTypeM [TypeParam] closeOverTypes substs tparams t = - case tparams of - [] -> fmap catMaybes $ mapM closeOver $ M.toList substs' - _ -> do mapM_ checkClosedOver $ M.toList substs' - return tparams - where substs' = M.filterWithKey (\k _ -> k `S.member` visible) substs + fmap ((tparams++) . catMaybes) $ mapM closeOver $ M.toList to_close_over + where to_close_over = M.filterWithKey (\k _ -> k `S.member` visible) substs visible = typeVars t - checkClosedOver (k, v) - | not (canBeClosedOver v) || - k `elem` map typeParamName tparams = return () - | otherwise = - typeError (srclocOf v) $ - unlines ["Type variable " ++ quote (prettyName k) ++ - " not closed over by type parameters " ++ - intercalate ", " (map pretty tparams) ++ ".", - "This is usually because a parameter needs a type annotation."] - - canBeClosedOver NoConstraint{} = True - canBeClosedOver _ = False - - closeOver (k, NoConstraint (Just Unlifted) loc) = return $ Just $ TypeParamType Unlifted k loc - closeOver (k, NoConstraint _ loc) = return $ Just $ TypeParamType Lifted k loc - closeOver (_, _) = return Nothing + -- Avoid duplicate type parameters. + closeOver (k, _) + | k `elem` map typeParamName tparams = + return Nothing + closeOver (k, NoConstraint (Just Unlifted) usage) = + return $ Just $ TypeParamType Unlifted k $ srclocOf usage + closeOver (k, NoConstraint _ usage) = + return $ Just $ TypeParamType Lifted k $ srclocOf usage + closeOver (k, ParamType l loc) = + return $ Just $ TypeParamType l k loc + closeOver (_, _) = + return Nothing --- Consumption @@ -1948,10 +2050,8 @@ -> TypeBase dim as -> ShapeDecl dim -> Uniqueness -> TermTypeM (TypeBase dim as) arrayOfM loc t shape u = do - zeroOrderType loc "used in array" t - maybe nope return $ arrayOf t shape u - where nope = typeError loc $ - "Cannot form an array with elements of type " ++ pretty t + zeroOrderType (mkUsage loc "use as array element") "used in array" t + return $ arrayOf t shape u -- | Perform substitutions of instantiated variables on the type -- annotations (including the instance lists) of an expression, or diff --git a/src/Language/Futhark/TypeChecker/Types.hs b/src/Language/Futhark/TypeChecker/Types.hs --- a/src/Language/Futhark/TypeChecker/Types.hs +++ b/src/Language/Futhark/TypeChecker/Types.hs @@ -10,6 +10,7 @@ , checkForDuplicateNames , checkForDuplicateNamesInType , checkTypeParams + , typeParamToArg , typeExpUses , checkShapeParamUses @@ -41,31 +42,40 @@ -- unification cannot happen, 'Nothing' is returned, otherwise a type -- that combines the aliasing of @t1@ and @t2@ is returned. -- Uniqueness is unified with @uf@. -unifyTypesU :: (Monoid als, Eq als, ArrayDim dim) => - (Uniqueness -> Uniqueness -> Maybe Uniqueness) - -> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als) -unifyTypesU _ (Prim t1) (Prim t2) +unifyTypesU :: (Monoid als, ArrayDim dim) => + (Uniqueness -> Uniqueness -> Maybe Uniqueness) + -> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als) +unifyTypesU uf (Array als1 u1 et1 shape1) (Array als2 u2 et2 shape2) = + Array (als1 <> als2) <$> uf u1 u2 + <*> unifyScalarTypes uf et1 et2 <*> unifyShapes shape1 shape2 +unifyTypesU uf (Scalar t1) (Scalar t2) = Scalar <$> unifyScalarTypes uf t1 t2 +unifyTypesU _ _ _ = Nothing + +unifyScalarTypes :: (Monoid als, ArrayDim dim) => + (Uniqueness -> Uniqueness -> Maybe Uniqueness) + -> ScalarTypeBase dim als -> ScalarTypeBase dim als -> Maybe (ScalarTypeBase dim als) +unifyScalarTypes _ (Prim t1) (Prim t2) | t1 == t2 = Just $ Prim t1 | otherwise = Nothing -unifyTypesU uf (TypeVar als1 u1 t1 targs1) (TypeVar als2 u2 t2 targs2) +unifyScalarTypes uf (TypeVar als1 u1 t1 targs1) (TypeVar als2 u2 t2 targs2) | t1 == t2 = do u3 <- uf u1 u2 targs3 <- zipWithM (unifyTypeArgs uf) targs1 targs2 Just $ TypeVar (als1 <> als2) u3 t1 targs3 | otherwise = Nothing -unifyTypesU uf (Array als1 u1 et1 shape1) (Array als2 u2 et2 shape2) = - Array (als1 <> als2) <$> uf u1 u2 - <*> unifyArrayElemTypes uf et1 et2 <*> unifyShapes shape1 shape2 -unifyTypesU uf (Record ts1) (Record ts2) +unifyScalarTypes uf (Record ts1) (Record ts2) | length ts1 == length ts2, sort (M.keys ts1) == sort (M.keys ts2) = Record <$> traverse (uncurry (unifyTypesU uf)) (M.intersectionWith (,) ts1 ts2) -unifyTypesU uf (Arrow as1 mn1 t1 t1') (Arrow as2 _ t2 t2') = +unifyScalarTypes uf (Arrow as1 mn1 t1 t1') (Arrow as2 _ t2 t2') = Arrow (as1 <> as2) mn1 <$> unifyTypesU (flip uf) t1 t2 <*> unifyTypesU uf t1' t2' -unifyTypesU _ e1@Enum{} e2@Enum{} - | e1 == e2 = Just e1 -unifyTypesU _ _ _ = Nothing +unifyScalarTypes uf (Sum cs1) (Sum cs2) + | length cs1 == length cs2, + sort (M.keys cs1) == sort (M.keys cs2) = + Sum <$> traverse (uncurry (zipWithM (unifyTypesU uf))) + (M.intersectionWith (,) cs1 cs2) +unifyScalarTypes _ _ _ = Nothing unifyTypeArgs :: (ArrayDim dim) => (Uniqueness -> Uniqueness -> Maybe Uniqueness) @@ -77,39 +87,6 @@ unifyTypeArgs _ _ _ = Nothing -unifyArrayElemTypes :: (ArrayDim dim) => - (Uniqueness -> Uniqueness -> Maybe Uniqueness) - -> ArrayElemTypeBase dim - -> ArrayElemTypeBase dim - -> Maybe (ArrayElemTypeBase dim) -unifyArrayElemTypes _ (ArrayPrimElem bt1) (ArrayPrimElem bt2) - | bt1 == bt2 = - Just $ ArrayPrimElem bt1 -unifyArrayElemTypes _ (ArrayPolyElem bt1 targs1) (ArrayPolyElem bt2 targs2) - | bt1 == bt2, targs1 == targs2 = - Just $ ArrayPolyElem bt1 targs1 -unifyArrayElemTypes uf (ArrayRecordElem et1) (ArrayRecordElem et2) - | sort (M.keys et1) == sort (M.keys et2) = - ArrayRecordElem <$> - traverse (uncurry $ unifyRecordArrayElemTypes uf) (M.intersectionWith (,) et1 et2) -unifyArrayElemTypes _ (ArrayEnumElem cs1) (ArrayEnumElem cs2) - | cs1 == cs2 = - Just $ ArrayEnumElem cs1 -unifyArrayElemTypes _ _ _ = - Nothing - -unifyRecordArrayElemTypes :: (ArrayDim dim) => - (Uniqueness -> Uniqueness -> Maybe Uniqueness) - -> RecordArrayElemTypeBase dim - -> RecordArrayElemTypeBase dim - -> Maybe (RecordArrayElemTypeBase dim) -unifyRecordArrayElemTypes uf (RecordArrayElem et1) (RecordArrayElem et2) = - RecordArrayElem <$> unifyArrayElemTypes uf et1 et2 -unifyRecordArrayElemTypes uf (RecordArrayArrayElem et1 shape1) (RecordArrayArrayElem et2 shape2) = - RecordArrayArrayElem <$> unifyArrayElemTypes uf et1 et2 <*> unifyShapes shape1 shape2 -unifyRecordArrayElemTypes _ _ _ = - Nothing - -- | @x \`subtypeOf\` y@ is true if @x@ is a subtype of @y@ (or equal to -- @y@), meaning @x@ is valid whenever @y@ is. subtypeOf :: ArrayDim dim => @@ -158,12 +135,14 @@ let fs' = fmap (\(x,_,_) -> x) fs_ts_ls ts_s = fmap (\(_,y,_) -> y) fs_ts_ls ls = fmap (\(_,_,z) -> z) fs_ts_ls - return (TERecord (M.toList fs') loc, Record ts_s, foldl' max Unlifted ls) + return (TERecord (M.toList fs') loc, + Scalar $ Record ts_s, + foldl' max Unlifted ls) checkTypeExp (TEArray t d loc) = do (t', st, l) <- checkTypeExp t d' <- checkDimDecl d case (l, arrayOf st (ShapeDecl [d']) Nonunique) of - (Unlifted, Just st') -> return (TEArray t' d' loc, st', Unlifted) + (Unlifted, st') -> return (TEArray t' d' loc, st', Unlifted) _ -> throwError $ TypeError loc $ "Cannot create array with elements of type `" ++ pretty st ++ "` (might be functional)." where checkDimDecl AnyDim = @@ -177,12 +156,12 @@ unless (mayContainArray st) $ warn loc $ "Declaring `" <> pretty st <> "` as unique has no effect." return (TEUnique t' loc, st `setUniqueness` Unique, l) - where mayContainArray Prim{} = False + where mayContainArray (Scalar Prim{}) = False mayContainArray Array{} = True - mayContainArray (Record fs) = any mayContainArray fs - mayContainArray TypeVar{} = True - mayContainArray Arrow{} = False - mayContainArray Enum{} = False + mayContainArray (Scalar (Record fs)) = any mayContainArray fs + mayContainArray (Scalar TypeVar{}) = True + mayContainArray (Scalar Arrow{}) = False + mayContainArray (Scalar (Sum cs)) = (any . any) mayContainArray cs checkTypeExp (TEArrow (Just v) t1 t2 loc) = do (t1', st1, _) <- checkTypeExp t1 bindSpaced [(Term, v)] $ do @@ -191,13 +170,13 @@ localEnv env $ do (t2', st2, _) <- checkTypeExp t2 return (TEArrow (Just v') t1' t2' loc, - Arrow mempty (Just v') st1 st2, + Scalar $ Arrow mempty (Named v') st1 st2, Lifted) checkTypeExp (TEArrow Nothing t1 t2 loc) = do (t1', st1, _) <- checkTypeExp t1 (t2', st2, _) <- checkTypeExp t2 return (TEArrow Nothing t1' t2' loc, - Arrow mempty Nothing st1 st2, + Scalar $ Arrow mempty Unnamed st1 st2, Lifted) checkTypeExp ote@TEApply{} = do (tname, tname_loc, targs) <- rootAndArgs ote @@ -240,13 +219,22 @@ throwError $ TypeError tloc $ "Type argument " ++ pretty a ++ " not valid for a type parameter " ++ pretty p -checkTypeExp t@(TEEnum names loc) = do - unless (sort names == sort (nub names)) $ +checkTypeExp t@(TESum cs loc) = do + let constructors = map fst cs + unless (sort constructors == sort (nub constructors)) $ throwError $ TypeError loc $ "Duplicate constructors in " ++ pretty t - unless (length names <= 256) $ - throwError $ TypeError loc "Enums must have 256 or fewer constructors." - return (TEEnum names loc, Enum names, Unlifted) + unless (length constructors <= 256) $ + throwError $ TypeError loc "Sum types must have 256 or fewer constructors." + + cs_ts_ls <- (traverse . traverse) checkTypeExp $ M.fromList cs + let cs' = (fmap . fmap) (\(x,_,_) -> x) cs_ts_ls + ts_s = (fmap . fmap) (\(_, y, _) -> y) cs_ts_ls + ls = (concatMap . fmap) (\(_, _, z) -> z) cs_ts_ls + return (TESum (M.toList cs') loc, + Scalar $ Sum ts_s, + foldl' max Unlifted ls) + -- | Check for duplication of names inside a pattern group. Produces -- a description of all names used in the pattern group. checkForDuplicateNames :: MonadTypeChecker m => @@ -259,13 +247,14 @@ check (RecordPattern fs _) = mapM_ (check . snd) fs check (PatternAscription p _ _) = check p check PatternLit{} = return () + check (PatternConstr _ _ ps _) = mapM_ check ps seen v loc = do already <- gets $ M.lookup v case already of Just prev_loc -> lift $ throwError $ TypeError loc $ - "Name " ++ pretty v ++ " also bound at " ++ locStr prev_loc + "Name " ++ quote (pretty v) ++ " also bound at " ++ locStr prev_loc Nothing -> modify $ M.insert v loc @@ -286,7 +275,7 @@ pats (TEApply t1 (TypeArgExpType t2) _) = pats t1 ++ pats t2 pats (TEApply t1 TypeArgExpDim{} _) = pats t1 pats TEVar{} = [] - pats TEEnum{} = [] + pats (TESum cs _) = concatMap (concatMap pats . snd) cs -- | Ensure that every shape parameter is used in positive position at -- least once before being used in negative position. @@ -329,7 +318,8 @@ case seen of Just prev -> throwError $ TypeError loc $ - "Type parameter " ++ pretty v ++ " previously defined at " ++ locStr prev + "Type parameter " ++ quote (pretty v) ++ + " previously defined at " ++ locStr prev ++ "." Nothing -> do modify $ M.insert (ns,v) loc lift $ checkName ns v loc @@ -339,6 +329,13 @@ checkTypeParam (TypeParamType l pv loc) = TypeParamType l <$> checkParamName Type pv loc <*> pure loc +-- | Construct a type argument corresponding to a type parameter. +typeParamToArg :: TypeParam -> StructTypeArg +typeParamToArg (TypeParamDim v ploc) = + TypeArgDim (NamedDim $ qualName v) ploc +typeParamToArg (TypeParamType _ v ploc) = + TypeArgType (Scalar (TypeVar () Nonunique (typeName v) [])) ploc + -- | Return the shapes used in a given type expression in positive and negative -- position, respectively. typeExpUses :: TypeExp VName -> ([VName], [VName]) @@ -353,7 +350,7 @@ typeExpUses (TEArrow _ t1 t2 _) = let (pos, neg) = typeExpUses t1 <> typeExpUses t2 in (mempty, pos <> neg) -typeExpUses TEEnum{} = mempty +typeExpUses (TESum cs _) = foldMap (mconcat . fmap typeExpUses . snd) cs dimDeclUses :: DimDecl VName -> ([VName], [VName]) dimDeclUses (NamedDim v) = ([qualLeaf v], []) @@ -368,37 +365,22 @@ substituteTypes :: Monoid als => TypeSubs -> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als substituteTypes substs ot = case ot of Array als u at shape -> - maybe nope (`addAliases` (<>als)) $ - arrayOf (substituteTypesInArrayElem at) (substituteInShape shape) u - Prim t -> Prim t - TypeVar als u v targs + arrayOf (substituteTypes substs (Scalar at) `setAliases` mempty) + (substituteInShape shape) u `addAliases` (<>als) + Scalar (Prim t) -> Scalar $ Prim t + Scalar (TypeVar als u v targs) | Just (TypeSub (TypeAbbr _ ps t)) <- M.lookup (qualLeaf (qualNameFromTypeName v)) substs -> applyType ps (t `setAliases` mempty) (map substituteInTypeArg targs) `setUniqueness` u `addAliases` (<>als) - | otherwise -> TypeVar als u v $ map substituteInTypeArg targs - Record ts -> - Record $ fmap (substituteTypes substs) ts - Arrow als v t1 t2 -> - Arrow als v (substituteTypes substs t1) (substituteTypes substs t2) - Enum cs -> Enum cs - where nope = error "substituteTypes: Cannot create array after substitution." - - substituteTypesInArrayElem (ArrayPrimElem t) = - Prim t - substituteTypesInArrayElem (ArrayPolyElem v targs) - | Just (TypeSub (TypeAbbr _ ps t)) <- - M.lookup (qualLeaf (qualNameFromTypeName v)) substs = - applyType ps (t `setAliases` mempty) (map substituteInTypeArg targs) - | otherwise = - TypeVar mempty Nonunique v (map substituteInTypeArg targs) - substituteTypesInArrayElem (ArrayRecordElem ts) = - Record ts' - where ts' = fmap (substituteTypes substs . recordArrayElemToType) ts - substituteTypesInArrayElem (ArrayEnumElem cs) = - Enum cs - - substituteInTypeArg (TypeArgDim d loc) = + | otherwise -> Scalar $ TypeVar als u v $ map substituteInTypeArg targs + Scalar (Record ts) -> + Scalar $ Record $ fmap (substituteTypes substs) ts + Scalar (Arrow als v t1 t2) -> + Scalar $ Arrow als v (substituteTypes substs t1) (substituteTypes substs t2) + Scalar (Sum cs) -> + Scalar $ Sum $ (fmap . fmap) (substituteTypes substs) cs + where substituteInTypeArg (TypeArgDim d loc) = TypeArgDim (substituteInDim d) loc substituteInTypeArg (TypeArgType t loc) = TypeArgType (substituteTypes substs t) loc @@ -460,37 +442,24 @@ (VName -> Maybe (Subst (TypeBase dim as))) -> TypeBase dim as -> TypeBase dim as substTypesAny lookupSubst ot = case ot of - Prim t -> Prim t Array als u et shape -> - maybe nope (`addAliases` (<>als)) $ - arrayOf (subsArrayElem et) shape u + arrayOf (substTypesAny lookupSubst' (Scalar et)) shape u `setAliases` als + Scalar (Prim t) -> Scalar $ Prim t -- We only substitute for a type variable with no arguments, since -- type parameters cannot have higher kind. - TypeVar als u v targs -> + Scalar (TypeVar als u v targs) -> case lookupSubst $ qualLeaf (qualNameFromTypeName v) of Just (Subst t) -> t `setUniqueness` u `addAliases` (<>als) - Just PrimSubst -> TypeVar mempty u v $ map subsTypeArg targs - Nothing -> TypeVar als u v $ map subsTypeArg targs - Record ts -> Record $ fmap (substTypesAny lookupSubst) ts - Arrow als v t1 t2 -> - Arrow als v (substTypesAny lookupSubst t1) (substTypesAny lookupSubst t2) - Enum names -> Enum names - - where nope = error "substTypesAny: Cannot create array after substitution." - - subsArrayElem (ArrayPrimElem t) = Prim t - subsArrayElem (ArrayPolyElem v targs) = - case lookupSubst $ qualLeaf $ qualNameFromTypeName v of - Just (Subst t) -> t - -- It is intentional that we do not handle PrimSubst - -- specially here, as we are inside an array, and that - -- gives the aliasing. - _ -> TypeVar mempty Nonunique v $ map subsTypeArg targs - subsArrayElem (ArrayRecordElem ts) = - Record $ substTypesAny lookupSubst . recordArrayElemToType <$> ts - subsArrayElem (ArrayEnumElem cs) = Enum cs + Just PrimSubst -> Scalar $ TypeVar mempty u v $ map subsTypeArg targs + Nothing -> Scalar $ TypeVar als u v $ map subsTypeArg targs + Scalar (Record ts) -> Scalar $ Record $ fmap (substTypesAny lookupSubst) ts + Scalar (Arrow als v t1 t2) -> + Scalar $ Arrow als v (substTypesAny lookupSubst t1) (substTypesAny lookupSubst t2) + Scalar (Sum ts) -> + Scalar $ Sum $ (fmap . fmap) (substTypesAny lookupSubst) ts - subsTypeArg (TypeArgType t loc) = + where subsTypeArg (TypeArgType t loc) = TypeArgType (substTypesAny lookupSubst' t) loc - where lookupSubst' = fmap (fmap $ bimap id (const ())) . lookupSubst subsTypeArg t = t + + lookupSubst' = fmap (fmap $ bimap id (const ())) . lookupSubst diff --git a/src/Language/Futhark/TypeChecker/Unify.hs b/src/Language/Futhark/TypeChecker/Unify.hs --- a/src/Language/Futhark/TypeChecker/Unify.hs +++ b/src/Language/Futhark/TypeChecker/Unify.hs @@ -2,6 +2,9 @@ {-# LANGUAGE GeneralizedNewtypeDeriving #-} module Language.Futhark.TypeChecker.Unify ( Constraint(..) + , Usage + , mkUsage + , mkUsage' , Constraints , lookupSubst , MonadUnify(..) @@ -40,23 +43,39 @@ -- a partial constraint on their type. type Constraints = M.Map VName Constraint -data Constraint = NoConstraint (Maybe Liftedness) SrcLoc +-- | A usage that caused a type constraint. +data Usage = Usage (Maybe String) SrcLoc + +mkUsage :: SrcLoc -> String -> Usage +mkUsage = flip (Usage . Just) + +mkUsage' :: SrcLoc -> Usage +mkUsage' = Usage Nothing + +instance Show Usage where + show (Usage Nothing loc) = "use at " ++ locStr loc + show (Usage (Just s) loc) = s ++ " at " ++ locStr loc + +instance Located Usage where + locOf (Usage _ loc) = locOf loc + +data Constraint = NoConstraint (Maybe Liftedness) Usage | ParamType Liftedness SrcLoc - | Constraint (TypeBase () ()) SrcLoc - | Overloaded [PrimType] SrcLoc - | HasFields (M.Map Name (TypeBase () ())) SrcLoc - | Equality SrcLoc - | HasConstrs [Name] SrcLoc + | Constraint (TypeBase () ()) Usage + | Overloaded [PrimType] Usage + | HasFields (M.Map Name (TypeBase () ())) Usage + | Equality Usage + | HasConstrs (M.Map Name [TypeBase () ()]) Usage deriving Show instance Located Constraint where - locOf (NoConstraint _ loc) = locOf loc + locOf (NoConstraint _ usage) = locOf usage locOf (ParamType _ loc) = locOf loc - locOf (Constraint _ loc) = locOf loc - locOf (Overloaded _ loc) = locOf loc - locOf (HasFields _ loc) = locOf loc - locOf (Equality loc) = locOf loc - locOf (HasConstrs _ loc) = locOf loc + locOf (Constraint _ usage) = locOf usage + locOf (Overloaded _ usage) = locOf usage + locOf (HasFields _ usage) = locOf usage + locOf (Equality usage) = locOf usage + locOf (HasConstrs _ usage) = locOf usage lookupSubst :: VName -> Constraints -> Maybe (Subst (TypeBase () ())) lookupSubst v constraints = case M.lookup v constraints of @@ -86,9 +105,26 @@ Just ParamType{} -> True _ -> False +unifySharedConstructors :: MonadUnify m => + Usage + -> M.Map Name [TypeBase () ()] + -> M.Map Name [TypeBase () ()] + -> m () +unifySharedConstructors usage cs1 cs2 = + forM_ (M.toList $ M.intersectionWith (,) cs1 cs2) $ \(c, (f1, f2)) -> + unifyConstructor c f1 f2 + where unifyConstructor c f1 f2 + | length f1 == length f2 = + zipWithM_ (unify usage) f1 f2 + | otherwise = typeError usage $ "Cannot unify constructor " ++ + quote (prettyName c) ++ "." + +indent :: String -> String +indent = intercalate "\n" . map (" "++) . lines + -- | Unifies two types. -unify :: MonadUnify m => SrcLoc -> TypeBase () () -> TypeBase () () -> m () -unify loc orig_t1 orig_t2 = do +unify :: MonadUnify m => Usage -> TypeBase () () -> TypeBase () () -> m () +unify usage orig_t1 orig_t2 = do orig_t1' <- normaliseType orig_t1 orig_t2' <- normaliseType orig_t2 breadCrumb (MatchingTypes orig_t1' orig_t2') $ subunify orig_t1 orig_t2 @@ -100,41 +136,46 @@ t1' = applySubst (`lookupSubst` constraints) t1 t2' = applySubst (`lookupSubst` constraints) t2 - failure = - typeError loc $ "Couldn't match expected type `" ++ - pretty t1' ++ "' with actual type `" ++ pretty t2' ++ "'." + failure + -- This case is to avoid repeating the types that are also + -- shown in the breadcrumb. + | t1 == orig_t1, t2 == orig_t2 = + typeError (srclocOf usage) "Types do not match." + | otherwise = + typeError (srclocOf usage) $ "Couldn't match expected type\n" ++ + indent (pretty t1') ++ "\nwith actual type\n" ++ indent (pretty t2') case (t1', t2') of _ | t1' == t2' -> return () - (Record fs, - Record arg_fs) + (Scalar (Record fs), + Scalar (Record arg_fs)) | M.keys fs == M.keys arg_fs -> forM_ (M.toList $ M.intersectionWith (,) fs arg_fs) $ \(k, (k_t1, k_t2)) -> breadCrumb (MatchingFields k) $ subunify k_t1 k_t2 - (TypeVar _ _ (TypeName _ tn) targs, - TypeVar _ _ (TypeName _ arg_tn) arg_targs) + (Scalar (TypeVar _ _ (TypeName _ tn) targs), + Scalar (TypeVar _ _ (TypeName _ arg_tn) arg_targs)) | tn == arg_tn, length targs == length arg_targs -> zipWithM_ unifyTypeArg targs arg_targs - (TypeVar _ _ (TypeName [] v1) [], - TypeVar _ _ (TypeName [] v2) []) -> + (Scalar (TypeVar _ _ (TypeName [] v1) []), + Scalar (TypeVar _ _ (TypeName [] v2) [])) -> case (isRigid' v1, isRigid' v2) of (True, True) -> failure - (True, False) -> linkVarToType loc v2 t1' - (False, True) -> linkVarToType loc v1 t2' - (False, False) -> linkVarToType loc v1 t2' + (True, False) -> linkVarToType usage v2 t1' + (False, True) -> linkVarToType usage v1 t2' + (False, False) -> linkVarToType usage v1 t2' - (TypeVar _ _ (TypeName [] v1) [], _) + (Scalar (TypeVar _ _ (TypeName [] v1) []), _) | not $ isRigid' v1 -> - linkVarToType loc v1 t2' - (_, TypeVar _ _ (TypeName [] v2) []) + linkVarToType usage v1 t2' + (_, Scalar (TypeVar _ _ (TypeName [] v2) [])) | not $ isRigid' v2 -> - linkVarToType loc v2 t1' + linkVarToType usage v2 t1' - (Arrow _ _ a1 b1, - Arrow _ _ a2 b2) -> do + (Scalar (Arrow _ _ a1 b1), + Scalar (Arrow _ _ a2 b2)) -> do subunify a1 a2 subunify b1 b2 @@ -143,12 +184,16 @@ Just t2'' <- peelArray 1 t2' -> subunify t1'' t2'' + (Scalar (Sum cs), + Scalar (Sum arg_cs)) + | M.keys cs == M.keys arg_cs -> + unifySharedConstructors usage cs arg_cs (_, _) -> failure where unifyTypeArg TypeArgDim{} TypeArgDim{} = return () unifyTypeArg (TypeArgType t _) (TypeArgType arg_t _) = subunify t arg_t - unifyTypeArg _ _ = typeError loc + unifyTypeArg _ _ = typeError usage "Cannot unify a type argument with a dimension argument (or vice versa)." applySubstInConstraint :: VName -> Subst (TypeBase () ()) -> Constraint -> Constraint @@ -157,205 +202,228 @@ applySubstInConstraint vn subst (HasFields fs loc) = HasFields (M.map (applySubst (flip M.lookup $ M.singleton vn subst)) fs) loc applySubstInConstraint _ _ (NoConstraint l loc) = NoConstraint l loc -applySubstInConstraint _ _ (Overloaded ts loc) = Overloaded ts loc +applySubstInConstraint _ _ (Overloaded ts usage) = Overloaded ts usage applySubstInConstraint _ _ (Equality loc) = Equality loc applySubstInConstraint _ _ (ParamType l loc) = ParamType l loc -applySubstInConstraint _ _ (HasConstrs ns loc) = HasConstrs ns loc +applySubstInConstraint vn subst (HasConstrs cs loc) = + HasConstrs (M.map (map (applySubst (flip M.lookup $ M.singleton vn subst))) cs) loc -linkVarToType :: MonadUnify m => SrcLoc -> VName -> TypeBase () () -> m () -linkVarToType loc vn tp = do +linkVarToType :: MonadUnify m => Usage -> VName -> TypeBase () () -> m () +linkVarToType usage vn tp = do constraints <- getConstraints if vn `S.member` typeVars tp - then typeError loc $ "Occurs check: cannot instantiate " ++ + then typeError usage $ "Occurs check: cannot instantiate " ++ prettyName vn ++ " with " ++ pretty tp' - else do modifyConstraints $ M.insert vn $ Constraint tp' loc + else do modifyConstraints $ M.insert vn $ Constraint tp' usage modifyConstraints $ M.map $ applySubstInConstraint vn $ Subst tp' case M.lookup vn constraints of - Just (NoConstraint (Just Unlifted) unlift_loc) -> - zeroOrderType loc ("used at " ++ locStr unlift_loc) tp' + + Just (NoConstraint (Just Unlifted) unlift_usage) -> + zeroOrderType usage (show unlift_usage) tp' + Just (Equality _) -> - equalityType loc tp' - Just (Overloaded ts old_loc) - | tp `notElem` map Prim ts -> + equalityType usage tp' + + Just (Overloaded ts old_usage) + | tp `notElem` map (Scalar . Prim) ts -> case tp' of - TypeVar _ _ (TypeName [] v) [] - | not $ isRigid v constraints -> linkVarToTypes loc v ts + Scalar (TypeVar _ _ (TypeName [] v) []) + | not $ isRigid v constraints -> linkVarToTypes usage v ts _ -> - typeError loc $ "Cannot unify `" ++ prettyName vn ++ "' with type `" ++ - pretty tp ++ "' (`" ++ prettyName vn ++ - "` must be one of " ++ intercalate ", " (map pretty ts) ++ - " due to use at " ++ locStr old_loc ++ ")." - Just (HasFields required_fields old_loc) -> + typeError usage $ "Cannot unify " ++ quote (prettyName vn) ++ + "' with type\n" ++ indent (pretty tp) ++ "\nas " ++ + quote (prettyName vn) ++ " must be one of " ++ + intercalate ", " (map pretty ts) ++ + " due to " ++ show old_usage ++ ")." + + Just (HasFields required_fields old_usage) -> case tp of - Record tp_fields + Scalar (Record tp_fields) | all (`M.member` tp_fields) $ M.keys required_fields -> - mapM_ (uncurry $ unify loc) $ M.elems $ + mapM_ (uncurry $ unify usage) $ M.elems $ M.intersectionWith (,) required_fields tp_fields - TypeVar _ _ (TypeName [] v) [] + Scalar (TypeVar _ _ (TypeName [] v) []) | not $ isRigid v constraints -> modifyConstraints $ M.insert v $ - HasFields required_fields old_loc + HasFields required_fields old_usage _ -> - let required_fields' = - intercalate ", " $ map field $ M.toList required_fields - field (l, t) = pretty l ++ ": " ++ pretty t - in typeError loc $ - "Cannot unify `" ++ prettyName vn ++ "' with type `" ++ - pretty tp ++ "' (must be a record with fields {" ++ - required_fields' ++ - "} due to use at " ++ locStr old_loc ++ ")." - Just (HasConstrs cs old_loc) -> + typeError usage $ + "Cannot unify " ++ quote (prettyName vn) ++ " with type\n" ++ + indent (pretty tp) ++ "\nas " ++ quote (prettyName vn) ++ + " must be a record with fields\n" ++ + pretty (Record required_fields) ++ + "\ndue to " ++ show old_usage ++ "." + + Just (HasConstrs required_cs old_usage) -> case tp of - Enum t_cs - | intersect cs t_cs == cs -> return () - | otherwise -> typeError loc $ - "Cannot unify `" ++ prettyName vn ++ "' with type `" - ++ pretty tp ++ "'" - TypeVar _ _ (TypeName [] v) [] - | not $ isRigid v constraints -> - let addConstrs (HasConstrs cs' loc') (HasConstrs cs'' _) = - HasConstrs (cs' `union` cs'') loc' - addConstrs c _ = c - in modifyConstraints $ M.insertWith addConstrs v $ - HasConstrs cs old_loc - _ -> typeError loc "Cannot unify." + Scalar (Sum ts) + | all (`M.member` ts) $ M.keys required_cs -> + unifySharedConstructors usage required_cs ts + Scalar (TypeVar _ _ (TypeName [] v) []) + | not $ isRigid v constraints -> do + case M.lookup v constraints of + Just (HasConstrs v_cs _) -> + unifySharedConstructors usage required_cs v_cs + _ -> return () + modifyConstraints $ M.insertWith combineConstrs v $ + HasConstrs required_cs old_usage + where combineConstrs (HasConstrs cs1 usage1) (HasConstrs cs2 _) = + HasConstrs (M.union cs1 cs2) usage1 + combineConstrs hasCs _ = hasCs + _ -> noSumType + _ -> return () + where tp' = removeUniqueness tp + noSumType = typeError usage "Cannot unify a sum type with a non-sum type" removeUniqueness :: TypeBase dim as -> TypeBase dim as -removeUniqueness (Record ets) = - Record $ fmap removeUniqueness ets -removeUniqueness (Arrow als p t1 t2) = - Arrow als p (removeUniqueness t1) (removeUniqueness t2) +removeUniqueness (Scalar (Record ets)) = + Scalar $ Record $ fmap removeUniqueness ets +removeUniqueness (Scalar (Arrow als p t1 t2)) = + Scalar $ Arrow als p (removeUniqueness t1) (removeUniqueness t2) +removeUniqueness (Scalar (Sum cs)) = + Scalar $ Sum $ (fmap . fmap) removeUniqueness cs removeUniqueness t = t `setUniqueness` Nonunique -mustBeOneOf :: MonadUnify m => [PrimType] -> SrcLoc -> TypeBase () () -> m () -mustBeOneOf [req_t] loc t = unify loc (Prim req_t) t +mustBeOneOf :: MonadUnify m => [PrimType] -> Usage -> TypeBase () () -> m () +mustBeOneOf [req_t] loc t = unify loc (Scalar (Prim req_t)) t mustBeOneOf ts loc t = do constraints <- getConstraints let t' = applySubst (`lookupSubst` constraints) t isRigid' v = isRigid v constraints case t' of - TypeVar _ _ (TypeName [] v) [] + Scalar (TypeVar _ _ (TypeName [] v) []) | not $ isRigid' v -> linkVarToTypes loc v ts - Prim pt | pt `elem` ts -> return () + Scalar (Prim pt) | pt `elem` ts -> return () _ -> failure where failure = typeError loc $ "Cannot unify type \"" ++ pretty t ++ "\" with any of " ++ intercalate "," (map pretty ts) ++ "." -linkVarToTypes :: MonadUnify m => SrcLoc -> VName -> [PrimType] -> m () -linkVarToTypes loc vn ts = do +linkVarToTypes :: MonadUnify m => Usage -> VName -> [PrimType] -> m () +linkVarToTypes usage vn ts = do vn_constraint <- M.lookup vn <$> getConstraints case vn_constraint of - Just (Overloaded vn_ts vn_loc) -> + Just (Overloaded vn_ts vn_usage) -> case ts `intersect` vn_ts of - [] -> typeError loc $ "Type constrained to one of " ++ + [] -> typeError usage $ "Type constrained to one of " ++ intercalate "," (map pretty ts) ++ " but also one of " ++ - intercalate "," (map pretty vn_ts) ++ " at " ++ locStr vn_loc ++ "." - ts' -> modifyConstraints $ M.insert vn $ Overloaded ts' loc + intercalate "," (map pretty vn_ts) ++ " due to " ++ show vn_usage ++ "." + ts' -> modifyConstraints $ M.insert vn $ Overloaded ts' usage - _ -> modifyConstraints $ M.insert vn $ Overloaded ts loc + _ -> modifyConstraints $ M.insert vn $ Overloaded ts usage equalityType :: (MonadUnify m, Pretty (ShapeDecl dim), Monoid as) => - SrcLoc -> TypeBase dim as -> m () -equalityType loc t = do + Usage -> TypeBase dim as -> m () +equalityType usage t = do unless (orderZero t) $ - typeError loc $ - "Type \"" ++ pretty t ++ "\" does not support equality." + typeError usage $ + "Type \"" ++ pretty t ++ "\" does not support equality (is higher-order)." mapM_ mustBeEquality $ typeVars t where mustBeEquality vn = do constraints <- getConstraints case M.lookup vn constraints of - Just (Constraint (TypeVar _ _ (TypeName [] vn') []) _) -> + Just (Constraint (Scalar (TypeVar _ _ (TypeName [] vn') [])) _) -> mustBeEquality vn' - Just (Constraint vn_t _) + Just (Constraint vn_t cusage) | not $ orderZero vn_t -> - typeError loc $ "Type \"" ++ pretty t ++ - "\" does not support equality." + typeError usage $ + unlines ["Type \"" ++ pretty t ++ "\" does not support equality.", + "Constrained to be higher-order due to " ++ show cusage ++ "."] | otherwise -> return () Just (NoConstraint _ _) -> - modifyConstraints $ M.insert vn (Equality loc) + modifyConstraints $ M.insert vn (Equality usage) Just (Overloaded _ _) -> return () -- All primtypes support equality. - Just HasConstrs{} -> - return () + Just (HasConstrs cs _) -> + mapM_ (equalityType usage) $ concat $ M.elems cs _ -> - typeError loc $ "Type " ++ pretty (prettyName vn) ++ + typeError usage $ "Type " ++ pretty (prettyName vn) ++ " does not support equality." zeroOrderType :: (MonadUnify m, Pretty (ShapeDecl dim), Monoid as) => - SrcLoc -> String -> TypeBase dim as -> m () -zeroOrderType loc desc t = do + Usage -> String -> TypeBase dim as -> m () +zeroOrderType usage desc t = do unless (orderZero t) $ - typeError loc $ "Type " ++ desc ++ - " must not be functional, but is " ++ pretty t ++ "." + typeError usage $ "Type " ++ desc ++ + " must not be functional, but is " ++ quote (pretty t) ++ "." mapM_ mustBeZeroOrder . S.toList . typeVars $ t where mustBeZeroOrder vn = do constraints <- getConstraints case M.lookup vn constraints of - Just (Constraint vn_t old_loc) + Just (Constraint vn_t old_usage) | not $ orderZero t -> - typeError loc $ "Type " ++ desc ++ + typeError usage $ "Type " ++ desc ++ " must be non-function, but inferred to be " ++ - pretty vn_t ++ " at " ++ locStr old_loc ++ "." + quote (pretty vn_t) ++ " due to " ++ show old_usage ++ "." Just (NoConstraint _ _) -> - modifyConstraints $ M.insert vn (NoConstraint (Just Unlifted) loc) + modifyConstraints $ M.insert vn (NoConstraint (Just Unlifted) usage) Just (ParamType Lifted ploc) -> - typeError loc $ "Type " ++ desc ++ - " must be non-function, but type parameter " ++ prettyName vn ++ " at " ++ + typeError usage $ "Type " ++ desc ++ + " must be non-function, but type parameter " ++ quote (prettyName vn) ++ " at " ++ locStr ploc ++ " may be a function." _ -> return () +-- | In @mustHaveConstr usage c t fs@, the type @t@ must have a +-- constructor named @c@ that takes arguments of types @ts@. mustHaveConstr :: MonadUnify m => - SrcLoc -> Name -> TypeBase dim as -> m () -mustHaveConstr loc c t = do + Usage -> Name -> TypeBase dim as -> [TypeBase () ()] -> m () +mustHaveConstr usage c t fs = do + let struct_f = toStructural <$> fs constraints <- getConstraints case t of - TypeVar _ _ (TypeName _ tn) [] + Scalar (TypeVar _ _ (TypeName _ tn) []) | Just NoConstraint{} <- M.lookup tn constraints -> - modifyConstraints $ M.insert tn $ HasConstrs [c] loc + modifyConstraints $ M.insert tn $ HasConstrs (M.singleton c struct_f) usage | Just (HasConstrs cs _) <- M.lookup tn constraints -> - if c `elem` cs - then return () - else modifyConstraints $ M.insert tn $ HasConstrs (c:cs) loc - Enum cs - | c `elem` cs -> return () - | otherwise -> throwError $ TypeError loc $ - "Type " ++ pretty (toStructural t) ++ - " does not have a " ++ pretty c ++ " constructor." - _ -> do unify loc (toStructural t) $ Enum [c] + case M.lookup c cs of + Nothing -> modifyConstraints $ M.insert tn $ HasConstrs (M.insert c fs cs) usage + Just fs' + | length fs == length fs' -> zipWithM_ (unify usage) fs fs' + | otherwise -> typeError usage $ "Different arity for constructor " + ++ quote (pretty c) ++ "." + + Scalar (Sum cs) -> + case M.lookup c cs of + Nothing -> typeError usage $ "Constuctor " ++ quote (pretty c) ++ " not present in type." + Just fs' + | length fs == length fs' -> zipWithM_ (unify usage) fs (toStructural <$> fs') + | otherwise -> typeError usage $ "Different arity for constructor " ++ + quote (pretty c) ++ "." + + _ -> do unify usage (toStructural t) $ Scalar $ Sum $ M.singleton c fs return () mustHaveField :: (MonadUnify m, Monoid as) => - SrcLoc -> Name -> TypeBase dim as -> m (TypeBase dim as) -mustHaveField loc l t = do + Usage -> Name -> TypeBase dim as -> m (TypeBase dim as) +mustHaveField usage l t = do constraints <- getConstraints - l_type <- newTypeVar loc "t" + l_type <- newTypeVar (srclocOf usage) "t" let l_type' = toStructural l_type case t of - TypeVar _ _ (TypeName _ tn) [] + Scalar (TypeVar _ _ (TypeName _ tn) []) | Just NoConstraint{} <- M.lookup tn constraints -> do - modifyConstraints $ M.insert tn $ HasFields (M.singleton l l_type') loc + modifyConstraints $ M.insert tn $ HasFields (M.singleton l l_type') usage return l_type | Just (HasFields fields _) <- M.lookup tn constraints -> do case M.lookup l fields of - Just t' -> unify loc l_type' t' + Just t' -> unify usage l_type' t' Nothing -> modifyConstraints $ M.insert tn $ - HasFields (M.insert l l_type' fields) loc + HasFields (M.insert l l_type' fields) usage return l_type - Record fields + Scalar (Record fields) | Just t' <- M.lookup l fields -> do - unify loc l_type' (toStructural t') + unify usage l_type' (toStructural t') return t' | otherwise -> - throwError $ TypeError loc $ - "Attempt to access field '" ++ pretty l ++ "' of value of type " ++ - pretty (toStructural t) ++ "." - _ -> do unify loc (toStructural t) $ Record $ M.singleton l l_type' + typeError usage $ + "Attempt to access field " ++ quote (pretty l) ++ "` of value of type " ++ + quote (pretty (toStructural t)) ++ "." + _ -> do unify usage (toStructural t) $ Scalar $ Record $ M.singleton l l_type' return l_type -- Simple MonadUnify implementation. @@ -376,8 +444,8 @@ put (x, i+1) return i let v = VName (mkTypeVarName desc i) 0 - modifyConstraints $ M.insert v $ NoConstraint Nothing loc - return $ TypeVar mempty Nonunique (typeName v) [] + modifyConstraints $ M.insert v $ NoConstraint Nothing $ Usage Nothing loc + return $ Scalar $ TypeVar mempty Nonunique (typeName v) [] -- | Construct a the name of a new type variable given a base -- description and a tag number (note that this is distinct from @@ -397,11 +465,11 @@ -> TypeBase () () -> TypeBase () () -> Either TypeError (TypeBase () ()) doUnification loc tparams t1 t2 = runUnifyM tparams $ do - unify loc t1 t2 + unify (Usage Nothing loc) t1 t2 normaliseType t2 runUnifyM :: [TypeParam] -> UnifyM a -> Either TypeError a runUnifyM tparams (UnifyM m) = runExcept $ evalStateT m (constraints, 0) where constraints = M.fromList $ mapMaybe f tparams f TypeParamDim{} = Nothing - f (TypeParamType l p loc) = Just (p, NoConstraint (Just l) loc) + f (TypeParamType l p loc) = Just (p, NoConstraint (Just l) $ Usage Nothing loc) diff --git a/src/futhark.hs b/src/futhark.hs --- a/src/futhark.hs +++ b/src/futhark.hs @@ -96,9 +96,6 @@ args <- getArgs prog <- getProgName case args of - -- The -t case here is for temporary compatibility with - -- futhark-mode. It will go away eventually. - "-t":args' -> Misc.mainCheck prog args' cmd:args' | Just (m, _) <- lookup cmd commands -> m (unwords [prog, cmd]) args' _ -> mainWithOptions () [] msg (const . const Nothing) prog args