[mlir][sparse] migrate sparse operations into new sparse tensor dialect

This is the very first step toward removing the glue and clutter from linalg and
replace it with proper sparse tensor types. This revision migrates the LinalgSparseOps
into SparseTensorOps of a sparse tensor dialect. This also provides a new home for
sparse tensor related transformation.

NOTE: the actual replacement with sparse tensor types (and removal of linalg glue/clutter)
will follow but I am trying to keep the amount of changes per revision manageable.

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

mlir/include/mlir/Dialect/CMakeLists.txt

@@ -17,6 +17,7 @@ add_subdirectory(PDLInterp)
 add_subdirectory(Quant)
 add_subdirectory(SCF)
 add_subdirectory(Shape)
+add_subdirectory(SparseTensor)
 add_subdirectory(SPIRV)
 add_subdirectory(StandardOps)
 add_subdirectory(Tensor)

mlir/include/mlir/Dialect/Linalg/IR/CMakeLists.txt

@@ -80,12 +80,6 @@ add_public_tablegen_target(MLIRLinalgStructuredOpsIncGen)
 add_dependencies(MLIRLinalgStructuredOpsIncGen LinalgOdsGen)
 add_dependencies(mlir-headers MLIRLinalgStructuredOpsIncGen)
 
-set(LLVM_TARGET_DEFINITIONS LinalgSparseOps.td)
-mlir_tablegen(LinalgSparseOps.h.inc -gen-op-decls)
-mlir_tablegen(LinalgSparseOps.cpp.inc -gen-op-defs)
-add_public_tablegen_target(MLIRLinalgSparseOpsIncGen)
-add_dependencies(mlir-headers MLIRLinalgSparseOpsIncGen)
-
 set(LLVM_TARGET_DEFINITIONS LinalgInterfaces.td)
 mlir_tablegen(LinalgInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(LinalgInterfaces.cpp.inc -gen-op-interface-defs)

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h

@@ -127,7 +127,4 @@ class IndexedGenericOp;
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.h.inc"
 
-#define GET_OP_CLASSES
-#include "mlir/Dialect/Linalg/IR/LinalgSparseOps.h.inc"
-
 #endif // MLIR_DIALECT_LINALG_LINALGOPS_H_

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

@@ -1129,9 +1129,6 @@ void populateSparsificationPatterns(
     RewritePatternSet &patterns,
     const SparsificationOptions &options = SparsificationOptions());
 
-/// Sets up sparsification conversion rules with the given options.
-void populateSparsificationConversionPatterns(RewritePatternSet &patterns);
-
 } // namespace linalg
 } // namespace mlir
 

mlir/include/mlir/Dialect/SparseTensor/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(IR)

mlir/include/mlir/Dialect/SparseTensor/IR/CMakeLists.txt

@@ -0,0 +1,2 @@
+add_mlir_dialect(SparseTensorOps sparse_tensor)
+add_mlir_doc(SparseTensorOps SparseTensorOps Dialects/ -gen-dialect-doc)

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

@@ -0,0 +1,23 @@
+//===- SparseTensor.h - Sparse tensor dialect -------------------*- C++ -*-===//
+//
+// 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_IR_SPARSETENSOR_H_
+#define MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSOR_H_
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorOps.h.inc"
+
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorOpsDialect.h.inc"
+
+#endif // MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSOR_H_

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorBase.td

@@ -0,0 +1,29 @@
+//===- SparseTensorBase.td - Sparse tensor dialect base ----*- 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 SPARSETENSOR_BASE
+#define SPARSETENSOR_BASE
+
+include "mlir/IR/OpBase.td"
+
+def SparseTensor_Dialect : Dialect {
+  let name = "sparse_tensor";
+  let cppNamespace = "::mlir::sparse_tensor";
+  let description = [{
+    The `sparse tensor` dialect is intended to hold primitives that
+    form a bridge between high-level operations on sparse tensors
+    and lower-level operations on the actual sparse storage schemes
+    consisting of pointers, indices, and values. This bridge
+    simplifies a `sparse compiler` pass by postponing actual
+    code generation for the supported primitives to a later phase,
+    either by generating calls into a runtime support library
+    or by further lowering the primitives into actual code.
+  }];
+}
+
+#endif // SPARSETENSOR_BASE

mlir/include/mlir/Dialect/Linalg/IR/LinalgSparseOps.td → mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

@@ -1,49 +1,27 @@
-//===- LinalgSparseOps.td - Linalg dialect sparse ops ------*- tablegen -*-===//
+//===- SparseTensorOps.td - Sparse tensor dialect ops ------*- 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
 //
 //===----------------------------------------------------------------------===//
-//
-// The following operations bootstrap working with sparse tensors solely
-// within the Linalg dialect. They provide temporary bridges between a
-// future SparseTensorType (now an opaque pointer), the actual TensorType,
-// and MemRef arrays underlying an actual sparse storage scheme in memory.
-//
-// Lacking a proper sparse tensor type, the 'sparse_tensor' operation
-// provides a bridge between an opaque pointer and a regular tensor type
-// just to simplify feeding the value into a Linalg op. The operation
-// simply disappears during lowering.
-//
-// The other operations form the bridge between the opaque pointer and
-// the actual storage of pointers, indices, and values. These operations
-// resemble 'buffer_cast' in the sense that they map tensors to
-// their bufferized memrefs, but they lower into actual calls since
-// sparse storage does not bufferize into a single memrefs, as dense
-// tensors do, but into a hierarchical storage scheme where pointers
-// access memrefs with indices and eventually into values.
-//
-// TODO: introduce SparseTensorType as first class citizen in MLIR
-//
-//===----------------------------------------------------------------------===//
 
-#ifndef LINALG_SPARSE_OPS
-#define LINALG_SPARSE_OPS
+#ifndef SPARSETENSOR_OPS
+#define SPARSETENSOR_OPS
 
-include "mlir/Dialect/Linalg/IR/LinalgBase.td"
+include "mlir/Dialect/SparseTensor/IR/SparseTensorBase.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 // Base class.
-class Linalg_SparseOp<string mnemonic, list<OpTrait> traits = []>
-  : Op<Linalg_Dialect, mnemonic, traits> {
+class SparseTensor_Op<string mnemonic, list<OpTrait> traits = []>
+  : Op<SparseTensor_Dialect, mnemonic, traits> {
   let printer = [{ return ::print(p, *this); }];
   let verifier = ?;
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
-def Linalg_SparseTensorFromPointerOp :
-    Linalg_SparseOp<"sparse_tensor">,
+// TODO: remove me
+def SparseTensor_FromPointerOp : SparseTensor_Op<"fromPtr">,
     Arguments<(ins AnyType:$ptr)>,
     Results<(outs AnyTensor:$result)> {
   let summary = "Views an opaque sparse tensor pointer as a tensor";
@@ -60,14 +38,13 @@ def Linalg_SparseTensorFromPointerOp :
     ```mlir
      !SparseTensor = type !llvm.ptr<i8>
 
-     %0 = linalg.sparse_tensor %arg0 : !SparseTensor to tensor<64x64xf64>
+     %0 = sparse_tensor.fromPtr %arg0 : !SparseTensor to tensor<64x64xf64>
     ```
   }];
   let assemblyFormat = "$ptr attr-dict `:` type($ptr) `to` type($result)";
 }
 
-def Linalg_SparseTensorToPointersMemRefOp :
-    Linalg_SparseOp<"sparse_pointers", [NoSideEffect]>,
+def SparseTensor_ToPointersOp : SparseTensor_Op<"pointers", [NoSideEffect]>,
     Arguments<(ins AnyTensor:$tensor, Index:$dim)>,
     Results<(outs AnyStridedMemRefOfRank<1>:$result)> {
   let summary = "Extract pointers array at given dimension from a tensor";
@@ -83,15 +60,14 @@ def Linalg_SparseTensorToPointersMemRefOp :
      Example:
 
     ```mlir
-    %1 = linalg.sparse_pointers %0, %c1 : tensor<64x64xf64> to memref<?xindex>
+    %1 = sparse_tensor.pointers %0, %c1 : tensor<64x64xf64> to memref<?xindex>
     ```
   }];
   let assemblyFormat = "$tensor `,` $dim attr-dict `:` type($tensor)"
       " `to` type($result)";
 }
 
-def Linalg_SparseTensorToIndicesMemRefOp :
-    Linalg_SparseOp<"sparse_indices", [NoSideEffect]>,
+def SparseTensor_ToIndicesOp : SparseTensor_Op<"indices", [NoSideEffect]>,
     Arguments<(ins AnyTensor:$tensor, Index:$dim)>,
     Results<(outs AnyStridedMemRefOfRank<1>:$result)> {
   let summary = "Extract indices array at given dimension from a tensor";
@@ -107,15 +83,14 @@ def Linalg_SparseTensorToIndicesMemRefOp :
      Example:
 
     ```mlir
-    %1 = linalg.sparse_indices %0, %c1 : tensor<64x64xf64> to memref<?xindex>
+    %1 = sparse_tensor.indices %0, %c1 : tensor<64x64xf64> to memref<?xindex>
     ```
   }];
   let assemblyFormat = "$tensor `,` $dim attr-dict `:` type($tensor)"
       " `to` type($result)";
 }
 
-def Linalg_SparseTensorToValuesMemRefOp :
-    Linalg_SparseOp<"sparse_values", [NoSideEffect]>,
+def SparseTensor_ToValuesOp : SparseTensor_Op<"values", [NoSideEffect]>,
     Arguments<(ins AnyTensor:$tensor)>,
     Results<(outs AnyStridedMemRefOfRank<1>:$result)> {
   let summary = "Extract numerical values array from a tensor";
@@ -131,10 +106,10 @@ def Linalg_SparseTensorToValuesMemRefOp :
      Example:
 
     ```mlir
-    %1 = linalg.sparse_values %0 : tensor<64x64xf64> to memref<?xf64>
+    %1 = sparse_tensor.values %0 : tensor<64x64xf64> to memref<?xf64>
     ```
   }];
   let assemblyFormat = "$tensor attr-dict `:` type($tensor) `to` type($result)";
 }
 
-#endif // LINALG_SPARSE_OPS
+#endif // SPARSETENSOR_OPS

mlir/include/mlir/Dialect/SparseTensor/Transforms/Transforms.h

@@ -0,0 +1,23 @@
+//===- Transforms.h - Sparse tensor transformations -------------*- C++ -*-===//
+//
+// 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_TRANSFORMS_H_
+#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_TRANSFORMS_H_
+
+#include "mlir/IR/PatternMatch.h"
+
+namespace mlir {
+namespace sparse_tensor {
+
+/// Sets up sparsification conversion rules with the given options.
+void populateSparsificationConversionPatterns(RewritePatternSet &patterns);
+
+} // namespace sparse_tensor
+} // namespace mlir
+
+#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_TRANSFORMS_H_

mlir/include/mlir/InitAllDialects.h

@@ -37,6 +37,7 @@
 #include "mlir/Dialect/SDBM/SDBMDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
@@ -74,6 +75,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
                   ROCDL::ROCDLDialect,
                   SDBMDialect,
                   shape::ShapeDialect,
+                  sparse_tensor::SparseTensorDialect,
                   tensor::TensorDialect,
                   tosa::TosaDialect,
                   x86vector::X86VectorDialect>();

mlir/lib/Dialect/CMakeLists.txt

@@ -18,6 +18,7 @@ add_subdirectory(Quant)
 add_subdirectory(SCF)
 add_subdirectory(SDBM)
 add_subdirectory(Shape)
+add_subdirectory(SparseTensor)
 add_subdirectory(SPIRV)
 add_subdirectory(StandardOps)
 add_subdirectory(Tensor)

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

@@ -2384,9 +2384,6 @@ struct FoldTensorCastOp;
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
 
-#define GET_OP_CLASSES
-#include "mlir/Dialect/Linalg/IR/LinalgSparseOps.cpp.inc"
-
 /// Return the dims that are `iteratorTypeName` loops in the LinalgOp `op`.
 /// Assumes `op` is a LinalgOp.
 void mlir::linalg::getDimsOfType(Operation *op, StringRef iteratorTypeName,

mlir/lib/Dialect/Linalg/IR/LinalgTypes.cpp

@@ -99,10 +99,6 @@ void mlir::linalg::LinalgDialect::initialize() {
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
       >();
-  addOperations<
-#define GET_OP_LIST
-#include "mlir/Dialect/Linalg/IR/LinalgSparseOps.cpp.inc"
-      >();
 
   // Fill the Linalg-specific OpName to RegionBuilder map.
   addNamedOpBuilders<

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

@@ -11,7 +11,6 @@ add_mlir_dialect_library(MLIRLinalgTransforms
   Interchange.cpp
   Loops.cpp
   Promotion.cpp
-  SparseLowering.cpp
   Sparsification.cpp
   Tiling.cpp
   Transforms.cpp
@@ -38,6 +37,7 @@ add_mlir_dialect_library(MLIRLinalgTransforms
   MLIRSCF
   MLIRSCFTransforms
   MLIRPass
+  MLIRSparseTensor
   MLIRStandard
   MLIRStandardOpsTransforms
   MLIRStandardToLLVM

mlir/lib/Dialect/Linalg/Transforms/Sparsification.cpp

@@ -45,6 +45,7 @@
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Matchers.h"
 
@@ -360,7 +361,7 @@ static void findSparseAnnotations(Merger &merger, linalg::GenericOp op) {
 /// Returns true if tensor was set up with sparse storage scheme.
 static bool linkedSparse(linalg::GenericOp op, unsigned tensor) {
   if (tensor < op.getNumInputs())
-    return isa_and_nonnull<linalg::SparseTensorFromPointerOp>(
+    return isa_and_nonnull<sparse_tensor::FromPointerOp>(
         op.getInput(tensor).getDefiningOp());
   return false;
 }
@@ -576,12 +577,10 @@ static void genBuffers(Merger &merger, CodeGen &codegen,
             dynShape, genIntType(rewriter, codegen.options.indType));
         Value dim = rewriter.create<ConstantIndexOp>(loc, d);
         // Generate sparse primitives to obtains pointer and indices.
-        codegen.pointers[t][i] =
-            rewriter.create<linalg::SparseTensorToPointersMemRefOp>(
-                loc, ptrTp, tensor, dim);
-        codegen.indices[t][i] =
-            rewriter.create<linalg::SparseTensorToIndicesMemRefOp>(loc, indTp,
-                                                                   tensor, dim);
+        codegen.pointers[t][i] = rewriter.create<sparse_tensor::ToPointersOp>(
+            loc, ptrTp, tensor, dim);
+        codegen.indices[t][i] = rewriter.create<sparse_tensor::ToIndicesOp>(
+            loc, indTp, tensor, dim);
       }
       // Find lower and upper bound in current dimension.
       Value up;
@@ -608,8 +607,7 @@ static void genBuffers(Merger &merger, CodeGen &codegen,
       auto dynShape = {ShapedType::kDynamicSize};
       auto sparseTp = MemRefType::get(dynShape, tensorType.getElementType());
       codegen.buffers[t] =
-          rewriter.create<linalg::SparseTensorToValuesMemRefOp>(loc, sparseTp,
-                                                                tensor);
+          rewriter.create<sparse_tensor::ToValuesOp>(loc, sparseTp, tensor);
     }
   }
 }

mlir/lib/Dialect/SparseTensor/CMakeLists.txt

@@ -0,0 +1,2 @@
+add_subdirectory(IR)
+add_subdirectory(Transforms)

mlir/lib/Dialect/SparseTensor/IR/CMakeLists.txt

@@ -0,0 +1,13 @@
+add_mlir_dialect_library(MLIRSparseTensor
+  SparseTensorDialect.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
+
+  DEPENDS
+  MLIRSparseTensorOpsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRDialect
+  MLIRIR
+  )

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

@@ -0,0 +1,25 @@
+//===- SparseTensorDialect.cpp - Sparse tensor dialect implementation -----===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/OpImplementation.h"
+
+using namespace mlir;
+using namespace mlir::sparse_tensor;
+
+void SparseTensorDialect::initialize() {
+  addOperations<
+#define GET_OP_LIST
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorOps.cpp.inc"
+      >();
+}
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorOps.cpp.inc"