Skip to content

Reshaping to support BSR and GCXS #63

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 68 additions & 0 deletions spec/latest/index.bs
Original file line number Diff line number Diff line change
Expand Up @@ -384,6 +384,74 @@ Special note: If the sparse level is the root level, the `pointers` array should
be ommitted, as its first value will be `0` and its last value will be the
length of any of the `indices` arrays in this level.

### Slicing and Dicing ### {#slice_and_dice}

Several sparse matrix formats, such as [BSR](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.bsr_array.html#scipy.sparse.bsr_array) or [GCXS](https://sparse.pydata.org/en/0.15.1/generated/sparse.GCXS.html), require multiple dimensions of the underlying storage to be split, transposed, and/or combined into other dimensions. For example, the BSR format stores a sparse matrix using dense, same-size tiles. If the original matrix `A` is `m` by `n`, the blocked matrix `B` is a sparse matrix of dense blocks, or a 4-tensor of size `m/b` by `n/b` by `b` by `b`. The relationship between the two could be described as `A[i, j] = B[floordiv(i, b), floordiv(j, b), mod(i, b), mod(j, b)]`.

As another example, the GCXS format stores N-dimensional tensors using 2-dimensional matrices, by combining dimensions. For example, if the original tensor `A` is `m` by `n` by `p`, the underlying matrix `B` might be `m` by `n*p`. The relationship between the two could be described as `A[i, j, k] = B[i, j * p + k]`.

In this section, we introduce an optional specification to split and combine dimensions.

Note that dimensions may not be able to be split or combined evenly. For example, if our original matrix is of size `5` by `7`, there is no way to use `2` by `2` blocks to tile the matrix evenly. In this case, we can pad our original matrix, decompose it into a tensor, and declare that the final matrix is a window into the full `6` by `8` matrix we would represent. For this reason, we introduce slicing operations into the spec.

The spec adds the following keys representing operations to be applied:

The `split_dims` key, when present, is a list of tuples of integer dimensions resulting from splitting the dimensions of the tensor. The dimensions in the `i`th tuple must multiply to the size of the `i`th dimension in the original tensor. The dimensions of the output tensor is defined to be the concatenation of the dimension tuples. The flattened output tensor should be equal to the flattened input tensor.

The `combine_dims` key, when present, is a list of tuples of integers describing the dimensions to combine, and in which order. The `i`th dimension of the output is the product of the sizes of the dimensions listed in the `i`th tuple. The flattened output should be equal to the flattened input tensor after transposing it to the order specified by concatenating the tuples.

The `slice` key, when present, is a list of tuples of integers describing the starting and ending index of each dimension. If the `i`th tuple is `(a, b)`, then the `i`th dimension of the output should contain indices starting at `a` and ending just before `b`.

The operations when present are to be applied in the order `split_dims`, `combine_dims`, `slice`, followed by the `transpose` key if present.
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Leave a note that points out that combine dims can transpose too actually.


Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If combine_dims, slice, split_dims, or transpose are present, a storage_shape` must be specified, to describe the shape before transformation

As an example, an `11` by `37` BCSR can be represented as:

```json
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

need outer brace here, also need commas, also there's no tuples in json

"shape": [3, 10, 4, 4]
"custom": {
"level": {
"level_desc": "dense",
"rank": 1,
"level": {
"level_desc": "sparse",
"rank": 1,
"level": {
"level_desc": "dense",
"rank": 1,
"level": {
"level_desc": "dense",
"rank": 1,
"level": {
"level_desc": "element",
}
}
}
}
}
}
"combine_dims"=[(0, 2), (1, 3)],
"slice"=[(0, 11), (0, 37)]
```

As another example, a `10` by `20` by `30` GCXS tensor can be represented as:

```json
"shape": [10, 600]
"custom": {
"level": {
"level_desc": "dense",
"rank": 1,
"level": {
"level_desc": "sparse",
"rank": 1,
"level": {
"level_desc": "element",
}
}
}
}
"split_dims"=[(10,), (20, 30)],
```

### Equivalent Formats ### {#equivalent_formats}

Expand Down