Overhaul of Index label usage - Githubissues

pnnl / COMET

Other

37 stars 7 forks source link

Overhaul of Index label usage #55

Closed pthomadakis closed 8 months ago

pthomadakis commented 8 months ago

This PR introduces the following significant changes to usage of index labels on different sections of the codebase.

Tensors are no longer "bound" to the index labels they were declared with in the DSL. This fixes the issue described in #47.
Index labels are no longer used to denote tensor sizes, except for convenience in the DSL. Thus, index labels are now treated as just symbols for the einsum operations. arith.index is used for tensor sizes instead.
Dialects below the TA now use the affine maps coming with each operation to infer index spaces.

Other changes include:

Completely removed ta.labeled_tensor
Introduction of the ta.dim operation that returns the size of a ta.dense/sparse tensor's dimension, similar to the standar tensor.dim
At the DSL level, one can also declare the size of a dense tensor directly, without using an index label e.g. Tensor<double> A([5][4], {Dense});

Bug fixes:

Fixed issue that prohibited spgemm in the style of C=A*A to work
Fixed issue in inferring output format of compound CSR*Dense operations in MLIRGen
Fixed bug in transpose operation for the DSL. Operation was dependent on labels alphabetical order.