32 Commits

Author SHA1 Message Date
River Riddle
3655069234 [mlir] Move the Builtin FuncOp to the Func dialect
This commit moves FuncOp out of the builtin dialect, and into the Func
dialect. This move has been planned in some capacity from the moment
we made FuncOp an operation (years ago). This commit handles the
functional aspects of the move, but various aspects are left untouched
to ease migration: func::FuncOp is re-exported into mlir to reduce
the actual API churn, the assembly format still accepts the unqualified
`func`. These temporary measures will remain for a little while to
simplify migration before being removed.

Differential Revision: https://reviews.llvm.org/D121266
2022-03-16 17:07:03 -07:00
Bixia Zheng
3580721a59 [mlir][sparse][taco] Support the use of index values in tensor expressions.
PyTACO DSL doesn't support the use of index values as in A[i] = B[i]+ i.
We extend the DSL to support such a use in MLIR-PyTACO.

Remove an obsolete unit test. Add unit tests and PyTACO tests.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D121716
2022-03-15 15:30:55 -07:00
Bixia Zheng
3a4229696d [mlir][sparse][taco] Reorder a class.
Define IndexExpr before IndexVar. This is to prepare for the next change
to support the use of index values in tensor expressions.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D121649
2022-03-15 08:51:22 -07:00
gysit
7294be2b8e [mlir][linalg] Replace linalg.fill by OpDSL variant.
The revision removes the linalg.fill operation and renames the OpDSL generated linalg.fill_tensor operation to replace it. After the change, all named structured operations are defined via OpDSL and there are no handwritten operations left.

A side-effect of the change is that the pretty printed form changes from:
```
%1 = linalg.fill(%cst, %0) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
```
changes to
```
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
```
Additionally, the builder signature now takes input and output value ranges as it is the case for all other OpDSL operations:
```
rewriter.create<linalg::FillOp>(loc, val, output)
```
changes to
```
rewriter.create<linalg::FillOp>(loc, ValueRange{val}, ValueRange{output})
```
All other changes remain minimal. In particular, the canonicalization patterns are the same and the `value()`, `output()`, and `result()` methods are now implemented by the FillOpInterface.

Depends On D120726

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D120728
2022-03-14 10:51:08 +00:00
Bixia Zheng
30c5269d93 [mlir][sparse][taco] Add a few unary operations.
Add operations -, abs, ceil and floor to the index notation.

Add test cases.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D121388
2022-03-11 08:08:55 -08:00
Bixia Zheng
5b87e0521d [mlir][sparse][taco] Split the evaluate method into compile and compute.
This is to align with the PyTACO API better.

Modify an existing unit test to test the new routines.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D121083
2022-03-07 16:58:41 -08:00
Bixia Zheng
c25f3dfff3 [mlir][sparse][taco] Support tensor dimension storage ordering and more general
sparsity values.

Previously, we can't properly handle input tensors with a dimension
ordering that is different from the natural ordering or with a mixed of
compressed and dense dimensions. This change fixes the problems by
passing the dimension ordering and sparsity values to the runtime
routine.

Modify an existing test to test the situation.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D120777
2022-03-01 15:36:38 -08:00
River Riddle
23aa5a7446 [mlir] Rename the Standard dialect to the Func dialect
The last remaining operations in the standard dialect all revolve around
FuncOp/function related constructs. This patch simply handles the initial
renaming (which by itself is already huge), but there are a large number
of cleanups unlocked/necessary afterwards:

* Removing a bunch of unnecessary dependencies on Func
* Cleaning up the From/ToStandard conversion passes
* Preparing for the move of FuncOp to the Func dialect

See the discussion at https://discourse.llvm.org/t/standard-dialect-the-final-chapter/6061

Differential Revision: https://reviews.llvm.org/D120624
2022-03-01 12:10:04 -08:00
Bixia Zheng
20eaa88fff [mlir][sparse] Extend convertToMLIRSparseTensor to support permutation and more general sparsity values.
Previously, convertToMLIRSparseTensor assumes identity storage ordering and all
compressed dimensions. This change extends the function with two parameters for
users to specify the storage ordering and the sparsity of each dimension.

Modify PyTACO to reflect this change.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D120643
2022-03-01 10:51:39 -08:00
Bixia Zheng
6f07191101 [mlir][sparse][taco] Support reduction to scalar tensors.
The PyTACO DSL doesn't support reduction to scalars. This change
enhances the MLIR-PyTACO implementation to support reduction to scalars.

Extend an existing test to show the syntax of reduction to scalars and
two methods to retrieve the scalar values.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D120572
2022-02-25 14:17:45 -08:00
Bixia Zheng
c601dfbcc2 [mlir][sparse][taco] Use np.array_equal to compare integer values.
Fix MLIR-PyTACO and some tests to use np.array_equal to compare integer
values.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D120526
2022-02-25 07:38:15 -08:00
gysit
cd2776b0d5 [mlir][OpDSL] Split arithmetic functions.
Split arithmetic function into unary and binary functions. The revision prepares the introduction of unary and binary function attributes that work similar to type function attributes.

Depends On D120108

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D120109
2022-02-25 15:27:42 +00:00
Bixia Zheng
90f22ab3ad [mlir][sparse][taco] Add support for scalar tensors.
This change allows the use of scalar tensors with index 0 in tensor index
expressions. In this case, the scalar value is broadcast to match the
dimensions of other tensors in the same expression.

Using scalar tensors as a destination in tensor index expressions is not
supported in the PyTACO DSL.

Add a PyTACO test to show the use of scalar tensors.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D120524
2022-02-25 07:20:15 -08:00
gysit
51fdd802c7 [mlir][OpDSL] Add type function attributes.
Previously, OpDSL operation used hardcoded type conversion operations (cast or cast_unsigned). Supporting signed and unsigned casts thus meant implementing two different operations. Type function attributes allow us to define a single operation that has a cast type function attribute which at operation instantiation time may be set to cast or cast_unsigned. We may for example, defina a matmul operation with a cast argument:

```
@linalg_structured_op
def matmul(A=TensorDef(T1, S.M, S.K), B=TensorDef(T2, S.K, S.N), C=TensorDef(U, S.M, S.N, output=True),
    cast=TypeFnAttrDef(default=TypeFn.cast)):
  C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
```

When instantiating the operation the attribute may be set to the desired cast function:

```
linalg.matmul(lhs, rhs, outs=[out], cast=TypeFn.cast_unsigned)
```

The revsion introduces a enum in the Linalg dialect that maps one-by-one to the type functions defined by OpDSL.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D119718
2022-02-25 08:25:23 +00:00
Bixia Zheng
c8ae8cfb5d [mlir][sparse][taco] Add support for float32.
Previously, we only support float64. We now support float32 and float64. When
constructing a tensor without providing a data type, the default is float32.

Fix the tests to data type consistency. All PyTACO application tests now use
float32 to match the default data type of TACO. Other tests may use float32 or
float64.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D120356
2022-02-23 18:24:22 -08:00
Aart Bik
6438783fda [mlir][sparse] provide more types for external to/from MLIR routines
These routines will need to be specialized a lot more based on value types,
index types, pointer types, and permutation/dimension ordering. This is a
careful first step, providing some functionality needed in PyTACO bridge.

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D120154
2022-02-18 13:36:52 -08:00
Shao-Ce SUN
21ac474392 [NFC] Correct typo interger to integer 2022-02-17 21:17:47 +08:00
Bixia Zheng
746c68eafd [mlir][sparse][taco] Handle tensor copy and trivial reduction expression.
Handle tensor copy, such as A[i, j] = B[i, j]. Also, handle trivial
reduction expression, such as A[i] = B[i, j].

Add unit tests.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D119867
2022-02-15 15:57:18 -08:00
Bixia Zheng
ad932a75f9 [mlir][sparse][taco] Support true dense tensors and all dense sparse tensors.
The only method to create a true dense tensor (i.e un-annotated) in MLIR-PyTACO
is through the from_array method. However, the annotated all dense tensors are
also implemented as true dense tensor currently. The PR fixes the
implementation to support annotated all dense sparse tensors.

Extend the tensor init method to support the construction of a tensor without
any sparsity annotation.

Change the tensor to_file method to only support writing unpacked sparse
tensors to file through the MLIR sparse tensor dialect.

Add unit tests for true dense tensors and all dense sparse tensors.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D119500
2022-02-14 15:35:01 -08:00
Aart Bik
719b865be2 [mlir][sparse][pytaco] add SDDMM test with two different ways of defining kernel
Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D119465
2022-02-10 13:33:06 -08:00
Aart Bik
8189a2b8bd [mlir][sparse][pytaco] migrate to sparse compiler pipeline
Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D119395
2022-02-10 07:42:54 -08:00
Aart Bik
6195a25487 [mlir][sparse][pytaco] test cleanup
removed obsoleted TODO
removed strange Fp precision for coordinates
lined up meta data testing code for readability

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D119377
2022-02-09 16:58:25 -08:00
Bixia Zheng
61a3dd70ff [mlir][taco] Use sparse_tensor.out to write sparse tensors to files.
Add a Python method, output_sparse_tensor, to use sparse_tensor.out to write
a sparse tensor value to a file.

Modify the method that evaluates a tensor expression to return a pointer of the
MLIR sparse tensor for the result to delay the extraction of the coordinates and
non-zero values.

Implement the Tensor to_file method to evaluate the tensor assignment and write
the result to a file.

Add unit tests. Modify test golden files to reflect the change that TNS outputs
now have a comment line and two meta data lines.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D118956
2022-02-08 08:47:05 -08:00
River Riddle
ace01605e0 [mlir] Split out a new ControlFlow dialect from Standard
This dialect is intended to model lower level/branch based control-flow constructs. The initial set
of operations are: AssertOp, BranchOp, CondBranchOp, SwitchOp; all split out from the current
standard dialect.

See https://discourse.llvm.org/t/standard-dialect-the-final-chapter/6061

Differential Revision: https://reviews.llvm.org/D118966
2022-02-06 14:51:16 -08:00
Bixia Zheng
93c81f44cc [mlir][taco] Uses sparse_tensor.new to read tensor input data from files.
Replace the Python implementation for reading tensor input data from files with
create_sparse_tensor that uses sparse_tensor.new.

The MLIR TNS format has two extra meta data lines. Add the extra meta data to a
test data file.

Implement TACO tensor methods evaluate and unpack.

Add unit tests.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D118803
2022-02-03 08:26:33 -08:00
River Riddle
cf70f7ce8b [mlir] Remove dangling reference to std-bufferize which got removed 2022-02-02 15:04:40 -08:00
Bixia Zheng
ae7ee655a9 [mlir][taco] Add a utility to create an MLIR sparse tensor from a file.
Move the functions that retrieve the supporting C library, compile an MLIR
module and build a JIT execution engine to mlir_pytaco_utils.

Add a function to create an MLIR sparse tensor from a file and return a pointer
to the MLIR sparse tensor as well as the shape of the sparse tensor.

Add unit tests.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D118496
2022-02-01 15:43:53 -08:00
Matthias Springer
ab47418df6 [mlir][bufferize] Merge tensor-constant-bufferize into arith-bufferize
The bufferization of arith.constant ops is also switched over to BufferizableOpInterface-based bufferization. The old implementation is deleted. Both implementations utilize GlobalCreator, now renamed to just `getGlobalFor`.

GlobalCreator no longer maintains a set of all created allocations to avoid duplicate allocations of the same constant. Instead, `getGlobalFor` scans the module to see if there is already a global allocation with the same constant value.

For compatibility reasons, it is still possible to create a pass that bufferizes only `arith.constant`. This pass (createConstantBufferizePass) could be deleted once all users were switched over to One-Shot bufferization.

Differential Revision: https://reviews.llvm.org/D118483
2022-01-30 21:37:48 +09:00
Bixia Zheng
91865cc027 [mlir][taco] Accept an integer list for the ordering when defining a tensor format.
The unit tests for PyTACO hasn't been upstreamed yet. A unit test for this
change will be added when we upstream all the unit tests for PyTACO.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D118417
2022-01-28 10:33:25 -08:00
Bixia Zheng
b7fd91c84b Upstream MLIR PyTACO implementation.
Add TACO tests to test/Integration/Dialect/SparseTensor/taco. Add the MLIR
PyTACO implementation as tools under the directory.

Reviewed By: aartbik, mehdi_amini

Differential Revision: https://reviews.llvm.org/D117260
2022-01-21 08:38:36 -08:00
Mehdi Amini
4e08ce7adb Revert "Upstream MLIR PyTACO implementation."
This reverts commit 778a264da9eba0c8523cdc10a10822fd3e458dd3.

This broke the bot: tests are failing at the moment.
2022-01-13 23:14:13 +00:00
Bixia Zheng
778a264da9 Upstream MLIR PyTACO implementation.
Add TACO tests to test/Integration/Dialect/SparseTensor/taco. Add the MLIR
PyTACO implementation as tools under the directory.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D117126
2022-01-13 14:50:28 -08:00