Replicate simple proofs on tensors

[bollu]

• 1D/2D/nD: (map f . map g = map (f . g)
• 2D: transpose . transpose = id
• 1D/2D: extract subview (fill const t) = fill const (extract subview t).

Doing the last one (extract and fill commutation) for ND holors would be a fun challenge.

[bollu]

Interesting thing to note: The equation map f . map g = map (f . g) cannot be encoded naively within MLIR, because within MLIR, the functions f, g correspond to regions. There is no way to describe f . g, since there is no nice way to write "perform the sequential composition of regions and return a new region". So this is a theorem that we can prove with no associated peephole rewrite. @tobiasgrosser @goens @ChrisHughes24 , Note that we we will not have an associated rewrite to the map (f . g) = map f . map g theorem.

[goens]

@bollu I guess the question is how do you describe a generic f in this? i.e. can you write either the LHS/RHS with a parametrized f and g as here? And we encode things with func then it should work, right?

Let me ask this differently, can you encode this in lambda calculus, without named functions? I think it's useful to understand this issue precisely

[bollu]

@goens In the tensor dialect, a generic f is represented as a region. Since regions are second-class citizens of the MLIR syntax, we cannot "talk about" the region itself. That is, we cannot store a region in an SSA variable. The closest encoding I can think of uses the rgn dialect to do something like:

%f = rgn.var { ... }
%g = rgn.var { ... }
%t1 = tensor.zero ...
%t2 = tensor.map %t1 { ^entry(%x : i32):
   %h = rgn.sequence %f, %g
   rgn.run %h %x
}

This is very awkward. It does provide "an encoding", just not what I would consider an idiomatic one.

[goens]

@bollu yeah that's the point with what I meant above. My understading is the following:

1. The func dialect should be there to talk about generic functions, right? So in a sense what you're suggesting there with rgn should also be expressible with func.
2. The problem is not that it's not expressible "as a peephole rewrite", the problem is that this statement is not expressible parametrically/polymorphically/universally quantified over f and g, because of what you call regions being second-class citizens. For any given (concrete) f and g, this would be readily expressible as a "peephole rewrite" (or similar), wouldn't it?
3. I'm not sure where second-class citizens issue comes from. In a sense, regions are first-class citizens in that they can also take regions as arguments, just like higher-order functions in lambda calculus. Thus, this theorem should be expressible in (untyped) lambda calculus, where f and g are just free variables (or bound all the way outside, whatever you prefer), right? Couldn't we do a similar trick to express this with regions?

Do you think this is a fair assesment and how does the situation look in untyped lambda calculus? If it's different, why would it?

[bollu]

For any given (concrete) f and g, this would be readily expressible as a "peephole rewrite" (or similar), wouldn't it?

Yes

---

The func dialect should be there to talk about generic functions, right? So in a sense what you're suggesting there with rgn should also be expressible with func.

Sort of. It will not express the same thing, but something similar. See that the difference between writing:

tensor.map %t1 {  (%x : i32) : <code> }

versus

%h = rgn.const { ... }
tensor.map %t1 { (%x : i32) : rgn.run %h x }

is the difference between tensor.map t (fun x => <code>) and writing let h x := fun x => <code> in tensor.map t (fun x => h x)

both of these programs are semantically equivalent, but syntactically they are not.

---

In a sense, regions are first-class citizens in that they can also take regions as arguments, just like higher-order functions in lambda calculus. Thus, this theorem should be expressible in (untyped) lambda calculus, where f and g are just free variables (or bound all the way outside, whatever you prefer), right? Couldn't we do a similar trick to express this with regions?

@goens Let me try to ask for help where I do not know a solution: Could you find an encoding of the proven theorem Tensor2d.map_functorial in our framework, such that upon calling simp, the proof state reduces to forall g f, t.map (g ∘ f) = (t.map f).map g?

[bollu]

@goens gentle nudge on this :)

[goens]

@bollu I haven't had the bandwidth to look into this, I'm afraid I won't have it while I'm at PLDI either, sorry!