[Mlir-commits] [mlir] bc5a023 - [mlir][sparse] cleanup of pass code
Aart Bik
llvmlistbot at llvm.org
Wed Aug 16 13:23:56 PDT 2023
Author: Aart Bik
Date: 2023-08-16T13:23:48-07:00
New Revision: bc5a023b81ec6ce7213f779e49cfc177a0efd8e9
URL: https://github.com/llvm/llvm-project/commit/bc5a023b81ec6ce7213f779e49cfc177a0efd8e9
DIFF: https://github.com/llvm/llvm-project/commit/bc5a023b81ec6ce7213f779e49cfc177a0efd8e9.diff
LOG: [mlir][sparse] cleanup of pass code
This reorders the pass declarations and definitions in a
consistent order and makes the layout a bit more clear.
Also makes the "minipipeline" sparsification-and-bufferization
available through the command line for individual testing.
Reviewed By: Peiming
Differential Revision: https://reviews.llvm.org/D158105
Added:
Modified:
mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 8d632bf7453adb..2b237842bc6ed5 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -8,10 +8,6 @@
//
// This header file defines prototypes of all sparse tensor passes.
//
-// In general, this file takes the approach of keeping "mechanism" (the
-// actual steps of applying a transformation) completely separate from
-// "policy" (heuristics for when and where to apply transformations).
-//
//===----------------------------------------------------------------------===//
#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES_H_
@@ -21,15 +17,15 @@
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"
+//===----------------------------------------------------------------------===//
+// Include the generated pass header (which needs some early definitions).
+//===----------------------------------------------------------------------===//
+
namespace mlir {
namespace bufferization {
struct OneShotBufferizationOptions;
} // namespace bufferization
-//===----------------------------------------------------------------------===//
-// The Sparsification pass.
-//===----------------------------------------------------------------------===//
-
/// Defines a parallelization strategy. Any independent loop is a candidate
/// for parallelization. The loop is made parallel if (1) allowed by the
/// strategy (e.g., AnyStorageOuterLoop considers either a dense or sparse
@@ -43,12 +39,25 @@ enum class SparseParallelizationStrategy {
kAnyStorageAnyLoop
};
-// TODO : Zero copy is disabled due to correctness bugs.Tracker #64316
+/// Defines data movement strategy between host and device for GPU.
+// TODO : Zero copy is disabled due to correctness bugs (tracker #64316)
enum class GPUDataTransferStrategy { kRegularDMA, kZeroCopy, kPinnedDMA };
#define GEN_PASS_DECL
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
+//===----------------------------------------------------------------------===//
+// The PreSparsificationRewriting pass.
+//===----------------------------------------------------------------------===//
+
+void populatePreSparsificationRewriting(RewritePatternSet &patterns);
+
+std::unique_ptr<Pass> createPreSparsificationRewritePass();
+
+//===----------------------------------------------------------------------===//
+// The Sparsification pass.
+//===----------------------------------------------------------------------===//
+
/// Options for the Sparsification pass.
struct SparsificationOptions {
SparsificationOptions(SparseParallelizationStrategy p,
@@ -77,6 +86,19 @@ std::unique_ptr<Pass> createSparsificationPass();
std::unique_ptr<Pass>
createSparsificationPass(const SparsificationOptions &options);
+//===----------------------------------------------------------------------===//
+// The PostSparsificationRewriting pass.
+//===----------------------------------------------------------------------===//
+
+void populatePostSparsificationRewriting(RewritePatternSet &patterns,
+ bool enableRT, bool enableForeach,
+ bool enableConvert);
+
+std::unique_ptr<Pass> createPostSparsificationRewritePass();
+std::unique_ptr<Pass>
+createPostSparsificationRewritePass(bool enableRT, bool enableForeach = true,
+ bool enableConvert = true);
+
//===----------------------------------------------------------------------===//
// The SparseTensorConversion pass.
//===----------------------------------------------------------------------===//
@@ -94,17 +116,6 @@ class SparseTensorTypeToPtrConverter : public TypeConverter {
/// directly via the algorithm in <https://arxiv.org/abs/2001.02609>;
/// however, beware that there are many formats not supported by this
/// conversion method.
-///
-/// The presence of the `kAuto` option violates our usual goal of keeping
-/// policy completely separated from mechanism. The reason it exists is
-/// because (at present) this strategy can only be specified on a per-file
-/// basis. To see why this is a problem, note that `kDirect` cannot
-/// support certain conversions; so if there is no `kAuto` setting,
-/// then whenever a file contains a single non-`kDirect`-able conversion
-/// the user would be forced to use `kViaCOO` for all conversions in
-/// that file! In the future, instead of using this enum as a `Pass`
-/// option, we could instead move it to being an attribute on the
-/// conversion op; at which point `kAuto` would no longer be necessary.
enum class SparseToSparseConversionStrategy { kAuto, kViaCOO, kDirect };
/// Converts command-line sparse2sparse flag to the strategy enum.
@@ -140,7 +151,7 @@ class SparseTensorTypeToBufferConverter : public TypeConverter {
SparseTensorTypeToBufferConverter();
};
-/// Sets up sparse tensor conversion rules.
+/// Sets up sparse tensor codegen rules.
void populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
RewritePatternSet &patterns,
bool createSparseDeallocs,
@@ -152,25 +163,42 @@ createSparseTensorCodegenPass(bool createSparseDeallocs,
bool enableBufferInitialization);
//===----------------------------------------------------------------------===//
-// The PreSparsificationRewriting pass.
+// The SparseBufferRewrite pass.
//===----------------------------------------------------------------------===//
-void populatePreSparsificationRewriting(RewritePatternSet &patterns);
+void populateSparseBufferRewriting(RewritePatternSet &patterns,
+ bool enableBufferInitialization);
-std::unique_ptr<Pass> createPreSparsificationRewritePass();
+std::unique_ptr<Pass> createSparseBufferRewritePass();
+std::unique_ptr<Pass>
+createSparseBufferRewritePass(bool enableBufferInitialization);
//===----------------------------------------------------------------------===//
-// The PostSparsificationRewriting pass.
+// The SparseVectorization pass.
//===----------------------------------------------------------------------===//
-void populatePostSparsificationRewriting(RewritePatternSet &patterns,
- bool enableRT, bool enableForeach,
- bool enableConvert);
+void populateSparseVectorizationPatterns(RewritePatternSet &patterns,
+ unsigned vectorLength,
+ bool enableVLAVectorization,
+ bool enableSIMDIndex32);
-std::unique_ptr<Pass> createPostSparsificationRewritePass();
-std::unique_ptr<Pass>
-createPostSparsificationRewritePass(bool enableRT, bool enableForeach = true,
- bool enableConvert = true);
+std::unique_ptr<Pass> createSparseVectorizationPass();
+std::unique_ptr<Pass> createSparseVectorizationPass(unsigned vectorLength,
+ bool enableVLAVectorization,
+ bool enableSIMDIndex32);
+
+//===----------------------------------------------------------------------===//
+// The SparseGPU pass.
+//===----------------------------------------------------------------------===//
+
+void populateSparseGPUCodegenPatterns(RewritePatternSet &patterns,
+ unsigned numThreads);
+
+void populateSparseGPULibgenPatterns(RewritePatternSet &patterns, bool enableRT,
+ GPUDataTransferStrategy gpuDataTransfer);
+
+std::unique_ptr<Pass> createSparseGPUCodegenPass();
+std::unique_ptr<Pass> createSparseGPUCodegenPass(unsigned numThreads);
//===----------------------------------------------------------------------===//
// The SparseStorageSpecifierToLLVM pass.
@@ -186,7 +214,11 @@ void populateStorageSpecifierToLLVMPatterns(TypeConverter &converter,
std::unique_ptr<Pass> createStorageSpecifierToLLVMPass();
//===----------------------------------------------------------------------===//
-// Other rewriting rules and passes.
+// The mini-pipeline for sparsification and bufferization.
+//
+// Note that this mini-pipeline is not defined through the tablegen pass
+// mechanism, and, thus, is not individually available through the command-line.
+// It is solely used as part of the full sparse compiler pipeline.
//===----------------------------------------------------------------------===//
std::unique_ptr<Pass> createSparsificationAndBufferizationPass(
@@ -197,32 +229,6 @@ std::unique_ptr<Pass> createSparsificationAndBufferizationPass(
bool enableBufferInitialization, unsigned vectorLength,
bool enableVLAVectorization, bool enableSIMDIndex32);
-void populateSparseBufferRewriting(RewritePatternSet &patterns,
- bool enableBufferInitialization);
-
-std::unique_ptr<Pass> createSparseBufferRewritePass();
-std::unique_ptr<Pass>
-createSparseBufferRewritePass(bool enableBufferInitialization);
-
-void populateSparseVectorizationPatterns(RewritePatternSet &patterns,
- unsigned vectorLength,
- bool enableVLAVectorization,
- bool enableSIMDIndex32);
-
-std::unique_ptr<Pass> createSparseVectorizationPass();
-std::unique_ptr<Pass> createSparseVectorizationPass(unsigned vectorLength,
- bool enableVLAVectorization,
- bool enableSIMDIndex32);
-
-void populateSparseGPUCodegenPatterns(RewritePatternSet &patterns,
- unsigned numThreads);
-
-void populateSparseGPULibgenPatterns(RewritePatternSet &patterns, bool enableRT,
- GPUDataTransferStrategy gpuDataTransfer);
-
-std::unique_ptr<Pass> createSparseGPUCodegenPass();
-std::unique_ptr<Pass> createSparseGPUCodegenPass(unsigned numThreads);
-
//===----------------------------------------------------------------------===//
// Registration.
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 95a4feb52256a2..f59785c0fcef9a 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -123,7 +123,6 @@ def SparsificationPass : Pass<"sparsification", "ModuleOp"> {
];
}
-
def PostSparsificationRewrite : Pass<"post-sparsification-rewrite", "ModuleOp"> {
let summary = "Applies sparse tensor rewriting rules after sparsification";
let description = [{
More information about the Mlir-commits
mailing list