gysit
commented
3 years ago in principle there are two approaches to implement subdomains. Either using a stencil.combine to combine multiple applies operation or by defining subdomains inside the stencil apply operator.
1) stencil.combine
%version1 = stencil.apply
%version2 = stencil.apply
%0 = stencil.combine 2 15 %version1 %version2 on ([0,0,0]:[64,64,60]) : (stencil.temp stencil.temp) -> stencil.temp
this example executes the stencil 1 if dimension k is < 15 and 2 otherwise. The operation is a normal shaped op and we need to extend the access analysis to infer the shape for the two inputs version1 and version2. Additionally we need a verifier that ensures the all both version1 and version2 have only one use. That way the stencil combines are tightly coupled to the stencil apply.
the advantage of the stencil.combine is that we can bubble the combines down in the chain of stencil apply operators. Similarly to hoisting the if / else in an if / else based approach and we can decide if we want to generate if/else or different kernels that run in streams. Additionally due to the bubbling we can use the existing stencil inlining pass at least in theory. We may run into trouble since a single evaluation of the output stencil may have to run multiple versions of the inlined stencil if the inlined stencil is a combination of two stencils.
2) introduce splits inside the stencil operator.
%0 = stencil.apply {
%res = stencil.split 0 in LB to 5 {
// version1
stencil.yield %version1
} and 5 to UB {
// version2
stencil.yield %version2
}
stencil.return %res
}
this example implements a similar split as above and the splits could be nested. This syntax does not change the stencil.apply on the course grained level and is similar to what we may use for the tridiagonal solves. However inlining and possibly onrolling again becomes complicating and potentially more complicating than for the stencil.combine. We may allow the stencil combine to have more or less than two bodies to get a cleaner result and less nesting. However, one split shall split only one dimension.
The placeholders LB and UB specify the lower and upper bounds of the parent stencil.apply
Introduce an explicit subdomain operation to split the domain. Besides implementing different stencils on subdomains the operation should also enable a number of additional fixes: