3fffffa882
This class represents a rewrite pattern list that has been frozen, and thus immutable. This replaces the uses of OwningRewritePatternList in pattern driver related API, such as dialect conversion. When PDL becomes more prevalent, this API will allow for optimizing a set of patterns once without the need to do this per run of a pass. Differential Revision: https://reviews.llvm.org/D89104
95 lines
3.1 KiB
C++
95 lines
3.1 KiB
C++
//===- ExpandAtomic.cpp - Code to perform expanding atomic ops ------------===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file implements expansion of AtomicRMWOp into GenericAtomicRMWOp.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "PassDetail.h"
|
|
#include "mlir/Dialect/StandardOps/IR/Ops.h"
|
|
#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Transforms/DialectConversion.h"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
|
|
/// Converts `atomic_rmw` that cannot be lowered to a simple atomic op with
|
|
/// AtomicRMWOpLowering pattern, e.g. with "minf" or "maxf" attributes, to
|
|
/// `generic_atomic_rmw` with the expanded code.
|
|
///
|
|
/// %x = atomic_rmw "maxf" %fval, %F[%i] : (f32, memref<10xf32>) -> f32
|
|
///
|
|
/// will be lowered to
|
|
///
|
|
/// %x = std.generic_atomic_rmw %F[%i] : memref<10xf32> {
|
|
/// ^bb0(%current: f32):
|
|
/// %cmp = cmpf "ogt", %current, %fval : f32
|
|
/// %new_value = select %cmp, %current, %fval : f32
|
|
/// atomic_yield %new_value : f32
|
|
/// }
|
|
struct AtomicRMWOpConverter : public OpRewritePattern<AtomicRMWOp> {
|
|
public:
|
|
using OpRewritePattern::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(AtomicRMWOp op,
|
|
PatternRewriter &rewriter) const final {
|
|
CmpFPredicate predicate;
|
|
switch (op.kind()) {
|
|
case AtomicRMWKind::maxf:
|
|
predicate = CmpFPredicate::OGT;
|
|
break;
|
|
case AtomicRMWKind::minf:
|
|
predicate = CmpFPredicate::OLT;
|
|
break;
|
|
default:
|
|
return failure();
|
|
}
|
|
|
|
auto loc = op.getLoc();
|
|
auto genericOp =
|
|
rewriter.create<GenericAtomicRMWOp>(loc, op.memref(), op.indices());
|
|
OpBuilder bodyBuilder = OpBuilder::atBlockEnd(genericOp.getBody());
|
|
|
|
Value lhs = genericOp.getCurrentValue();
|
|
Value rhs = op.value();
|
|
Value cmp = bodyBuilder.create<CmpFOp>(loc, predicate, lhs, rhs);
|
|
Value select = bodyBuilder.create<SelectOp>(loc, cmp, lhs, rhs);
|
|
bodyBuilder.create<AtomicYieldOp>(loc, select);
|
|
|
|
rewriter.replaceOp(op, genericOp.getResult());
|
|
return success();
|
|
}
|
|
};
|
|
|
|
struct ExpandAtomic : public ExpandAtomicBase<ExpandAtomic> {
|
|
void runOnFunction() override {
|
|
OwningRewritePatternList patterns;
|
|
patterns.insert<AtomicRMWOpConverter>(&getContext());
|
|
|
|
ConversionTarget target(getContext());
|
|
target.addLegalOp<GenericAtomicRMWOp>();
|
|
target.addDynamicallyLegalOp<AtomicRMWOp>([](AtomicRMWOp op) {
|
|
return op.kind() != AtomicRMWKind::maxf &&
|
|
op.kind() != AtomicRMWKind::minf;
|
|
});
|
|
if (failed(mlir::applyPartialConversion(getFunction(), target,
|
|
std::move(patterns))))
|
|
signalPassFailure();
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
std::unique_ptr<Pass> mlir::createExpandAtomicPass() {
|
|
return std::make_unique<ExpandAtomic>();
|
|
}
|