Feature request: Rank polymorphism thru size and type lists

[leonardschneider]

I propose to implement rank polymorphism by introducing size and types lists type arguments.

Inspired by the blog in https://futhark-lang.org/blog/2022-10-03-futhark-on-arraycast.html

A size or type list can be expanded thanks a .. suffix in a function declaration. This suffix be followed by indices, in case there is an ambiguity on the expansion to be performed, like so a[t]..t

For example

-- array expansion
[dim..] becomes [dim[0]][dim[1]]...
-- tuple expansion
(type..) becomes (type[0], type[1], ...)

The .. suffix can also be used in function types. For example,

-- function argument type expansion
type.. -> i64 becomes type[0] -> type[1] -> ... -> i64

Single size can also be expanded like so

-- function argument repetition expansion
i64..n -> i64 becomes i64 -> i64 -> ... <n times> -> i64

Functional function can then be generalized as follows. Note that variadic functions are now possible.

-- generic curry
val curry '^[t]as : (f: as.. -> c) -> (as: as).. -> c

-- generic uncurrry
val uncurry '^[t]as: (f: (as..)) -> as.. -> c

-- generic const
val const '^a '^[t]bs: (x: a) -> bs.. -> a

SOAC functions can then be generalized as follows

-- generic heterogeneous map
-- def map2 '[2]as = map, etc
val map '[t]as [n] 'x : (f: as.. -> x) -> (as: [n](as[t]))..t -> *[n]x

-- homogeneous map
-- translates to: map (map (map ... <n times> f)) as
val hmap [n]dim 't 'x: (f: t -> x) -> (as: [dim..]t) -> x

-- generic reduce
-- def reduce_1d [1]dim 'a = reduce 0
val reduce [n]dim 'a: (axis: i64)(op: a -> a -> a) -> (ne: a) -> (as: [dim..]a) -> [dim[0:axis]++dim[axis+1:n]]a

-- generic reduce_by_index
-- def reduce_index_2d 'a [2]dim = reduce_index, etc
val reduce_by_index 'a [k]dim [n]: (dest: *[dim..]a) -> (f: a -> a -> a) -> (ne: a) -> (is: [n](i64..k)) -> (as: [n]a) -> *[dim..]a

-- generic scatter
-- def scatter_2d 't [2]dim = scatter
val scatter 't [k]dim [n]: (dest: *[dim..]t) -> (is: [n](i64..k)) -> (vs: [n]t) -> *[dim..]t

Array functions can then be generalized as follows

-- generic concat
val ++ [n[dim]] 't: (xs: [dim[n]]t)..n -> *[sum dim]t
-- where: sum dim == reduce (+) 0 dim

-- generic iota
val iota: (dim: [n]i64).. -> *[prod dim](i64..n)

-- generic flatten
-- def flatten_2d 't [2]dim = flatten
val flatten 't [n]dim: (xs: [dim..]t) -> [prod dim]t
-- where prod dim == reduce (*) 1 dim

-- generic unflatten
-- def unflatten_2d 't [2]dim = unflatten, etc
val unflatten 't n[dim]: (xs: [prod dim]t) -> [dim..]t

-- generic transpose
-- can take any permutation
-- def transpose_2d [m][o] = transpose [1, 0] 
val transpose 't [n]dim: (perm: [n]i64) -> (a: [dim..]t) -> [(select perm dim)..]t
-- where: select perm dim == map (\i -> dim[perm[i]]) (iota (length perm))

-- generic tabulate
-- def tabulate_2d 'a: (dim: [2]i64).. -> (f: [2]i64.. -> a) -> *[[2]dim]a  = tabulate., etc
val tabulate 'a: (dim: i64)..n -> (f: i64..n -> a) -> *[dim..]a

Zip functions can then be generalized as follows

-- generic zip
val zip [n] '[t]a: (as: [n](a[t]))..t -> *[n](a..)

-- generic unzip
val unzip [n] '[t]a: (xs: [n](a..)) -> (([n](a[t]))..t)

TBH, I have no idea about the feasibility or complexity on implementing this. I've encountered the need for more rank polymorphism when exploring how to implement a futhark onnx runtime. The only way to go around with current the current futhark implementation would be to template everything, which could work, but be quite ugly. Ideally, I would much prefer have some independent generic futhark code base to translate an onnx graph in a pretty direct manner.

[zfnmxt]

What exactly does a[t]..t mean? Why does t appear twice? E.g., in

-- generic heterogeneous map
-- def map2 '[2]as = map, etc
val map '[t]as [n] 'x : (f: as.. -> x) -> (as: [n](as[t]))..t -> *[n]x

Does (as: [n](as[t]))..t mean you accept t arrays of length [n], where the mth array has type as[m]? What's the precise meaning of as[t] in that type?

Why not a definition of

val map '[t]as [n]dim 'x : (f: as.. -> x) -> (as: [dim...](as[t]))..t -> *[n]x

so that the map is both variadic and rank polymorphic?

I think @athas's vision with this stuff is to support some APL-like syntax; i.e., rank polymorphism via function application as mentioned in the post. So, rather than having (with the above variadic/rank polymorphic map) having to write

map (+) xss yss

(and at this point maybe it'd be better to call map lift or something...)

for some arrays xss and yss, say xss : [l][m][n]i32 and yss : [l][m][n]i32, you'd simply write the more natural

xss + yss

which the compiler then translates to

map2 (map2 (+)) xss yss

We actually looked into rank polymorphic function application and I hacked up a couple of prototypes. As mentioned in the post, the type checking is pretty non-trivial (whole PhD dissertations have been written on this issue and haven't solved it). I'd still like
to return to it, but I'm busy with other things right now. The crux of the issue lies in figuring out two things:

1. Whether it's even possible (both in terms of decidability and performance) to do type checking in general with the current design (roughly, arbitrary rank polymorphism for function application, including higher-order functions).
2. If 1. is "no", whether we can come up with a restricted design that's still useful, can be simply explained in terms of a few simple rules, and isn't confusing to users. For these sorts of features, I think it's important that there's a very simple/understandable "calculus", otherwise it's too confusing.

As for variadic functions, maybe @athas can comment on that. I haven't thought much about it but obviously a generic zip is nicer that having zipn for a number of n.

[athas]

I don't think variadic functions are possible to combine with type inference. It's not even clear how you as a programmer would write them (your examples are all just types, not definitions).

Rank polymorphism is certainly something we (sort of) actively investigate, and @zfnmxt did a good job layout out the current vision. In particular, it is rooted in not complicating the function type system itself. There's discussion and work here: #1747