Skip to content

Commit 96aee79

Browse files
kumarakkumarak
committed
transform: update argument specification for handling structured format
1 parent 3044e3e commit 96aee79

7 files changed

Lines changed: 283 additions & 55 deletions

File tree

‎docs/GettingStarted/patch_specifications.md‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,7 @@ arguments:
191191
| `argument` | Function call argument by position | `index` | Access arguments of matched function calls |
192192
| `variable` | Local variable or symbol by name | `symbol` | Access local variables in scope |
193193
| `constant` | Literal constant value | `value` | Pass fixed values to patch functions |
194+
| `return_value` | Return value of function or operation | None | Access return value (for ApplyAfter mode) |
194195

195196
### Argument Examples
196197

@@ -224,6 +225,17 @@ arguments:
224225
value: "4294967295"
225226
```
226227

228+
#### Return Value Handling
229+
```yaml
230+
# Access function return value (ApplyAfter mode)
231+
arguments:
232+
- name: "result"
233+
source: "return_value"
234+
- name: "success_code"
235+
source: "constant"
236+
value: "0"
237+
```
238+
227239
#### Mixed Argument Types
228240
```yaml
229241
# Comprehensive validation

‎include/patchestry/Passes/PatchSpec.hpp‎

Lines changed: 51 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,23 @@ namespace patchestry::passes {
2727

2828
enum class MatchKind : uint8_t { NONE = 0, OPERATION, FUNCTION };
2929

30+
enum class ArgumentSourceType : uint8_t {
31+
OPERAND = 0, // Reference to operation operand by index
32+
VARIABLE, // Reference to variable by name
33+
SYMBOL, // Reference to symbol by name
34+
CONSTANT, // Literal constant value
35+
RETURN_VALUE // Return value of function or operation
36+
};
37+
38+
struct ArgumentSource
39+
{
40+
ArgumentSourceType source;
41+
std::string name; // Descriptive name for the argument
42+
std::optional< unsigned > index; // Operand/argument index (required for OPERAND type)
43+
std::optional< std::string > symbol; // Symbol name (required for VARIABLE/SYMBOL type)
44+
std::optional< std::string > value; // Constant value (required for CONSTANT type)
45+
};
46+
3047
struct ArgumentMatch
3148
{
3249
unsigned index;
@@ -68,7 +85,8 @@ namespace patchestry::passes {
6885
std::string patch_file;
6986
std::string patch_function;
7087
std::optional< std::string > patch_module;
71-
std::vector< std::string > arguments;
88+
std::vector< ArgumentSource >
89+
argument_sources; // New: structured argument specifications
7290
};
7391

7492
struct PatchSpec
@@ -122,6 +140,7 @@ namespace patchestry::passes {
122140
LLVM_YAML_IS_SEQUENCE_VECTOR(patchestry::passes::PatchSpec)
123141
LLVM_YAML_IS_SEQUENCE_VECTOR(patchestry::passes::VariableMatch)
124142
LLVM_YAML_IS_SEQUENCE_VECTOR(patchestry::passes::ArgumentMatch)
143+
LLVM_YAML_IS_SEQUENCE_VECTOR(patchestry::passes::ArgumentSource)
125144
LLVM_YAML_IS_SEQUENCE_VECTOR(patchestry::passes::FunctionContext)
126145

127146
class PatchSpecContext
@@ -156,6 +175,36 @@ class PatchSpecContext
156175
};
157176

158177
namespace llvm::yaml {
178+
// Parse ArgumentSource
179+
template<>
180+
struct MappingTraits< patchestry::passes::ArgumentSource >
181+
{
182+
static void mapping(IO &io, patchestry::passes::ArgumentSource &arg) {
183+
std::string source_str;
184+
io.mapRequired("source", source_str);
185+
186+
if (source_str == "operand") {
187+
arg.source = patchestry::passes::ArgumentSourceType::OPERAND;
188+
} else if (source_str == "argument") {
189+
arg.source = patchestry::passes::ArgumentSourceType::OPERAND; // Treat argument
190+
// same as operand
191+
} else if (source_str == "variable") {
192+
arg.source = patchestry::passes::ArgumentSourceType::VARIABLE;
193+
} else if (source_str == "symbol") {
194+
arg.source = patchestry::passes::ArgumentSourceType::SYMBOL;
195+
} else if (source_str == "constant") {
196+
arg.source = patchestry::passes::ArgumentSourceType::CONSTANT;
197+
} else if (source_str == "return_value") {
198+
arg.source = patchestry::passes::ArgumentSourceType::RETURN_VALUE;
199+
}
200+
201+
io.mapRequired("name", arg.name);
202+
io.mapOptional("index", arg.index);
203+
io.mapOptional("symbol", arg.symbol);
204+
io.mapOptional("value", arg.value);
205+
}
206+
};
207+
159208
// Prase ArgumentMatch
160209
template<>
161210
struct MappingTraits< patchestry::passes::ArgumentMatch >
@@ -219,7 +268,7 @@ namespace llvm::yaml {
219268
io.mapOptional("code", patch.code);
220269
io.mapOptional("patch_file", patch.patch_file);
221270
io.mapOptional("patch_function", patch.patch_function);
222-
io.mapOptional("arguments", patch.arguments);
271+
io.mapOptional("arguments", patch.argument_sources);
223272

224273
std::string mode_str;
225274
io.mapRequired("mode", mode_str);

‎lib/patchestry/Passes/InstrumentationPass.cpp‎

Lines changed: 207 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -519,7 +519,7 @@ namespace patchestry::passes {
519519
*
520520
* @param op The operation to be instrumented
521521
*/
522-
void InstrumentationPass::instrument_operation(mlir::Operation * op) {
522+
void InstrumentationPass::instrument_operation(mlir::Operation *op) {
523523
if (!config || config->patches.empty()) {
524524
LOG(ERROR) << "No patch configuration found. Skipping...\n";
525525
return;
@@ -560,7 +560,8 @@ namespace patchestry::passes {
560560

561561
/**
562562
* @brief Prepares the arguments for a function call based on the patch information.
563-
* This function handles argument type casting and argument matching.
563+
* This function handles argument type casting and argument matching using
564+
* the new structured ArgumentSource specifications.
564565
*
565566
* @param builder The MLIR operation builder.
566567
* @param op The call operation to be instrumented.
@@ -572,41 +573,6 @@ namespace patchestry::passes {
572573
mlir::OpBuilder &builder, mlir::Operation *call_op, cir::FuncOp patch_func,
573574
const PatchInfo &patch, llvm::SmallVector< mlir::Value > &args
574575
) {
575-
if (call_op->getNumOperands() == 0) {
576-
assert(patch.arguments.empty() && "Patch arguments are not empty");
577-
assert(
578-
patch_func.getNumArguments() == 0 && "Patch function arguments are not empty"
579-
);
580-
return;
581-
}
582-
583-
if (patch_func.getNumArguments() > call_op->getNumOperands()) {
584-
LOG(ERROR) << "Number of arguments in patch is greater than the number of "
585-
"arguments in the call operation\n";
586-
return;
587-
}
588-
589-
// Check if patch function argument is taking return value
590-
if (patch.arguments.size() == 1 && patch.arguments[0] == "return_value") {
591-
if (call_op->getResultTypes().size() == 0) {
592-
LOG(ERROR) << "Call operation does not have a return value\n";
593-
return;
594-
}
595-
auto patch_argument_type = patch_func.getArguments().front().getType();
596-
auto call_return_type = call_op->getResultTypes().front();
597-
if (patch_argument_type != call_return_type) {
598-
auto cast_op = builder.create< cir::CastOp >(
599-
call_op->getLoc(), patch_argument_type,
600-
getCastKind(call_return_type, patch_argument_type),
601-
call_op->getResults().front()
602-
);
603-
args.append(cast_op->getResults().begin(), cast_op->getResults().end());
604-
return;
605-
}
606-
args.append(call_op->getResults().begin(), call_op->getResults().end());
607-
return;
608-
}
609-
610576
auto create_cast = [&](mlir::Value value, mlir::Type type) -> mlir::Value {
611577
if (value.getType() == type) {
612578
return value;
@@ -618,18 +584,212 @@ namespace patchestry::passes {
618584
return cast_op->getResults().front();
619585
};
620586

621-
if (auto orig_call_op = mlir::dyn_cast< cir::CallOp >(call_op)) {
622-
for (unsigned i = 0; i < patch.arguments.size(); i++) {
623-
auto patch_arg_type = patch_func.getArgumentTypes()[i];
624-
args.push_back(create_cast(orig_call_op.getArgOperands()[i], patch_arg_type));
587+
// Handle structured argument specifications
588+
for (size_t i = 0;
589+
i < patch.argument_sources.size() && i < patch_func.getNumArguments(); ++i)
590+
{
591+
const auto &arg_spec = patch.argument_sources[i];
592+
auto patch_arg_type = patch_func.getArgumentTypes()[i];
593+
mlir::Value arg_value;
594+
595+
switch (arg_spec.source) {
596+
case ArgumentSourceType::OPERAND: {
597+
// Get operand by index
598+
if (!arg_spec.index.has_value()) {
599+
LOG(ERROR) << "OPERAND source requires index field\n";
600+
continue;
601+
}
602+
unsigned idx = arg_spec.index.value();
603+
604+
if (auto orig_call_op = mlir::dyn_cast< cir::CallOp >(call_op)) {
605+
if (idx >= orig_call_op.getArgOperands().size()) {
606+
LOG(ERROR) << "Operand index " << idx << " out of range\n";
607+
continue;
608+
}
609+
arg_value = orig_call_op.getArgOperands()[idx];
610+
} else {
611+
if (idx >= call_op->getNumOperands()) {
612+
LOG(ERROR) << "Operand index " << idx << " out of range\n";
613+
continue;
614+
}
615+
arg_value = call_op->getOperand(idx);
616+
}
617+
break;
618+
}
619+
case ArgumentSourceType::VARIABLE: {
620+
// Handle local variables only
621+
if (!arg_spec.symbol.has_value()) {
622+
LOG(ERROR) << "VARIABLE source requires symbol field\n";
623+
continue;
624+
}
625+
626+
const std::string &var_name = arg_spec.symbol.value();
627+
mlir::Value var_value;
628+
bool found = false;
629+
630+
// Look for local variables in function scope only
631+
auto func = call_op->getParentOfType< cir::FuncOp >();
632+
if (!func) {
633+
LOG(ERROR) << "Cannot find parent function for local variable lookup\n";
634+
continue;
635+
}
636+
637+
// Search for local variables in function scope
638+
func.walk([&](mlir::Operation *op) {
639+
if (auto alloca_op = mlir::dyn_cast< cir::AllocaOp >(op)) {
640+
if (auto name_attr = op->getAttrOfType< mlir::StringAttr >("name"))
641+
{
642+
if (name_attr.getValue() == var_name) {
643+
var_value = alloca_op.getResult();
644+
found = true;
645+
return mlir::WalkResult::interrupt();
646+
}
647+
}
648+
}
649+
return mlir::WalkResult::advance();
650+
});
651+
652+
if (!found) {
653+
LOG(WARNING) << "Local variable '" << var_name << "' not found\n";
654+
continue;
655+
}
656+
arg_value = builder.create< cir::LoadOp >(call_op->getLoc(), var_value);
657+
break;
658+
}
659+
case ArgumentSourceType::SYMBOL: {
660+
// Handle global variables, functions, and any symbol in symbol table
661+
if (!arg_spec.symbol.has_value()) {
662+
LOG(ERROR) << "SYMBOL source requires symbol field\n";
663+
continue;
664+
}
665+
666+
const std::string &symbol_name = arg_spec.symbol.value();
667+
mlir::Value symbol_value;
668+
bool found = false;
669+
670+
auto module = call_op->getParentOfType< mlir::ModuleOp >();
671+
if (!module) {
672+
LOG(ERROR) << "Cannot find parent module for symbol lookup\n";
673+
continue;
674+
}
675+
676+
// Look for global variables
677+
if (auto global_op = module.lookupSymbol< cir::GlobalOp >(symbol_name)) {
678+
// Create a GetGlobal operation to access the global variable
679+
auto global_type = global_op.getSymType();
680+
if (auto global_ptr_type =
681+
mlir::dyn_cast< cir::PointerType >(global_type))
682+
{
683+
symbol_value = builder.create< cir::GetGlobalOp >(
684+
call_op->getLoc(), global_ptr_type, symbol_name
685+
);
686+
found = true;
687+
} else {
688+
// For non-pointer globals, create a pointer type
689+
auto ptr_type =
690+
cir::PointerType::get(builder.getContext(), global_type);
691+
symbol_value = builder.create< cir::GetGlobalOp >(
692+
call_op->getLoc(), ptr_type, symbol_name
693+
);
694+
found = true;
695+
}
696+
}
697+
698+
// Look for functions
699+
if (!found) {
700+
if (auto func_op = module.lookupSymbol< cir::FuncOp >(symbol_name)) {
701+
// Create a function reference
702+
auto func_type = func_op.getFunctionType();
703+
auto func_ptr_type =
704+
cir::PointerType::get(builder.getContext(), func_type);
705+
auto symbol_ref =
706+
mlir::FlatSymbolRefAttr::get(builder.getContext(), symbol_name);
707+
708+
// Create a constant operation for the function pointer
709+
symbol_value = builder.create< cir::GetGlobalOp >(
710+
call_op->getLoc(), func_ptr_type, symbol_ref
711+
);
712+
found = true;
713+
}
714+
}
715+
716+
if (!found) {
717+
LOG(WARNING)
718+
<< "Symbol '" << symbol_name << "' not found in symbol table\n";
719+
continue;
720+
}
721+
arg_value = symbol_value;
722+
break;
723+
}
724+
case ArgumentSourceType::RETURN_VALUE: {
725+
// Handle return value of function or operation
726+
if (call_op->getNumResults() == 0) {
727+
LOG(ERROR) << "Operation/function does not have a return value\n";
728+
continue;
729+
}
730+
731+
// Get the first result (most common case)
732+
arg_value = call_op->getResult(0);
733+
break;
734+
}
735+
case ArgumentSourceType::CONSTANT: {
736+
// Create constant value
737+
if (!arg_spec.value.has_value()) {
738+
LOG(ERROR) << "CONSTANT source requires value field\n";
739+
continue;
740+
}
741+
742+
const std::string &const_value = arg_spec.value.value();
743+
744+
// Parse constant based on patch function argument type
745+
if (auto int_type = mlir::dyn_cast< cir::IntType >(patch_arg_type)) {
746+
try {
747+
// Parse integer constant
748+
int64_t int_val =
749+
std::stoll(const_value, nullptr, 0); // Support hex, oct, dec
750+
auto attr = mlir::IntegerAttr::get(patch_arg_type, int_val);
751+
arg_value = builder.create< cir::ConstantOp >(
752+
call_op->getLoc(), patch_arg_type, attr
753+
);
754+
} catch (const std::exception &e) {
755+
LOG(ERROR) << "Failed to parse integer constant '" << const_value
756+
<< "': " << e.what() << "\n";
757+
continue;
758+
}
759+
} else if (auto ptr_type =
760+
mlir::dyn_cast< cir::PointerType >(patch_arg_type))
761+
{
762+
try {
763+
// Parse pointer constant (usually hex address)
764+
uint64_t ptr_val = std::stoull(const_value, nullptr, 0);
765+
auto int_type = cir::IntType::get(builder.getContext(), 64, false);
766+
auto int_attr = mlir::IntegerAttr::get(
767+
int_type, static_cast< int64_t >(ptr_val)
768+
);
769+
auto int_const = builder.create< cir::ConstantOp >(
770+
call_op->getLoc(), int_type, int_attr
771+
);
772+
arg_value = builder.create< cir::CastOp >(
773+
call_op->getLoc(), patch_arg_type, cir::CastKind::int_to_ptr,
774+
int_const
775+
);
776+
} catch (const std::exception &e) {
777+
LOG(ERROR) << "Failed to parse pointer constant '" << const_value
778+
<< "': " << e.what() << "\n";
779+
continue;
780+
}
781+
} else {
782+
LOG(ERROR)
783+
<< "Unsupported constant type for value '" << const_value << "'\n";
784+
continue;
785+
}
786+
break;
787+
}
625788
}
626-
return;
627-
}
628789

629-
// llvm::SmallVector< mlir::Value, 4 > argument_vec;
630-
for (unsigned i = 0; i < patch.arguments.size(); i++) {
631-
auto patch_arg_type = patch_func.getArgumentTypes()[i];
632-
args.push_back(create_cast(call_op->getOperands()[i], patch_arg_type));
790+
if (arg_value) {
791+
args.push_back(create_cast(arg_value, patch_arg_type));
792+
}
633793
}
634794
}
635795

0 commit comments

Comments
 (0)