| # RUN: env PYTHONPATH=%ir2vec_python_dir %python %s %S/../Inputs/input.ll %ir2vec_test_vocab_dir/dummy_3D_nonzero_opc_vocab.json | FileCheck %s | |
| import sys | |
| import ir2vec | |
| ll_file = sys.argv[1] | |
| vocab_path = sys.argv[2] | |
| vocab = ir2vec.loadVocab(vocab_path) | |
| tool = ir2vec.initEmbedding( | |
| filename=ll_file, mode=ir2vec.IR2VecKind.Symbolic, vocab=vocab | |
| ) | |
| # Success case | |
| inst_map = tool.getInstEmbMap("add") | |
| for inst in sorted(inst_map.keys()): | |
| print(f"INST: {inst}") | |
| print(f" EMB: {inst_map[inst].tolist()}") | |
| # CHECK: INST: %sum = add i32 %a, %b | |
| # CHECK: EMB: [37.0, 38.0, 39.0] | |
| # CHECK: INST: ret i32 %sum | |
| # CHECK: EMB: [1.0, 2.0, 3.0] | |
| # Error: Function not found | |
| try: | |
| tool.getInstEmbMap("nonexistent") | |
| except ValueError: | |
| print("ERROR: Function not found") | |
| # CHECK: ERROR: Function not found |