[mlir][sparse] Add SparseTensorStorageExpansion Pass to expand compounded sparse tensor tuples

[Peiming]

Summary

This patch adds SparseTensorStorageExpansion pass, it flattens the tuple used to store a sparse
tensor handle.

Right now, it only set up the skeleton for the pass, more lowering rules for sparse tensor storage
operation need to be added.

[Peiming]

add missing period in comments...

[aartbik]

very good progress!

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

154

period at end here too

160

can you break the multi line a bit

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp

162

I had a fix for that in my revision that looks a bit more idiomatic

197

I definitely like having this as a new pass for the time being, since it nicely progressively lowers and can be tested in isolation!

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageExpansion.cpp

10

nested -> compound?

mlir/test/Dialect/SparseTensor/sparse_tensor_storage.mlir

10

can you indent so that the tuples "line up"?

12

can you also add a call to this test already?

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

152

One complication is that we actually expand all tuples, not just storage formats. I don't think that is an immediate problem, since tuples seem mostly unused. But for safety perhaps we should indeed consider having our own extension of tuple as a sparse_storage_tuple or so, so that we only rewrite what we introduced....

[Peiming]

address comments from Aart

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp

162

Reverted!

mlir/test/Dialect/SparseTensor/sparse_tensor_storage.mlir

12

Not yet.

[aartbik]

Few last comments, please address, but good to merge in.

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

168

This comment seems copy and paste? How about

/// Sparse tensor type converter from compound to expanded form.

Or something like that.

174

Similar

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp

161

maybe remove 161 as well to avoid merge conflicts now

[Peiming]

address comments from Aart

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

 
 /// Sets up sparse tensor conversion rules.
 void populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
                                          RewritePatternSet &patterns);
 
 std::unique_ptr<Pass> createSparseTensorCodegenPass();
 
 //===----------------------------------------------------------------------===//
+// The SparseTensorStorageExpansion pass.
+//===----------------------------------------------------------------------===//
+
+/// Sparse tensor storage type converter from compound to expanded form.
+class SparseTensorStorageTupleExpander : public TypeConverter {
+public:
+  SparseTensorStorageTupleExpander();
+};
+
+/// Sets up sparse tensor storage expansion rules.
+void populateSparseTensorStorageExpansionPatterns(TypeConverter &typeConverter,
+                                                  RewritePatternSet &patterns);
+
+std::unique_ptr<Pass> createSparseTensorStorageExpansionPass();
+
+//===----------------------------------------------------------------------===//
 // Other rewriting rules and passes.
 //===----------------------------------------------------------------------===//
 
 void populateSparseTensorRewriting(RewritePatternSet &patterns);
 
 std::unique_ptr<Pass> createDenseBufferizationPass(
     const bufferization::OneShotBufferizationOptions &options);
 

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

let description = [{
     ```
   }];
   let constructor = "mlir::createSparseTensorCodegenPass()";
   let dependentDialects = [
     "sparse_tensor::SparseTensorDialect",
   ];
 }
 
+def SparseTensorStorageExpansion : Pass<"sparse-tensor-storage-expansion", "ModuleOp"> {
+  let summary = "Expand compounded sparse tensor storage into individual SSA values";
+  let description = [{
+    A pass that expands sparse tensor storage (aggregated by tuple) into
+    individual SSA values. It also lowers sparse tensor storage operations,
+    e.g., sparse_tensor.storage_get and sparse_tensor.storage_set.
+
+    Example of the conversion:
+
+    ```mlir
+    Before:
+      func.func @sparse_storage_set(%arg0: tuple<memref<?xf64>,
+                                                 memref<?xf64>,
+                                                 f64>)
+                                        -> tuple<memref<?xf64>,
+                                                 memref<?xf64>,
+                                                 f64> {
+        return %arg0 : tuple<memref<?xf64>, memref<?xf64>, f64>
+      }
+    After:
+      func.func @sparse_storage_set(%arg0: memref<?xf64>,
+                                    %arg1: memref<?xf64>,
+                                    %arg2: f64)
+                                    -> (memref<?xf64>, memref<?xf64>, f64) {
+        return %arg0, %arg1, %arg2 : memref<?xf64>, memref<?xf64>, f64
+      }
+    ```
+  }];
+  let constructor = "mlir::createSparseTensorStorageExpansionPass()";
+  let dependentDialects = [
+    "sparse_tensor::SparseTensorDialect",
+  ];
+}
+
+
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES

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

 add_mlir_dialect_library(MLIRSparseTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   CodegenUtils.cpp
   DenseBufferizationPass.cpp
   Sparsification.cpp
   SparseTensorCodegen.cpp
   SparseTensorConversion.cpp
   SparseTensorPasses.cpp
   SparseTensorRewriting.cpp
+  SparseTensorStorageExpansion.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
 
   DEPENDS
   MLIRSparseTensorPassIncGen
 
   LINK_LIBS PUBLIC

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp

 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 namespace mlir {
 #define GEN_PASS_DEF_SPARSIFICATIONPASS
 #define GEN_PASS_DEF_SPARSETENSORCONVERSIONPASS
 #define GEN_PASS_DEF_SPARSETENSORCODEGEN
+#define GEN_PASS_DEF_SPARSETENSORSTORAGEEXPANSION
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 namespace {
 
struct SparseTensorCodegenPass
   SparseTensorCodegenPass(const SparseTensorCodegenPass &pass) = default;
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
     SparseTensorTypeToBufferConverter converter;
     ConversionTarget target(*ctx);
     // Everything in the sparse dialect must go!
     target.addIllegalDialect<SparseTensorDialect>();
     // All dynamic rules below accept new function, call, return.
     target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {
       return converter.isSignatureLegal(op.getFunctionType());
     });
     target.addDynamicallyLegalOp<func::CallOp>([&](func::CallOp op) {
       return converter.isSignatureLegal(op.getCalleeType());
     });
     target.addDynamicallyLegalOp<func::ReturnOp>([&](func::ReturnOp op) {
       return converter.isLegal(op.getOperandTypes());
scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,
                                                          target);
     populateSparseTensorCodegenPatterns(converter, patterns);
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 
+struct SparseTensorStorageExpansionPass
+    : public impl::SparseTensorStorageExpansionBase<
+          SparseTensorStorageExpansionPass> {
+
+  SparseTensorStorageExpansionPass() = default;
+  SparseTensorStorageExpansionPass(
+      const SparseTensorStorageExpansionPass &pass) = default;
+
+  void runOnOperation() override {
+    auto *ctx = &getContext();
+    RewritePatternSet patterns(ctx);
+    SparseTensorStorageTupleExpander converter;
+    ConversionTarget target(*ctx);
+    // Now, everything in the sparse dialect must go!
+    target.addIllegalDialect<SparseTensorDialect>();
+    // All dynamic rules below accept new function, call, return.
+    target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {
+      return converter.isSignatureLegal(op.getFunctionType());
+    });
+    target.addDynamicallyLegalOp<func::CallOp>([&](func::CallOp op) {
+      return converter.isSignatureLegal(op.getCalleeType());
+    });
+    target.addDynamicallyLegalOp<func::ReturnOp>([&](func::ReturnOp op) {
+      return converter.isLegal(op.getOperandTypes());
+    });
+    // Populate with rules and apply rewriting rules.
+    populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
+                                                                   converter);
+    populateCallOpTypeConversionPattern(patterns, converter);
+    scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,
+                                                         target);
+    populateSparseTensorStorageExpansionPatterns(converter, patterns);
+    if (failed(applyPartialConversion(getOperation(), target,
+                                      std::move(patterns))))
+      signalPassFailure();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Strategy flag methods.
 //===----------------------------------------------------------------------===//
 
 SparseParallelizationStrategy
 mlir::sparseParallelizationStrategy(int32_t flag) {

 std::unique_ptr<Pass> mlir::createSparseTensorConversionPass(
     const SparseTensorConversionOptions &options) {
   return std::make_unique<SparseTensorConversionPass>(options);
 }
 
 std::unique_ptr<Pass> mlir::createSparseTensorCodegenPass() {
   return std::make_unique<SparseTensorCodegenPass>();
 }
+
+std::unique_ptr<Pass> mlir::createSparseTensorStorageExpansionPass() {
+  return std::make_unique<SparseTensorStorageExpansionPass>();
+}

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageExpansion.cpp

+//===- SparseTensorStorageExpansion.cpp - Sparse tensor storage expansion ===//
+//
+// 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 sparse tensor storage expansion pass expands the compound storage for
+// sparse tensors (using tuple) to flattened SSA values.
+//
+//===----------------------------------------------------------------------===//
+
+#include "CodegenUtils.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+using namespace mlir::sparse_tensor;
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// Helper methods.
+//===----------------------------------------------------------------------===//
+
+/// Expands sparse tensor storage tuple.
+static Optional<LogicalResult>
+convertSparseTensorStorageTuple(Type t, SmallVectorImpl<Type> &result) {
+  if (auto tuple = t.dyn_cast<TupleType>()) {
+    // Note that it does not handle nest tuples, but it is fine
+    // for sparse compiler as they will not be generated.
+    result.append(tuple.getTypes().begin(), tuple.getTypes().end());
+    return success();
+  }
+  return llvm::None;
+}
+
+//===----------------------------------------------------------------------===//
+// Conversion rules.
+//===----------------------------------------------------------------------===//
+
+/// Sparse tensor storage conversion rule for returns.
+class SparseStorageReturnConverter
+    : public OpConversionPattern<func::ReturnOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(func::ReturnOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    SmallVector<Value, 8> flattened;
+    for (auto operand : adaptor.getOperands()) {
+      if (auto cast =
+              dyn_cast<UnrealizedConversionCastOp>(operand.getDefiningOp());
+          cast && cast->getResultTypes()[0].isa<TupleType>())
+        // An unrealized_conversion_cast will be inserted by type converter to
+        // inter-mix the gap between 1:N conversion between tuple and types.
+        // In this case, take the operands in the cast and replace the tuple
+        // output with the flattened type array.
+        flattened.append(cast.getOperands().begin(), cast.getOperands().end());
+      else
+        flattened.push_back(operand);
+    }
+    // Create a return with the flattened value extracted from tuple.
+    rewriter.replaceOpWithNewOp<func::ReturnOp>(op, flattened);
+    return success();
+  }
+};
+
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Sparse tensor storage expansion
+//===----------------------------------------------------------------------===//
+
+mlir::SparseTensorStorageTupleExpander::SparseTensorStorageTupleExpander() {
+  addConversion([](Type type) { return type; });
+  addConversion(convertSparseTensorStorageTuple);
+}
+
+//===----------------------------------------------------------------------===//
+// Public method for populating conversion rules.
+//===----------------------------------------------------------------------===//
+
+/// Populates the given patterns list with conversion rules required
+/// to expand compounded sparse tensor tuples.
+void mlir::populateSparseTensorStorageExpansionPatterns(
+    TypeConverter &typeConverter, RewritePatternSet &patterns) {
+  patterns.add<SparseStorageReturnConverter>(typeConverter,
+                                             patterns.getContext());
+}

mlir/test/Dialect/SparseTensor/sparse_tensor_storage.mlir

+// RUN: mlir-opt %s -sparse-tensor-storage-expansion | FileCheck %s
+
+// CHECK-LABEL:  func @sparse_storage_expand(
+// CHECK-SAME:     %[[TMP_arg0:.*0]]: memref<?xf64>,
+// CHECK-SAME:     %[[TMP_arg1:.*1]]: memref<?xf64>,
+// CHECK-SAME:     %[[TMP_arg2:.*]]: f64
+// CHECK           return %[[TMP_arg0]], %[[TMP_arg1]], %[[TMP_arg2]]
+func.func @sparse_storage_expand(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>)
+                                     -> tuple<memref<?xf64>, memref<?xf64>, f64> {
+  return %arg0 : tuple<memref<?xf64>, memref<?xf64>, f64>
+}