402 lines · python
1# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.2# See https://llvm.org/LICENSE.txt for license information.3# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception4"""IR2Vec/MIR2Vec Triplet Generator5 6Generates IR2Vec or MIR2Vec triplets by applying random optimization levels to7LLVM IR files (or processing MIR files) and extracting triplets using llvm-ir2vec.8Automatically generates preprocessed files (entity2id.txt, relation2id.txt, and9train2id.txt) necessary for training IR2Vec or MIR2Vec vocabularies.10 11Usage:12 For LLVM IR:13 python generateTriplets.py <llvm_build_dir> <num_optimizations> <ll_file_list> <output_dir>14 15 For Machine IR:16 python generateTriplets.py --mode=mir <llvm_build_dir> <mir_file_list> <output_dir>17"""18 19import argparse20import logging21import os22import random23import subprocess24from concurrent.futures import ThreadPoolExecutor, as_completed25from pathlib import Path26from typing import List, Set, Tuple27 28# Configuration29OPT_LEVELS = ["O0", "O1", "O2", "O3", "Os", "Oz"]30DEFAULT_MAX_WORKERS = 10031 32logger = logging.getLogger(__name__)33 34 35# TODO: Change this to a dataclass with slots36# when Python 3.10+ is the minimum version37# https://docs.python.org/3/library/dataclasses.html#dataclasses.dataclass38class TripletResult:39 """Result from processing a single LLVM IR file"""40 41 __slots__ = ["triplets", "max_relation"]42 43 def __init__(self, triplets: Set[str], max_relation: int):44 self.triplets = triplets45 self.max_relation = max_relation46 47 48class IR2VecTripletGenerator:49 """Main class for generating IR2Vec or MIR2Vec triplets"""50 51 def __init__(52 self,53 llvm_build_dir: Path,54 num_optimizations: int,55 output_dir: Path,56 max_workers: int = DEFAULT_MAX_WORKERS,57 mode: str = "llvm",58 ):59 self.llvm_build_dir = llvm_build_dir60 self.num_optimizations = num_optimizations61 self.output_dir = output_dir62 self.max_workers = max_workers63 self.mode = mode # "llvm" or "mir"64 65 # Tool paths66 self.opt_binary = os.path.join(llvm_build_dir, "bin", "opt")67 self.ir2vec_binary = os.path.join(llvm_build_dir, "bin", "llvm-ir2vec")68 69 self._validate_setup()70 71 # Create output directory if it doesn't exist72 self.output_dir.mkdir(parents=True, exist_ok=True)73 74 def _validate_setup(self):75 """Validate that all required tools and paths exist"""76 if not self.llvm_build_dir.exists():77 raise FileNotFoundError(78 f"LLVM build directory not found: {self.llvm_build_dir}"79 )80 81 if not os.path.isfile(self.opt_binary) or not os.access(82 self.opt_binary, os.X_OK83 ):84 raise FileNotFoundError(85 f"opt binary not found or not executable: {self.opt_binary}"86 )87 88 if not os.path.isfile(self.ir2vec_binary) or not os.access(89 self.ir2vec_binary, os.X_OK90 ):91 raise FileNotFoundError(92 f"llvm-ir2vec binary not found or not executable: {self.ir2vec_binary}"93 )94 95 if self.mode not in ["llvm", "mir"]:96 raise ValueError(f"Mode must be 'llvm' or 'mir', got: {self.mode}")97 98 # For LLVM IR mode, validate optimization count99 if self.mode == "llvm" and not (1 <= self.num_optimizations <= len(OPT_LEVELS)):100 raise ValueError(101 f"Number of optimizations must be between 1-{len(OPT_LEVELS)}"102 )103 104 def _select_optimization_levels(self) -> List[str]:105 """Select unique random optimization levels"""106 return random.sample(OPT_LEVELS, self.num_optimizations)107 108 def _process_single_file(self, input_file: Path) -> TripletResult:109 """Process a single LLVM IR or MIR file"""110 all_triplets = set()111 max_relation = 1112 113 if self.mode == "mir":114 # For MIR files, process directly without optimization115 triplets, file_max_relation = self._run_mir_pipeline(input_file)116 if triplets:117 all_triplets.update(triplets)118 max_relation = max(max_relation, file_max_relation)119 logger.debug(f"Generated {len(triplets)} triplets for {input_file}")120 else:121 # For LLVM IR files, apply multiple optimization levels122 opt_levels = self._select_optimization_levels()123 for opt_level in opt_levels:124 triplets, file_max_relation = self._run_pipeline(input_file, opt_level)125 if triplets:126 all_triplets.update(triplets)127 max_relation = max(max_relation, file_max_relation)128 logger.debug(129 f"Generated {len(triplets)} triplets for {input_file} with {opt_level}"130 )131 132 return TripletResult(all_triplets, max_relation)133 134 def _run_pipeline(self, input_file: Path, opt_level: str) -> Tuple[Set[str], int]:135 """Run opt | llvm-ir2vec pipeline using subprocess pipes."""136 try:137 # Run opt first138 opt_proc = subprocess.Popen(139 [self.opt_binary, f"-{opt_level}", str(input_file), "-o", "-"],140 stdout=subprocess.PIPE,141 stderr=subprocess.PIPE,142 text=True,143 )144 145 # Run llvm-ir2vec with opt's output as input146 ir2vec_proc = subprocess.Popen(147 [self.ir2vec_binary, "triplets", "--mode=llvm", "-", "-o", "-"],148 stdin=opt_proc.stdout,149 stdout=subprocess.PIPE,150 stderr=subprocess.PIPE,151 text=True,152 )153 154 opt_proc.stdout.close()155 stdout, _ = ir2vec_proc.communicate()156 opt_proc.wait()157 158 # Check if either process failed159 if opt_proc.returncode != 0 or ir2vec_proc.returncode != 0:160 return set(), 1161 162 return self._parse_triplet_output(stdout)163 except (subprocess.SubprocessError, OSError):164 return set(), 1165 166 def _run_mir_pipeline(self, input_file: Path) -> Tuple[Set[str], int]:167 """Run llvm-ir2vec pipeline for MIR files."""168 try:169 # Run llvm-ir2vec directly on MIR file170 result = subprocess.run(171 [172 self.ir2vec_binary,173 "triplets",174 "--mode=mir",175 str(input_file),176 "-o",177 "-",178 ],179 stdout=subprocess.PIPE,180 stderr=subprocess.PIPE,181 text=True,182 check=False,183 )184 185 if result.returncode != 0:186 return set(), 1187 188 return self._parse_triplet_output(result.stdout)189 except (subprocess.SubprocessError, OSError):190 return set(), 1191 192 def _parse_triplet_output(self, output: str) -> Tuple[Set[str], int]:193 """Parse triplet output and extract max relation"""194 if not output.strip():195 return set(), 1196 197 lines = output.strip().split("\n")198 max_relation = 1199 200 # Extract max relation from metadata line201 if lines and lines[0].startswith("MAX_RELATION="):202 max_relation = int(lines[0].split("=")[1])203 lines = lines[1:]204 205 # Remove duplicate triplets by converting to a set206 return set(lines), max_relation207 208 def generate_triplets(self, file_list: Path) -> None:209 """Main method to generate triplets from a list of LLVM IR or MIR files"""210 # Store file_list_path for later use in entity generation211 self.file_list_path = file_list212 213 input_files = self._read_file_list(file_list)214 215 if self.mode == "mir":216 logger.info(217 f"Processing {len(input_files)} MIR files using {self.max_workers} workers"218 )219 else:220 logger.info(221 f"Processing {len(input_files)} files with {self.num_optimizations} "222 f"optimization levels using {self.max_workers} workers"223 )224 225 all_triplets = set()226 global_max_relation = 1227 228 with ThreadPoolExecutor(max_workers=self.max_workers) as executor:229 future_to_file = {230 executor.submit(self._process_single_file, file): file231 for file in input_files232 }233 234 for future in as_completed(future_to_file):235 try:236 result = future.result()237 all_triplets.update(result.triplets)238 global_max_relation = max(global_max_relation, result.max_relation)239 except (subprocess.SubprocessError, OSError, ValueError) as e:240 file_path = future_to_file[future]241 logger.error(f"Error processing {file_path}: {e}")242 243 self._generate_output_files(all_triplets, global_max_relation)244 logger.info("Processing completed successfully")245 246 def _read_file_list(self, file_list: Path) -> List[Path]:247 """Read and validate the list of input files"""248 input_files = []249 with open(file_list, "r") as f:250 for line_num, line in enumerate(f, 1):251 if line := line.strip():252 file_path = Path(line)253 if file_path.exists():254 input_files.append(file_path)255 else:256 logger.warning(f"File not found (line {line_num}): {file_path}")257 258 if not input_files:259 raise ValueError("No valid input files found")260 return input_files261 262 def _generate_output_files(self, all_triplets: Set[str], max_relation: int) -> None:263 """Generate the final output files"""264 logger.info(f"Generating output files with {len(all_triplets)} unique triplets")265 266 # Write all output files -- train2id.txt, entity2id.txt, relation2id.txt267 train2id_file = os.path.join(self.output_dir, "train2id.txt")268 entity2id_file = os.path.join(self.output_dir, "entity2id.txt")269 relation2id_file = os.path.join(self.output_dir, "relation2id.txt")270 271 with open(train2id_file, "w") as f:272 f.write(f"{len(all_triplets)}\n")273 f.writelines(f"{triplet}\n" for triplet in all_triplets)274 275 self._generate_entity2id(entity2id_file)276 self._generate_relation2id(relation2id_file, max_relation)277 278 def _generate_entity2id(self, output_file: Path) -> None:279 """Generate entity2id.txt using llvm-ir2vec"""280 if self.mode == "mir":281 # For MIR mode, we need to provide a sample MIR file to determine target282 # Use the first file from the processed list283 input_files = self._read_file_list(self.file_list_path)284 if not input_files:285 raise ValueError("No input files available for entity generation")286 287 subprocess.run(288 [289 str(self.ir2vec_binary),290 "entities",291 "--mode=mir",292 str(input_files[0]),293 "-o",294 str(output_file),295 ],296 check=True,297 capture_output=True,298 )299 else:300 subprocess.run(301 [302 str(self.ir2vec_binary),303 "entities",304 "--mode=llvm",305 "-o",306 str(output_file),307 ],308 check=True,309 capture_output=True,310 )311 312 def _generate_relation2id(self, output_file: Path, max_relation: int) -> None:313 """Generate relation2id.txt from max relation"""314 max_relation = max(max_relation, 1) # At least Next relation315 num_relations = max_relation + 1316 317 with open(output_file, "w") as f:318 f.write(f"{num_relations}\n")319 if self.mode == "llvm":320 # LLVM IR has Type relation at 0321 f.write("Type\t0\n")322 f.write("Next\t1\n")323 f.writelines(f"Arg{i-2}\t{i}\n" for i in range(2, num_relations))324 else:325 # MIR doesn't have Type relation, starts with Next at 0326 f.write("Next\t0\n")327 f.writelines(f"Arg{i-1}\t{i}\n" for i in range(1, num_relations))328 329 330def main():331 """Main entry point"""332 parser = argparse.ArgumentParser(333 description="Generate IR2Vec or MIR2Vec triplets from LLVM IR or Machine IR files",334 formatter_class=argparse.RawDescriptionHelpFormatter,335 )336 337 parser.add_argument(338 "llvm_build_dir", type=Path, help="Path to LLVM build directory"339 )340 parser.add_argument(341 "num_optimizations",342 type=int,343 nargs="?",344 default=1,345 help="Number of optimization levels to apply (1-6) for LLVM IR mode",346 )347 parser.add_argument(348 "input_file_list",349 type=Path,350 help="File containing list of LLVM IR or MIR files to process",351 )352 parser.add_argument(353 "output_dir", type=Path, help="Output directory for generated files"354 )355 parser.add_argument(356 "--mode",357 type=str,358 choices=["llvm", "mir"],359 default="llvm",360 help="Operation mode: 'llvm' for LLVM IR (default) or 'mir' for Machine IR",361 )362 parser.add_argument(363 "-j",364 "--max-workers",365 type=int,366 default=DEFAULT_MAX_WORKERS,367 help=f"Maximum number of parallel workers (default: {DEFAULT_MAX_WORKERS})",368 )369 parser.add_argument(370 "-v", "--verbose", action="store_true", help="Enable debug logging"371 )372 parser.add_argument(373 "-q", "--quiet", action="store_true", help="Suppress all output except errors"374 )375 376 args = parser.parse_args()377 378 # Configure logging379 level = (380 logging.ERROR381 if args.quiet382 else (logging.DEBUG if args.verbose else logging.INFO)383 )384 logging.basicConfig(385 level=level,386 format="[%(asctime)s] %(levelname)s: %(message)s",387 datefmt="%H:%M:%S",388 )389 390 generator = IR2VecTripletGenerator(391 args.llvm_build_dir,392 args.num_optimizations,393 args.output_dir,394 args.max_workers,395 args.mode,396 )397 generator.generate_triplets(args.input_file_list)398 399 400if __name__ == "__main__":401 main()402