shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

[mlir] adapt TransformEachOpTrait to parameter values

Browse Source 
Adapt the implementation of TransformEachOpTrait to the existence of
parameter values recently introduced into the transform dialect. In
particular, allow `applyToOne` hooks to return a list containing a mix
of `Operation *` that will be associated with handles and `Attribute`
that will be associated with parameter values by the trait
implementation of the transform interface's `apply` method.

Disentangle the "transposition" of the list of per-payload op partial
results to decrease its overall complexity and detemplatize the code
that doesn't really need templates. This removes the poorly documented
special handling for single-result ops with TransformEachOpTrait that
could have assigned null pointer values to handles.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D140979
This commit is contained in:
Alex Zinenko 2023-01-03 16:01:07 +00:00
parent 054ec47c91
commit 4b455a71b7
18 changed files with 355 additions and 171 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
View File

@ -93,7 +93,7 @@ def EmptyTensorToAllocTensorOp
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::tensor::EmptyOp target,
::llvm::SmallVector<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}

6
mlir/include/mlir/Dialect/GPU/TransformOps/GPUTransformOps.td
View File

@ -109,7 +109,7 @@ def MapNestedForeachToThreads :
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::Operation *target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -130,7 +130,7 @@ def MapForeachToBlocks :
The operation searches top level `scf.foreach_thread` ops under
`gpu_launch` and maps each such op to GPU blocks. Mapping is
one-to-one and the induction variables of `scf.foreach_thread` are
rewritten to gpu.block_id according to the `thread_dim_apping` attribute.
rewritten to gpu.block_id according to the `thread_dim_mapping` attribute.
Dynamic, `scf.foreach_thread` trip counts are currently not supported.
Dynamic block dim sizes are currently not supported.
@ -167,7 +167,7 @@ def MapForeachToBlocks :
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::Operation *target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}

1
mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h
View File

@ -23,6 +23,7 @@ class LinalgOp;
} // namespace linalg
namespace transform {
class TransformHandleTypeInterface;
// Types needed for builders.
struct TileSizesSpec {};
struct NumThreadsSpec {};

22
mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
View File

@ -47,7 +47,7 @@ def DecomposeOp : Op<Transform_Dialect, "structured.decompose",
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::linalg::LinalgOp target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -160,7 +160,7 @@ def GeneralizeOp : Op<Transform_Dialect, "structured.generalize",
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::linalg::LinalgOp target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -202,7 +202,7 @@ def InterchangeOp : Op<Transform_Dialect, "structured.interchange",
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::linalg::GenericOp target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -336,7 +336,7 @@ def MultiTileSizesOp : Op<Transform_Dialect, "structured.multitile_sizes",
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::linalg::LinalgOp target,
::llvm::SmallVector<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
TransformState &state);
}];
}
@ -380,7 +380,7 @@ def PadOp : Op<Transform_Dialect, "structured.pad",
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::linalg::LinalgOp target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -423,7 +423,7 @@ def PromoteOp : Op<Transform_Dialect, "structured.promote",
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::linalg::LinalgOp target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -492,7 +492,7 @@ def ScalarizeOp : Op<Transform_Dialect, "structured.scalarize",
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::linalg::LinalgOp target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -700,7 +700,7 @@ def SplitReductionOp : Op<Transform_Dialect, "structured.split_reduction",
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::linalg::LinalgOp target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -803,7 +803,7 @@ def TileReductionUsingScfOp : Op<Transform_Dialect, "structured.tile_reduction_u
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::linalg::LinalgOp target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -911,7 +911,7 @@ def TileReductionUsingForeachThreadOp :
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::linalg::LinalgOp target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
@ -1228,7 +1228,7 @@ def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::Operation *target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}

2
mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td
View File

@ -44,7 +44,7 @@ def MemRefMultiBufferOp : Op<Transform_Dialect, "memref.multibuffer",
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
memref::AllocOp target,
::llvm::SmallVector<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}

6
mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
View File

@ -109,7 +109,7 @@ def LoopPeelOp : Op<Transform_Dialect, "loop.peel",
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::scf::ForOp target,
::llvm::SmallVector<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -151,7 +151,7 @@ def LoopPipelineOp : Op<Transform_Dialect, "loop.pipeline",
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::scf::ForOp target,
::llvm::SmallVector<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -184,7 +184,7 @@ def LoopUnrollOp : Op<Transform_Dialect, "loop.unroll",
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::Operation *target,
::llvm::SmallVector<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}

224
mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h
View File

@ -542,8 +542,8 @@ public:
};
/// Trait implementing the TransformOpInterface for operations applying a
/// transformation to a single operation handle and producing zero, one or
/// multiple operation handles.
/// transformation to a single operation handle and producing an arbitrary
/// number of handles and parameter values.
/// The op must implement a method with the following signature:
/// - DiagnosedSilenceableFailure applyToOne(OpTy,
/// SmallVector<Operation*> &results, state)
@ -732,7 +732,82 @@ public:
namespace mlir {
namespace transform {
/// A single result of applying a transform op with `ApplyEachOpTrait` to a
/// single payload operation.
using ApplyToEachResult = llvm::PointerUnion<Operation *, Attribute>;
/// A list of results of applying a transform op with `ApplyEachOpTrait` to a
/// single payload operation, co-indexed with the results of the transform op.
class ApplyToEachResultList {
public:
ApplyToEachResultList() = default;
explicit ApplyToEachResultList(unsigned size) : results(size) {}
/// Sets the list of results to `size` null pointers.
void assign(unsigned size, std::nullptr_t) { results.assign(size, nullptr); }
/// Sets the list of results to the given range of values.
template <typename Range>
void assign(Range &&range) {
// This is roughly the implementation of SmallVectorImpl::assign.
// Dispatching to it with map_range and template type inference would result
// in more complex code here.
results.clear();
results.reserve(llvm::size(range));
for (auto element : range) {
if constexpr (std::is_convertible_v<decltype(*std::begin(range)),
Operation *>) {
results.push_back(static_cast<Operation *>(element));
} else {
results.push_back(static_cast<Attribute>(element));
}
}
}
/// Appends an element to the list.
void push_back(Operation *op) { results.push_back(op); }
void push_back(Attribute attr) { results.push_back(attr); }
/// Reserves space for `size` elements in the list.
void reserve(unsigned size) { results.reserve(size); }
/// Iterators over the list.
auto begin() { return results.begin(); }
auto end() { return results.end(); }
auto begin() const { return results.begin(); }
auto end() const { return results.end(); }
/// Returns the number of elements in the list.
size_t size() const { return results.size(); }
/// Element access. Expects the index to be in bounds.
ApplyToEachResult &operator[](size_t index) { return results[index]; }
const ApplyToEachResult &operator[](size_t index) const {
return results[index];
}
private:
/// Underlying storage.
SmallVector<ApplyToEachResult> results;
};
namespace detail {
/// Check that the contents of `partialResult` matches the number, kind (payload
/// op or parameter) and nullity (either all or none) requirements of
/// `transformOp`. Report errors and return failure otherwise.
LogicalResult checkApplyToOne(Operation *transformOp, Location payloadOpLoc,
const ApplyToEachResultList &partialResult);
/// "Transpose" the results produced by individual applications, arranging them
/// per result value of the transform op, and populate `transformResults` with
/// that. The number, kind and nullity of per-application results are assumed to
/// have been verified.
void setApplyToOneResults(Operation *transformOp,
TransformResults &transformResults,
ArrayRef<ApplyToEachResultList> results);
/// Applies a one-to-one or a one-to-many transform to each of the given
/// targets. Puts the results of transforms, if any, in `results` in the same
/// order. Fails if any of the application fails. Individual transforms must be
@ -744,22 +819,28 @@ namespace detail {
/// - a concrete Op class, in which case a check is performed whether
/// `targets` contains operations of the same class and a silenceable failure
/// is reported if it does not.
template <typename FnTy>
DiagnosedSilenceableFailure applyTransformToEach(
Location loc, int expectedNumResults, ArrayRef<Operation *> targets,
SmallVectorImpl<SmallVector<Operation *>> &results, FnTy transform) {
SmallVector<Diagnostic> silenceableStack;
using OpTy = typename llvm::function_traits<FnTy>::template arg_t<0>;
template <typename TransformOpTy>
DiagnosedSilenceableFailure
applyTransformToEach(TransformOpTy transformOp, ArrayRef<Operation *> targets,
SmallVectorImpl<ApplyToEachResultList> &results,
TransformState &state) {
using OpTy = typename llvm::function_traits<
decltype(&TransformOpTy::applyToOne)>::template arg_t<0>;
static_assert(std::is_convertible<OpTy, Operation *>::value,
"expected transform function to take an operation");
SmallVector<Diagnostic> silenceableStack;
unsigned expectedNumResults = transformOp->getNumResults();
for (Operation *target : targets) {
// Emplace back a placeholder for the returned new ops.
// Emplace back a placeholder for the returned new ops and params.
// This is filled with `expectedNumResults` if the op fails to apply.
results.push_back(SmallVector<Operation *>());
ApplyToEachResultList placeholder;
placeholder.reserve(expectedNumResults);
results.push_back(std::move(placeholder));
auto specificOp = dyn_cast<OpTy>(target);
if (!specificOp) {
Diagnostic diag(loc, DiagnosticSeverity::Error);
Diagnostic diag(transformOp->getLoc(), DiagnosticSeverity::Error);
diag << "transform applied to the wrong op kind";
diag.attachNote(target->getLoc()) << "when applied to this op";
// Producing `expectedNumResults` nullptr is a silenceableFailure mode.
@ -770,11 +851,16 @@ DiagnosedSilenceableFailure applyTransformToEach(
continue;
}
DiagnosedSilenceableFailure result = transform(specificOp, results.back());
if (result.isDefiniteFailure())
return result;
if (result.isSilenceableFailure())
result.takeDiagnostics(silenceableStack);
DiagnosedSilenceableFailure res =
transformOp.applyToOne(specificOp, results.back(), state);
if (res.isDefiniteFailure() ||
failed(detail::checkApplyToOne(transformOp, specificOp->getLoc(),
results.back()))) {
return DiagnosedSilenceableFailure::definiteFailure();
}
if (res.isSilenceableFailure())
res.takeDiagnostics(silenceableStack);
}
if (!silenceableStack.empty()) {
return DiagnosedSilenceableFailure::silenceableFailure(
@ -783,23 +869,6 @@ DiagnosedSilenceableFailure applyTransformToEach(
return DiagnosedSilenceableFailure::success();
}
/// Helper function: transpose MxN into NxM; assumes that the input is a valid.
static inline SmallVector<SmallVector<Operation *, 1>>
transposeResults(const SmallVector<SmallVector<Operation *>, 1> &m) {
SmallVector<SmallVector<Operation *, 1>> res;
if (m.empty())
return res;
int64_t rows = m.size(), cols = m[0].size();
for (int64_t j = 0; j < cols; ++j)
res.push_back(SmallVector<Operation *, 1>(rows, nullptr));
for (int64_t i = 0; i < rows; ++i) {
assert(static_cast<int64_t>(m[i].size()) == cols);
for (int64_t j = 0; j < cols; ++j) {
res[j][i] = m[i][j];
}
}
return res;
}
} // namespace detail
} // namespace transform
} // namespace mlir
@ -808,8 +877,6 @@ template <typename OpTy>
mlir::DiagnosedSilenceableFailure
mlir::transform::TransformEachOpTrait<OpTy>::apply(
TransformResults &transformResults, TransformState &state) {
using TransformOpType = typename llvm::function_traits<
decltype(&OpTy::applyToOne)>::template arg_t<0>;
ArrayRef<Operation *> targets =
state.getPayloadOps(this->getOperation()->getOperand(0));
@ -818,88 +885,35 @@ mlir::transform::TransformEachOpTrait<OpTy>::apply(
// propagate gracefully.
// In this case, we fill all results with an empty vector.
if (targets.empty()) {
SmallVector<Operation *> empty;
for (auto r : this->getOperation()->getResults())
transformResults.set(r.template cast<OpResult>(), empty);
SmallVector<Operation *> emptyPayload;
SmallVector<Attribute> emptyParams;
for (OpResult r : this->getOperation()->getResults()) {
if (r.getType().isa<TransformParamTypeInterface>())
transformResults.setParams(r, emptyParams);
else
transformResults.set(r, emptyPayload);
}
return DiagnosedSilenceableFailure::success();
}
// Step 2. Call applyToOne on each target and record newly produced ops in its
// corresponding results entry.
int expectedNumResults = this->getOperation()->getNumResults();
SmallVector<SmallVector<Operation *>, 1> results;
SmallVector<ApplyToEachResultList, 1> results;
results.reserve(targets.size());
DiagnosedSilenceableFailure result = detail::applyTransformToEach(
this->getOperation()->getLoc(), expectedNumResults, targets, results,
[&](TransformOpType specificOp, SmallVector<Operation *> &partialResult) {
auto res = static_cast<OpTy *>(this)->applyToOne(specificOp,
partialResult, state);
if (res.isDefiniteFailure())
return res;
// TODO: encode this implicit must always produce `expectedNumResults`
// and nullptr is fine with a proper trait.
if (static_cast<int>(partialResult.size()) != expectedNumResults) {
auto loc = this->getOperation()->getLoc();
auto diag = mlir::emitError(loc, "applications of ")
<< OpTy::getOperationName() << " expected to produce "
<< expectedNumResults << " results (actually produced "
<< partialResult.size() << ").";
diag.attachNote(loc)
<< "If you need variadic results, consider a generic `apply` "
<< "instead of the specialized `applyToOne`.";
diag.attachNote(loc)
<< "Producing " << expectedNumResults << " null results is "
<< "allowed if the use case warrants it.";
diag.attachNote(specificOp->getLoc()) << "when applied to this op";
return DiagnosedSilenceableFailure::definiteFailure();
}
// Check that all is null or none is null
// TODO: relax this behavior and encode with a proper trait.
if (llvm::any_of(partialResult, [](Operation *op) { return op; }) &&
llvm::any_of(partialResult, [](Operation *op) { return !op; })) {
auto loc = this->getOperation()->getLoc();
auto diag = mlir::emitError(loc, "unexpected application of ")
<< OpTy::getOperationName()
<< " produces both null and non null results.";
diag.attachNote(specificOp->getLoc()) << "when applied to this op";
return DiagnosedSilenceableFailure::definiteFailure();
}
return res;
});
cast<OpTy>(this->getOperation()), targets, results, state);
// Step 3. Propagate the definite failure if any and bail out.
if (result.isDefiniteFailure())
return result;
// Step 4. If there are no results, return early.
if (OpTy::template hasTrait<OpTrait::ZeroResults>())
return result;
// Step 4. "Transpose" the results produced by individual applications,
// arranging them per result value of the transform op. The number, kind and
// nullity of per-application results have been verified by the callback
// above.
detail::setApplyToOneResults(this->getOperation(), transformResults, results);
// Step 5. Perform transposition of M applications producing N results each
// into N results for each of the M applications.
SmallVector<SmallVector<Operation *, 1>> transposedResults =
detail::transposeResults(results);
// Step 6. Single result applies to M ops produces one single M-result.
if (OpTy::template hasTrait<OpTrait::OneResult>()) {
assert(transposedResults.size() == 1 && "Expected single result");
transformResults.set(
this->getOperation()->getResult(0).template cast<OpResult>(),
transposedResults[0]);
// ApplyToOne may have returned silenceableFailure, propagate it.
return result;
}
// Step 7. Filter out empty results and set the transformResults.
for (const auto &it :
llvm::zip(this->getOperation()->getResults(), transposedResults)) {
SmallVector<Operation *, 1> filtered;
llvm::copy_if(std::get<1>(it), std::back_inserter(filtered),
[](Operation *op) { return op; });
transformResults.set(std::get<0>(it).template cast<OpResult>(), filtered);
}
// Step 8. ApplyToOne may have returned silenceableFailure, propagate it.
// Step 5. ApplyToOne may have returned silenceableFailure, propagate it.
return result;
}

2
mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
View File

@ -109,7 +109,7 @@ def CastOp : TransformDialectOp<"cast",
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::Operation *target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}

2
mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
View File

@ -75,7 +75,7 @@ void transform::OneShotBufferizeOp::getEffects(
DiagnosedSilenceableFailure
EmptyTensorToAllocTensorOp::applyToOne(tensor::EmptyOp target,
SmallVector<Operation *> &results,
ApplyToEachResultList &results,
transform::TransformState &state) {
IRRewriter rewriter(target->getContext());
rewriter.setInsertionPoint(target);

20
mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
View File

@ -291,14 +291,14 @@ static void generateGpuBlockIds(RewriterBase &rewriter,
DiagnosedSilenceableFailure
transform::MapForeachToBlocks::applyToOne(Operation *target,
SmallVectorImpl<Operation *> &results,
ApplyToEachResultList &results,
transform::TransformState &state) {
LaunchOp gpuLaunch = dyn_cast<LaunchOp>(target);
TrivialPatternRewriter rewriter(getContext());
auto transformOp = cast<TransformOpInterface>(getOperation());
if (!getGenerateGpuLaunch() && !gpuLaunch) {
results.assign({target});
results.push_back(target);
DiagnosedSilenceableFailure diag =
emitSilenceableError()
<< "Given target is not gpu.launch, set `generate_gpu_launch` "
@ -312,7 +312,7 @@ transform::MapForeachToBlocks::applyToOne(Operation *target,
mlir::transform::gpu::findTopLevelForeachThreadOp(
target, topLevelForeachThreadOp, transformOp);
if (!diag.succeeded()) {
results.assign({target});
results.push_back(target);
diag.attachNote(target->getLoc()) << "when applied to this payload op";
return diag;
}
@ -325,7 +325,7 @@ transform::MapForeachToBlocks::applyToOne(Operation *target,
DiagnosedSilenceableFailure diag =
createGpuLaunch(rewriter, target->getLoc(), transformOp, gpuLaunch);
if (!diag.succeeded()) {
results.assign({target});
results.push_back(target);
return diag;
}
rewriter.setInsertionPointToStart(&gpuLaunch.getBody().front());
@ -352,7 +352,7 @@ transform::MapForeachToBlocks::applyToOne(Operation *target,
gridDim[0], gridDim[1], gridDim[2]);
}
results.assign({gpuLaunch});
results.push_back(gpuLaunch);
return diag;
}
@ -520,14 +520,12 @@ DiagnosedSilenceableFailure mlir::transform::gpu::mapNestedForeachToThreadsImpl(
}
DiagnosedSilenceableFailure transform::MapNestedForeachToThreads::applyToOne(
::mlir::Operation *target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::TransformState &state) {
Operation *target, ApplyToEachResultList &results, TransformState &state) {
LaunchOp gpuLaunch = dyn_cast<LaunchOp>(target);
auto transformOp = cast<TransformOpInterface>(getOperation());
if (!gpuLaunch) {
results.assign({target});
results.push_back(target);
return emitSilenceableError() << "Given target is not gpu.launch";
}
@ -538,7 +536,7 @@ DiagnosedSilenceableFailure transform::MapNestedForeachToThreads::applyToOne(
checkGpuLimits(transformOp, std::nullopt, std::nullopt, std::nullopt,
blockDim[0], blockDim[1], blockDim[2]);
if (diag.isSilenceableFailure()) {
results.assign({target});
results.push_back(target);
diag.attachNote(getLoc()) << getBlockDimAttrName() << " is very large";
return diag;
}
@ -562,7 +560,7 @@ DiagnosedSilenceableFailure transform::MapNestedForeachToThreads::applyToOne(
blockDim[2]);
}
results.assign({gpuLaunch});
results.push_back(gpuLaunch.getOperation());
return diag;
}

33
mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
View File

@ -66,7 +66,7 @@ static FailureOr<LinalgOp> tryApply(Operation *operation, Args &&...args) {
DiagnosedSilenceableFailure
transform::DecomposeOp::applyToOne(linalg::LinalgOp target,
SmallVectorImpl<Operation *> &results,
transform::ApplyToEachResultList &results,
transform::TransformState &state) {
#define DOWNSCALE(trans) \
{ \
@ -577,7 +577,7 @@ transform::FuseIntoContainingOp::apply(transform::TransformResults &results,
DiagnosedSilenceableFailure
transform::GeneralizeOp::applyToOne(linalg::LinalgOp target,
SmallVectorImpl<Operation *> &results,
transform::ApplyToEachResultList &results,
transform::TransformState &state) {
// Exit early if no transformation is needed.
if (isa<GenericOp>(target)) {
@ -599,7 +599,7 @@ transform::GeneralizeOp::applyToOne(linalg::LinalgOp target,
DiagnosedSilenceableFailure
transform::InterchangeOp::applyToOne(linalg::GenericOp target,
SmallVectorImpl<Operation *> &results,
transform::ApplyToEachResultList &results,
transform::TransformState &state) {
ArrayRef<int64_t> interchangeVector = getIteratorInterchange();
// Exit early if no transformation is needed.
@ -708,7 +708,8 @@ transform::MatchOp::apply(transform::TransformResults &results,
//===---------------------------------------------------------------------===//
DiagnosedSilenceableFailure transform::MultiTileSizesOp::applyToOne(
LinalgOp target, SmallVector<Operation *> &results, TransformState &state) {
LinalgOp target, transform::ApplyToEachResultList &results,
TransformState &state) {
OpBuilder builder(target.getContext());
builder.setInsertionPoint(target);
OpFoldResult targetSize = builder.getIndexAttr(getTargetSize());
@ -748,7 +749,7 @@ void transform::MultiTileSizesOp::getEffects(
DiagnosedSilenceableFailure
transform::PadOp::applyToOne(linalg::LinalgOp target,
SmallVectorImpl<Operation *> &results,
transform::ApplyToEachResultList &results,
transform::TransformState &state) {
// Convert the integer packing flags to booleans.
SmallVector<bool> packPaddings;
@ -861,7 +862,7 @@ LogicalResult transform::PadOp::verify() {
DiagnosedSilenceableFailure
transform::PromoteOp::applyToOne(linalg::LinalgOp target,
SmallVectorImpl<Operation *> &results,
transform::ApplyToEachResultList &results,
transform::TransformState &state) {
LinalgPromotionOptions promotionOptions;
if (!getOperandsToPromote().empty())
@ -955,7 +956,7 @@ LogicalResult transform::ReplaceOp::verify() {
DiagnosedSilenceableFailure
transform::ScalarizeOp::applyToOne(linalg::LinalgOp target,
SmallVectorImpl<Operation *> &results,
transform::ApplyToEachResultList &results,
transform::TransformState &state) {
scf::SCFTilingOptions tilingOptions;
tilingOptions.setTileSizeComputationFunction([&](OpBuilder &b, Operation *) {
@ -991,7 +992,10 @@ transform::ScalarizeOp::applyToOne(linalg::LinalgOp target,
rewriter.replaceOp(target, maybeTilingResult->replacements);
else
rewriter.eraseOp(target);
results.append(maybeTilingResult->tiledOps);
results.reserve(maybeTilingResult->tiledOps.size());
for (Operation *tiled : maybeTilingResult->tiledOps)
results.push_back(tiled);
return DiagnosedSilenceableFailure::success();
}
@ -1172,10 +1176,9 @@ void transform::SplitReductionOp::build(
result.addTypes({resultType, resultType, resultType, resultType});
}
DiagnosedSilenceableFailure
transform::SplitReductionOp::applyToOne(linalg::LinalgOp target,
SmallVectorImpl<Operation *> &results,
transform::TransformState &state) {
DiagnosedSilenceableFailure transform::SplitReductionOp::applyToOne(
linalg::LinalgOp target, transform::ApplyToEachResultList &results,
transform::TransformState &state) {
ControlSplitReductionFn splitFn = [&](LinalgOp) {
return linalg::SplitReductionOptions{int64_t(getSplitFactor()),
unsigned(getInsertSplitDimension()),
@ -1219,7 +1222,7 @@ void transform::TileReductionUsingScfOp::build(
}
DiagnosedSilenceableFailure transform::TileReductionUsingScfOp::applyToOne(
linalg::LinalgOp target, SmallVectorImpl<Operation *> &results,
linalg::LinalgOp target, transform::ApplyToEachResultList &results,
transform::TransformState &state) {
TrivialPatternRewriter rewriter(getContext());
rewriter.setInsertionPoint(target);
@ -1263,7 +1266,7 @@ void transform::TileReductionUsingForeachThreadOp::build(
DiagnosedSilenceableFailure
transform::TileReductionUsingForeachThreadOp::applyToOne(
linalg::LinalgOp target, SmallVectorImpl<Operation *> &results,
linalg::LinalgOp target, transform::ApplyToEachResultList &results,
transform::TransformState &state) {
TrivialPatternRewriter rewriter(getContext());
rewriter.setInsertionPoint(target);
@ -1952,7 +1955,7 @@ private:
DiagnosedSilenceableFailure
transform::VectorizeOp::applyToOne(Operation *target,
SmallVectorImpl<Operation *> &results,
transform::ApplyToEachResultList &results,
transform::TransformState &state) {
if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
auto diag = this->emitOpError("requires isolated-from-above targets");

7
mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp
View File

@ -21,10 +21,9 @@ using namespace mlir;
// MemRefMultiBufferOp
//===----------------------------------------------------------------------===//
DiagnosedSilenceableFailure
transform::MemRefMultiBufferOp::applyToOne(memref::AllocOp target,
SmallVector<Operation *> &results,
transform::TransformState &state) {
DiagnosedSilenceableFailure transform::MemRefMultiBufferOp::applyToOne(
memref::AllocOp target, transform::ApplyToEachResultList &results,
transform::TransformState &state) {
auto newBuffer = memref::multiBuffer(target, getFactor());
if (failed(newBuffer)) {
Diagnostic diag(target->getLoc(), DiagnosticSeverity::Note);

6
mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
View File

@ -125,7 +125,7 @@ transform::LoopOutlineOp::apply(transform::TransformResults &results,
DiagnosedSilenceableFailure
transform::LoopPeelOp::applyToOne(scf::ForOp target,
SmallVector<Operation *> &results,
transform::ApplyToEachResultList &results,
transform::TransformState &state) {
scf::ForOp result;
IRRewriter rewriter(target->getContext());
@ -182,7 +182,7 @@ loopScheduling(scf::ForOp forOp,
DiagnosedSilenceableFailure
transform::LoopPipelineOp::applyToOne(scf::ForOp target,
SmallVector<Operation *> &results,
transform::ApplyToEachResultList &results,
transform::TransformState &state) {
scf::PipeliningOption options;
options.getScheduleFn =
@ -210,7 +210,7 @@ transform::LoopPipelineOp::applyToOne(scf::ForOp target,
DiagnosedSilenceableFailure
transform::LoopUnrollOp::applyToOne(Operation *op,
SmallVector<Operation *> &results,
transform::ApplyToEachResultList &results,
transform::TransformState &state) {
LogicalResult result(failure());
if (scf::ForOp scfFor = dyn_cast<scf::ForOp>(op))

87
mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp
View File

@ -407,6 +407,93 @@ bool transform::TransformResults::isParam(unsigned resultNumber) const {
return paramSegments[resultNumber].data() != nullptr;
}
//===----------------------------------------------------------------------===//
// Utilities for TransformEachOpTrait.
//===----------------------------------------------------------------------===//
LogicalResult
transform::detail::checkApplyToOne(Operation *transformOp,
Location payloadOpLoc,
const ApplyToEachResultList &partialResult) {
Location transformOpLoc = transformOp->getLoc();
StringRef transformOpName = transformOp->getName().getStringRef();
unsigned expectedNumResults = transformOp->getNumResults();
// TODO: encode this implicit must always produce `expectedNumResults`
// and nullptr is fine with a proper trait.
if (partialResult.size() != expectedNumResults) {
auto diag = mlir::emitError(transformOpLoc, "applications of ")
<< transformOpName << " expected to produce "
<< expectedNumResults << " results (actually produced "
<< partialResult.size() << ").";
diag.attachNote(transformOpLoc)
<< "If you need variadic results, consider a generic `apply` "
<< "instead of the specialized `applyToOne`.";
diag.attachNote(transformOpLoc)
<< "Producing " << expectedNumResults << " null results is "
<< "allowed if the use case warrants it.";
diag.attachNote(payloadOpLoc) << "when applied to this op";
return failure();
}
// Check that all is null or none is null
// TODO: relax this behavior and encode with a proper trait.
if (llvm::any_of(
partialResult,
[](llvm::PointerUnion<Operation *, Attribute> ptr) { return ptr; }) &&
llvm::any_of(partialResult,
[](llvm::PointerUnion<Operation *, Attribute> ptr) {
return !ptr;
})) {
auto diag = mlir::emitError(transformOpLoc, "unexpected application of ")
<< transformOpName
<< " produces both null and non null results.";
diag.attachNote(payloadOpLoc) << "when applied to this op";
return failure();
}
// Check that the right kind of value was produced.
for (const auto &[ptr, res] :
llvm::zip(partialResult, transformOp->getResults())) {
if (ptr.is<Operation *>() &&
!res.getType().template isa<TransformHandleTypeInterface>()) {
mlir::emitError(transformOpLoc)
<< "applications of " << transformOpName
<< " expected to produce an Attribute for result #"
<< res.getResultNumber();
return failure();
}
if (ptr.is<Attribute>() &&
!res.getType().template isa<TransformParamTypeInterface>()) {
mlir::emitError(transformOpLoc)
<< "applications of " << transformOpName
<< " expected to produce an Operation * for result #"
<< res.getResultNumber();
return failure();
}
}
return success();
}
void transform::detail::setApplyToOneResults(
Operation *transformOp, TransformResults &transformResults,
ArrayRef<ApplyToEachResultList> results) {
for (OpResult r : transformOp->getResults()) {
if (r.getType().isa<TransformParamTypeInterface>()) {
auto params = llvm::to_vector(
llvm::map_range(results, [r](const ApplyToEachResultList &oneResult) {
return oneResult[r.getResultNumber()].get<Attribute>();
}));
transformResults.setParams(r, params);
} else {
auto payloads = llvm::to_vector(
llvm::map_range(results, [r](const ApplyToEachResultList &oneResult) {
return oneResult[r.getResultNumber()].get<Operation *>();
}));
transformResults.set(r, payloads);
}
}
}
//===----------------------------------------------------------------------===//
// Utilities for PossibleTopLevelTransformOpTrait.
//===----------------------------------------------------------------------===//

3
mlir/lib/Dialect/Transform/IR/TransformOps.cpp
View File

@ -261,8 +261,7 @@ LogicalResult transform::AlternativesOp::verify() {
//===----------------------------------------------------------------------===//
DiagnosedSilenceableFailure
transform::CastOp::applyToOne(Operation *target,
SmallVectorImpl<Operation *> &results,
transform::CastOp::applyToOne(Operation *target, ApplyToEachResultList &results,
transform::TransformState &state) {
results.push_back(target);
return DiagnosedSilenceableFailure::success();

33
mlir/test/Dialect/Transform/test-interpreter.mlir
View File

@ -991,3 +991,36 @@ func.func private @three_test_ops(%arg0: i32) {
"test.op_a"(%arg0) { attr = 0 : i32} : (i32) -> i32
return
}
// -----
transform.sequence failures(propagate) {
^bb0(%arg0: !transform.any_op):
// expected-error @below {{expected to produce an Operation * for result #0}}
transform.test_produce_transform_param_or_forward_operand %arg0
{ first_result_is_param }
: (!transform.any_op) -> (!transform.any_op, !transform.param<i64>)
}
// -----
// expected-note @below {{when applied to this op}}
module {
transform.sequence failures(propagate) {
^bb0(%arg0: !transform.any_op):
// expected-error @below {{produces both null and non null results}}
transform.test_produce_transform_param_or_forward_operand %arg0
{ first_result_is_null }
: (!transform.any_op) -> (!transform.any_op, !transform.param<i64>)
}
}
// -----
transform.sequence failures(propagate) {
^bb0(%arg0: !transform.any_op):
// expected-error @below {{expected to produce an Attribute for result #1}}
transform.test_produce_transform_param_or_forward_operand %arg0
{ second_result_is_handle }
: (!transform.any_op) -> (!transform.any_op, !transform.param<i64>)
}

39
mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp
View File

@ -243,7 +243,7 @@ DiagnosedSilenceableFailure mlir::test::TestEmitRemarkAndEraseOperandOp::apply(
}
DiagnosedSilenceableFailure mlir::test::TestWrongNumberOfResultsOp::applyToOne(
Operation *target, SmallVectorImpl<Operation *> &results,
Operation *target, transform::ApplyToEachResultList &results,
transform::TransformState &state) {
OperationState opState(target->getLoc(), "foo");
results.push_back(OpBuilder(target).create(opState));
@ -252,7 +252,7 @@ DiagnosedSilenceableFailure mlir::test::TestWrongNumberOfResultsOp::applyToOne(
DiagnosedSilenceableFailure
mlir::test::TestWrongNumberOfMultiResultsOp::applyToOne(
Operation *target, SmallVectorImpl<Operation *> &results,
Operation *target, transform::ApplyToEachResultList &results,
transform::TransformState &state) {
static int count = 0;
if (count++ == 0) {
@ -264,7 +264,7 @@ mlir::test::TestWrongNumberOfMultiResultsOp::applyToOne(
DiagnosedSilenceableFailure
mlir::test::TestCorrectNumberOfMultiResultsOp::applyToOne(
Operation *target, SmallVectorImpl<Operation *> &results,
Operation *target, transform::ApplyToEachResultList &results,
transform::TransformState &state) {
OperationState opState(target->getLoc(), "foo");
results.push_back(OpBuilder(target).create(opState));
@ -274,7 +274,7 @@ mlir::test::TestCorrectNumberOfMultiResultsOp::applyToOne(
DiagnosedSilenceableFailure
mlir::test::TestMixedNullAndNonNullResultsOp::applyToOne(
Operation *target, SmallVectorImpl<Operation *> &results,
Operation *target, transform::ApplyToEachResultList &results,
transform::TransformState &state) {
OperationState opState(target->getLoc(), "foo");
results.push_back(nullptr);
@ -284,7 +284,7 @@ mlir::test::TestMixedNullAndNonNullResultsOp::applyToOne(
DiagnosedSilenceableFailure
mlir::test::TestMixedSuccessAndSilenceableOp::applyToOne(
Operation *target, SmallVectorImpl<Operation *> &results,
Operation *target, transform::ApplyToEachResultList &results,
transform::TransformState &state) {
if (target->hasAttr("target_me"))
return DiagnosedSilenceableFailure::success();
@ -429,6 +429,35 @@ LogicalResult mlir::test::TestProduceIntegerParamWithTypeOp::verify() {
return success();
}
void mlir::test::TestProduceTransformParamOrForwardOperandOp::getEffects(
SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
transform::onlyReadsHandle(getIn(), effects);
transform::producesHandle(getOut(), effects);
transform::producesHandle(getParam(), effects);
}
DiagnosedSilenceableFailure
mlir::test::TestProduceTransformParamOrForwardOperandOp::applyToOne(
Operation *target, ::transform::ApplyToEachResultList &results,
::transform::TransformState &state) {
Builder builder(getContext());
if (getFirstResultIsParam()) {
results.push_back(builder.getI64IntegerAttr(0));
} else if (getFirstResultIsNull()) {
results.push_back(nullptr);
} else {
results.push_back(state.getPayloadOps(getIn()).front());
}
if (getSecondResultIsHandle()) {
results.push_back(state.getPayloadOps(getIn()).front());
} else {
results.push_back(builder.getI64IntegerAttr(42));
}
return DiagnosedSilenceableFailure::success();
}
namespace {
/// Test extension of the Transform dialect. Registers additional ops and
/// declares PDL as dependent dialect since the additional ops are using PDL

31
mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.td
View File

@ -173,7 +173,7 @@ def TestWrongNumberOfResultsOp
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::Operation * target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -189,7 +189,7 @@ def TestWrongNumberOfMultiResultsOp
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::Operation * target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -206,7 +206,7 @@ def TestCorrectNumberOfMultiResultsOp
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::Operation * target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -223,7 +223,7 @@ def TestMixedNullAndNonNullResultsOp
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::Operation * target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -239,7 +239,7 @@ def TestMixedSuccessAndSilenceableOp
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::Operation * target,
::llvm::SmallVectorImpl<::mlir::Operation *> &results,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
@ -313,4 +313,25 @@ def TestProduceIntegerParamWithTypeOp
let hasVerifier = 1;
}
def TestProduceTransformParamOrForwardOperandOp
: Op<Transform_Dialect, "test_produce_transform_param_or_forward_operand",
[DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
TransformEachOpTrait, TransformOpInterface]> {
let arguments = (ins TransformHandleTypeInterface:$in,
UnitAttr:$first_result_is_param,
UnitAttr:$first_result_is_null,
UnitAttr:$second_result_is_handle);
let results = (outs TransformHandleTypeInterface:$out,
TransformParamTypeInterface:$param);
let assemblyFormat = "$in attr-dict `:` functional-type(operands, results)";
let cppNamespace = "::mlir::test";
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::Operation *target,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
#endif // MLIR_TESTTRANSFORMDIALECTEXTENSION_TD