[all-commits] [llvm/llvm-project] c66303: [mlir][sparse] Switch to One-Shot Bufferize

[Matthias Springer via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: c66303c2870c9a77a0f2a8aa16fd0ea87b0358e6
      https://github.com/llvm/llvm-project/commit/c66303c2870c9a77a0f2a8aa16fd0ea87b0358e6
  Author: Matthias Springer <springerm at google.com>
  Date:   2022-07-14 (Thu, 14 Jul 2022)

  Changed paths:
    M mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
    M mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
    M mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
    M mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
    M mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp
    M mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt
    A mlir/lib/Dialect/SparseTensor/Transforms/DenseBufferizationPass.cpp
    M mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
    M mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
    M mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
    M mlir/test/Dialect/SparseTensor/dense.mlir
    M mlir/test/Dialect/SparseTensor/one_shot_bufferize_tensor_copy_insertion.mlir
    M mlir/test/Dialect/SparseTensor/sparse_1d.mlir
    M mlir/test/Dialect/SparseTensor/sparse_2d.mlir
    M mlir/test/Dialect/SparseTensor/sparse_3d.mlir
    M mlir/test/Dialect/SparseTensor/sparse_affine.mlir
    M mlir/test/Dialect/SparseTensor/sparse_fp_ops.mlir
    M mlir/test/Dialect/SparseTensor/sparse_int_ops.mlir
    M mlir/test/Dialect/SparseTensor/sparse_kernels.mlir
    M mlir/test/Dialect/SparseTensor/sparse_lower.mlir
    M mlir/test/Dialect/SparseTensor/sparse_lower_col.mlir
    M mlir/test/Dialect/SparseTensor/sparse_lower_inplace.mlir
    M mlir/test/Dialect/SparseTensor/sparse_nd.mlir
    M mlir/test/Dialect/SparseTensor/sparse_out.mlir
    M mlir/test/Dialect/SparseTensor/sparse_outbuf.mlir
    M mlir/test/Dialect/SparseTensor/sparse_perm.mlir
    M mlir/test/Dialect/SparseTensor/sparse_perm_lower.mlir
    M mlir/test/Dialect/SparseTensor/sparse_scalars.mlir
    M mlir/test/Dialect/SparseTensor/sparse_vector.mlir
    M mlir/test/Dialect/SparseTensor/sparse_vector_chain.mlir
    M mlir/test/Dialect/SparseTensor/sparse_vector_index.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_binary.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_sparse2dense.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion_sparse2sparse.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_dot.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reshape.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scale.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_c32.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_f16.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tanh.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_transpose.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_unary.mlir
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_vector_ops.mlir

  Log Message:
  -----------
  [mlir][sparse] Switch to One-Shot Bufferize

This change removes the partial bufferization passes from the sparse compilation pipeline and replaces them with One-Shot Bufferize. One-Shot Analysis (and TensorCopyInsertion) is used to resolve all out-of-place bufferizations, dense and sparse. Dense ops are then bufferized with BufferizableOpInterface. Sparse ops are still bufferized in the Sparsification pass.

Details:
* Dense allocations are automatically deallocated, unless they are yielded from a block. (In that case the alloc would leak.) All test cases are modified accordingly. E.g., some funcs now have an "out" tensor argument that is returned from the function. (That way, the allocation happens at the call site.)
* Sparse allocations are *not* automatically deallocated. They must be "released" manually. (No change, this will be addressed in a future change.)
* Sparse tensor copies are not supported yet. (Future change)
* Sparsification no longer has to consider inplacability. If necessary, allocations and/or copies are inserted during TensorCopyInsertion. All tensors are inplaceable by the time Sparsification is running. Instead of marking a tensor as "not inplaceable", it can be marked as "not writable", which will trigger an allocation and/or copy during TensorCopyInsertion.

Differential Revision: https://reviews.llvm.org/D129356