diff --git a/doc/AD.pdf b/doc/AD.pdf index be78d6d3d..9a24068af 100644 Binary files a/doc/AD.pdf and b/doc/AD.pdf differ diff --git a/doc/AD.tex b/doc/AD.tex index 593e7ee37..e8e70bb1b 100644 --- a/doc/AD.tex +++ b/doc/AD.tex @@ -3,6 +3,7 @@ \settopmatter{printfolios=true,printccs=false,printacmref=false} \usepackage{color} +\usepackage{proof} \usepackage{url} \usepackage{stmaryrd} % double brackets \usepackage{listings} @@ -15,6 +16,21 @@ \renewcommand{\arraystretch}{1.2} +% The cateorical language +\newcommand{\CatLang}{$\mathcal CL$} +\newcommand{\KC}[2]{{\mathcal C}\lb \, #1 \, \rb\,#2} +\newcommand{\clcomp}{\, ; \,} % CL composition +\newcommand{\ck}[1]{{\mathcal K}(#1)} % CL constant +\newcommand{\cid}{\mathcal{I}} % CL function +\newcommand{\cprune}[2]{\mathcal{P}[#1/#2]} % CL function +\newcommand{\cfun}[1]{{\mathcal F}(#1)} % CL function +\newcommand{\clet}[3]{{\mathcal L}(#1,#2,#3)} % CL let +\newcommand{\cbuild}[3]{{\mathcal B}(#1,#2,#3)} % CL build +\newcommand{\cif}[3]{\mathcal{IF}(#1,#2,#3)} % CL conditional +\newcommand{\ctupx}{\gtrdot} % CL tuple extension (terms) +\newcommand{\ctupxt}{\gtrdot} % CL tuple extension (types) +\newcommand{\capply}{\diamond} % CL application + \newcommand{\deriv}{\partial} % The symbol we use for ``derivative'' % Was \nabla\!, but mathematicans hate that @@ -48,8 +64,14 @@ \newcommand{\olmvcat}{\;\overline{\lmvcat}\;} % Infix pairing \newcommand{\olmhcat}{\;\overline{\lmhcat}\;} % Infix pairing -\newcommand{\fwdDf}[1]{#1'} % Forward derivative, f' -\newcommand{\revDf}[1]{#1'_R} % Reverse derivative, f` +\newcommand{\fwdDf}[1]{\mathit{fwd\$}#1} % Forward derivative, f' +\newcommand{\revDf}[1]{\mathit{rev\$}#1} % Reverse derivative, f` + +% With splitting +\newcommand{\revfDf}[1]{\mathit{revf\$}#1} % Reverse derivative, f` +\newcommand{\revrDf}[1]{\mathit{revr\$}#1} % Reverse derivative, f` + + \newcommand{\lb}{\llbracket} \newcommand{\rb}{\rrbracket} \newcommand{\sel}[2]{\pi_{#1,#2}} @@ -64,7 +86,7 @@ \newcommand{\indexfunname}{\mathit{ixR}} %% Reversed from current code \newcommand{\indexfun}[2]{\indexfunname(#2,#1)} % Vector indexing \indexfun{v}{i} \renewcommand{\d}[1]{\delta_{\mathit{#1}}} % delta_x etc -\renewcommand{\vector}[2]{\mathit{Vec}\;#1\;#2} % the type (Vec n t) +\renewcommand{\vector}[1]{\mathit{Vec}\;#1} % the type (Vec t) \newcommand{\tcolon}{\!:\!} \newcommand{\typ}[2]{#1 \tcolon #2} % x:S, with less horizontal whitespace @@ -72,6 +94,7 @@ \renewcommand{\dot}{.\,} % dot with some space after \newcommand{\real}{\mathbb{R}} % R, the reals11 \newcommand{\nat}{\mathbb{N}} % N, the natural numbers +\newcommand{\bool}{\mathbb{B}} % B, the booleans \newcommand{\darrow}{\Rightarrow} % => % Linear maps @@ -135,7 +158,7 @@ \section{The language} \tau & ::= & \nat & \mbox{Natural numbers} \\ & | & \real & \mbox{Real numbers} \\ & | & (\tau_1, \tau_2) & \mbox{Pairs} \\ - & | & \vector{n}{\tau} & \mbox{Vectors} \\ + & | & \vector{\tau} & \mbox{Vectors} \\ & | & \tau_1 \to \tau_2 & \mbox{Functions} \\ & | & \tau_1 \linto \tau_2 & \mbox{Linear maps} \\ \end{array} @@ -190,10 +213,10 @@ \subsection{Built in functions} \sel{1}{2} & :: & (t_1,t_2) \to t_1 & \mbox{Selection} \\ \sel{2}{2} & :: & (t_1,t_2) \to t_2 & \mbox{..ditto..} \\ % \deltafun & :: & Field \; t \darrow (\nat,\; \nat) \to t & \mbox{Delta-function} \\ - \buildfun & :: & (n::\nat,\; \nat \to t) \to \vector{n}{t} & \mbox{Vector build} \\ - \indexfunname & :: & (\nat, \; \vector{n}{t}) \to t & \mbox {Indexing (NB arg order)} \\ - \sumfunname & :: & \vector{n}{t} \to t & \mbox{Sum a vector} \\ - \sizefun & :: & \vector{n}{t} \to \nat & \mbox{Size of a vector} \\ + \buildfun & :: & (n::\nat,\; \nat \to t) \to \vector{t} & \mbox{Vector build} \\ + \indexfunname & :: & (\nat, \; \vector{t}) \to t & \mbox {Indexing (NB arg order)} \\ + \sumfunname & :: & \vector{t} \to t & \mbox{Sum a vector} \\ + \sizefun & :: & \vector{t} \to \nat & \mbox{Size of a vector} \\ \end{array} $$ \\[3mm] @@ -206,11 +229,11 @@ \subsection{Built in functions} \gradf{*}(x,y) & = & \lmscalar{y} \lmhcat \lmscalar{x} \\[2mm] \gradf{\sel{1}{2}} & :: & (t,t) \to ((t,t) \linto t) \\ \gradf{\sel{1}{2}}(x) & = & \lmone \lmhcat \lmzero \\[2mm] - \gradf{\indexfunname} & :: & (\nat, \;\vector{n}{t}) \to ((\nat, \; \vector{n}{t}) \linto t) \\ + \gradf{\indexfunname} & :: & (\nat, \;\vector{t}) \to ((\nat, \; \vector{t}) \linto t) \\ \gradf{\indexfunname}(i,v) & = & \lmzero \lmhcat \lmhcatv{\buildfun( \sizefun(v), \lambda j. \mbox{\lstinline|if|}\, i=j \,\mbox{\lstinline|then|}\, \lmone \, \mbox{\lstinline|else|}\, \lmzero )} \\[2mm] - \gradf{\sumfunname} & :: & \vector{n}{\real} \to (\vector{n}{\real} \linto \real) \\ + \gradf{\sumfunname} & :: & \vector{\real} \to (\vector{\real} \linto \real) \\ \gradf{\sumfun{v}} & = & lmhcatv{\buildfun( \sizefun(v),\; \lambda i. \lmone )} \\[2mm] \ldots \end{array} @@ -242,19 +265,19 @@ \subsection{Built in functions} \subsection{Vectors} -The language supports one-dimensional vectors, of type $\vector{n}{T}$, +The language supports one-dimensional vectors, of type $\vector{T}$, whose elements have type $T$ (Figure~\ref{fig:syntax}). A matrix can be represented as a vector of vectors. Vectors are supported by the following built-in functions (Figure~\ref{fig:built-in}): \begin{itemize} -\item $\buildfun :: (\nat,\; \nat \to t) \to \vector{n}{t}$ for vector construction. -\item $\indexfunname :: (\nat,\; \vector{n}{t}) \to t$ for indexing. Informally we allow ourselves to write +\item $\buildfun :: (\nat,\; \nat \to t) \to \vector{t}$ for vector construction. +\item $\indexfunname :: (\nat,\; \vector{t}) \to t$ for indexing. Informally we allow ourselves to write $v[i]$ instead of $\indexfun{v}{i}$. -\item $\sumfunname :: \vector{n}{\real} \to \real$ to add up the elements of a vector. +\item $\sumfunname :: \vector{\real} \to \real$ to add up the elements of a vector. We specifically do not have a general, higher order, fold operator; we say why in Section~\ref{sec:gen-fold}. -\item $\sizefun :: \vector{n}{t} \to \nat$ takes the size of a vector. +\item $\sizefun :: \vector{t} \to \nat$ takes the size of a vector. \item Arithmetic functions $(*), (+)$ etc are overloaded to work over vectors, always elementwise. \end{itemize} @@ -309,11 +332,11 @@ \subsection{Linear maps} \label{sec:lin-maps} \mbox{Scale} & \lmscalar{\cdot} & : & \real \to (s \linto s) \\ \mbox{VCat} & (\lmvcat) & : & (s \linto t_1,\; s \linto t_2) \to (s \linto (t_1,t_2)) & \mbox{Vertical juxtaposition} \\ - \mbox{VCatV} & \lmvcatv{\cdot} & : & \vector{n}{(s \linto t)} \to (s \linto \vector{n}{t}) + \mbox{VCatV} & \lmvcatv{\cdot} & : & \vector{(s \linto t)} \to (s \linto \vector{t}) & \mbox{...vector version} \\ \mbox{HCat} & (\lmhcat) & : & (t_1 \linto s,\; t_2 \linto s) \to ((t_1,t_2) \linto s) & \mbox{Horizontal juxtaposition} \\ - \mbox{HCatV} & \lmhcatv{\cdot} & : & \vector{n}{(t \linto s)} \to (\vector{n}{t} \linto s) + \mbox{HCatV} & \lmhcatv{\cdot} & : & \vector{(t \linto s)} \to (\vector{t} \linto s) & \mbox{...vector version} \\ \mbox{Transpose} & \lmtrans{\cdot} & : & (s \linto t) \to (t \linto s) & \mbox{Matrix transpose} \\[2mm] @@ -453,13 +476,13 @@ \subsubsection{Building vector spaces} What types are vector spaces? Look the syntax of types in Figure~\ref{fig:syntax}. \begin{itemize} \item The real numbers $\real$ is a vector space, using the standard $+$ and $*$ for reals; and $\dotprod_{\real} = *$. - \item If $V$ is a vector space then $\vector{n}{V}$ is a vector space, with + \item If $V$ is a vector space then $\vector{V}$ is a vector space, with \begin{itemize} \item $v_1 + v_2$ is vector addittion \item $r * v$ multiplies each element of the vector $v$ by the real $r$. \item $v_1 \dotprod v_2$ is a the usual vector dot-product. \end{itemize} - We often write $\vector{n}{\real}$ as $\real^N$. + We often write $\vector{\real}$ as $\real^N$. \item If $V_1$ and $V_2$ are vector spaces, then the product space $(V_1, V_2)$ is a vector space \begin{itemize} \item $(v_1,v_2) +(w_1,w_2) = (v_1 + w_1, v_2 + w_2)$. @@ -823,11 +846,11 @@ \section{AD for vectors} \label{sec:ad-vectors} \subsection{General folds} \label{sec:gen-fold} -We have $\sumfunname :: \vector{n}{\real} \to \real$. What is $\gradf{\sumfunname}$? +We have $\sumfunname :: \vector{\real} \to \real$. What is $\gradf{\sumfunname}$? One way to define its semantics is by applying it: $$ \begin{array}{rcl} - \gradf{\sumfunname} & :: & \vector{n}{\real} \to (\vector{n}{\real} \linto \real) \\ + \gradf{\sumfunname} & :: & \vector{\real} \to (\vector{\real} \linto \real) \\ \gradf{\sumfunname}(v) \lmapply dv & = & \sumfun{dv} \end{array} $$ @@ -945,9 +968,9 @@ \subsection{Tupling: basic version} (\olmhcat) & : & ((t_1, t_1 \linto s),\; (t_2, t_2 \linto s)) \to ((t_1,t_2),\; (t_1,t_2) \linto s) \\ (t_1,m_1) \olmhcat (t_2,m_2) & = & ((t_1,t_2), m_1 \lmhcat m_2) \\[2mm] - \swap & : & \vector{n}{(a,b)} \to (\vector{n}{a}, \vector{n}{b}) \\[2mm] + \swap & : & \vector{(a,b)} \to (\vector{a}, \vector{b}) \\[2mm] \olmtrans{\cdot} & : & (r, s \linto t) \to (r, t \linto s) -% B & : & (N, N \to (s \linto t, t)) \to (s \linto \vector{n}{t}, \vector{n}{t}) +% B & : & (N, N \to (s \linto t, t)) \to (s \linto \vector{t}, \vector{t}) \end{array} $$ \\ @@ -960,10 +983,10 @@ \subsection{Tupling: basic version} \tgradf{*}(x,y) & = & (\lmscalar{y} \lmhcat \lmscalar{x}, \, x*y) \\[2mm] % \gradf{\sel{1}{2}} & :: & (t,t) \to ((t,t) \linto t) \\ % \gradf{\sel{1}{2}}(x) & = & \lmone \lmhcat \lmzero \\[2mm] -% \gradf{\indexfunname} & :: & (\nat, \;\vector{n}{t}) \to ((\nat, \; \vector{n}{t}) \linto t) \\ +% \gradf{\indexfunname} & :: & (\nat, \;\vector{t}) \to ((\nat, \; \vector{t}) \linto t) \\ % \gradf{\indexfunname}(i,v) & = & \lmzero \lmhcat \lmbuildt( \sizefun(v), % \lambda j. \mbox{\lstinline|if|}\, i=j \,\mbox{\lstinline|then|}\, \lmone \, \mbox{\lstinline|else|}\, \lmzero ) \\[2mm] -% \gradf{\sumfunname} & :: & Field \; t \darrow \vector{n}{t} \to (\vector{n}{t} \linto t) \\ +% \gradf{\sumfunname} & :: & Field \; t \darrow \vector{t} \to (\vector{t} \linto t) \\ % \gradf{\sumfun{v}} & = & \lmbuildt( \sizefun(v),\; \lambda i. \lmone ) \\[2mm] % \ldots \end{array} @@ -1028,7 +1051,7 @@ \subsection{Polymorphic tupling: forward mode} (\olmhcat) & : & ((t_1, t_1 \linto s),\; (t_2, t_2 \linto s)) \to ((t_1,t_2),\; (t_1,t_2) \linto s) \\ (t_1,m_1) \olmhcat (t_2,m_2) & = & (t_1 + t_2, m_1 \lmhcat m_2) \\[2mm] \olmtrans{\cdot} & : & (t, s \linto t) \to (t, t \linto s) -% B & : & (N, N \to (s \linto t, t)) \to (s \linto \vector{n}{t}, \vector{n}{t}) +% B & : & (N, N \to (s \linto t, t)) \to (s \linto \vector{t}, \vector{t}) \end{array} $$ \\ @@ -1041,10 +1064,10 @@ \subsection{Polymorphic tupling: forward mode} \ogradf{*}((x,y), m) & = & (x*y,\, (\lmscalar{y} \lmhcat \lmscalar{x}) \lmcomp m) \\[2mm] % \gradf{\sel{1}{2}} & :: & (t,t) \to ((t,t) \linto t) \\ % \gradf{\sel{1}{2}}(x) & = & \lmone \lmhcat \lmzero \\[2mm] -% \gradf{\indexfunname} & :: & (\nat, \;\vector{n}{t}) \to ((\nat, \; \vector{n}{t}) \linto t) \\ +% \gradf{\indexfunname} & :: & (\nat, \;\vector{t}) \to ((\nat, \; \vector{t}) \linto t) \\ % \gradf{\indexfunname}(i,v) & = & \lmzero \lmhcat \lmbuildt( \sizefun(v), % \lambda j. \mbox{\lstinline|if|}\, i=j \,\mbox{\lstinline|then|}\, \lmone \, \mbox{\lstinline|else|}\, \lmzero ) \\[2mm] -% \gradf{\sumfunname} & :: & Field \; t \darrow \vector{n}{t} \to (\vector{n}{t} \linto t) \\ +% \gradf{\sumfunname} & :: & Field \; t \darrow \vector{t} \to (\vector{t} \linto t) \\ % \gradf{\sumfun{v}} & = & \lmbuildt( \sizefun(v),\; \lambda i. \lmone ) \\[2mm] % \ldots \end{array} @@ -1191,6 +1214,189 @@ \subsection{Polymorphic tupling: reverse mode} all intermediate linear maps. +\section{Compiling through categories} + + +\begin{figure} + \fbox{\begin{minipage}{\columnwidth} +$$ + \begin{array}{rcll} + \multicolumn{4}{l}{\mbox{\bf Atoms}} \\ + f,g,h & ::= & \multicolumn{2}{l}{\mbox{Function}} \\ + k & ::= & \multicolumn{2}{l}{\mbox{Literal constants}} \\ + \multicolumn{4}{l}{\mbox{\bf Terms}} \\ + \mathit{pgm} & ::= & \mathit{def}_1 \ldots \mathit{def}_n \\ + \mathit{def} & ::= & f {:} S \darrow T = c \\ + c & ::= & \cid & \mbox{Identity} \\ + & | & \ck{k} & \mbox{Constant} \\ + & | & \cprune{i_1,\ldots,i_m}{n}& \mbox{Pruning} (0 \leq m \leq n) \\ + & | & \cfun{f} & \mbox{Function constant} \\ + & | & c_1 \clcomp c_2 & \mbox{Composition} \\ + & | & (c_1,\ldots, c_n) & \mbox{Tuple} \\ + & | & \cif{c_1}{c_2}{c_3} & \mbox{Conditional} \\ + & | & \clet{x}{c_r}{c_b} & \mbox{Let} \\ + & | & \cbuild{c_s}{i}{c_e} & \mbox{Build} + \end{array} +$$ +\end{minipage}} +\caption{Syntax of \CatLang} \label{fig:cl-syntax} +\end{figure} + +% --------------------- Semantics +\begin{figure} +\fbox{\begin{minipage}{\columnwidth} +\mbox{{\bf Semantics (aka conversion from \CatLang):}\; \fbox{$e \capply c = e$}} +$$ +\begin{array}{r@{\hspace{1mm}}c@{\hspace{1mm}}l} + t \capply \cid & = & t \\ + t \capply \cprune{i_1,\ldots,i_m}{n} & = & (\sel{i_1}{n}(t), \ldots, \sel{i_m}{n}(t)) \\ + t \capply \ck{k} & = & k \\ + t \capply \cfun{f} & = & f(t) \\ + t \capply (c_1 \clcomp c_2) & = & (t \capply c_1) \capply c_2 \\ + t \capply (c_1,\ldots, c_n) & = & (t \capply c_1,\ldots, t \capply c_n) \\ + t \capply \cif{c_1}{c_2}{c_3} & = & \lstinline|if|\; (t \capply c_1) \; (t \capply c_2)\; (t \capply c_3) \\ + t \capply \clet{x}{c_r}{c_b} & = & \lstinline|let|\;x\, = \, t \capply c_r \; \lstinline{in}\; (t \ctupx x) \capply c_b \\ + t \capply \cbuild{c_s}{i}{c_e} & = & \lstinline|build|\;(t \capply c_x)\;(\lambda i. (t \ctupx i) \capply c_e) +\end{array} +$$ +\mbox{\bf Conversion to \CatLang} +$$ +\begin{array}{l} +\begin{array}{rcl} +\Gamma & ::= & (x_1\tcolon\tau_1, \ldots, x_n\tcolon\tau_n) \\[1mm] +\end{array} \\ +\begin{array}{rcl} +\phi((x_1\tcolon\tau_1,\ldots, x_n\tcolon\tau_n),\,x_i) & = & i \\ +T(x_1\tcolon\tau_1, \ldots, x_n\tcolon\tau_n) & = & (\tau_1, \ldots, \tau_n) \\ +\end{array} +\\[5mm] +\begin{array}{l} + \KC{f(x_1\tcolon\tau_1,\ldots,x_n\tcolon\tau_n) = e}{} \\ + \quad = \quad \cfun{f} = \KC{e}(x_1\tcolon\tau_1,\ldots,x_n\tcolon\tau_n) + \end{array} \\ \\ +\begin{array}{r@{\hspace{1mm}}c@{\hspace{1mm}}l} +\multicolumn{3}{c}{\fbox{If $\Gamma \vdash e:\tau$ then $\KC{e}\Gamma : T(\Gamma) \darrow \tau$}} \\ + \KC{k}{\Gamma} & = & \ck{k} \\ + \KC{x}{\Gamma} & = & \cfun{\pi(\Gamma,x)} \\ + \KC{f(e)}{\Gamma} & = & \cfun{f} \clcomp \KC{e}{\Gamma} \\ + \KC{\lstinline|if|\;e_1\;e_2\;e_3}{\Gamma} \\ + \multicolumn{3}{l}{\hspace{2cm} = \cif{\KC{e_1}{\Gamma}}{\KC{e_2}{\Gamma}}{\KC{e_3}{\Gamma}}} \\ + \KC{(e_1, \ldots, e_n)}{\Gamma} & = & (\KC{e_1}{\Gamma}, \ldots, \KC{e_n}{\Gamma}) \\ + \KC{\lstinline|let|\;x{:}\tau = e_r \;\lstinline|in|\; e_b}{\Gamma} & = & \clet{x}{\KC{e_r}{\Gamma}}{\KC{e_b}{(\Gamma,x{:}\tau)}} \\ + \KC{\lstinline|build|\;e_s\;(\lambda i.e_e)}{\Gamma} & = & \cbuild{\KC{e_s}{\Gamma}}{i}{\KC{e_e}{(\Gamma,i)}} +\end{array} +\\[1mm] +\mbox{\bf Pruning} \\ +\begin{array}{r@{\hspace{1mm}}c@{\hspace{1mm}}l} + \KC{e}\Gamma & = & \cprune{\phi(\Gamma,v_1), \ldots, \phi(\Gamma, v_m)}{\sizefun(\Gamma)}(\KC{e}{\Gamma'}) \\ + \multicolumn{3}{l}{\quad \mbox{where} + $\begin{array}[t]{r@{\hspace{1mm}}c@{\hspace{1mm}}l} + \{v_1, \ldots, v_m\} & = & \mathit{fv}(e) \\ + \Gamma' & = & (v_1\tcolon\Gamma(v_1), \ldots, v_n\tcolon\Gamma(v_n)) + \end{array}$} \\ +\end{array} +\end{array} +$$ +\end{minipage}} +\caption{Semantics of \CatLang} \label{fig:cl-semantics} +\end{figure} + +\begin{figure} + \fbox{\begin{minipage}{\columnwidth} + \begin{center}{\large \fbox{$\Gamma \vdash c : S \darrow T$}}\end{center} + + \begin{gather*} + \infer{\Gamma \vdash \cid : S \darrow S} + {} \qquad + \\[1mm] + \infer{\Gamma \vdash \cprune{i_1, \ldots, i_m}{n} : (s_1, \ldots, s_n) \darrow (s_{i_1}, \ldots, s_{i_m})} + {} + \\[1mm] + \infer{\Gamma \vdash \cfun{f} : S \darrow T} + {f : S \to T \in \Gamma} \qquad + \infer{\Gamma \vdash \ck{k} : () \darrow \real} + {} + \\[1mm] + \infer{\Gamma \vdash c_1 \clcomp c_2 : S \darrow T} + {\Gamma \vdash c_1 : S \darrow R \quad \Gamma \vdash c_2 : R \darrow T} + \\[1mm] + \infer{\Gamma \vdash (c_1, \ldots, c_n) : S \darrow (T_1, \ldots, T_n)} + {\Gamma \vdash c_1:S \darrow T_1 \quad \ldots \quad \Gamma \vdash c_n:S \darrow T_n} + \\[1mm] + \infer{\Gamma \vdash \cif{c_1}{c_2}{c_3} : S \darrow T} + {\Gamma \vdash c_1 : S \darrow \bool \quad \Gamma \vdash c_2:S \darrow T \quad \Gamma \vdash c_3:S \darrow T} + \\[1mm] + \infer{\Gamma \vdash \clet{x}{c_r}{c_b} : S \darrow T} + {\Gamma \vdash c_r : S \darrow R \quad \Gamma \vdash c_b: (S \ctupxt R) \darrow T} + \\[1mm] + \infer{\Gamma \vdash \cbuild{c_s}{i}{c_e} : S \darrow \vector{T}} + {\Gamma \vdash c_s : S \darrow \nat \quad \Gamma \vdash c_e: (S \ctupxt \nat) \darrow T} + \end{gather*} +\end{minipage}} + \caption{Type system for \CatLang} \label{fig:cl-types} +\end{figure} + +\subsection{Splitting for reverse mode} + +Supppose $f$ is defined in terms of $g$ and $h$: +$$ +\lstinline|f(x) = g(h(x))| +$$ +Here are the types: +$$ +\begin{array}{rcl} +f & : & R \to T \\ +g & : & S \to T \\ +h & : & R \to S \\ +\revDf{f} & : & (R, \tangent{T}) \to (T, \tangent{R}) \\ +\revDf{g} & : & (S, \tangent{T}) \to (T, \tangent{S}) \\ +\revDf{h} & : & (R, \tangent{S}) \to (S, \tangent{R}) \\ +\end{array} +$$ +How can we define $\revDf{f}$ by calling $\revDf{g}$ and $\revDf{h}$? It would +have to look something like this +$$ +\begin{array}{rcl} +\revDf{f}(r,dt) & = & \lstinline|letrec| \begin{array}[t]{l} + (t, ds) = \revDf{g}( s, dt ) \\ + (s, dr) = \revDf{h}( r, ds ) \\ + \end{array} \\ + && \lstinline|in| \; (t, dr) +\end{array} +$$ +We can't call $\revDf{g}$ before $\revDf{h}$, nor the other way around. +That's why there is a letrec! Even leaving aside how we generate this code, +We'd need lazy evaluation to execute it. + +The key idea for splitting is this. Given $f : S \to T$, produce two functions +$$ +\begin{array}{rcl} + \revfDf{f} & : & S \to (T,X) \\ + \revrDf{f} & : & (X,\tangent{T}) \to \tangent{S} +\end{array} +$$ +where the type $X$ depends on the details of f's definition. The idea is that +$X$ records all the stuff that $f$ computed when running forward that is necessary +for it to run backward. Now we can write +$$ +\begin{array}{rcl} +\revDf{f}(s,dt) & = & \lstinline|letrec| \begin{array}[t]{l} + (t, \mathit{xf}) = \revfDf{f}( s ) \\ + ds = \revfDf{f}( \mathit{xf}, dt ) \\ + \end{array} \\ +&& \lstinline|in| \; (t, ds) +\\[1mm] + +\revfDf{f}(r) & = & \lstinline|letrec| \begin{array}[t]{l} + (s, xh) = \revfDf{h}( r ) \\ + (t, xg) = \revfDf{g}( r ) \\ + \end{array} \\ + && \lstinline|in| \; (t, (xh,xg)) +\\[1mm] +\revrDf{f}((xh,xg),dt) & = & \revrDf{h}(dh, \revrDf{g}(dg, gt)) +\end{array} +$$ + \section{Implementation} The implementation differs from this document as follows: @@ -1505,6 +1711,9 @@ \section{Old stuff about Build and lambda} \label{sec:build-lam-fail} the solution I adopted was not to have a new language construct, but rather to differentiate only $\buildfun(n, \lambda i.e)$. Any other use of $\buildfun$ will fail.} + + + \end{document} do C( s -o t) -> s -o C(t) diff --git a/knossos.cabal b/knossos.cabal index ed7b03f4f..69f863acd 100644 --- a/knossos.cabal +++ b/knossos.cabal @@ -12,6 +12,7 @@ executable ksc CSE Cgen KMonad + Ksc.CatLang Ksc.Futhark Ksc.Pipeline Lang diff --git a/src/ksc/AD.hs b/src/ksc/AD.hs index 0cb90053e..1aaf0e469 100644 --- a/src/ksc/AD.hs +++ b/src/ksc/AD.hs @@ -119,6 +119,8 @@ gradBuild TupleAD s n ti body grad_body = mkLet (gradTVar TupleAD s ti) (Tuple [Var ti, lmZero s (Var ti)]) $ gradE TupleAD s body +gradBuild SplitAD _ _ _ _ = error "gradBuild:SplitAD" -- Doesn't make sense + --------------- gradFold :: ADPlan -> Shape -> TVar -> TExpr -> TExpr -> TExpr -> TExpr gradFold BasicAD s ti body acc v = @@ -156,6 +158,8 @@ gradFold TupleAD s _ti _body acc _v = lmDummyFold (TypeTuple [t_acc, TypeLM (typeof s) t_acc]) where t_acc = typeof acc +gradFold SplitAD _ _ _ _ _ = error "gradFold:SplitAD" -- Doesn't make sense + --------------- gradCall :: ADPlan -> Shape -> TFun -> TExpr -> TExpr gradCall BasicAD s f args @@ -194,6 +198,8 @@ gradCall TupleAD s f args grad_arg = gradE TupleAD s arg arg = args +gradCall SplitAD _ _ _ = error "gradCall:SplitAD" -- Doesn't make sense + ---------------------- gradLet :: HasCallStack => ADPlan -> Shape -> TVar -> TExpr -> TExpr -> TExpr gradLet BasicAD s v e1 e2 @@ -212,6 +218,8 @@ gradLet TupleAD s v e1 e2 gradE TupleAD s e2 where dv = gradTVar TupleAD s v +gradLet SplitAD _ _ _ _ = error "gradLet:SplitAD" -- Doesn't make sense + {- Note [Shadowing after grad] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ We transform let x = @@ -234,7 +242,7 @@ lmVCat_AD :: ADPlan -> [TExpr] -> TExpr lmVCat_AD BasicAD ms = lmVCat ms lmVCat_AD TupleAD ms = Tuple [ Tuple (map pFst ms) , lmVCat (map pSnd ms) ] - +lmVCat_AD SplitAD _ = error "lmVCat_AD:SplitAD" -- Doesn't make sense --------------------------------- diff --git a/src/ksc/Annotate.hs b/src/ksc/Annotate.hs index fcb6e5b5a..8f1b5673e 100644 --- a/src/ksc/Annotate.hs +++ b/src/ksc/Annotate.hs @@ -4,7 +4,7 @@ ScopedTypeVariables, TypeApplications #-} module Annotate ( - annotDecls, lintDefs + annotDecls, lintDefs, callResultTy_maybe, callResultTy ) where import Lang @@ -214,28 +214,34 @@ tcVar var mb_ty -- callResultTy_maybe is given a (global) function and the type of its -- argument, and returns the type of its result. -- --- It has special cases for a bunch opuilt-in functions with polymorphic +-- It has special cases for a bunch built-in functions with polymorphic -- types; that is, where the result type is a function of the argument types -- Otherwise it just looks in the global symbol table. callResultTy_maybe :: GblSymTab -> Fun -> Type -> Either SDoc Type -callResultTy_maybe env fun args +callResultTy_maybe env fun arg_ty | is_user_fun fun - = userCallResultTy_maybe fun env args + = userCallResultTy_maybe fun env arg_ty | otherwise - = primCallResultTy_maybe fun (typeof args) + = primCallResultTy_maybe fun arg_ty where is_user_fun = isUserFun . funIdOfFun +callResultTy :: GblSymTab -> Fun -> Type -> Type +callResultTy gst fun arg_ty + = case callResultTy_maybe gst fun arg_ty of + Left err -> pprTrace "callResultType:" err TypeUnknown + Right res_ty -> res_ty + ----------------------------------------------- -- The typecheck monad ----------------------------------------------- userCallResultTy_maybe :: HasCallStack => Fun -> GblSymTab -> Type -> Either SDoc Type -userCallResultTy_maybe fn env args - = case lookupGblST (fn, typeof args) env of - Just def -> userCallResultTy_help def args +userCallResultTy_maybe fn env arg_ty + = case lookupGblST (fn, arg_ty) env of + Just def -> userCallResultTy_help def arg_ty Nothing -> Left (text "Not in scope: userCall:" <+> ppr fn $$ ppr env) userCallResultTy_help :: HasCallStack @@ -243,13 +249,12 @@ userCallResultTy_help :: HasCallStack userCallResultTy_help (Def { def_fun = fn , def_res_ty = ret_ty , def_args = params }) - args - = case check_args bndr_tys arg_tys of + arg_ty + = case check_args bndr_ty arg_ty of Just err -> Left err Nothing -> Right ret_ty where - bndr_tys = tVarType params - arg_tys = typeof args + bndr_ty = tVarType params check_args :: Type -> Type -> Maybe SDoc -- Return (Just err) if there's a wrong-ness @@ -402,18 +407,18 @@ lookupLclTc v Just ty -> return ty } lookupGblTc :: Fun -> TypedExpr -> TcM Type -lookupGblTc fun args +lookupGblTc fun arg = do { st <- getSymTabTc - ; case callResultTy_maybe (gblST st) fun (typeof args) of + ; case callResultTy_maybe (gblST st) fun (typeof arg) of Left err -> do { addErr $ hang err 2 (mk_extra st) ; return TypeUnknown } Right res_ty -> return res_ty } where mk_extra st = vcat [ text "In a call of:" <+> ppr fun <+> parens (text (show fun)) - , text " Arg types:" <+> ppr (typeof args) - , text " Args:" <+> ppr (exprOf args) - , text "ST lookup:" <+> ppr (Map.lookup (fun, typeof args) (gblST st)) + , text " Arg types:" <+> ppr (typeof arg) + , text " Args:" <+> ppr (exprOf arg) + , text "ST lookup:" <+> ppr (Map.lookup (fun, typeof arg) (gblST st)) -- This is very verbose, and obscures error messages, but can be useful for typos. -- Perhaps think about printing it only for failed lookup of userfun -- , text "ST keys:" <+> gblDoc st diff --git a/src/ksc/Cgen.hs b/src/ksc/Cgen.hs index a0675855a..60b1e7293 100644 --- a/src/ksc/Cgen.hs +++ b/src/ksc/Cgen.hs @@ -199,7 +199,7 @@ params_withPackedParams param = case typeof param of TypeTuple tys -> let params = zipWith mkParam [1..] tys mkParam i ty = TVar ty (Simple name) - where name = nameOfVar (tVarVar param) ++ "arg" ++ show i + where name = tVarName param ++ "arg" ++ show i packParams = Let param (Tuple (map Var params)) in (params, ensureDon'tReuseParams params . packParams) _ -> ([param], id) @@ -525,6 +525,8 @@ cgenUserFun (f, ty) = case f of DrvFun s (AD BasicAD Rev) -> "rev$" ++ cgenFunId (s, ty) DrvFun s (AD TupleAD Fwd) -> "fwdt$" ++ cgenFunId (s, ty) DrvFun s (AD TupleAD Rev) -> "revt$" ++ cgenFunId (s, ty) + DrvFun s (AD SplitAD Fwd) -> "revsf$" ++ cgenFunId (s, ty) + DrvFun s (AD SplitAD Rev) -> "revsr$" ++ cgenFunId (s, ty) cgenAnyFun :: HasCallStack => (TFun, Type) -> CType -> String cgenAnyFun (tf, ty) cftype = case tf of diff --git a/src/ksc/Ksc/CatLang.hs b/src/ksc/Ksc/CatLang.hs new file mode 100644 index 000000000..c83a7342c --- /dev/null +++ b/src/ksc/Ksc/CatLang.hs @@ -0,0 +1,949 @@ +{-# LANGUAGE MonoLocalBinds #-} + +{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-} + +-- Ksc.CatLang: Language for compilation to categories +module Ksc.CatLang( CLDef, toCLDefs, fromCLDefs, toCLDef_maybe, fromCLDef +-- , fwdAdDefs, revAdDefs, fsAdDefs + ) where + +import Prelude hiding( (<>) ) +import Lang +import LangUtils +import Prim +import OptLet +import Annotate( callResultTy ) +import qualified Data.Set as S +import Data.Maybe( mapMaybe ) +import Data.List( mapAccumL ) + +data CLExpr + = CLId + | CLPrune [Int] Int + | CLKonst Konst + | CLCall Type FunId -- The Type is the result type + | CLComp CLExpr CLExpr -- Reverse composition, c1;c2 + | CLTuple [CLExpr] -- Tuple + | CLIf CLExpr CLExpr CLExpr -- If + | CLLet TVar CLExpr CLExpr -- Let $var = $rhs in $body + | CLBuild CLExpr TVar CLExpr -- Build $size (Lam ($var : Integer) $body) + | CLFold TVar CLExpr CLExpr CLExpr + -- ^ Fold (Lam $t body) $acc $vector + +data CLDef = CLDef { cldef_fun :: FunId + , cldef_arg :: TVar -- Arg type S + , cldef_rhs :: CLExpr + , cldef_res_ty :: Type } -- Result type T + +mkCLComp :: CLExpr -> CLExpr -> CLExpr +-- Optimise away the identity function to reduce clutter +mkCLComp CLId e = e +mkCLComp e CLId = e +mkCLComp e1 e2 = CLComp e1 e2 + +clResultType :: CLExpr -> [Type] -> Type +clResultType CLId tys = mkTupleTy tys +clResultType (CLPrune ts _) tys = mkTupleTy (pick ts tys) +clResultType (CLKonst k) _ = typeofKonst k +clResultType (CLComp c1 c2) tys = clResultType c2 [clResultType c1 tys] +clResultType (CLTuple cs) tys = mkTupleTy [clResultType c tys | c <- cs] +clResultType (CLLet tv _ b) tys = clResultType b (typeof tv : tys) +clResultType (CLBuild _ _ b) tys = TypeVec (clResultType b (TypeInteger : tys)) +clResultType (CLCall ty _) _ = ty +clResultType (CLIf _ c _) tys = clResultType c tys +clResultType (CLFold t l _ _) tys = clResultType l (typeof t : tys) + +foldUnimplemented :: String -> error +foldUnimplemented s = error (s ++ ": CLFold not yet implemented") + +----------------------------------------------- +-- Pretty printing +----------------------------------------------- + +instance Pretty CLDef where + ppr (CLDef { cldef_fun = f + , cldef_arg = arg + , cldef_rhs = rhs + , cldef_res_ty = res_ty }) + = sep [ hang (text "def" <+> ppr f <+> pprParendType res_ty) + 2 (parens (pprTVar arg)) + , nest 2 (text "=" <+> ppr rhs) ] + +instance Pretty CLExpr where + pprPrec = pprCLExpr + +pprCLExpr :: Prec -> CLExpr -> SDoc +pprCLExpr p c@(CLComp {}) + = parensIf p precOne $ + sep [ pprCLExpr precTwo c1 + , vcat [ text ";" <+> pprCLExpr precTwo c + | c <- cs ] ] + where + (c1,cs) = case collect c of + (c1:cs) -> (c1,cs) + [] -> pprPanic "pprCLExpr" (text "") + + collect (CLComp a b) = collect a ++ collect b + collect e = [e] + +pprCLExpr _ CLId = text "Id" +pprCLExpr _ (CLKonst k) = ppr k +pprCLExpr _ (CLCall _ f) = ppr f +pprCLExpr p (CLPrune ts n) = sep [ text "Prune" <> char '[' + <> cat (punctuate comma (map ppr ts)) + <> char '/' <> int n <> char ']' ] +pprCLExpr _ (CLTuple cs) + | [] <- cs = text "[]" + | [c] <- cs = char '[' <+> pprCLExpr precZero c <+> char ']' + | (c1:rest) <- cs = cat ( [ char '[' <+> pprCLExpr precZero c1 ] ++ + [ char ',' <+> pprCLExpr precZero c | c <- rest ] ) <+> char ']' + +pprCLExpr p (CLIf b t e) = parensIf p precZero $ + sep [ text "if" <+> ppr b + , text "then" <+> ppr t + , text "else" <+> ppr e] +pprCLExpr p (CLLet x r b) + = parensIf p precZero $ + vcat [ text "let" <+> (bracesSp $ sep [ ppr x <+> char '=', nest 2 (pprCLExpr precZero r) ]) + , pprCLExpr precZero b ] + +pprCLExpr p (CLBuild es x ee) + = parensIf p precZero $ + text "build" <+> sep [ pprCLExpr precOne es + , parens $ char '\\' <> ppr x <> char '.' <+> pprCLExpr precZero ee ] + +pprCLExpr p (CLFold t elam eacc ev) + = parensIf p precZero $ + text "fold" <+> sep [ parens $ char '\\' <> ppr t <> char '.' <+> pprCLExpr precZero elam + , pprCLExpr precOne eacc + , pprCLExpr precOne ev ] + +----------------------------------------------- +-- Convert to CLDedf +----------------------------------------------- + +toCLDefs :: [TDef] -> [CLDef] +toCLDefs defs = mapMaybe toCLDef_maybe defs + +toCLDef_maybe :: TDef -> Maybe CLDef +toCLDef_maybe (Def { def_fun = fun + , def_args = arg + , def_res_ty = res_ty + , def_rhs = rhs }) + | Fun f <- fun + , UserRhs e <- rhs + = Just CLDef { cldef_fun = f + , cldef_arg = arg + , cldef_rhs = toCLExpr [arg] e + , cldef_res_ty = res_ty } + | otherwise + = Nothing + +data EnvPruned = Pruned | NotPruned + +toCLExpr :: [TVar] -> TExpr -> CLExpr +toCLExpr = to_cl_expr NotPruned + +to_cl_expr :: EnvPruned -> [TVar] -> TExpr -> CLExpr + +to_cl_expr Pruned [w] (Var v) | v == w = CLId +to_cl_expr Pruned env (Var v) = pprPanic "toCLExpr:var" (ppr v $$ ppr env) +to_cl_expr Pruned [] (Konst k) = CLKonst k +to_cl_expr Pruned env (Konst k) = pprPanic "toCLExpr:konst" (ppr k $$ ppr env) +to_cl_expr Pruned env (Tuple es) = CLTuple (map (toCLExpr env) es) +to_cl_expr pruned env (Call f e) = to_cl_call pruned env f e + +-- We shouldn't do this because asserts can be essential to good +-- optimisation, but we will do it for now, to make progress with +-- CatLang +to_cl_expr pruned env (Assert _ e) = to_cl_expr pruned env e + +to_cl_expr Pruned env (If e1 e2 e3) + = CLIf (toCLExpr env e1) (toCLExpr env e2) (toCLExpr env e3) +to_cl_expr Pruned env (Let tv rhs body) + = CLLet tv (toCLExpr env rhs) (toCLExpr (tv:env) body) + +to_cl_expr _ _ e@(Lam {}) = pprPanic "toCLExpr Lam" (ppr e) +to_cl_expr _ _ e@(App {}) = pprPanic "toCLExpr App" (ppr e) +to_cl_expr _ _ e@(Dummy {}) = pprPanic "toCLExpr Dummy" (ppr e) + +to_cl_expr NotPruned env e = prune env e + +--------------- +to_cl_call :: EnvPruned -> [TVar] -> TFun -> TExpr -> CLExpr +-- Calls: +-- * for build, prune before the build +-- * for other calls, prune in the argument + +to_cl_call pruned env f e + | f `isThePrimFun` "build" + , Tuple [n, Lam tvi body] <- e + = case pruned of + NotPruned -> prune env call + Pruned -> CLBuild (toCLExpr env n) tvi (toCLExpr (tvi:env) body) + + | f `isThePrimFun` "sumbuild" + , Tuple [n, lam] <- e + = to_cl_expr pruned env (pSum (pBuild n lam)) + + | f `isThePrimFun` "fold" + , Tuple [Lam t body, acc, v] <- e + = case pruned of + NotPruned -> prune env call + Pruned -> + CLFold t (toCLExpr (t:env) body) (toCLExpr env acc) (toCLExpr env v) + + | TFun ty (Fun fun_id) <- f + = to_cl_expr pruned env e`mkCLComp` CLCall ty fun_id + + | TFun _ (GradFun _ _) <- f + = pprPanic "toCLExpr Call of GradFun" (ppr call) + | TFun _ (DrvFun _ _) <- f + = pprPanic "toCLExpr Call of DrvFun" (ppr call) + where + call = Call f e + +--------------- +prune :: [TVar] -> TExpr -> CLExpr +-- See Note [Pruning] +prune env e + | no_prune_reqd = to_cl_expr Pruned env e + | otherwise = CLPrune trim_indices (length env) + `mkCLComp` to_cl_expr Pruned trimmed_env e + where + fvs = freeVarsOf e + no_prune_reqd = all (\tv -> tv `S.member` fvs) env + + trimmed_prs :: [(TVar,Int)] + trimmed_prs = trim fvs 1 env + + (trimmed_env, trim_indices) = unzip trimmed_prs + + trim :: S.Set TVar -> Int -> [TVar] -> [(TVar,Int)] + trim _ _ [] = [] + trim fvs n (tv:tvs) + | tv `S.member` fvs = (tv,n) : trim (tv `S.delete` fvs) (n+1) tvs + | otherwise = trim fvs (n+1) tvs + +find :: Eq a => a -> [a] -> Maybe Int +-- Return 1-indexed position of the item in the list +find x xs + = go 1 xs + where + go j (x1:xs) | x == x1 = Just j + | otherwise = go (j+1) xs + go _ [] = Nothing + + +{- Note [Pruning] +~~~~~~~~~~~~~~~~~ +We need to be careful when we have shadowing. Then +we might have a call + to_cl_expr [x,v,x,w] e +where the first 'x' in the environment shadows the second. +So when pruning we must be careful to pick the first element of +the tuple, not the third, let alone both! +-} + +----------------------------------------------- +-- Convert from CLDef +----------------------------------------------- + +fromCLDefs :: [CLDef] -> [TDef] +fromCLDefs cldefs = map fromCLDef cldefs + +fromCLDef :: CLDef -> TDef +fromCLDef (CLDef { cldef_fun = f + , cldef_arg = arg + , cldef_rhs = rhs + , cldef_res_ty = res_ty }) + = Def { def_fun = Fun f + , def_args = arg + , def_res_ty = res_ty + , def_rhs = UserRhs rhs' } + where + rhs' = fromCLExpr (mkInScopeSet [arg]) (Var arg) rhs + +fromCLExpr :: InScopeSet -> TExpr -> CLExpr -> TExpr +-- (fromCLExpr is arg c) +-- We may freely duplicate 'arg', so make sure that all +-- of the TExprs in 'arg' are trivial -- usually variables -- +-- and hence can be duplicated without dupliating work +fromCLExpr _ _ (CLKonst k) = Konst k +fromCLExpr _ arg CLId = arg +fromCLExpr is arg (CLPrune ts n) = mkTuple (pick ts (splitTuple arg n)) +fromCLExpr is arg (CLTuple es) = makeArgDupable is arg $ \is arg -> + Tuple (map (fromCLExpr is arg) es) +fromCLExpr is arg (CLIf b t e) = makeArgDupable is arg $ \ is arg -> + If (fromCLExpr is arg b) + (fromCLExpr is arg t) + (fromCLExpr is arg e) + +fromCLExpr _ arg (CLCall ty f) + = Call (TFun ty (Fun f)) arg + +fromCLExpr is arg (CLComp e1 e2) + = fromCLExpr is (fromCLExpr is arg e1) e2 + +fromCLExpr is arg (CLLet tv rhs body) + = Let tv' rhs' (fromCLExpr is' (pTupCons (Var tv') arg) body) + where + rhs' = fromCLExpr is arg rhs + (is', tv') = notInScopeTV is tv + +fromCLExpr is arg (CLBuild size tv elt) + = makeArgDupable is arg $ \is arg -> + pBuild (fromCLExpr is arg size) + (Lam tv' (fromCLExpr is' (pTupCons (Var tv') arg) elt)) + where + (is', tv') = notInScopeTV is tv + +fromCLExpr is arg (CLFold t lam acc v) + = makeArgDupable is arg $ \is arg -> + mkPrimCall3 "fold" + (Lam t' (fromCLExpr is' (pTupCons (Var t') arg) lam)) + (fromCLExpr is arg acc) + (fromCLExpr is arg v) + where + (is', t') = notInScopeTV is t + +------------------------------------------- +-- Forward AD, tupled +-- S => T = (S, dS) -> (T, dT) +------------------------------------------- +{- + +fwdAdDefs :: GblSymTab -> [CLDef] -> (GblSymTab, [TDef]) +fwdAdDefs gst cldefs + = let defs = map fwdAdDef cldefs + in (extendGblST gst defs, defs) + +fwdAdDef :: CLDef -> TDef +fwdAdDef (CLDef { cldef_fun = f + , cldef_arg = TVar ty v + , cldef_rhs = rhs + , cldef_res_ty = res_ty }) + = Def { def_fun = DrvFun f (AD { adPlan = TupleAD, adDir = Fwd }) + , def_args = arg' + , def_res_ty = tangentPair res_ty + , def_rhs = UserRhs rhs' } + where + arg' = TVar (tangentPair ty) v + rhs' = mkPairLet (mkInScopeSet [arg']) "x" "dx" (Var arg') $ \ is x dx -> + fwdAdExpr is [x] [dx] rhs + +fwdAdExpr :: InScopeSet -> [TExpr] -> [TExpr] -> CLExpr -> TExpr +fwdAdExpr _ _ _ (CLKonst k) = Tuple [ Konst k, mkTangentZero (Konst k) ] +fwdAdExpr is s ds (CLIf a b c) = If (fromCLExpr is s a) + (fwdAdExpr is s ds b) + (fwdAdExpr is s ds c) + +-- (.) :: (b => c) -> (a => b) -> (a => c) +-- (f . g) s = f (g s) +fwdAdExpr is s ds (CLComp g f) + = mkPairLet is "gr" "dgr" (fwdAdExpr is s ds g) $ \is a da -> + fwdAdExpr is [a] [da] f + +fwdAdExpr _ s ds CLId + = Tuple [mkTuple s, mkTuple ds] + +fwdAdExpr is s ds (CLPrune ts _ c) + = fwdAdExpr is (pick ts s) (pick ts ds) c + +fwdAdExpr _ s ds (CLCall ty f) + = Call (TFun (tangentPair ty) fwd_fun) + (Tuple [mkTuple s, mkTuple ds]) + where + fwd_fun = DrvFun f ad_mode + ad_mode = AD { adPlan = TupleAD, adDir = Fwd } + +-- (,) :: (s => t1) -> (s => t2) -> (s => (t1,t2)) +-- (c1,c2) (s,ds) = let (ar,dar) = c1 (s,ds) +-- (br,dbr) = c2 (s,ds) +-- in ((ar,br),(dar,dbr)) +fwdAdExpr is s ds (CLTuple cs) + = mkTempLets is "t" es $ \vs -> + Tuple [ mkTuple (map pFst vs), mkTuple (map pSnd vs) ] + where + es = map (fwdAdExpr is s ds) cs + +-- (let x = c in b) (s,ds) = let (x,dx) = c (s,ds) +-- in b (x:s, dx:ds) +fwdAdExpr is s ds (CLLet tv rhs body) + = mkPairLet is nm ("d" ++ nm) (fwdAdExpr is s ds rhs) $ \ is x dx -> + fwdAdExpr is (x:s) (dx:ds) body + where + nm = tVarName tv + +-- (build s i e) <> (s,ds) = let (n,_) = c (s,ds) +-- in unzip (build n (\i. e (i:s, ():ds)) +fwdAdExpr is s ds (CLBuild n tvi elt) + = mkTempLet is "np" (fwdAdExpr is s ds n) $ \ is np -> + let (is', tvi') = notInScopeTV is tvi + in pUnzip $ pBuild (pFst np) $ + Lam tvi' (fwdAdExpr is' (Var tvi' : s) (unitExpr:ds) elt) + +fwdAdExpr _ _ _ CLFold{} = foldUnimplemented "fwdAdExpr" + + +tangentPair :: Type -> Type +tangentPair ty = TypeTuple [ty, tangentType ty] + + +------------------------------------------- +-- Reverse AD, not tupled +-- S => T = (S, dT) -> dS +------------------------------------------- + +revAdDefs :: GblSymTab -> [CLDef] -> (GblSymTab, [TDef]) +revAdDefs gst cldefs + = let defs = map revAdDef cldefs + in (extendGblST gst defs, defs) + +revAdDef :: CLDef -> TDef +revAdDef (CLDef { cldef_fun = f + , cldef_arg = TVar arg_ty v + , cldef_rhs = rhs + , cldef_res_ty = res_ty }) + = Def { def_fun = DrvFun f (AD { adPlan = BasicAD, adDir = Rev }) + , def_args = arg' + , def_res_ty = tangentType arg_ty + , def_rhs = UserRhs rhs' } + where + arg' = TVar (TypeTuple [arg_ty, tangentType res_ty]) (Simple "arg") + nm = varName v + rhs' = mkPairLet (mkInScopeSet [arg']) nm "dr" (Var arg') $ \ is x dr -> + revAdExpr is [x] dr rhs + +revAdExpr :: InScopeSet -> [TExpr] -> TExpr -> CLExpr -> TExpr +revAdExpr _ s _ (CLKonst _) = mkTangentZero (mkTuple s) +revAdExpr is s dt (CLIf a b c) = If (fromCLExpr is s a) + (revAdExpr is s dt b) + (revAdExpr is s dt c) + +-- (.) :: (b => c) -> (a => b) -> (a => c) +-- (f . g) (a,dc) = let b = g <> a +-- db = f (b,dc) +-- in +-- g (a,db) +revAdExpr is s dt (CLComp g f) + = mkTempLet is "b" (fromCLExpr is s g) $ \ is b -> + mkTempLet is "db" (revAdExpr is [b] dt f) $ \ is db -> + revAdExpr is s db g + +revAdExpr _ _ dt CLId + = dt + +revAdExpr is s dt (CLPrune ts _ c) + = mkTempLet is "dps" (revAdExpr is (pick ts s) dt c) $ \ _ dps -> + let do_one (s,i) + | Just j <- find i ts = pSel j n dps + | otherwise = mkTangentZero s + in mkTuple (map do_one (s `zip` [1..])) + where + n = length ts + +revAdExpr _ s dt (CLCall _ f) + = Call (TFun res_ty rev_fun) (Tuple [mkTuple s, dt]) + where + res_ty = tangentType (typeof (mkTuple s)) + rev_fun = DrvFun f ad_mode + ad_mode = AD { adPlan = BasicAD, adDir = Rev } + +-- (,) :: (s => t1) -> (s => t2) -> (s => (t1,t2)) +-- (c1,c2) (s,dt) = let (dt1,dt2) = dt +-- ds1 = c1 (s,dt1) +-- ds2 = c2 (s,dt2) +-- in (ds1 + ds2) +revAdExpr is s dt (CLTuple cs) + = foldr1 pAdd $ + [ revAdExpr is s (pSel i n dt) c + | (c,i) <- cs `zip` [1..] ] + where + n = length cs + +-- (build sz (\i.e) (s,dt) +-- = sumbuild (sz <> s) (\i. (e (i:s, dt[i]))) +revAdExpr is s dt (CLBuild sz tvi elt) + = pSumBuild (fromCLExpr is s sz) $ + Lam tvi' $ pTupTail $ + revAdExpr is' (Var tvi' : s) + (pIndex (Var tvi') dt) elt + where + (is', tvi') = notInScopeTV is tvi + +-- (let x = c in b) (s,dt) = let x = c <> s +-- (dxb:dsb) = b (x:s, dt) +-- dsx = c (s,dxb) +-- in +-- dsb + dsx +revAdExpr is s dt (CLLet tv rhs body) + = mkTempLet is (tVarName tv) (fromCLExpr is s rhs) $ \is x -> + mkTempLet is "dx" (revAdExpr is (x : s) dt body) $ \ is dx -> + pAdd (pTupTail dx) (revAdExpr is s (pTupHead dx) rhs) + +revAdExpr _ _ _ CLFold{} = foldUnimplemented "revAdExpr" + +-------------------------------------------- +-- Spilt AD +-- fwds: S => T = S -> (T, X) +-- revs: S => T = (dT,X) -> dS +-------------------------------------------- + +type FwdSplit = InScopeSet -> [TExpr] -> TExpr +type RevSplit = InScopeSet -> TExpr -> TExpr -> TExpr + +data SplitResult = SR { sr_fwd :: FwdSplit + , sr_rev :: RevSplit + , sr_empx :: XShape } +-- For an expression of type S => T +-- if XShape = XNonEmpty, then +-- fwds: S -> (T, X) +-- revs: (dT,X) -> dS +-- if XShape = XEmpty, then +-- fwds: S -> X +-- revs: dT -> dS +-- +-- The *top-level* fwds$f is always S -> (T, X) +-- even if X = (), so that we can tell what X is +-- But the top-level revs$g is dT -> dS if X = () + +data XShape = XEmpty | XNonEmpty + +instance Pretty XShape where + ppr XEmpty = text "XEmpty" + ppr XNonEmpty = text "XEmpty" + +andXS :: XShape -> XShape -> XShape +andXS XEmpty XEmpty = XEmpty +andXS _ _ = XNonEmpty + +isNonEmpty :: XShape -> Bool +isNonEmpty XNonEmpty = True +isNonEmpty XEmpty = False + +------------------- +fsAdDefs :: GblSymTab -> [CLDef] -> (GblSymTab, [TDef]) +fsAdDefs gst cldefs = (gst', concat defs) + where + (gst', defs) = mapAccumL fsAdDef gst cldefs + +------------------- +fsAdDef :: GblSymTab -> CLDef -> (GblSymTab, [TDef]) +fsAdDef gst (CLDef { cldef_fun = f + , cldef_arg = arg + , cldef_rhs = rhs + , cldef_res_ty = res_ty }) + = ( extendGblST gst [fwd_def, rev_def] + , [fwd_def, rev_def]) + where + arg_ty = typeof arg + + SR { sr_fwd = fwd, sr_rev = rev, sr_empx = emp } = fsAdExpr gst [arg_ty] rhs + + fwd_def = Def { def_fun = DrvFun f (AD { adPlan = SplitAD, adDir = Fwd }) + , def_args = arg + , def_res_ty = fwd_rhs_ty + , def_rhs = UserRhs final_fwd_rhs } + fwd_in_scope = mkInScopeSet [arg] + fwd_rhs = fwd fwd_in_scope [Var arg] + final_fwd_rhs = case emp of + XEmpty -> Tuple [fwd_rhs, unitExpr] + XNonEmpty -> fwd_rhs + + fwd_rhs_ty = typeof final_fwd_rhs + x_ty = case fwd_rhs_ty of + TypeTuple [_, x_ty] -> x_ty + _ -> pprPanic "fsAdDef" (ppr fwd_rhs_ty) + + rev_def = Def { def_fun = DrvFun f (AD { adPlan = SplitAD, adDir = Rev }) + , def_args = rev_arg + , def_res_ty = tangentType (typeof arg) + , def_rhs = UserRhs rev_rhs } + + (rev_arg, rev_rhs) + = case emp of + XNonEmpty -> (dtx, mkPairLet is_dtx "dt" "x" (Var dtx) $ \is dt x -> + rev is dt x) + XEmpty -> (dt, rev is_dt (Var dt) unitExpr) + + dtx = TVar (TypeTuple [tangentType res_ty, x_ty]) (Simple "dtx") + dt = TVar (tangentType res_ty) (Simple "dt") + is_dtx = mkInScopeSet [dtx] + is_dt = mkInScopeSet [dt] + +------------------- +fsAdExpr :: GblSymTab -> [Type] -> CLExpr -> SplitResult +-- This is the main workhorse +-- The [Type] is the type of the argument tuple at that point + +fsAdExpr _ _ (CLKonst k) + = SR { sr_fwd = \ _ _ -> Konst k + , sr_rev = \ _ _ _ -> mkTangentZero (Konst k) + , sr_empx = XEmpty } + +fsAdExpr _ _ CLId + = SR { sr_fwd = \ _ s -> mkTuple s + , sr_rev = \ _ dt _ -> dt + , sr_empx = XEmpty } + +-- prune ts c s = let (rc, xc) = c (pick is s) +-- in (rc, (xc,zeros)) +-- prune ts c (dt,x) = let (xc,zeros) = x +-- ds = c (dt,xc) +-- in spread ts ds zeros + +fsAdExpr gst tys (CLPrune ts _ c) + = SR{ sr_fwd = fwd, sr_rev = rev, sr_empx = c_emp `andXS` zero_emp } + where + fwd is s = mkXPairLet is "rc" "xc" (c_emp, c_fwd is (pick ts s)) $ \_ rc xc -> + mkXFwdPair rc [ (c_emp, xc), (zero_emp, xzero) ] + where + xzero = mkTuple [ (s!!(i-1)) | i <- zero_ts ] + + rev is dt x = mkTempLet is "ds" (c_rev is dt xc) $ \_ ds -> + mkTuple (map (do_one ds xzero) (tys `zip` [1..])) + where + (xc:xzero:_) = splitXTuple x [c_emp, zero_emp] + + m = length ts + z = length zero_ts + t_prs = ts `zip` [1..] + z_prs = zero_ts `zip` [1..] + + do_one ds xzero (ty, i_n) + | Just j_m <- findIndex t_prs i_n = pSel j_m m ds + | Just k_m <- findIndex z_prs i_n = pSel k_m z xzero + | otherwise = mkTangentZeroFromType ty + + + SR { sr_fwd = c_fwd, sr_rev = c_rev, sr_empx = c_emp } + = fsAdExpr gst (pick ts tys) c + + zero_emp | null zero_ts = XEmpty + | otherwise = XNonEmpty + + zero_ts :: [Int] -- Positions in [1..n] for which I need a zero + zero_ts = [ i | (ty,i) <- tys `zip` [1..] + , not (i `elem` ts) + , needValueForTangentZero ty ] + +{- +fsAdExpr gst is s (CLIf a b c) = If (fromCLExpr is s a) + (fsAdExpr gst is s b) + (fsAdExpr gst is s c) +-} +fsAdExpr _ _ e@(CLIf {}) = pprPanic "fsAdExpr:If" (ppr e) + +-- (.) :: (b => c) -> (a => b) -> (a => c) +-- (f . g) s = let (rg, xg) = g s +-- (rf, xf) = f rg +-- in (rf, (xf,xg)) +-- (f . g) (dt,x) = let (xf,xg) = x +-- dtf = f (dt,xf) +-- dtg = g (dtf,xg) +-- in dtg +fsAdExpr gst tys (CLComp g f) + = SR { sr_fwd = \ is s -> + mkXPairLet is "rg" "xg" (g_empx, fwd_g is s) $ \is rg xg -> + mkXPairLet is "rf" "xf" (f_empx, fwd_f is [rg]) $ \_ rf xf -> + mkXFwdPair rf [(f_empx,xf), (g_empx, xg)] + , sr_rev = \ is dt x -> + let (xf:xg:_) = splitXTuple x [f_empx, g_empx] in + mkTempLet is "dtf" (rev_f is dt xf) $ \is dtf -> + rev_g is dtf xg + , sr_empx = empx } + where + empx = f_empx `andXS` g_empx + SR { sr_fwd = fwd_f, sr_rev = rev_f, sr_empx = f_empx } = fsAdExpr gst [g_res_ty] f + SR { sr_fwd = fwd_g, sr_rev = rev_g, sr_empx = g_empx } = fsAdExpr gst tys g + g_res_ty = clResultType g tys + +fsAdExpr _ [ty] (CLCall _ (SelFun i n)) + = SR{ sr_fwd = fwd, sr_rev = rev, sr_empx = zero_emp } + where + fwd _ [s] = mkXFwdPair (pSel i n s) [ (zero_emp, xzero) ] + where + xzero = mkTuple [ mkTangentZero (pSel i n s) | i <- zero_ts ] + fwd _ s = pprPanic "fsAdExpr:SelFun" (ppr s) + + rev _ dt x = mkTuple (map (do_one dt x) (tys `zip` [1..])) + tys = case ty of + TypeTuple tys -> tys + _ -> pprPanic "fsAdExpr:SelFuny" (ppr ty) + + z = length zero_ts + z_prs = zero_ts `zip` [1..] + + do_one dt xzero (ty, i_n) + | i_n == i = dt + | Just k_m <- findIndex z_prs i_n = pSel k_m z xzero + | otherwise = mkTangentZeroFromType ty + + zero_emp | null zero_ts = XEmpty + | otherwise = XNonEmpty + + zero_ts :: [Int] -- Positions in [1..n] for which I need a zero + zero_ts = [ j | (ty,j) <- tys `zip` [1..] + , j /= i + , needValueForTangentZero ty ] + +fsAdExpr gst tys (CLCall _ f) + = SR { sr_fwd = fwd, sr_rev = rev, sr_empx = emp } + where + fs_fun = DrvFun f (AD { adPlan = SplitAD, adDir = Fwd }) + fr_fun = DrvFun f (AD { adPlan = SplitAD, adDir = Rev }) + + fwd _ [s] = case emp of + XEmpty -> pFst (mk_call fs_fun s) + XNonEmpty -> mk_call fs_fun s + fwd _ s = pprPanic "fsAdExpr:CLCall" (ppr s) + + rev _ dt x = mk_call fr_fun $ + case emp of { XEmpty -> dt; XNonEmpty -> Tuple [dt,x] } + + emp = funXShape gst f (mkTupleTy tys) + + mk_call fun arg = Call (TFun res_ty fun) arg + where + res_ty = callResultTy gst fun (typeof arg) + +-- (,) :: (s => t1) -> (s => t2) -> (s => (t1,t2)) +-- (c1,c2) s = let (ar,x1) = c1 s +-- (br,x2) = c2 s +-- in ((ar,br),(x1,x2)) +-- (c1,c2) (dt,x) = let (x1, x2) = x +-- (dt1, dt2) = dt +-- dr1 = c1 (dt1, x1) +-- dr2 = c2 (dt2, x2) +-- in dt1 + dt2 +fsAdExpr gst tys (CLTuple cs) + = SR { sr_fwd = fwd, sr_rev = rev, sr_empx = all_emp } + where + srs = map (fsAdExpr gst tys) cs + emps = map sr_empx srs + all_emp = foldr andXS XEmpty emps + + fwd is s = mkXPairLets is "r" "x" [ (sr_empx sr, sr_fwd sr is s) + | sr <- srs] $ \_ ts xs -> + mkXFwdPair (mkTuple ts) (emps `zip` xs) + + rev is dt x + = foldr1 pAdd [ sr_rev sr is dt x + | (sr,(dt,x)) <- srs `zip` (dts `zip` xs) ] + where + -- We should probably use lets for these + dts = splitTuple dt (length cs) + xs = splitXTuple x emps + + +-- (let v = c in b) s = let (rc,xc) = c s +-- let (rb,xb) = b (rc:s) +-- in (rb, (xc,xb)) +-- (let v = c in b) (dt,x) = let (xc,xb) = x +-- rb = b (dt,xb) +-- (dc:dr1) = rb +-- dr2 = c (dc,xc) +-- in dr1 + dr2 +fsAdExpr gst tys (CLLet tv rhs body) + = SR { sr_fwd = fwd, sr_rev = rev, sr_empx = rhs_emp `andXS` body_emp } + where + fwd is s = mkXPairLet is "rc" "xc" (rhs_emp, fwd_rhs is s) $ \ is rc xc -> + mkXPairLet is "rb" "xb" (body_emp, fwd_body is (rc:s)) $ \ _ rb xb -> + mkXFwdPair rb [(rhs_emp, xc), (body_emp, xb)] + rev is dt x + = mkTempLet is "rb" (rev_body is dt xbody) $ \ is rb -> + pAdd (pTupTail rb) (rev_rhs is (pSel 1 (tupLen rb) rb) xrhs) + where + (xrhs:xbody:_) = splitXTuple x [rhs_emp, body_emp] + + SR { sr_fwd = fwd_rhs, sr_rev = rev_rhs, sr_empx = rhs_emp } + = fsAdExpr gst tys rhs + SR { sr_fwd = fwd_body, sr_rev = rev_body, sr_empx = body_emp } + = fsAdExpr gst (typeof tv:tys) body + +-- (build n i e) s = unzip (build n (\i. e (i:s))) +-- (build n i e) (dt,x) = sumbuild n (\i. e (dt[i], x[i])) +fsAdExpr gst tys (CLBuild sz tvi elt) + = SR { sr_fwd = fwd, sr_rev = rev, sr_empx = elt_emp } + where + fwd is s = let (is', tvi') = notInScopeTV is tvi in + pUnzip $ pBuild (fromCLExpr is s sz) $ + Lam tvi' $ fwd_elt is' (Var tvi' : s) + rev is dt x = let (is', tvi') = notInScopeTV is tvi in + pSumBuild (pSize dt) $ Lam tvi' $ + rev_elt is' (pIndex (Var tvi') dt) (pIndex (Var tvi') x) + + SR { sr_fwd = fwd_elt, sr_rev = rev_elt, sr_empx = elt_emp } + = fsAdExpr gst (TypeInteger:tys) elt + +fsAdExpr _ _ CLFold{} = foldUnimplemented "fsAdEXpr CLFold not yet implemented" + + +funXShape :: GblSymTab -> FunId -> Type -> XShape +funXShape gst fn arg_ty + = case lookupGblST (fs_fn, arg_ty) gst of + Just (Def { def_res_ty = res_ty }) + | TypeTuple [_, x] <- res_ty + -> case x of + TypeTuple [] -> XEmpty + _ -> XNonEmpty + _ -> pprPanic "hasEmptyX" (ppr fn) + -- Not in GST, or bad shape + where + fs_fn = DrvFun fn (AD { adPlan = SplitAD, adDir = Fwd }) + + +mkXFwdPair :: TExpr -> [(XShape,TExpr)] -> TExpr +mkXFwdPair r prs + | null non_emp_xs = r -- Ignore xs + | otherwise = Tuple [r, mkTuple non_emp_xs ] + where + non_emp_xs = [ x | (XNonEmpty,x) <- prs ] + +mkXPairLet :: InScopeSet -> String -> String -> (XShape, TExpr) + -> (InScopeSet -> TExpr -> TExpr -> TExpr) + -> TExpr +mkXPairLet is sf ss (xshape, e) thing_inside + = case xshape of + XNonEmpty -> mkPairLet is sf ss e thing_inside + XEmpty -> mkTempLet is sf e $ \is ve -> + thing_inside is ve unitExpr + +mkXPairLets :: InScopeSet -> String -> String -> [(XShape,TExpr)] + -> (InScopeSet -> [TExpr] -> [TExpr] -> TExpr) + -> TExpr +mkXPairLets is sf ss prs thing_inside + = go is prs thing_inside + where + go is [] thing_inside = thing_inside is [] [] + go is (pr:prs) thing_inside = mkXPairLet is sf ss pr $ \is t x -> + go is prs $ \is ts xs -> + thing_inside is (t:ts) (x:xs) + +splitXTuple :: TExpr -> [XShape] -> [TExpr] +-- splitXTuple x [XEmpty, XEmpty] = [(), ()] +-- splitXTuple x [XNonEmpty, XEmpty] = [x, ()] +-- splitXTuple x [XNonEmpty, XNonEmpty] = [fst x, snd x] +-- NB: the all-empty case does not mention x at all +splitXTuple x bs + | nval == 1 = [ case b of { XEmpty -> unitExpr; XNonEmpty -> x } + | b <- bs ] + | otherwise = snd $ mapAccumL do_one 1 bs + where + nval = count isNonEmpty bs + + do_one :: Int -> XShape -> (Int, TExpr) + do_one i XNonEmpty = (i+1, pSel i nval x) + do_one i XEmpty = (i, unitExpr) +-} + +----------------------------------------------- +-- Utilities +----------------------------------------------- + + +splitTuple :: TExpr -> Int -> [TExpr] +-- Expects e to be a tuple-typed expression; +-- returns its n components. +-- May duplicate e +splitTuple _ 0 = [] +splitTuple e 1 = [e] +splitTuple (Tuple es) n + | n == length es = es + | otherwise = pprPanic "splitTuple" (ppr n $$ ppr es) +splitTuple (Call f a) n + | f `isThePrimFun` "tupCons" + , Tuple [e1,es] <- a + = e1 : splitTuple es (n-1) +splitTuple e n = [ pSel i n e | i <- [1..n] ] + +splitTupleK :: InScopeSet -> TExpr -> Int + -> (InScopeSet -> [TExpr] -> TExpr) + -> TExpr +-- Expects e to be a tuple-typed expression; +-- returns its n components. +-- May duplicate e +splitTupleK is _ 0 thing_inside = thing_inside is [] +splitTupleK is e 1 thing_inside = thing_inside is [e] +splitTupleK is (Tuple es) n thing_inside + | n == length es = thing_inside is es + | otherwise = pprPanic "splitTuple" (ppr n $$ ppr es) + +splitTupleK is (Call f a) n thing_inside + | f `isThePrimFun` "tupCons" + , Tuple [e1,es] <- a + = splitTupleK is es (n-1) $ \ is es -> + thing_inside is (e1 : es) + +splitTupleK is e n thing_inside + = mkTempLet is "tx" e $ \ is v -> + thing_inside is [ pSel i n v | i <- [1..n] ] + +pick :: Pretty a => [Int] -> [a] -> [a] +-- Pick the specifed items from the list +pick ts es = [ get t | t <- ts ] + where + get t | t > length es = pprTrace "pick" (ppr ts $$ ppr es) (head es) + | otherwise = es !! (t-1) + +findIndex :: Eq a => [(a,b)] -> a -> Maybe b +findIndex prs i + = go prs + where + go [] = Nothing + go ((a,b):prs) | a ==i = Just b + | otherwise = go prs + +tupLen :: TExpr -> Int +tupLen e = case typeof e of + TypeTuple tys -> length tys + x -> pprTrace "fsAdExpr:CLLet" (ppr x $$ ppr e) 0 + +makeArgDupable :: InScopeSet -> TExpr + -> (InScopeSet -> TExpr -> TExpr) + -> TExpr +makeArgDupable is e thing_inside + | isDupable e + = thing_inside is e + | Let v r b <- e + = Let v r $ + makeArgDupable (v `extendInScopeSet` is) b $ \ is b -> + thing_inside is b + | otherwise + = mkTempLet is "ax" e thing_inside + +isDupable (Tuple es) = all isTrivial es +isDupable (Call f a) = f `isThePrimFun` "tupCons" + && isDupable a +isDupable e = False + +mkTempLets :: InScopeSet -> String -> [TExpr] + -> ([TExpr] -> TExpr) -> TExpr +mkTempLets _ _ [] thing_inside + = thing_inside [] +mkTempLets is s (e:es) thing_inside + = mkTempLet is s e $ \ is' ve -> + mkTempLets is' s es $ \ ves -> + thing_inside (ve : ves) + +mkPairLet :: InScopeSet -> String -> String -> TExpr + -> (InScopeSet -> TExpr -> TExpr -> TExpr) + -> TExpr +mkPairLet is sf ss e thing_inside + = mkTempLet is "t" e $ \is ve -> + mkTempLet is sf (pFst ve) $ \is vf -> + mkTempLet is ss (pSnd ve) $ \is vs -> + thing_inside is vf vs + +mkTempLet :: InScopeSet -> String -> TExpr + -> (InScopeSet -> TExpr -> TExpr) -> TExpr +mkTempLet is s e thing_inside + | isTrivial e = thing_inside is e + | otherwise = Let tv' e $ + thing_inside is' (Var tv') + where + (is', tv') = notInScopeTV is tv + tv = TVar (typeof e) (Simple s) + + diff --git a/src/ksc/Ksc/Pipeline.hs b/src/ksc/Ksc/Pipeline.hs index 9d3474a8d..862f351ab 100644 --- a/src/ksc/Ksc/Pipeline.hs +++ b/src/ksc/Ksc/Pipeline.hs @@ -8,9 +8,10 @@ import ANF (anfDefs) import qualified Cgen import CSE (cseDefs) import KMonad (KM, KMT, runKM, banner, liftIO) +import Ksc.CatLang import Lang (ADDir(Rev, Fwd), ADPlan(BasicAD, TupleAD), Decl, DeclX(DefDecl), DefX(Def), Fun(Fun), - FunId(UserFun), TDef, + FunId(UserFun), TDef, Pretty, def_fun, displayN, partitionDecls, pps, ppr, renderSexp) import LangUtils (GblSymTab, emptyGblST, extendGblST) @@ -26,10 +27,63 @@ import GHC.Stack (HasCallStack) -- The demo driver ------------------------------------- +demoCL :: String -> IO () +demoCL file + = do { pr_decls <- parseF "src/runtime/prelude.ks" + ; my_decls <- parseF ("test/ksc/" ++ file ++ ".ks") + ; runKM $ + + do { banner "Original decls" + ; displayN my_decls + + ; (env1, pr_tc) <- annotDecls emptyGblST pr_decls + ; (env2, my_tc) <- annotDecls env1 my_decls + ; let (_pr_rules, _pr_defs) = partitionDecls pr_tc + (_my_rules, my_defs) = partitionDecls my_tc + rulebase = mkRuleBase (_pr_rules ++ _my_rules) + + ; displayPassM veryVerbose "Typechecked declarations" env2 my_defs + + ; banner "toCLDefs" + ; let cl_defs = toCLDefs my_defs + ; displayN cl_defs + + ; displayPassM veryVerbose "fromCLDefs" env2 + (fromCLDefs cl_defs) + +{- + ; let (env3, ad_defs) = fwdAdDefs env2 cl_defs + ; displayPassM veryVerbose "Forward tupled AD" env3 ad_defs + + ; (env4, opt_ad_defs) <- optDefs rulebase env3 ad_defs + ; displayPassM veryVerbose "Optimized forward tupled AD " env4 opt_ad_defs + + ; let (env5, rev_defs) = revAdDefs env4 cl_defs + ; displayPassM veryVerbose "Reverse AD" env5 rev_defs + + ; (env6, opt_rev_defs) <- optDefs rulebase env5 rev_defs + ; displayPassM veryVerbose "Optimized reverse AD " env5 opt_rev_defs + + ; (env7, cse_rev_defs) <- cseDefs rulebase env6 opt_rev_defs + ; displayPassM veryVerbose "Optimized (CSE'd) reverse AD " env7 cse_rev_defs + + ; let (env8, fs_defs) = fsAdDefs env7 cl_defs + ; displayPassM veryVerbose "Split AD" env8 fs_defs + + ; (env9, opt_fs_defs) <- optDefs rulebase env8 fs_defs + ; (env10, cse_fs_defs) <- cseDefs rulebase env9 opt_fs_defs + ; displayPassM veryVerbose "Optimized (CSE'd) split AD" env10 cse_fs_defs +-} + } } + + +veryVerbose :: Maybe Int +veryVerbose = Just 999 + demoF :: ADPlan -> [String] -> IO () -- Read source code from specified input files, optimise, -- differentiate, optimise, and display results of each step -demoF = demoFFilter (Just 999) id +demoF = demoFFilter veryVerbose id demoFFilter :: Maybe Int -> ([Decl] -> [Decl]) -> ADPlan -> [String] -> IO () demoFFilter verbosity theFilter adp files = do @@ -41,7 +95,10 @@ demo d = runKM (demoN (Just 999) BasicAD [d]) demoN :: Maybe Int -> ADPlan -> [Decl] -> KM () demoN verbosity adp decls - = let disp = displayPassM verbosity in + = let disp = displayPassM verbosity + dispNoLint :: Pretty def => String -> [def] -> KMT IO () + dispNoLint = displayPassMNoLint verbosity + in do { flip mapM_ verbosity $ \v -> do banner "Original declarations" displayN (take v decls) @@ -52,6 +109,11 @@ demoN verbosity adp decls ; disp "Typechecked declarations" env defs + ; let cl_defs = toCLDefs defs + ; dispNoLint "toCLDefs" cl_defs + + ; disp "fromCLDefs" env (fromCLDefs cl_defs) + ; (env1, opt_defs) <- optDefs rulebase env defs ; disp "Optimized original definition" env1 opt_defs @@ -87,11 +149,15 @@ demoN verbosity adp decls type DisplayLintT m a = String -> GblSymTab -> [TDef] -> KMT m a type DisplayLint a = DisplayLintT IO a +displayPassMNoLint :: Pretty def => Maybe Int -> String -> [def] -> KMT IO () +displayPassMNoLint mverbosity what decls + = do flip mapM_ mverbosity $ \verbosity -> do + banner what + displayN (take verbosity decls) + displayPassM :: Maybe Int -> DisplayLint () displayPassM mverbosity what env decls - = do { flip mapM_ mverbosity $ \verbosity -> do - banner what - displayN (take verbosity decls) + = do { displayPassMNoLint mverbosity what decls ; lintDefs what env decls } @@ -120,6 +186,17 @@ theDefs display decls = do { ; return (defs, env, rulebase) } +theDefsViaCatLang :: DisplayLint a + -> [Decl] -> KMT IO ([TDef], GblSymTab, RuleBase) +theDefsViaCatLang display decls = do { + (defs, env, rulebase) <- theDefs display decls + ; let defsViaCL = flip map defs $ \x -> case toCLDef_maybe x of + Nothing -> x + Just d -> fromCLDef d + ; display "Typechecked defs via CatLang" env defsViaCL + ; return (defsViaCL, env, rulebase) + } + theDiffs :: DisplayLint a -> [TDef] -> GblSymTab @@ -169,6 +246,7 @@ defsAndDiffs display decls = do { ; (defs, env, rulebase) <- theDefs display decls ; theDiffs display defs env rulebase } + anfOptAndCse :: DisplayLint a -> RuleBase -> GblSymTab -> [TDef] -> KM [TDef] anfOptAndCse display rulebase env4 alldefs = @@ -185,13 +263,17 @@ anfOptAndCse display rulebase env4 alldefs = ; return cse } -displayCppGenDiffs :: (DisplayLint () - -> [TDef] - -> GblSymTab - -> RuleBase - -> KMT IO (GblSymTab, [TDef], [TDef], RuleBase)) - -> Maybe Int -> [String] -> String -> String -> IO () -displayCppGenDiffs generateDiffs verbosity ksFiles ksofile cppfile = +displayCppGenDefsDiffs :: + (DisplayLint () + -> [Decl] + -> KMT IO ([TDef], GblSymTab, RuleBase)) + -> (DisplayLint () + -> [TDef] + -> GblSymTab + -> RuleBase + -> KMT IO (GblSymTab, [TDef], [TDef], RuleBase)) + -> Maybe Int -> [String] -> String -> String -> IO () +displayCppGenDefsDiffs generateDefs generateDiffs verbosity ksFiles ksofile cppfile = let dd defs = mapM_ (liftIO . putStrLn . ("...\n" ++) . pps . flip take defs) verbosity display = displayPassM verbosity in @@ -203,7 +285,7 @@ displayCppGenDiffs generateDiffs verbosity ksFiles ksofile cppfile = ; let (main, decls) = moveMain decls0 ; dd main - ; (defs, env, rulebase) <- theDefs display decls + ; (defs, env, rulebase) <- generateDefs display decls ; (env3, defs, optdiffs, rulebase) <- generateDiffs display defs env rulebase ; (env4, ann_main) <- annotDecls env3 main @@ -217,12 +299,34 @@ displayCppGenDiffs generateDiffs verbosity ksFiles ksofile cppfile = ; liftIO (Cgen.cppGenWithFiles ksofile cppfile cse) } +displayCppGenDiffs :: (DisplayLint () + -> [TDef] + -> GblSymTab + -> RuleBase + -> KMT IO (GblSymTab, [TDef], [TDef], RuleBase)) + -> Maybe Int -> [String] -> String -> String -> IO () +displayCppGenDiffs = displayCppGenDefsDiffs theDefs + displayCppGenNoDiffs :: Maybe Int -> [String] -> String -> String -> IO () displayCppGenNoDiffs = displayCppGenDiffs (\_ defs env rulebase -> pure (env, defs, [], rulebase)) -displayCppGenAndCompile :: HasCallStack => (String -> String -> IO String) -> String -> Maybe Int -> [String] -> String -> IO String -displayCppGenAndCompile compile ext verbosity files file = do { +displayCppGenAndCompileDefsDiffs + :: HasCallStack + => (DisplayLint () -> [Decl] -> KMT IO ([TDef], GblSymTab, RuleBase)) + -> (DisplayLint () + -> [TDef] + -> GblSymTab + -> RuleBase + -> KMT IO (GblSymTab, [TDef], [TDef], RuleBase)) + -> (String -> String -> IO String) + -> String + -> Maybe Int + -> [String] + -> String + -> IO String +displayCppGenAndCompileDefsDiffs + generateDefs generateDiffs compile ext verbosity files file = do { ; let ksFile = file ++ ".ks" ; let ksFiles = map (++ ".ks") files ; let compiler = compile @@ -230,16 +334,32 @@ displayCppGenAndCompile compile ext verbosity files file = do { ; let exefile = "obj/" ++ file ++ ext ; let ksofile = outfile ++ ".kso" ; let cppfile = outfile ++ ".cpp" - ; displayCppGenDiffs theDiffs verbosity (ksFiles ++ [ksFile]) ksofile cppfile + ; displayCppGenDefsDiffs generateDefs generateDiffs verbosity (ksFiles ++ [ksFile]) ksofile cppfile ; compiler cppfile exefile } +displayCppGenAndCompile :: HasCallStack => (String -> String -> IO String) -> String -> Maybe Int -> [String] -> String -> IO String +displayCppGenAndCompile = displayCppGenAndCompileDefsDiffs theDefs theDiffs + displayCppGenCompileAndRun :: HasCallStack => String -> Maybe Int -> [String] -> String -> IO String displayCppGenCompileAndRun compilername verbosity file files = do { exefile <- displayCppGenAndCompile (Cgen.compile compilername) ".exe" verbosity file files ; Cgen.runExe exefile } +displayCppGenCompileAndRunViaCatLang :: HasCallStack + => String + -> Maybe Int + -> [String] + -> String + -> IO String +displayCppGenCompileAndRunViaCatLang compilername verbosity file files = do + { exefile <- displayCppGenAndCompileDefsDiffs + theDefsViaCatLang theDiffs + (Cgen.compile compilername) ".exe" verbosity file files + ; Cgen.runExe exefile + } + displayCppGenCompileAndRunWithOutput :: HasCallStack => String -> Maybe Int -> [String] -> String -> IO () displayCppGenCompileAndRunWithOutput compilername verbosity files file = do { output <- displayCppGenCompileAndRun compilername verbosity files file diff --git a/src/ksc/Lang.hs b/src/ksc/Lang.hs index cf5da08f6..6591373b7 100644 --- a/src/ksc/Lang.hs +++ b/src/ksc/Lang.hs @@ -24,11 +24,13 @@ import Test.Hspec mkGradType :: ADPlan -> Type -> Type -> Type mkGradType BasicAD s ty = TypeLM s ty mkGradType TupleAD s ty = TypeTuple [ty, TypeLM s ty] +mkGradType SplitAD _ _ = error "mkGradType:SplitAD" -- Doesn't make sense -- For TupleAD, mkGradType s t = (t, s -o t) mkGradTuple :: ADPlan -> TExpr -> TExpr -> TExpr mkGradTuple BasicAD _ lm = lm mkGradTuple TupleAD p lm = Tuple [p, lm] +mkGradTuple SplitAD _ _ = error "mkGradTuple:SplitAD" -- Doesn't make sense data Phase = Parsed | Typed | OccAnald @@ -143,6 +145,9 @@ data ExprX p | Assert (ExprX p) (ExprX p) | Dummy (MTypeX p) +unitExpr :: ExprX p +unitExpr = Tuple [] + instance InPhase p => Show (ExprX p) where show e = pps e @@ -176,6 +181,9 @@ data TypeX | TypeUnknown +unitType :: TypeX +unitType = TypeTuple [] + deriving instance Eq Type deriving instance Ord Type @@ -234,9 +242,9 @@ tangentType :: HasCallStack => Type -> Type tangentType TypeFloat = TypeFloat tangentType (TypeVec t) = TypeVec (tangentType t) tangentType (TypeTuple ts) = TypeTuple (map tangentType ts) -tangentType TypeInteger = TypeTuple [] -tangentType TypeBool = TypeTuple [] -tangentType TypeString = TypeTuple [] +tangentType TypeInteger = unitType +tangentType TypeBool = unitType +tangentType TypeString = unitType tangentType TypeUnknown = TypeUnknown tangentType t = pprPanic "tangentType" (ppr t) -- TypeLM, TypeLam @@ -267,7 +275,7 @@ data FunId = UserFun String -- UserFuns have a Def Int -- Arity deriving( Eq, Ord, Show ) -data Fun = Fun FunId -- The function f(x) +data Fun = Fun FunId -- The function f(x) | GradFun FunId ADPlan -- Full Jacobian Df(x) | DrvFun FunId ADMode -- Derivative derivative f'(x,dx) -- Rev <=> reverse mode f`(x,dr) @@ -294,7 +302,7 @@ funIdOfFun = \case data ADMode = AD { adPlan :: ADPlan, adDir :: ADDir } deriving( Eq, Ord, Show ) -data ADPlan = BasicAD | TupleAD +data ADPlan = BasicAD | TupleAD | SplitAD deriving( Eq, Ord, Show ) data ADDir = Fwd | Rev @@ -311,12 +319,15 @@ data Var | Grad String ADPlan -- Derivative of x deriving( Eq, Ord, Show ) -nameOfVar :: Var -> String -nameOfVar = \case +varName :: Var -> String +varName = \case Simple s -> s Delta s -> s Grad s _ -> s +tVarName :: TVarX -> String +tVarName tv = varName (tVarVar tv) + data Konst = KInteger Integer -- :: TypeInteger | KFloat Double -- :: TypeFloat | KBool Bool -- :: TypeBool @@ -601,6 +612,9 @@ punctuate (SDoc p) ss = comma :: SDoc comma = text "," +dcolon :: SDoc +dcolon = text "::" + empty :: SDoc empty = SDoc (\_ -> PP.empty) @@ -706,6 +720,7 @@ instance Pretty ADDir where instance Pretty ADPlan where ppr BasicAD = empty ppr TupleAD = char 't' + ppr SplitAD = char 's' instance Pretty Var where ppr (Simple s) = text s @@ -822,9 +837,11 @@ pprCall prec f e = mode (case (e, isInfix @p f) of (Tuple [e1, e2], Just prec') -> parensIf prec prec' $ - sep [pprExpr prec' e1, pprFunOcc @p f <+> pprExpr prec' e2] + sep [pprExpr prec' e1 + , char '`'<> pprFunOcc @p f <> char '`' + <+> pprExpr prec' e2] _ -> parensIf prec precCall $ - cat [pprFunOcc @p f, nest 2 (parensSp pp_args)] + cat [pprFunOcc @p f, nest 2 (case e of { Tuple {} -> ppr e; _ -> parens (ppr e) })] ) where pp_args = ppr e @@ -858,8 +875,8 @@ isInfixFun (Fun (PrimFun s)) isInfixFun _ = Nothing parensIf :: Prec -> Prec -> SDoc -> SDoc -parensIf ctxt inner doc | ctxt >= inner = parens doc - | otherwise = doc +parensIf ctxt inner doc | ctxt >= inner = parens doc + | otherwise = doc instance InPhase p => Pretty (DeclX p) where ppr (DefDecl d) = ppr d @@ -882,7 +899,7 @@ pprDef (Def { def_fun = f, def_args = vs, def_res_ty = res_ty, def_rhs = rhs }) , ppr rhs]) (sep [ hang (text "def" <+> pprFun f <+> pprParendType res_ty) 2 (parens (pprTVar vs)) - , nest 2 (text "=" <+> ppr rhs) ]) + , nest 2 (text "=" <+> pprExpr precZero rhs) ]) StubRhs -> text "<>" @@ -935,8 +952,8 @@ hspec = do e2 = Call (Fun (UserFun "f")) (Tuple [e, var "_t1", kInt 5]) describe "Pretty" $ do - test e "g( i )" - test e2 "f( (g( i ), _t1, 5) )" + test e "g(i)" + test e2 "f(g(i), _t1, 5)" describe "eqType" $ it "doesn't truncate" (eqType (TypeTuple []) (TypeTuple [TypeFloat]) `shouldBe` False) diff --git a/src/ksc/LangUtils.hs b/src/ksc/LangUtils.hs index ec64ec7ed..78f8a8d21 100644 --- a/src/ksc/LangUtils.hs +++ b/src/ksc/LangUtils.hs @@ -26,7 +26,10 @@ module LangUtils ( GblSymTab, extendGblST, lookupGblST, emptyGblST, modifyGblST, stInsertFun, LclSymTab, extendLclST, - SymTab(..), newSymTab, emptySymTab + SymTab(..), newSymTab, emptySymTab, + + -- Other utility functions + isSingleton, count ) where @@ -35,6 +38,18 @@ import qualified Data.Map as M import qualified Data.Set as S import Test.Hspec + +----------------------------------------------- +-- Utility functions +----------------------------------------------- + +isSingleton :: [a] -> Bool +isSingleton [x] = True +isSingleton _ = False + +count :: (a->Bool) -> [a] -> Int +count f = length . filter f + ----------------------------------------------- -- Functions over expressions ----------------------------------------------- diff --git a/src/ksc/Main.hs b/src/ksc/Main.hs index e90181b7e..3596c1417 100644 --- a/src/ksc/Main.hs +++ b/src/ksc/Main.hs @@ -10,6 +10,7 @@ import Parse (parseE) import Opt import Ksc.Pipeline (demoFFilter, displayCppGenCompileAndRun, + displayCppGenCompileAndRunViaCatLang, displayCppGenAndCompile, displayCppGenNoDiffs, genFuthark, ignoreMain) @@ -242,10 +243,11 @@ demoFOnTestPrograms ksTests = do dropWhile1 :: (a -> Bool) -> [a] -> [a] dropWhile1 f = tail . dropWhile f -testRunKS :: String -> String -> IO () -testRunKS compiler ksFile = do +testRunKSVia :: (String -> Maybe int -> [String] -> FilePath -> IO String) + -> String -> [Char] -> IO () +testRunKSVia via_ compiler ksFile = do let ksTest = System.FilePath.dropExtension ksFile - output <- displayCppGenCompileAndRun compiler Nothing ["src/runtime/prelude"] ksTest + output <- via_ compiler Nothing ["src/runtime/prelude"] ksTest let testResults = dropWhile1 (/= "TESTS FOLLOW") (lines output) @@ -273,6 +275,12 @@ testRunKS compiler ksFile = do putStrLn (unlines (reverse (take 30 (reverse (lines output))))) error ("These tests failed:\n" ++ unlines failures) +testRunKS :: String -> String -> IO () +testRunKS = testRunKSVia displayCppGenCompileAndRun + +testRunKSViaCatLang :: String -> String -> IO () +testRunKSViaCatLang = testRunKSVia displayCppGenCompileAndRunViaCatLang + testHspec :: IO () testHspec = do summary <- runSpec Main.hspec defaultConfig @@ -287,6 +295,8 @@ testC compiler fsTestKs = do compileKscPrograms compiler ksTestFiles_ testRunKS compiler "test/ksc/gmm.ks" testRunKS compiler "test/ksc/fold.ks" + testRunKSViaCatLang compiler "test/ksc/gmm.ks" + testRunKSViaCatLang compiler "test/ksc/fold.ks" compileKscPrograms compiler fsTestKs profileArgs :: String -> FilePath -> FilePath -> FilePath -> IO () diff --git a/src/ksc/Opt.hs b/src/ksc/Opt.hs index 9562a2f98..0e7962b71 100644 --- a/src/ksc/Opt.hs +++ b/src/ksc/Opt.hs @@ -183,24 +183,30 @@ rewriteCall _ fun (If e1 e2 e3) rewriteCall env (TFun _ (Fun fun)) arg = optFun env fun arg -rewriteCall env (TFun ty (GradFun f adm)) arg - = optGradFun (optEnvInScope env) adm ty f arg +rewriteCall env (TFun ty (GradFun f adp)) arg + = optGradFun (optEnvInScope env) adp ty f arg -rewriteCall _ f@(TFun (TypeLM _ _) _) _ - = trace ("NOTE: Unmatched LM call {" ++ pps f ++ "}") Nothing - -rewriteCall _ _ _ - = Nothing +rewriteCall env (TFun ty (DrvFun f adm)) arg + = optDrvFun (optEnvInScope env) adm ty f arg ----------------------- optFun :: OptEnv -> FunId -> TExpr -> Maybe TExpr -- RULE: sel_i_n (..., ei, ...) ==> ei -optFun _ (SelFun i _) arg +optFun _ (SelFun i n) arg + | n == 1, i == 1 + = Just arg + | Tuple es <- arg , i <= length es = Just (es !! (i-1)) + | Call tupCons (Tuple [e,es]) <- arg + , tupCons `isThePrimFun` "tupCons" + = case i of + 1 -> Just e + _ -> Just (pSel (i-1) (n-1) es) + | otherwise = Nothing @@ -238,6 +244,10 @@ optPrimFun _ op (Tuple [Konst (KFloat k1), Konst (KFloat k2)]) , " https://github.com/microsoft/knossos-ksc/pull/61/commits/29c2ab04568e17b953d3fe942aba5881ab15e1f8#r309892713" ] +optPrimFun _ "tupTail" (Tuple es) + | [_,e2] <- es = Just e2 + | _:es2 <- es = Just (Tuple es2) + -- RULE: (e1 : ()) + (e2 : ()) = () -- The type () contains only one value (), which is a zero of the type -- We use () as the tangent type for non-differentiatable types @@ -333,6 +343,20 @@ optPrimFun _ "lmAdd" (Tuple [ Call hcat1 (Tuple ps) optPrimFun _ _ _ = Nothing +---------------------- +optDrvFun :: HasCallStack => InScopeSet -> ADMode + -> Type -> FunId -> TExpr -> Maybe TExpr +optDrvFun _ (AD BasicAD Fwd) _ (SelFun i n) (Tuple [_,ds]) + = Just (pSel i n ds) +optDrvFun _ (AD TupleAD Fwd) _ (SelFun i n) (Tuple [s,ds]) + = Just (Tuple [pSel i n s, pSel i n ds]) +optDrvFun _ (AD BasicAD Rev) _ (SelFun i n) (Tuple [s,dt]) + = Just (mkTuple (map mk_dr [1..n])) -- Returns (0,0,dt,0) + where + mk_dr j | i==j = dt + | otherwise = mkTangentZero (pSel j n s) +optDrvFun _ _ _ _ _ = Nothing + ---------------------- optLMCompose :: TExpr -> TExpr -> Maybe TExpr optLMCompose f g @@ -553,6 +577,7 @@ optGradFun env TupleAD ty f args optGradFun _ BasicAD _ (SelFun i n) args = optGradSel i n args optGradFun _ BasicAD ty (PrimFun f) args = optGradPrim ty f args +optGradFun _ SplitAD _ _ _ = error "optGradFun:SplitAD" -- Doesn't make sense type TBinds = [(TVar, TExpr)] diff --git a/src/ksc/OptLet.hs b/src/ksc/OptLet.hs index db49d1117..7853858f5 100644 --- a/src/ksc/OptLet.hs +++ b/src/ksc/OptLet.hs @@ -5,11 +5,11 @@ ScopedTypeVariables #-} module OptLet( optLets - , Subst, InScopeSet, emptyInScopeSet + , Subst, InScopeSet, emptyInScopeSet, mkInScopeSet , mkEmptySubst, lookupSubst , substInScope, extendInScopeSet , substBndr, extendSubstMap, zapSubst - , substExpr, substVar, notInScope ) + , substExpr, substVar, notInScope, notInScopeTV ) where import Lang @@ -223,12 +223,12 @@ zapSubst (S { s_in_scope = in_scope }) -- * It clones the binder if it is already in scope -- * Extends the substitution and the in-scope set as appropriate substBndr :: TVar -> Subst -> (TVar, Subst) -substBndr (TVar ty v) (S { s_in_scope = in_scope, s_env = env }) +substBndr tv (S { s_in_scope = in_scope, s_env = env }) = (tv', S { s_env = env' , s_in_scope = is' }) where - (is', tv') = notInScopeTV in_scope v ty - env' = M.insert v (Var tv') env + (is', tv') = notInScopeTV in_scope tv + env' = M.insert (tVarVar tv) (Var tv') env substVar :: Subst -> TVar -> TExpr substVar subst tv = case lookupSubst (tVarVar tv) subst of @@ -254,8 +254,8 @@ substExpr subst e where (v', subst') = substBndr v subst -notInScopeTV :: InScopeSet -> Var -> Type -> (InScopeSet, TVar) -notInScopeTV is v ty +notInScopeTV :: InScopeSet -> TVar -> (InScopeSet, TVar) +notInScopeTV is (TVar ty v) = (v' `S.insert` is, TVar ty v') where v' = notInScope v is @@ -263,7 +263,7 @@ notInScopeTV is v ty notInScope :: Var -> InScopeSet -> Var -- Find a variant of the input Var that is not in the in-scope set -- --- Do this by adding "_1", "_2" etc +-- Do this by adding "1", "2" etc notInScope v in_scope | not (v `S.member` in_scope) = v @@ -280,23 +280,19 @@ notInScope v in_scope | otherwise = var' where var' = rebuild str' - str' = prefix ++ '_' : show n + str' = prefix ++ show n (prefix, _n) = parse_suffix [] (reverse str) parse_suffix :: String -- Digits parsed from RH end (in order) -> String -- String being parsed (reversed) - -> (String, Int) -- String before "_", plus number found after - -- E.g. parse_suffix "foo_23" = ("foo", 23) + -> (String, Int) -- String before number, plus number found after + -- E.g. parse_suffix "foo23" = ("foo", 23) -- parse_suffix "wombat" = ("wombat", 0) parse_suffix ds (c:cs) - | c == '_' - , not (null ds) - = (reverse cs, read ds) - | isDigit c - = parse_suffix (c:ds) cs - parse_suffix ds cs - = (reverse cs ++ ds, 0) + | isDigit c = parse_suffix (c:ds) cs + | not (null ds) = (reverse (c:cs), read ds) + parse_suffix ds cs = (reverse cs ++ ds, 0) optLetsE :: Subst -> ExprX OccAnald -> TExpr -- This function inline let-bindings that are only used once @@ -405,6 +401,6 @@ inline_me_help rhs | isTrivial rhs = True -- RHS is trivial, see isTrivial for what that means | isKZero rhs = True -- Inline zeros, as they will very likely disappear | TypeLM {} <- typeof rhs = True -- Always inline linear maprs (might not do this in future) - | Tuple ts <- rhs -- Always inline tuples whose fields are all trivial - , all inline_me_help ts = True -- See Note [Inline tuples] + | Tuple ts <- rhs -- Always inline tuples whose fields are all trivial + , all inline_me_help ts = True -- See Note [Inline tuples] | otherwise = False diff --git a/src/ksc/Prim.hs b/src/ksc/Prim.hs index 983b0fe12..d4ca9d7d7 100644 --- a/src/ksc/Prim.hs +++ b/src/ksc/Prim.hs @@ -5,7 +5,7 @@ module Prim where import Lang -import LangUtils (isTrivial) +import LangUtils import GHC.Stack import Data.Maybe @@ -67,6 +67,8 @@ mk_fun f = case find_dollar f of Just ("fwdt", s) -> DrvFun (mk_fun_id s) (AD TupleAD Fwd) Just ("rev", s) -> DrvFun (mk_fun_id s) (AD BasicAD Rev) Just ("revt", s) -> DrvFun (mk_fun_id s) (AD TupleAD Rev) + Just ("fwds", s) -> DrvFun (mk_fun_id s) (AD SplitAD Fwd) + Just ("revs", s) -> DrvFun (mk_fun_id s) (AD SplitAD Rev) Just ("get", s) -> Fun (mk_sel_fun s) _ -> Fun (mk_fun_id f) where @@ -88,11 +90,32 @@ mkZero, mkTangentZero :: HasCallStack => TExpr -> TExpr mkZero = getZero id mkTangentZero = getZero tangentType +mkTangentZeroFromType :: Type -> TExpr +-- Only works for types with no vectors, +-- for which needValueForTangentZero says True +mkTangentZeroFromType TypeInteger = unitExpr +mkTangentZeroFromType TypeFloat = Konst (KFloat 0.0) +mkTangentZeroFromType TypeString = unitExpr +mkTangentZeroFromType TypeBool = unitExpr +mkTangentZeroFromType (TypeTuple ts) = Tuple (map mkTangentZeroFromType ts) +mkTangentZeroFromType ty = pprTrace "mkTangentZeroFromType" (ppr ty) + Konst (KInteger 99) -- A distinctive value + +needValueForTangentZero :: Type -> Bool +needValueForTangentZero TypeInteger = False +needValueForTangentZero TypeFloat = False +needValueForTangentZero TypeString = False +needValueForTangentZero TypeBool = False +needValueForTangentZero (TypeTuple ts) = any needValueForTangentZero ts +needValueForTangentZero TypeUnknown = False +needValueForTangentZero (TypeVec {}) = True +needValueForTangentZero ty = pprPanic "needValueForTangentZero" (ppr ty) + getZero :: HasCallStack => (Type -> Type) -> TExpr -> TExpr getZero tangent_type e = go e where - go e = case tangent_type e_ty of + go e = case tangent_type (typeof e) of TypeInteger -> Konst (KInteger 0) TypeFloat -> Konst (KFloat 0.0) TypeString -> Konst (KString "") @@ -107,9 +130,7 @@ getZero tangent_type e -> mkAtomicNoFVs e $ \e -> Tuple $ map go $ [ pSel i n e | i <- [1..n] ] - _ -> pprPanic "mkZero" (ppr e_ty $$ ppr e) - where - e_ty = typeof e + _ -> pprPanic "mkZero" (ppr (typeof e) $$ ppr e) -- (mkAtomicNoFVs e body) returns the expression (let a = e in body a) -- where body :: TExpr -> TExpr is a function expecting an expression @@ -214,6 +235,7 @@ lmApplyR = mkPrimCall2 "lmApplyR" lmApply_AD :: HasCallStack => ADMode -> TExpr -> TExpr -> TExpr lmApply_AD (AD BasicAD dir) = lmApply_Dir dir lmApply_AD (AD TupleAD dir) = lmApplyT_Dir dir +lmApply_AD (AD SplitAD _) = error "lmApply_AD:SplitAD" -- Doesn't make sense lmApply_Dir :: HasCallStack => ADDir -> TExpr -> TExpr -> TExpr lmApply_Dir Fwd e ds = lmApply e ds @@ -344,6 +366,28 @@ pIndex = mkPrimCall2 "index" pSum :: TExpr -> TExpr pSum = mkPrimCall1 "sum" +pTupCons :: TExpr -> TExpr -> TExpr +-- Put arg1 on the front of the tuple arg2 +pTupCons arg1 arg2 + | Tuple args <- arg2 = Tuple (arg1:args) + | TypeTuple {} <- typeof arg2 = mkPrimCall2 "tupCons" arg1 arg2 + | otherwise = Tuple [arg1,arg2] + +pTupHead :: TExpr -> TExpr +pTupHead e | Tuple (e1:_) <- e = e1 + | TypeTuple es <- typeof e = pSel 1 (length es) e + | otherwise = pprTrace "pTupHead" (ppr e <+> dcolon <+> ppr (typeof e)) $ + Konst (KInteger 99) + +pTupTail :: TExpr -> TExpr +pTupTail e | Tuple (_:es) <- e = mkTuple es + | TypeTuple [_,_] <- e_ty = pSel 2 2 e + | TypeTuple {} <- e_ty = mkPrimCall1 "tupTail" e + | otherwise = pprTrace "pTupTail" (ppr e <+> dcolon <+> ppr e_ty) $ + Konst (KInteger 88) + where + e_ty = typeof e + pSumBuild :: TExpr -> TExpr -> TExpr pSumBuild = mkPrimCall2 "sumbuild" @@ -354,12 +398,22 @@ pSize :: TExpr -> TExpr pSize e = mkPrimCall1 "size" e pSel :: Int -> Int -> TExpr -> TExpr -pSel i n e = Call (TFun el_ty - (Fun (SelFun i n))) e - where - el_ty = case typeof e of - TypeTuple ts -> ts !! (i-1) - _ -> TypeUnknown -- Better error from Lint +pSel i n e + | n == 1 -- A 1-tuple is the same as the thing itself + -- So sel_1_1 e = e + = assert (text "pSel1") (i==1) e + + | Tuple es <- e -- Reduce clutter by optimising right here + -- sel_1_2 (e1,e2) --> e1 + = assert (text "pSel" <+> int i <+> int n <+> int (length es) $$ ppr es) (n == length es) + es !! (i-1) + + | otherwise + = Call (TFun el_ty (Fun (SelFun i n))) e + where + el_ty = case typeof e of + TypeTuple ts -> ts !! (i-1) + _ -> TypeUnknown -- Better error from Lint pFst,pSnd :: TExpr -> TExpr pFst = pSel 1 2 @@ -378,7 +432,7 @@ pMulff x1 x2 = userCall "mul" TypeFloat (Tuple [x1, x2]) --------------------------------------------- -- Types of primitive functions -- --- For each primitve, we give its type +-- For each primitive, we give its type -- And this is the /only/ place we do this --------------------------------------------- @@ -413,14 +467,26 @@ primCallResultTy_maybe fun arg_ty | AD BasicAD Rev <- adm -- f :: S1 -> T, then rev$f :: (S1, T_t) -> S1_t , TypeTuple [s, _dt] <- arg_ty -> Right (tangentType s) + + | AD SplitAD Fwd <- adm + , Right t_ty <- primCallResultTy_maybe (Fun f) arg_ty + -> Right (TypeTuple [t_ty, unitType]) -- Not clear this is always right + -- What is the X type for a primitive? + + | AD SplitAD Rev <- adm + -> Right arg_ty + | otherwise -> Left (text "Ill-typed call to:" <+> ppr fun) Fun (UserFun _) -> Left (text "Not in scope: user fun:" <+> ppr fun) selCallResultTy_maybe :: Int -> Int -> Type -> Either SDoc Type -selCallResultTy_maybe i n (TypeTuple arg_tys) - | i <= length arg_tys +selCallResultTy_maybe i n arg_ty + | i == 1, n == 1 + = Right arg_ty + | TypeTuple arg_tys <- arg_ty + , i <= length arg_tys , n == length arg_tys = Right (arg_tys !! (i - 1)) selCallResultTy_maybe _ _ _ = Left (text "Bad argument to selector") @@ -570,6 +636,11 @@ primFunCallResultTy_maybe fun args ("delta" , TypeTuple [TypeInteger, TypeInteger, t]) -> Just t + ("tupCons" , TypeTuple [t, TypeTuple ts] ) -> Just (TypeTuple (t:ts)) + ("tupHead" , TypeTuple (t:_) ) -> Just t + ("tupTail" , TypeTuple (_:ts) ) | length ts >= 2 + -> Just (TypeTuple ts) + _ -> Nothing isPrimFun :: String -> Bool @@ -582,6 +653,7 @@ isPrimFun f = f `elem` [ "$inline" -- ($inline f args...) Force inline f , "sumbuild" -- (sumbuild N f) (sum (build N f)) , "fold" -- (fold f z v) (Left) fold over v , "index" + , "tupCons", "tupTail" , "size" , "sum" , "unzip" -- Takes a vector of tuples to a tuple of vectors diff --git a/src/runtime/prelude.ks b/src/runtime/prelude.ks index 4d24ccc15..c6272cf22 100644 --- a/src/runtime/prelude.ks +++ b/src/runtime/prelude.ks @@ -1,5 +1,8 @@ +;; ---------------- Addition -------------- ;; add :: Number x Number -> Number ;; add (x, y) = x + y + +;; ------- Addition: Float ------ (edef add Float (Float Float)) (edef D$add (LM (Tuple Float Float) Float) (Float Float)) (edef Dt$add (Tuple Float (LM (Tuple Float Float) Float)) (Float Float)) @@ -10,6 +13,18 @@ ((dx1 (get$1$2 dxt)) (dx2 (get$2$2 dxt))) (add dx1 dx2))) + +; Tupled forward add +(def + fwdt$add (Tuple Float Float) + ((xt : (Tuple Float Float)) (dxt : (Tuple Float Float))) + (let + ((dx1 (get$1$2 dxt)) + (dx2 (get$2$2 dxt))) + (tuple (add (get$1$2 xt) (get$2$2 xt)) + (add dx1 dx2)))) + +; Reverse add (def rev$add (Tuple Float Float) ((xt : (Tuple Float Float)) (drt : Float)) @@ -17,6 +32,24 @@ ((d_dadd drt)) (tuple d_dadd d_dadd))) +; Split add +; fwds$add :: (F,F) -> (F, ()) +; revs$add :: (dF, ()) -> (dF,dF) + +(def + fwds$add (Tuple Float (Tuple)) + (xt : (Tuple Float Float)) + (let + ((x1 (get$1$2 xt)) + (x2 (get$2$2 xt))) + (tuple (add x1 x2) (tuple)))) + +(def + revs$add (Tuple Float Float) + (dr : Float) + (tuple dr dr)) + +;; ------- Addition: Integer ------ (edef add Integer (Integer Integer)) (edef D$add (LM (Tuple Integer Integer) Integer) (Integer Integer)) (edef Dt$add (Tuple Integer (LM (Tuple Integer Integer) Integer)) (Integer Integer)) @@ -29,6 +62,7 @@ ((xt : (Tuple Integer Integer)) (drt : (Tuple))) (tuple (tuple) (tuple))) +;; ---------------- Subtraction -------------- ;; sub :: Number x Number -> Number ;; sub (x, y) = x - y (edef sub Float (Float Float)) @@ -60,8 +94,11 @@ ((xt : (Tuple Integer Integer)) (drt : (Tuple))) (tuple (tuple) (tuple))) +;; ---------------- Division -------------- ;; div :: Number x Number -> Number ;; div (x, y) = x / y + +;; ------- Division: Float ------ (edef div Float (Float Float)) (edef D$div (LM (Tuple Float Float) Float) (Float Float)) (edef Dt$div (Tuple Float (LM (Tuple Float Float) Float)) (Float Float)) @@ -76,6 +113,21 @@ (div (sub (mul x2 dx1) (mul x1 dx2)) (mul x2 x2)))) + +; Tupled forward divide +(def + fwdt$div (Tuple Float Float) + ((xt : (Tuple Float Float)) (dxt : (Tuple Float Float))) + (let + ((x1 (get$1$2 xt)) + (x2 (get$2$2 xt)) + (dx1 (get$1$2 dxt)) + (dx2 (get$2$2 dxt))) + (tuple (div x1 x2) + (div (sub (mul x2 dx1) + (mul x1 dx2)) + (mul x2 x2))))) + (def rev$div (Tuple Float Float) ((xt : (Tuple Float Float)) (drt : Float)) @@ -86,6 +138,8 @@ (tuple (div d_ddiv x2) (neg (div (mul x1 d_ddiv) (mul x2 x2)))))) + +;; ------- Division: Integer ------ (edef div Integer (Integer Integer)) (edef D$div (LM (Tuple Integer Integer) Integer) (Integer Integer)) (edef Dt$div (Tuple Integer (LM (Tuple Integer Integer) Integer)) (Integer Integer)) @@ -93,13 +147,20 @@ fwd$div (Tuple) ((xt : (Tuple Integer Integer)) (dxt : (Tuple (Tuple) (Tuple)))) (tuple)) +(def + fwdt$div (Tuple Integer (Tuple)) + ((xt : (Tuple Integer Integer)) (dxt : (Tuple (Tuple) (Tuple)))) + (tuple (div xt) (tuple))) (def rev$div (Tuple (Tuple) (Tuple)) ((xt : (Tuple Integer Integer)) (drt : (Tuple))) (tuple (tuple) (tuple))) +;; ---------------- Multiplication -------------- ;; mul :: Number x Number -> Number ;; mul (x, y) = x * y + +;; ------- Multiplication: Float ------ (edef mul Float (Float Float)) (edef D$mul (LM (Tuple Float Float) Float) (Float Float)) (edef Dt$mul (Tuple Float (LM (Tuple Float Float) Float)) (Float Float)) @@ -112,6 +173,19 @@ (dx1 (get$1$2 dxt)) (dx2 (get$2$2 dxt))) (add (mul x2 dx1) (mul x1 dx2)))) + +; Tupled forward multiply +(def + fwdt$mul (Tuple Float Float) + ((xt : (Tuple Float Float)) (dxt : (Tuple Float Float))) + (let + ((x1 (get$1$2 xt)) + (x2 (get$2$2 xt)) + (dx1 (get$1$2 dxt)) + (dx2 (get$2$2 dxt))) + (tuple (mul x1 x2) (add (mul x2 dx1) (mul x1 dx2))))) + +; Reverse multiply (def rev$mul (Tuple Float Float) ((xt : (Tuple Float Float)) (drt : Float)) @@ -121,6 +195,28 @@ (d_dmul drt)) (tuple (mul d_dmul x2) (mul d_dmul x1)))) +; Split forward multiply +; fwds$mul :: (F,F) -> (F, (F,F)) +; revs$mul :: (dF, (F,F)) -> (dF,dF) +(def + fwds$mul (Tuple Float (Tuple Float Float)) + (xt : (Tuple Float Float)) + (let + ((x1 (get$1$2 xt)) + (x2 (get$2$2 xt))) + (tuple (mul x1 x2) (tuple x1 x2)))) + +(def + revs$mul (Tuple Float Float) + (xt : (Tuple Float (Tuple Float Float))) + (let + ((dr (get$1$2 xt)) + (x (get$2$2 xt)) + (x1 (get$1$2 x)) + (x2 (get$2$2 x))) + (tuple (mul dr x2) (mul dr x1)))) + +;; ------- Multiplication: Integer ------ (edef mul Integer (Integer Integer)) (edef D$mul (LM (Tuple Integer Integer) Integer) (Integer Integer)) (edef Dt$mul (Tuple Integer (LM (Tuple Integer Integer) Integer)) (Integer Integer)) @@ -128,11 +224,16 @@ fwd$mul (Tuple) ((xt : (Tuple Integer Integer)) (dxt : (Tuple (Tuple) (Tuple)))) (tuple)) +(def + fwdt$mul (Tuple Integer (Tuple)) + ((xt : (Tuple Integer Integer)) (dxt : (Tuple (Tuple) (Tuple)))) + (tuple (mul xt) (tuple))) (def rev$mul (Tuple (Tuple) (Tuple)) ((xt : (Tuple Integer Integer)) (drt : (Tuple))) (tuple (tuple) (tuple))) +;; ----------------- Negation -------------------------- ;; neg :: Number -> Number ;; neg x = -x (edef neg Float (Float)) @@ -166,6 +267,7 @@ ((xt : (Tuple Float Float)) (drt : (Tuple))) (tuple 0.0 0.0)) +;; ----------------- Relational operators -------------------------- (edef gt Bool (Integer Integer)) (edef D$gt (LM (Tuple Integer Integer) Bool) (Integer Integer)) (edef Dt$gt (Tuple Bool (LM (Tuple Integer Integer) Bool)) (Integer Integer)) @@ -314,6 +416,8 @@ (edef to_float Float (Integer)) (edef D$to_float (LM Integer Float) (Integer)) (def fwd$to_float Float ((x : Integer) (dx : (Tuple))) 0.0) +(def fwdt$to_float (Tuple Float Float) ((x : Integer) (dx : (Tuple))) + (tuple (to_float x) 0.0)) (def rev$to_float (Tuple) ((x : Integer) (d_dto_float : Float)) (tuple)) (edef Dt$to_float (Tuple Float (LM Integer Float)) (Integer)) diff --git a/test/ksc/cat1.ks b/test/ksc/cat1.ks new file mode 100644 index 000000000..efe278efc --- /dev/null +++ b/test/ksc/cat1.ks @@ -0,0 +1,34 @@ +; Copyright (c) Microsoft Corporation. +; Licensed under the MIT license. +;; (def g Float ( (x : Float) (y : Float) ) +;; x) ;; (mul@ff y x)) + + +(def g1 Float (x : Float) + (let (y (add x x)) + (add y y))) + +(def g2 Float (x : Float) + (let (y (add x x)) + (add y x))) + +(def g3 Float (x : Float) + (let (y (add x x)) + (mul y x))) + +(def g4 Float (x : Float) + (let (y (mul x x)) + (mul y x))) + +(def g5 Float (x : Float) + (let (y (add x x)) + (let (z (mul y x)) + z))) + +(def g6 Float( (x : Float) (y : Float) ) + x) +(def g7 Float( (x : Float) (y : Float) ) + y) + +(def g8 Float( (y : Vec Float) ) + 3.0) diff --git a/test/ksc/ex10.ks b/test/ksc/ex10.ks new file mode 100644 index 000000000..d8b408ca1 --- /dev/null +++ b/test/ksc/ex10.ks @@ -0,0 +1,15 @@ +; Copyright (c) Microsoft Corporation. +; Licensed under the MIT license. + +; This program demonstrates how Conalisation helps +; In reverse mode, in rev$h, we get one, rather than two, calls to rev$f + +(def f Float ( x : Float ) + (mul x x)) + +(def g Float ( x : Float ) + (add x x)) + +(def h Float ( x : Float ) + (let (z (f x)) + (add z (g z))))