From 719ea6a2a749caf4bc5d5420c405ced2ca044637 Mon Sep 17 00:00:00 2001 From: Nikolaj Bjorner Date: Thu, 13 Feb 2025 21:11:58 -0800 Subject: [PATCH] added ai scripts --- genaisrc/codeupdate.genai.mts | 76 ++++++++++++++++ genaisrc/mergeopt.genai.mts | 55 +----------- genaisrc/myai.genai.mts | 159 ++++++++++++++++++++++++++++++++++ genaisrc/myopt.genai.mts | 29 +------ genaisrc/myopttool.genai.mts | 18 ++++ 5 files changed, 260 insertions(+), 77 deletions(-) create mode 100644 genaisrc/codeupdate.genai.mts create mode 100644 genaisrc/myai.genai.mts create mode 100644 genaisrc/myopttool.genai.mts diff --git a/genaisrc/codeupdate.genai.mts b/genaisrc/codeupdate.genai.mts new file mode 100644 index 000000000..ae74331e9 --- /dev/null +++ b/genaisrc/codeupdate.genai.mts @@ -0,0 +1,76 @@ + +script({ + title: "Invoke LLM code update", +}) + + +async function runCodePrompt(role, message, code) { + const answer = await runPrompt( + (_) => { + _.def("ROLE", role); + _.def("REQUEST", message); + _.def("CODE", code); + _.$`Your role is . + The request is given by + original code: + .` + } + ) + console.log(answer.text); + return answer.text; +} + +async function invokeLLMUpdate(code, inputFile) { + + let role = `You are a highly experienced compiler engineer with over 20 years of expertise, + specializing in C and C++ programming. Your deep knowledge of best coding practices + and software engineering principles enables you to produce robust, efficient, and + maintainable code in any scenario.`; + + let userMessage = `Please modify the original code to ensure that it enforces the following: + - do not use pointer arithmetic for the updates. + - do not introduce uses of std::vector. + - only make replacements that are compatible with the ones listed below. + - add white space between operators: + For example: + i=0 + by + i = 0 + For example + a+b + by + a + b + - remove brackets around single statements: + For example: + { break; } + by + break; + - replaces uses of for loops using begin(), end() iterator patterns by C++21 style for loops + For example replace + for (auto it = x.begin(), end = x.end(); it != end; ++it) + by + for (auto & e : x) + + For example, replace + for (unsigned i = 0; i < a->get_num_args(); ++i) { + expr* arg = a->get_arg(i); + ... + } + by + for (auto arg : *a) { + ... + } + `; + + return runCodePrompt(role, userMessage, code); +} + + +const inputFile = env.files[0]; +const file = await workspace.readText(inputFile); +const answer = await invokeLLMUpdate(file.content, inputFile); +// Extract the code from the answer by removing ```cpp and ```: +let code = answer.replace(/```cpp/g, "").replace(/```/g, ""); +const outputFile = inputFile.filename + ".patch"; +await workspace.writeText(outputFile, code); + diff --git a/genaisrc/mergeopt.genai.mts b/genaisrc/mergeopt.genai.mts index 92b95b205..a68d511bd 100644 --- a/genaisrc/mergeopt.genai.mts +++ b/genaisrc/mergeopt.genai.mts @@ -2,56 +2,7 @@ script({ title: "Merge optimizations function changes for a C++ file" }) -// given a source file -// list files in code_slice directory based on that name with extension opt -// replace functions in src by the corresponding ones in the opt files. -// Save into - - -import * as fs from 'fs'; - - -function get_functions(captures : QueryCapture[], code : string) { - return captures.map(({ name, node }) => ({ - code : node.text, - start : node.startIndex, - end : node.endIndex, - name : node.text.split('(')[0].trim() - })) -} - - +import { mergeFunctions } from "./myai.genai.mts"; const inputFile = env.files[0]; - -const { captures: functions } = await parsers.code( - inputFile, - `(function_definition) @function` -); - - -let funs = get_functions(functions, inputFile.content); - -const modifiedFunctions = "slice_" + path.basename(inputFile.filename) + "*.opt"; - -let inputCode = inputFile.content; -console.log(modifiedFunctions); -const directory_path = path.join("code_slices", modifiedFunctions); -const files = await workspace.findFiles(directory_path); -for (const file of files) { - console.log(file.filename); - const code = file.content.match(/```cpp([\s\S]*?)```/); - if (!code) { - continue; - } - const modifiedFunction = code[1]; - const name = modifiedFunction.split('(')[0].trim(); - console.log(name); - const fun = funs.find(f => f.name === name); - if (fun) { - console.log("Updated function: " + name); - inputCode = inputCode.replace(fun.code, modifiedFunction); - } -} - -console.log(inputCode); -await workspace.writeText(inputFile.filename + ".opt.cpp", inputCode); \ No newline at end of file +let new_code = await mergeFunctions(inputFile); +await workspace.writeText(inputFile.filename + ".opt.cpp", new_code); \ No newline at end of file diff --git a/genaisrc/myai.genai.mts b/genaisrc/myai.genai.mts new file mode 100644 index 000000000..68daeac3c --- /dev/null +++ b/genaisrc/myai.genai.mts @@ -0,0 +1,159 @@ +function function_name_from_code(code: string) { + let name = code.split('(')[0].trim(); + name = name + .replace(/::/g, '_') + .replace(/ /g, '_') + .replace(/\*/g, ''); + return name; +} + +function tree_sitter_get_functions(captures: QueryCapture[], code: string) { + return captures.map(({ name, node }) => ({ + code: node.text, + start: node.startIndex, + end: node.endIndex, + name: function_name_from_code(node.text) + })); +} + +function output_name_from_file(inputFile: WorkspaceFile, name: string) { + let outputFile = "slice_" + path.basename(inputFile.filename) + name + ".cpp"; + return path.join("code_slices", outputFile); +} + +export async function splitFunctions(inputFile: WorkspaceFile) { + const { captures: functions } = await parsers.code( + inputFile, + `(function_definition) @function` + ); + return tree_sitter_get_functions(functions, inputFile.content); +} + + +export async function saveFunctions(inputFile: WorkspaceFile, funs: { code: string, name: string }[]) { + for (const fun of funs) { + let name = function_name_from_code(fun.code); + console.log(name); + let outputFile = output_name_from_file(inputFile, name); + await workspace.writeText(outputFile, `//Extracted ${name} in ${inputFile.filename}\n${fun.code}\n\n`); + } +} + +// given a source file +// list files in code_slice directory based on that name with extension opt +// replace functions in src by the corresponding ones in the opt files. +// Save into + +async function parseOptFunctions(inputFile: WorkspaceFile) { + const modifiedFunctions = "slice_" + path.basename(inputFile.filename) + "*.opt"; + console.log(modifiedFunctions); + const directory_path = path.join("code_slices", modifiedFunctions); + const files = await workspace.findFiles(directory_path); + let modifiedFunctionsList = []; + for (const file of files) { + console.log(file.filename); + const code = file.content.match(/```cpp([\s\S]*?)```/); + if (!code) { + continue; + } + const modifiedFunction = code[1]; + modifiedFunctionsList.push(modifiedFunction); + } + return modifiedFunctionsList; +} + +import * as fs from 'fs'; + +export async function mergeModifiedFunction(code :string, funs : { code: string, name: string }[], new_code : string) { + let name = function_name_from_code(new_code); + let fun = funs.find(f => f.name === name); + if (fun) { + console.log("Updated function: " + name); + code = code.replace(fun.code, new_code); + } + return code; +} + +async function canCompileCode(inputFile : WorkspaceFile, code : string) { + + // move input file to a temp file + // move code to the inputFile.filename + // invoke ninja in the build directory: ninja -b build + // move the temp file back to the original file + // return true iff it succeeded + + let tempFile = inputFile.filename + ".tmp"; + let original_content = inputFile.content; + await workspace.writeText(tempFile, inputFile.content); + await workspace.writeText(inputFile.filename, code); + let result = await host.exec(`cmd /k "C:\Program\ Files/Microsoft\ Visual\ Studio/2022/Enterprise/Common7/Tools/VsDevCmd.bat" -arch=x64 & ninja`, { cwd: "build" }); + + // await fs.delete(tempFile); + if (result.exitCode !== 0) { + await workspace.writeText(inputFile.filename, original_content); + console.log(result.stderr); + return false; + } + return true; +} + +export async function mergeCompileFunction(inputFile : WorkspaceFile, code : string, funs : { code: string, name: string }[], new_code_input : string) { + let match_new_code = new_code_input.match(/```cpp([\s\S]*?)```/); + if (!match_new_code) { + console.log("Invalid new code"); + return code; + } + let new_code = match_new_code[1]; + + let name = function_name_from_code(new_code); + let fun = funs.find(f => f.name == name); + + if (!fun) { + console.log(`Function name '${name}' not found`); + for (const fun of funs) + console.log("'" + fun.name + "'"); + return code; + } + console.log("Updated function: " + name); + let modified_code = code.replace(fun.code, new_code); + if (code == modified_code) { + console.log("No change in function: " + name); + return code; + } + let canCompile = await canCompileCode(inputFile, modified_code); + console.log("Can compile: " + canCompile); + if (canCompile) + return modified_code; + return code; +} + +export async function mergeFunctionsFromList(inputCode: string, funs: { code: string, name: string }[], modifiedFunctionList: string[]) { + let code = inputCode; + for (const new_code of modifiedFunctionList) + code = await mergeModifiedFunction(code, funs, new_code); + return code; +} + +export async function mergeFunctions(inputFile: WorkspaceFile) { + let funs = await splitFunctions(inputFile); + let modifiedFunctionList = await parseOptFunctions(inputFile); + return mergeFunctionsFromList(inputFile.content, funs, modifiedFunctionList); +} + +export async function invokeLLMOpt(code : string) { + const answer = await runPrompt( + (_) => { + _.def("CODE", code); + _.$`You are a highly experienced compiler engineer with over 20 years of expertise, + specializing in C and C++ programming. Your deep knowledge of best coding practices + and software engineering principles enables you to produce robust, efficient, and + maintainable code in any scenario. + + Please modify the original code in to ensure that it uses best practices for optimal code execution.' ` + }, { + system: [], + systemSafety: false + } + ); + return answer.text; +} \ No newline at end of file diff --git a/genaisrc/myopt.genai.mts b/genaisrc/myopt.genai.mts index 53b0ee184..576ffa10a 100644 --- a/genaisrc/myopt.genai.mts +++ b/genaisrc/myopt.genai.mts @@ -3,30 +3,9 @@ script({ files: "code_slices/muz/spacer/orig_spacer_antiunify.cpp_anti_unifier.cpp" }) +import { invokeLLMOpt } from "./myai.genai.mts"; -async function invokeLLMUpdate(code) { - const answer = await runPrompt( - (_) => { - _.def("CODE", code); - _.$` - You are a highly experienced compiler engineer with over 20 years of expertise, - specializing in C and C++ programming. Your deep knowledge of best coding practices - and software engineering principles enables you to produce robust, efficient, and - maintainable code in any scenario. - - Please modify the original code in to ensure that it uses best practices for optimal code execution.' ` - }, { - system: [], - systemSafety: false - } - ) - console.log(answer.text); - return answer.text; -} - - -const inputFile = env.files[0]; -const file = await workspace.readText(inputFile); -const answer = await invokeLLMUpdate(file.content); -const outputFile = inputFile.filename + ".opt"; +let file = env.files[0]; +let answer = await invokeLLMOpt(file); +const outputFile = file.filename + ".opt"; await workspace.writeText(outputFile, answer); diff --git a/genaisrc/myopttool.genai.mts b/genaisrc/myopttool.genai.mts new file mode 100644 index 000000000..5faa52719 --- /dev/null +++ b/genaisrc/myopttool.genai.mts @@ -0,0 +1,18 @@ + +script({ + title: "optimize functions in a file", + files: "src/muz/spacer/spacer_qe_project.cpp" +}) + +import { splitFunctions, invokeLLMOpt, mergeFunctionsFromList, mergeCompileFunction} from "./myai.genai.mts"; + +const inputFile = env.files[0]; +let funs = await splitFunctions(inputFile); +let new_code = inputFile.content; +for (const fun of funs) { + let answer = await invokeLLMOpt(fun.code); + if (answer) + new_code = await mergeCompileFunction(inputFile, new_code, funs, answer); +} + +await workspace.writeText(inputFile.filename + "opt.cpp", new_code);