xref: /llvm-project/mlir/test/lib/Dialect/Transform/TestTransformStateExtension.cpp (revision 7b4a7552719b7720b9c8ccb4bc04a9e6fa1ec0b6)
1 //===- TestTransformStateExtension.cpp ------------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "TestTransformStateExtension.h"
10 
11 using namespace mlir;
12 
13 LogicalResult
updateMapping(Operation * previous,Operation * updated)14 test::TestTransformStateExtension::updateMapping(Operation *previous,
15                                                  Operation *updated) {
16   // Update value handles. The new ops should have at least as many results as
17   // the replacement op. Fewer results are acceptable, if those results are not
18   // mapped to any handle.
19   for (auto r = updated->getNumResults(); r < previous->getNumResults(); ++r) {
20     SmallVector<Value> handles;
21     (void)getTransformState().getHandlesForPayloadValue(previous->getResult(r),
22                                                         handles);
23     if (!handles.empty())
24       return emitError(previous->getLoc())
25              << "cannot replace an op with another op producing fewer results "
26                 "while tracking handles";
27   }
28 
29   for (auto [oldValue, newValue] :
30        llvm::zip(previous->getResults(), updated->getResults()))
31     if (failed(replacePayloadValue(oldValue, newValue)))
32       return failure();
33 
34   // Update op handle.
35   return replacePayloadOp(previous, updated);
36 }
37