Skip to content

Commit d5ee6b6

Browse files
kumarakkumarak
committed
transform: update function matching to use same methods
1 parent 1b28f53 commit d5ee6b6

9 files changed

Lines changed: 254 additions & 53 deletions

‎include/patchestry/Passes/OperationMatcher.hpp‎

Lines changed: 61 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,8 @@ namespace mlir {
2222

2323
namespace cir {
2424
class FuncOp;
25-
}
25+
class CallOp;
26+
} // namespace cir
2627

2728
namespace patchestry::passes {
2829

@@ -37,6 +38,8 @@ namespace patchestry::passes {
3738
class OperationMatcher
3839
{
3940
public:
41+
enum Mode : uint8_t { OPERATION, FUNCTION };
42+
4043
/**
4144
* @brief Checks if an operation matches the given patch specification.
4245
*
@@ -48,7 +51,8 @@ namespace patchestry::passes {
4851
* @param spec The patch specification to match against
4952
* @return true if the operation matches the specification
5053
*/
51-
static bool matches(mlir::Operation *op, cir::FuncOp func, const PatchSpec &spec);
54+
static bool
55+
matches(mlir::Operation *op, cir::FuncOp func, const PatchSpec &spec, Mode mode);
5256

5357
/**
5458
* @brief Checks if an operation name matches the specified operation pattern.
@@ -79,6 +83,17 @@ namespace patchestry::passes {
7983
* @param argument_matches The argument match criteria
8084
* @return true if the arguments match the criteria
8185
*/
86+
static bool matches_arguments( // NOLINT
87+
mlir::Operation *op, const std::vector< ArgumentMatch > &argument_matches
88+
);
89+
90+
/**
91+
* @brief Checks if operation operands match the specified operand patterns.
92+
*
93+
* @param op The operation whose operands to check
94+
* @param operand_matches The operand match criteria
95+
* @return true if the operands match the criteria
96+
*/
8297
static bool matches_operands( // NOLINT
8398
mlir::Operation *op, const std::vector< OperandMatch > &operand_matches
8499
);
@@ -93,10 +108,53 @@ namespace patchestry::passes {
93108
* @param variable_matches The variable match criteria
94109
* @return true if the variables match the criteria
95110
*/
96-
static bool matches_symbols( // NOLINT
111+
static bool matches_variables( // NOLINT
97112
mlir::Operation *op, const std::vector< VariableMatch > &variable_matches
98113
);
99114

115+
/**
116+
* @brief Checks if symbols match the specified symbol patterns.
117+
*
118+
* @param op The operation to check
119+
* @param symbol_matches The symbol match criteria
120+
* @return true if the symbols match the criteria
121+
*/
122+
static bool matches_symbols( // NOLINT
123+
mlir::Operation *op, const std::vector< SymbolMatch > &symbol_matches
124+
);
125+
126+
/**
127+
* @brief Matches operation-based criteria.
128+
*
129+
* @param op The operation to check
130+
* @param func The function containing the operation
131+
* @param match The operation match criteria
132+
* @return true if the operation matches the criteria
133+
*/
134+
static bool matches_operation( // NOLINT
135+
mlir::Operation *op, cir::FuncOp func, const PatchMatch &match
136+
);
137+
138+
/**
139+
* @brief Matches function call-based criteria.
140+
*
141+
* @param op The operation to check (should be a call operation)
142+
* @param func The function containing the call
143+
* @param match The function match criteria
144+
* @return true if the function call matches the criteria
145+
*/
146+
static bool matches_function_call( // NOLINT
147+
mlir::Operation *op, cir::FuncOp func, const PatchMatch &match
148+
);
149+
150+
/**
151+
* @brief Extracts the called function name from a call operation.
152+
*
153+
* @param call_op The call operation
154+
* @return The name of the called function
155+
*/
156+
static std::string extract_callee_name(cir::CallOp call_op); // NOLINT
157+
100158
private:
101159
/**
102160
* @brief Performs pattern matching with support for regex.

‎include/patchestry/Passes/PatchSpec.hpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,8 @@ namespace llvm::yaml {
196196
io.mapOptional("function_context", match.function_context);
197197
io.mapOptional("argument_matches", match.argument_matches);
198198
io.mapOptional("variable_matches", match.variable_matches);
199+
io.mapOptional("symbol_matches", match.symbol_matches);
200+
io.mapOptional("operand_matches", match.operand_matches);
199201

200202
std::string kind_str;
201203
io.mapRequired("kind", kind_str);

‎lib/patchestry/Passes/InstrumentationPass.cpp‎

Lines changed: 35 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -468,41 +468,43 @@ namespace patchestry::passes {
468468
auto callee_name = op.getCallee()->str();
469469
assert(!callee_name.empty() && "Callee name is empty");
470470

471-
std::set< std::string > seen_functions;
471+
auto func = op->getParentOfType< cir::FuncOp >();
472+
if (!func) {
473+
LOG(ERROR) << "Call operation is not in a function. Skipping...\n";
474+
return;
475+
}
476+
472477
for (const auto &spec : config->patches) {
473-
const auto &match = spec.match;
474-
if (match.name != callee_name
475-
|| seen_functions.find(callee_name) != seen_functions.end()
476-
|| exclude_from_patching(func, spec))
478+
if (OperationMatcher::matches(op, func, spec, OperationMatcher::Mode::FUNCTION)
479+
&& !exclude_from_patching(func, spec))
477480
{
478-
continue;
479-
}
480-
481-
seen_functions.emplace(callee_name);
482-
483-
const auto &patch = spec.patch;
484-
// Create module from the patch file mlir representation
485-
auto patch_module = load_patch_module(*op->getContext(), *patch.patch_module);
486-
if (!patch_module) {
487-
LOG(ERROR) << "Failed to load patch module for function: " << callee_name
488-
<< "\n ";
489-
continue;
490-
}
491-
492-
switch (patch.mode) {
493-
case PatchInfoMode::APPLY_BEFORE: {
494-
apply_before_patch(op, match, patch, patch_module.get());
495-
break;
481+
const auto &patch = spec.patch;
482+
const auto &match = spec.match;
483+
// Create module from the patch file mlir representation
484+
auto patch_module =
485+
load_patch_module(*op->getContext(), *patch.patch_module);
486+
if (!patch_module) {
487+
LOG(ERROR)
488+
<< "Failed to load patch module for function: " << callee_name
489+
<< "\n ";
490+
continue;
496491
}
497-
case PatchInfoMode::APPLY_AFTER: {
498-
apply_after_patch(op, match, patch, patch_module.get());
499-
break;
492+
493+
switch (patch.mode) {
494+
case PatchInfoMode::APPLY_BEFORE: {
495+
apply_before_patch(op, match, patch, patch_module.get());
496+
break;
497+
}
498+
case PatchInfoMode::APPLY_AFTER: {
499+
apply_after_patch(op, match, patch, patch_module.get());
500+
break;
501+
}
502+
case PatchInfoMode::REPLACE:
503+
replace_call(op, match, patch, patch_module.get());
504+
break;
505+
default:
506+
break;
500507
}
501-
case PatchInfoMode::REPLACE:
502-
replace_call(op, match, patch, patch_module.get());
503-
break;
504-
default:
505-
break;
506508
}
507509
}
508510
});
@@ -529,7 +531,8 @@ namespace patchestry::passes {
529531
}
530532

531533
for (const auto &spec : config->patches) {
532-
if (OperationMatcher::matches(op, func, spec) && !exclude_from_patching(func, spec))
534+
if (OperationMatcher::matches(op, func, spec, OperationMatcher::Mode::OPERATION)
535+
&& !exclude_from_patching(func, spec))
533536
{
534537
const auto &patch = spec.patch;
535538
const auto &match = spec.match;

‎lib/patchestry/Passes/OperationMatcher.cpp‎

Lines changed: 146 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,28 @@
1919

2020
namespace patchestry::passes {
2121

22-
bool
23-
OperationMatcher::matches(mlir::Operation *op, cir::FuncOp func, const PatchSpec &spec) {
22+
bool OperationMatcher::matches(
23+
mlir::Operation *op, cir::FuncOp func, const PatchSpec &spec,
24+
OperationMatcher::Mode mode // NOLINT
25+
) {
2426
const auto &match = spec.match;
2527

26-
// If kind is not operation, or name is empty, return false
27-
if (match.kind != MatchKind::OPERATION || match.name.empty()) {
28+
// Handle different match kinds
29+
switch (mode) {
30+
case OperationMatcher::Mode::OPERATION:
31+
return matches_operation(op, func, match);
32+
case OperationMatcher::Mode::FUNCTION:
33+
return matches_function_call(op, func, match);
34+
}
35+
36+
return false;
37+
}
38+
39+
bool OperationMatcher::matches_operation(
40+
mlir::Operation *op, cir::FuncOp func, const PatchMatch &match
41+
) {
42+
// If the match kind is not operation, return false
43+
if (match.name.empty() || match.kind != MatchKind::OPERATION) {
2844
return false;
2945
}
3046

@@ -51,6 +67,46 @@ namespace patchestry::passes {
5167
return true;
5268
}
5369

70+
bool OperationMatcher::matches_function_call(
71+
mlir::Operation *op, cir::FuncOp func, const PatchMatch &match
72+
) {
73+
// If the match kind is not function, return false
74+
if (match.name.empty() || match.kind != MatchKind::FUNCTION) {
75+
return false;
76+
}
77+
78+
// For function-based matching, we expect a cir.call operation
79+
auto call_op = mlir::dyn_cast< cir::CallOp >(op);
80+
if (!call_op) {
81+
return false;
82+
}
83+
84+
// Check if the called function name matches
85+
if (!match.name.empty()) {
86+
std::string callee_name = extract_callee_name(call_op);
87+
if (!matches_pattern(callee_name, match.name)) {
88+
return false;
89+
}
90+
}
91+
92+
// Check function context match (the function containing the call)
93+
if (!matches_function_context(func, match.function_context)) {
94+
return false;
95+
}
96+
97+
// Check argument matches for function calls
98+
if (!matches_arguments(op, match.argument_matches)) {
99+
return false;
100+
}
101+
102+
// Check variable matches as one of the arguments
103+
if (!matches_variables(op, match.variable_matches)) {
104+
return false;
105+
}
106+
107+
return true;
108+
}
109+
54110
bool OperationMatcher::matches_operation_name(
55111
mlir::Operation *op, const std::string &operation_pattern
56112
) {
@@ -178,6 +234,92 @@ namespace patchestry::passes {
178234
return false;
179235
}
180236

237+
bool OperationMatcher::matches_arguments(
238+
mlir::Operation *op, const std::vector< ArgumentMatch > &argument_matches
239+
) {
240+
// If no argument matches specified, consider it a match
241+
if (argument_matches.empty()) {
242+
return true;
243+
}
244+
245+
auto operands = op->getOperands();
246+
247+
for (const auto &arg_match : argument_matches) {
248+
// Check if the argument index is valid
249+
if (arg_match.index >= operands.size()) {
250+
return false;
251+
}
252+
253+
auto operand = operands[arg_match.index];
254+
255+
// Check argument name if specified
256+
if (!arg_match.name.empty()) {
257+
std::string var_name = extract_variable_name(op, arg_match.index);
258+
if (!matches_pattern(var_name, arg_match.name)) {
259+
return false;
260+
}
261+
}
262+
263+
// Check argument type if specified
264+
if (!arg_match.type.empty()) {
265+
if (!matches_type(operand.getType(), arg_match.type)) {
266+
return false;
267+
}
268+
}
269+
}
270+
271+
return true;
272+
}
273+
274+
bool OperationMatcher::matches_variables(
275+
mlir::Operation *op, const std::vector< VariableMatch > &variable_matches
276+
) {
277+
// If no variable matches specified, consider it a match
278+
if (variable_matches.empty()) {
279+
return true;
280+
}
281+
282+
// Check operands for variable matches
283+
auto operands = op->getOperands();
284+
for (unsigned i = 0; i < operands.size(); ++i) {
285+
auto operand = operands[i];
286+
std::string var_name = extract_variable_name(op, i);
287+
std::string var_type = type_to_string(operand.getType());
288+
289+
for (const auto &var_match : variable_matches) {
290+
bool name_matches =
291+
var_match.name.empty() || matches_pattern(var_name, var_match.name);
292+
bool type_matches =
293+
var_match.type.empty() || matches_pattern(var_type, var_match.type);
294+
295+
if (name_matches && type_matches) {
296+
return true;
297+
}
298+
}
299+
}
300+
301+
return false;
302+
}
303+
304+
std::string OperationMatcher::extract_callee_name(cir::CallOp call_op) {
305+
// Extract the called function name from the call operation
306+
if (auto callee = call_op.getCalleeAttr()) {
307+
return callee.getLeafReference().str();
308+
}
309+
310+
// If no direct callee attribute, try to extract from operands
311+
auto operands = call_op.getOperands();
312+
if (!operands.empty()) {
313+
auto first_operand = operands[0];
314+
std::string name = extract_ssa_value_name(first_operand);
315+
if (!name.empty()) {
316+
return name;
317+
}
318+
}
319+
320+
return "";
321+
}
322+
181323
bool
182324
OperationMatcher::matches_pattern(const std::string &text, const std::string &pattern) {
183325
// if no pattern, match all

‎test/patchir-transform/measurement_update_after_patch.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ arch: "ARM:LE:32"
22
patches:
33
- name: "apply_after_spo2_lookup"
44
match:
5-
symbol: "spo2_lookup"
5+
name: "spo2_lookup"
66
kind: "function"
77
patch:
88
mode: "ApplyAfter"

‎test/patchir-transform/measurement_update_before_operation.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ arch: "ARM:LE:32"
22
patches:
33
- name: "apply_before_spo2_lookup_operation"
44
match:
5-
operation: "cir.call"
5+
name: "cir.call"
66
kind: "operation"
77
function_context:
88
- name: "/.*measurement.*/" # Functions containing "measurement"

0 commit comments

Comments
 (0)