-
Notifications
You must be signed in to change notification settings - Fork 74k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[XLA:GPU] Construct constraints for destructured summations in `Symbo…
…licTile` derivation. Take the summation expression `expr = sum(map(lambda [size, stride]: stride * size, sizes_and_strides))`. In order to assign a single stride for the summation expression, we need to ensure that the parameters (sizes) involved in the expression are such that the gap between them is always the same. Concretely, given a list of sizes `[s0, s1, ..., s{n}]` ordered in descending order of associated strides, we expect that each size `s{k}` is either: 1. `1` (and the corresponding stride is irrelevant); 2. fully captured---i.e. `s{k} = upper_bound(s{k})`. Assume `s{k}` is the leftmost fully captured dimension. In that case, `for i in {0, ..., n-k-1}`, `s{k+i+1}` is allowed to be fully captured if `s{k+i}` is also fully captured. Otherwise, `s{k+i+1} = 1`. The resulting stride is the smallest stride associated with a fully captured dimension; 3. partially captured---i.e. `1 < s{k} < upper_bound(s{k})`. In that case, `for i in {0, ..., k-1}`, `s{i} = 1`. `s{k+1}` is allowed to be fully captured (and thus the leftmost fully captured dimension), in which case we do as in `2.`. If `s{k+1}` is not fully captured, then `for i in {k+1, ..., n}`, `s{i} = 1`, and the stride of the expression is the stride associated with `s{k}`. As a regex-like summary, we expect the sizes to be as follows in row-major order (i.e. strictly decreasing order of strides): `(1*, partial_dim?, full_dims*, 1*)` and construct constraints as such. PiperOrigin-RevId: 646742678
- Loading branch information
1 parent
0a4261b
commit 0c1a3e5
Showing
2 changed files
with
199 additions
and
35 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters