Skip to content

Commit a73761c

Browse files
author
Jessica Shi
committed
add first draft of scheduling docs
1 parent 9f25b63 commit a73761c

File tree

2 files changed

+91
-0
lines changed

2 files changed

+91
-0
lines changed

documentation/docs/scheduling.md

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
The scheduling language enables users to specify and compose transformations to further optimize the code generated by taco.
2+
3+
Consider the following SpMV computation and associated code, which we will transform below:
4+
```c++
5+
Format csr({Dense,Sparse});
6+
Tensor<double> A("A", {512, 64}, csr);
7+
Tensor<double> x("x", {64}, {Dense});
8+
Tensor<double> y("y", {512}, {Dense});
9+
10+
IndexVar i, j;
11+
y(i) = A(i, j) * x(j);
12+
IndexStmt stmt = y.getAssignment().concretize();
13+
```
14+
```c
15+
for (int32_t i = 0; i < A1_dimension; i++) {
16+
double y_val = 0.0;
17+
for (int32_t jA = A2_pos[i]; jA < A2_pos[(i + 1)]; jA++) {
18+
int32_t j = A2_crd[jA];
19+
y_val += A_vals[jA] * x_vals[j];
20+
}
21+
y_vals[i] = y_val;
22+
}
23+
```
24+
# Pos
25+
26+
The `pos(i, ipos, access)` transformation takes in an index variable `i` that operates over the coordinate space of `access` and replaces it with a derived index variable `ipos` that operates over the same iteration range, but with respect to the the position space.
27+
28+
Since the `pos` transformation is not valid for dense level formats, for the SpMV example, the following would result in an error:
29+
```c++
30+
stmt = stmt.pos(i, IndexVar("ipos"), A);
31+
```
32+
33+
We could instead have:
34+
```c++
35+
stmt = stmt.pos(j, IndexVar("jpos"), A);
36+
```
37+
```c
38+
for (int32_t i = 0; i < A1_dimension; i++) {
39+
for (int32_t jposA = A2_pos[i]; jposA < A2_pos[(i + 1)]; jposA++) {
40+
if (jposA < A2_pos[i] || jposA >= A2_pos[(i + 1)])
41+
continue;
42+
43+
int32_t j = A2_crd[jposA];
44+
y_vals[i] = y_vals[i] + A_vals[jposA] * x_vals[j];
45+
}
46+
}
47+
```
48+
49+
# Split
50+
51+
The `split(i, i0, i1, splitFactor)` transformation splits (strip-mines) an index variable `i` into two nested index variables `i0` and `i1`. The size of the inner index variable `i1` is then held constant at `splitFactor`, which must be a positive integer.
52+
53+
For the SpMV example, we could have:
54+
```c++
55+
stmt = stmt.split(j, IndexVar("i0"), IndexVar("i1"), 16);
56+
```
57+
```c
58+
for (int32_t i0 = 0; i0 < ((A1_dimension + 15) / 16); i0++) {
59+
for (int32_t i1 = 0; i1 < 16; i1++) {
60+
int32_t i = i0 * 16 + i1;
61+
if (i >= A1_dimension)
62+
continue;
63+
64+
for (int32_t jA = A2_pos[i]; jA < A2_pos[(i + 1)]; jA++) {
65+
int32_t j = A2_crd[jA];
66+
y_vals[i] = y_vals[i] + A_vals[jA] * x_vals[j];
67+
}
68+
}
69+
}
70+
```
71+
72+
# Reorder
73+
74+
The `reorder(vars)` transformation takes in a new ordering for a set of index variables in the expression that are directly nested in the iteration order.
75+
76+
For the SpMV example, we could have:
77+
```c++
78+
stmt = stmt.reorder({j, i});
79+
```
80+
```c
81+
for (int32_t jA = A2_pos[iA]; jA < A2_pos[(iA + 1)]; jA++) {
82+
int32_t j = A2_crd[jA];
83+
for (int32_t i = 0; i < A1_dimension; i++) {
84+
y_vals[i] = y_vals[i] + A_vals[jA] * x_vals[j];
85+
}
86+
}
87+
```
88+
89+
90+

documentation/mkdocs.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ pages:
1515
- C++ Library:
1616
- 'Defining Tensors': 'tensors.md'
1717
- 'Computing on Tensors': 'computations.md'
18+
- 'Providing a Schedule': 'scheduling.md'
1819
- Python Library:
1920
- 'Tutorial': 'tutorial.md'
2021
- 'Defining Tensors': 'pytensors.md'

0 commit comments

Comments
 (0)