//===-- Passes.td - Sparse tensor pass definition file -----*- tablegen -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES
#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES

include "mlir/Pass/PassBase.td"

def Sparsification : Pass<"sparsification", "ModuleOp"> {
  let summary = "Automatically generate sparse tensor code from sparse tensor types";
  let description = [{
    A pass that implements the core functionality of a **sparse compiler**.
    Each Linalg operation (MLIR's tensor index notation) that operates on
    sparse tensor types is converted into code in which the sparsity is
    explicit both in terms of co-iterating looping logic as well as
    selected sparse storage schemes.

    See the `SparseTensor` dialect documentation for more background.

    Example input:

    ```mlir
    #matvec = {
      indexing_maps = [
        affine_map<(i,j) -> (i,j)>, // A
        affine_map<(i,j) -> (j)>,   // b
        affine_map<(i,j) -> (i)>    // x (out)
      ],
      iterator_types = ["parallel", "reduction"],
      doc = "X(i) += A(i,j) * B(j)"
    }

    // Multiply a sparse matrix A with a dense vector b into a dense vector x.
    func @kernel_matvec(%arga: tensor<?x?xf64, #SparseMatrix>,
                        %argb: tensor<?xf64>,
                        %argx: tensor<?xf64>) -> tensor<?xf64> {
      %0 = linalg.generic #matvec
        ins(%arga, %argb: tensor<?x?xf64, #SparseMatrix>, tensor<?xf64>)
        outs(%argx: tensor<?xf64>) {
        ^bb(%a: f64, %b: f64, %x: f64):
          %0 = arith.mulf %a, %b : f64
          %1 = arith.addf %x, %0 : f64
          linalg.yield %1 : f64
      } -> tensor<?xf64>
      return %0 : tensor<?xf64>
    }
    ```
  }];
  let constructor = "mlir::createSparsificationPass()";
  let dependentDialects = [
    "AffineDialect",
    "arith::ArithmeticDialect",
    "bufferization::BufferizationDialect",
    "LLVM::LLVMDialect",
    "memref::MemRefDialect",
    "scf::SCFDialect",
    "sparse_tensor::SparseTensorDialect",
    "vector::VectorDialect",
  ];
  let options = [
    Option<"parallelization", "parallelization-strategy", "int32_t", "0",
           "Set the parallelization strategy">,
    Option<"vectorization", "vectorization-strategy", "int32_t", "0",
           "Set the vectorization strategy">,
    Option<"vectorLength", "vl", "int32_t", "1",
           "Set the vector length">,
    Option<"enableSIMDIndex32", "enable-simd-index32", "bool", "false",
           "Enable i32 indexing into vectors (for efficiency)">
  ];
}

def SparseTensorConversion : Pass<"sparse-tensor-conversion", "ModuleOp"> {
  let summary = "Apply conversion rules to sparse tensor primitives and types";
  let description = [{
    A pass that converts sparse tensor primitives to calls into a runtime
    support library. All sparse tensor types are converted into opaque
    pointers to the underlying sparse storage schemes.

    Note that this is a current implementation choice to keep the conversion
    relatively simple. In principle, these primitives could also be
    converted to actual elaborate IR code that implements the primitives
    on the selected sparse tensor storage schemes.

    Example of the conversion:

    ```mlir
      Before:
        %c1 = arith.constant 1 : index
        %0 = sparse_tensor.pointers %arg0, %c1
          : tensor<8x8xf32, #sparse_tensor.encoding<{
              dimLevelType = [ "dense", "compressed" ],
              pointerBitWidth = 0,
              indexBitWidth = 0
            }>> to memref<?xindex>

      After:
        %c1 = arith.constant 1 : index
        %0 = call @sparsePointers(%arg0, %c1) : (!llvm.ptr<i8>, index) -> memref<?xindex>
    ```
  }];
  let constructor = "mlir::createSparseTensorConversionPass()";
  let dependentDialects = [
    "arith::ArithmeticDialect",
    "bufferization::BufferizationDialect",
    "LLVM::LLVMDialect",
    "linalg::LinalgDialect",
    "memref::MemRefDialect",
    "scf::SCFDialect",
    "sparse_tensor::SparseTensorDialect",
    "vector::VectorDialect",
  ];
}

#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES