f4f158b2f8
Similar to the parallelization strategies, the vectorization strategies provide control on what loops should be vectorize. Unlike the parallel strategies, only innermost loops are considered, but including reductions, with the control of vectorizing dense loops only or dense and sparse loops. The vectorized loops are always controlled by a vector mask to avoid overrunning the iterations, but subsequent vector operation folding removes redundant masks and replaces the operations with more efficient counterparts. Similarly, we will rely on subsequent loop optimizations to further optimize masking, e.g. using an unconditional full vector loop and scalar cleanup loop. The current strategy already demonstrates a nice interaction between the sparse compiler and all prior optimizations that went into the vector dialect. Ongoing discussion at: https://llvm.discourse.group/t/mlir-support-for-sparse-tensors/2020/10 Reviewed By: penpornk Differential Revision: https://reviews.llvm.org/D94551
116 lines
3.8 KiB
C++
116 lines
3.8 KiB
C++
//===- TestSparsification.cpp - Test sparsification of tensors ------------===//
|
|
//
|
|
// 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/Linalg/Transforms/Transforms.h"
|
|
#include "mlir/Dialect/Vector/VectorOps.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
|
|
struct TestSparsification
|
|
: public PassWrapper<TestSparsification, FunctionPass> {
|
|
|
|
TestSparsification() = default;
|
|
TestSparsification(const TestSparsification &pass) {}
|
|
|
|
Option<int32_t> parallelization{
|
|
*this, "parallelization-strategy",
|
|
llvm::cl::desc("Set the parallelization strategy"), llvm::cl::init(0)};
|
|
|
|
Option<int32_t> vectorization{
|
|
*this, "vectorization-strategy",
|
|
llvm::cl::desc("Set the vectorization strategy"), llvm::cl::init(0)};
|
|
|
|
Option<int32_t> vectorLength{
|
|
*this, "vl", llvm::cl::desc("Set the vector length"), llvm::cl::init(1)};
|
|
|
|
Option<int32_t> ptrType{*this, "ptr-type",
|
|
llvm::cl::desc("Set the pointer type"),
|
|
llvm::cl::init(0)};
|
|
|
|
Option<int32_t> indType{*this, "ind-type",
|
|
llvm::cl::desc("Set the index type"),
|
|
llvm::cl::init(0)};
|
|
|
|
/// Registers all dialects required by testing.
|
|
void getDependentDialects(DialectRegistry ®istry) const override {
|
|
registry.insert<scf::SCFDialect, vector::VectorDialect>();
|
|
}
|
|
|
|
/// Returns parallelization strategy given on command line.
|
|
linalg::SparseParallelizationStrategy parallelOption() {
|
|
switch (parallelization) {
|
|
default:
|
|
return linalg::SparseParallelizationStrategy::kNone;
|
|
case 1:
|
|
return linalg::SparseParallelizationStrategy::kDenseOuterLoop;
|
|
case 2:
|
|
return linalg::SparseParallelizationStrategy::kAnyStorageOuterLoop;
|
|
case 3:
|
|
return linalg::SparseParallelizationStrategy::kDenseAnyLoop;
|
|
case 4:
|
|
return linalg::SparseParallelizationStrategy::kAnyStorageAnyLoop;
|
|
}
|
|
}
|
|
|
|
/// Returns vectorization strategy given on command line.
|
|
linalg::SparseVectorizationStrategy vectorOption() {
|
|
switch (vectorization) {
|
|
default:
|
|
return linalg::SparseVectorizationStrategy::kNone;
|
|
case 1:
|
|
return linalg::SparseVectorizationStrategy::kDenseInnerLoop;
|
|
case 2:
|
|
return linalg::SparseVectorizationStrategy::kAnyStorageInnerLoop;
|
|
}
|
|
}
|
|
|
|
/// Returns the requested integer type.
|
|
linalg::SparseIntType typeOption(int32_t option) {
|
|
switch (option) {
|
|
default:
|
|
return linalg::SparseIntType::kNative;
|
|
case 1:
|
|
return linalg::SparseIntType::kI64;
|
|
case 2:
|
|
return linalg::SparseIntType::kI32;
|
|
}
|
|
}
|
|
|
|
/// Runs the test on a function.
|
|
void runOnFunction() override {
|
|
auto *ctx = &getContext();
|
|
OwningRewritePatternList patterns;
|
|
// Translate strategy flags to strategy options.
|
|
linalg::SparsificationOptions options(parallelOption(), vectorOption(),
|
|
vectorLength, typeOption(ptrType),
|
|
typeOption(indType));
|
|
// Apply rewriting.
|
|
linalg::populateSparsificationPatterns(ctx, patterns, options);
|
|
vector::populateVectorToVectorCanonicalizationPatterns(patterns, ctx);
|
|
applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
|
|
}
|
|
};
|
|
|
|
} // end anonymous namespace
|
|
|
|
namespace mlir {
|
|
namespace test {
|
|
|
|
void registerTestSparsification() {
|
|
PassRegistration<TestSparsification> sparsificationPass(
|
|
"test-sparsification",
|
|
"Test automatic generation of sparse tensor code");
|
|
}
|
|
|
|
} // namespace test
|
|
} // namespace mlir
|