ASM 코드 분할 및 병합 코드

ASM 코드 분할 Pass

// CallSplitterPass.cpp
#include "llvm/ADT/StringRef.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PassManager.h"
#include "llvm/Passes/PassPlugin.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/Path.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/IR/IntrinsicInst.h"

#include <filesystem>
#include <vector>
#include <string>
#include <fstream>

using namespace llvm;
namespace fs = std::filesystem;

static cl::opt<std::string> OutputDir("split-dir",
    cl::desc("Directory to place chunk files"),
    cl::init("splitted"));

static cl::opt<std::string> JsonOut("split-json",
    cl::desc("JSONL output file"),
    cl::init("chunks.jsonl"));

namespace {

// --- 유틸들 -----------------------------------------------------

static std::string getSafeModuleName(Module &M) {
    std::string moduleName = M.getSourceFileName();
    if (moduleName.empty()) {
        StringRef mid = M.getModuleIdentifier();
        moduleName = sys::path::stem(mid).str();
    }
    if (moduleName.empty())
        moduleName = "module";
    return moduleName;
}

static fs::path prepareModuleDir(const std::string &baseDirStr, const std::string &moduleName) {
    fs::path baseDir = baseDirStr;
    std::error_code ec;
    fs::create_directories(baseDir, ec);
    if (ec) {
        errs() << "Warning: cannot create dir '" << baseDir.string() << "': " << ec.message() << "\n";
    }

    fs::path moduleDir = baseDir / moduleName;
    fs::create_directories(moduleDir, ec);
    if (ec) {
        errs() << "Warning: cannot create module dir '" << moduleDir.string() << "': " << ec.message() << "\n";
    }

    return moduleDir;
}

static std::string instToString(const Instruction &I) {
    std::string s;
    raw_string_ostream rso(s);
    I.print(rso);
    rso.flush();
    return s;
}

static bool writeChunkFile(const fs::path &moduleDir, const std::string &chunkFilename,
                           const std::vector<std::string> &lines) {
    fs::path chunkPath = moduleDir / chunkFilename;
    std::error_code ec;
    raw_fd_ostream ofs(chunkPath.string(), ec, sys::fs::OF_Text);
    if (ec) {
        errs() << "Warning: cannot write chunk file " << chunkPath.string() << ": " << ec.message() << "\n";
        return false;
    }
    for (const auto &L : lines) ofs << L << "\n";
    return true;
}

static void writeChunkJson(std::ofstream &jsonl, const std::string &chunkName,
                           const std::set<std::string> &labels) {
    int li = 0;
    jsonl << "{\"id\":\"" << chunkName << "\",";
    jsonl << "\"labels\":[";
    for (std::string label : labels) {
        jsonl << "\"" << label << "\"";
        if (++li < labels.size()) jsonl << ",";
    }
    jsonl << "]}";
    jsonl << "\n";
}

static void flushChunk(std::ofstream &jsonl, const fs::path &moduleDir,
                       unsigned &chunk_id,
                       std::vector<std::string> &currentChunkLines,
                       std::set<std::string> &currentLabels) {
                        
    std::ostringstream oss;
    oss << std::setw(4) << std::setfill('0') << chunk_id;
    std::string chunkName = "chunk_" + oss.str();
    std::string chunkFilename = chunkName + ".s";
    writeChunkFile(moduleDir, chunkFilename, currentChunkLines);
    writeChunkJson(jsonl, chunkName, currentLabels);
    ++chunk_id;
    currentChunkLines.clear();
    currentLabels.clear();
}

// --- 함수별 처리 ----------------
static void processFunction(Function &F, std::ofstream &jsonl, const fs::path &moduleDir, unsigned &chunk_id) {
    std::vector<std::string> currentChunkLines;
    std::set<std::string> currentLabels;
    unsigned isInDispatch = 0;

    for (BasicBlock &BB : F) {
        for (Instruction &I : BB) {
            if (I.getParent() == &F.getEntryBlock() && &I == &*F.getEntryBlock().begin()) {
                currentLabels.insert("FUNC_ENTRY");
            }

            if (isInDispatch) {
                currentLabels.insert("DISPATCH_HANDLER"); // 라벨 추가
            } 
            else {
                currentLabels.insert("NO_DISPATCH"); // 라벨 추가
            }

            // 인스트럭션 문자열로 변환
            currentChunkLines.push_back(instToString(I));

            if (auto *CB = dyn_cast<CallBase>(&I)) {
                // 청크 저장 (라벨 포함)
                Function *calledFunc = CB->getCalledFunction();
                // 2. Dummy function 확인

                if (calledFunc && calledFunc->getName() == "dummy_function_succ") {
                    currentLabels.insert("DISPATCH_ENTRY"); // 라벨 추가
                    ++isInDispatch;
                }
                else if (calledFunc && calledFunc->getName() == "dummy_function_pred") {
                    currentLabels.insert("DISPATCH_EXIT"); // 라벨 추가
                    --isInDispatch;
                }
                else {
                    // currentLabels.insert("Call_BASE"); // 라벨 추가
                }
                
                flushChunk(jsonl, moduleDir, chunk_id, currentChunkLines, currentLabels);
            }
            else if (isa<BranchInst>(I) || isa<SwitchInst>(I) ||
                isa<IndirectBrInst>(I) || isa<ReturnInst>(I)) {
                // currentLabels.insert("BRANCH");
                flushChunk(jsonl, moduleDir, chunk_id, currentChunkLines, currentLabels);
            }

            // branch 기준 split
            /*
            else if (auto *BI = dyn_cast<BranchInst>(&I)) {
                currentLabels.insert("BRANCH");
                flushChunk(jsonl, moduleDir, chunk_id, currentChunkLines, currentLabels);
            }
            else if (auto *SI = dyn_cast<SwitchInst>(&I)) {
                currentLabels.insert("SWITCH");
                flushChunk(jsonl, moduleDir, chunk_id, currentChunkLines, currentLabels);
            }
            else if (auto *IBI = dyn_cast<IndirectBrInst>(&I)) {
                currentLabels.insert("INDIRECT_BRANCH");
                flushChunk(jsonl, moduleDir, chunk_id, currentChunkLines, currentLabels);
            }
            else if (auto *RI = dyn_cast<ReturnInst>(&I)) {
                currentLabels.insert("RETURN");
                flushChunk(jsonl, moduleDir, chunk_id, currentChunkLines, currentLabels);
            }
            */
        }
    }

    // 함수 끝에 남은 청크 저장
    if (!currentChunkLines.empty()) {
        flushChunk(jsonl, moduleDir, chunk_id, currentChunkLines, currentLabels);
    }
}

// --- Pass 구현 ------------------------------------------------
struct SplitterPass : public PassInfoMixin<SplitterPass> {
    PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
        std::string moduleName = getSafeModuleName(M);
        fs::path moduleDir = prepareModuleDir(OutputDir.getValue(), moduleName);

        fs::path jsonPath = moduleDir / JsonOut.getValue();
        std::ofstream jsonl(jsonPath, std::ios::out | std::ios::trunc);
        if (!jsonl.is_open()) {
            errs() << "Error: cannot open JSONL file: " << JsonOut << "\n";
            return PreservedAnalyses::all();
        }

        unsigned chunk_id = 1;
        for (Function &F : M) {
            if (F.isDeclaration()) continue;
            processFunction(F, jsonl, moduleDir, chunk_id);
        }

        jsonl.close();
        errs() << "[SplitterPass] wrote " << chunk_id-1 << " chunks to directory '" << moduleDir.string()
               << "' and JSONL '" << JsonOut.getValue() << "'\n";

        return PreservedAnalyses::all();
    }
};

} // namespace

// plugin registration
extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo llvmGetPassPluginInfo() {
    return {
        LLVM_PLUGIN_API_VERSION, "SplitterPass", LLVM_VERSION_STRING,
        [](PassBuilder &PB) {
            PB.registerPipelineParsingCallback(
                [](StringRef Name, ModulePassManager &MPM,
                   ArrayRef<PassBuilder::PipelineElement>) -> bool {
                    if (Name == "splitter") {
                        MPM.addPass(SplitterPass());
                        return true;
                    }
                    return false;
                });
        }
    };
}

 

주요 코드 :

 

• 호출(Call) 또는 분기(Branch, Switch, IndirectBranch, Return) 명령어를 만날 때마다 청크를 끊고 저장
• 함수 시작 시 "FUNC_ENTRY" 라벨 부여
• "dummy_function_succ" 호출 시 "DISPATCH_ENTRY" 부여 및 디스패치 구간(DISPATCH_HANDLER) 진입
• "dummy_function_pred" 호출 시 "DISPATCH_EXIT" 부여 및 디스패치 구간(DISPATCH_HANDLER)  종료
• 디스패치 구간 라벨은 "DISPATCH_HANDLER", 구간 외 라벨은 "NO_DISPATCH" 으로 부여

 

 

---

 

Chunk 합병 코드 1 (mergeChunk1.py)

import os
import json

# ===== 설정 =====
INPUT_DIR = "./splitted"
OUTPUT_DIR = "./merged1"
h = 20
MAX_SIZE = h * 1024  # 20KB

# ===== 폴더 선택 =====
subdirs = [d for d in os.listdir(INPUT_DIR) if os.path.isdir(os.path.join(INPUT_DIR, d))]

if not subdirs:
    print("❌ 'splitted' 폴더 안에 하위 폴더가 없습니다.")
    exit(1)

print("=== 병합 대상 코드 목록 ===")
for i, d in enumerate(subdirs, 1):
    print(f"{i}. {d}")

target = input("\n병합할 TARGET_DIR 이름 또는 번호 입력 (Enter 입력 시 전체 처리): ").strip()

# 전체 처리 여부
if target == "":
    targets = subdirs
    print("\n[!] 전체 폴더 자동 처리 모드 실행 중...\n")
else:
    if target.isdigit():
        idx = int(target) - 1
        if 0 <= idx < len(subdirs):
            targets = [subdirs[idx]]
        else:
            print(f"❌ 번호 '{target}'는 유효하지 않습니다.")
            exit(1)
    else:
        if target not in subdirs:
            print(f"❌ '{target}' 폴더를 찾을 수 없습니다. 경로를 확인하세요.")
            exit(1)
        targets = [target]


# ===== 병합 함수 정의 =====
def merge_target_code(TARGET_DIR):
    base_dir = os.path.join(INPUT_DIR, TARGET_DIR)
    output_dir = os.path.join(OUTPUT_DIR, TARGET_DIR)
    os.makedirs(output_dir, exist_ok=True)

    chunks = sorted(
        [f for f in os.listdir(base_dir) if f.endswith(".s")],
        key=lambda x: int(x.split("_")[1].split(".")[0])
    )

    jsonl_path = os.path.join(base_dir, "chunks.jsonl")
    if not os.path.exists(jsonl_path):
        print(f"⚠️ {TARGET_DIR}: chunks.jsonl 파일이 없습니다. 건너뜀.")
        return

    with open(jsonl_path, "r") as f:
        jsonl_entries = [json.loads(line) for line in f if line.strip()]

    if len(jsonl_entries) != len(chunks):
        raise ValueError(f"{TARGET_DIR}: .s 파일 개수와 jsonl 엔트리 수가 다릅니다!")

    merged_id = 1
    current_lines = []
    current_labels = set()
    current_size = 0
    merged_jsonl = []

    def flush_merged():
        nonlocal merged_id, current_lines, current_labels, current_size
        if not current_lines:
            return
        merged_name = f"chunk_{merged_id:04d}"
        merged_s_path = os.path.join(output_dir, f"{merged_name}.s")
        with open(merged_s_path, "w") as f:
            f.write("\n".join(current_lines))
        merged_jsonl.append({
            "id": merged_name,
            "labels": sorted(list(current_labels))
        })
        print(f"[+] {TARGET_DIR}: Saved {merged_name}.s ({current_size/1024:.1f} KB, labels={list(current_labels)})")
        merged_id += 1
        current_lines.clear()
        current_labels.clear()
        current_size = 0

    # ===== 병합 진행 =====
    for entry, s_file in zip(jsonl_entries, chunks):
        s_path = os.path.join(base_dir, s_file)
        with open(s_path, "r") as f:
            content = f.read()

        labels = entry.get("labels", [])
        size = len(content.encode())

        # DISPATCH 라벨 단독 처리
        if "DISPATCH_ENTRY" in labels or "DISPATCH_EXIT" in labels:
            flush_merged()
            merged_name = f"chunk_{merged_id:04d}"
            merged_s_path = os.path.join(output_dir, f"{merged_name}.s")
            with open(merged_s_path, "w") as f:
                f.write(content)
            merged_jsonl.append({
                "id": merged_name,
                "labels": sorted(labels)
            })
            print(f"[+] {TARGET_DIR} → {merged_name}.s (Dispatch 단독 처리)")
            merged_id += 1
            continue

        # 🔥 다음 chunk 추가 시 용량 초과 예측 → 미리 flush
        if current_size + size > MAX_SIZE:
            flush_merged()

        # 새 묶음에 현재 chunk 추가
        current_lines.append(content)
        current_labels.update(labels)
        current_size += size

    # 마지막 남은 조각 flush
    flush_merged()

    merged_jsonl_path = os.path.join(output_dir, "chunks.jsonl")
    with open(merged_jsonl_path, "w") as f:
        for entry in merged_jsonl:
            f.write(json.dumps(entry) + "\n")

    print(f"✅ {TARGET_DIR}: 병합 완료 — 총 {merged_id - 1}개 묶음 생성됨\n")


# ===== 전체/단일 폴더 실행 =====
for t in targets:
    merge_target_code(t)

print("🎉 모든 병합 작업 완료!")

 

병합 규칙

이 코드는 splitted 폴더 내 각 타깃 폴더(TARGET_DIR)를 대상으로
.s(어셈블리 청크) 파일들을 chunks.jsonl 메타데이터 정보를 이용해 병합을 수행한다.
병합된 결과는 merged1/<TARGET_DIR> 폴더에 저장된다.

 

병합 로직

1. 용량 기준 (h KB)

• h는 용량 기준을 나타내는 변수
• MAX_SIZE = h * 1024 (h KB)
• 현재 병합 중인 청크의 크기(current_size)와 새로 추가할 청크의 크기의 합이 h KB를 초과하는지 검사
• 초과 시, 현재 병합 중이던 청크를 flush
• 문제점 : 입력이 애초의 h KB 를 넘는 크기의 파일일 경우 기준을 충족하지 못함

 

if current_size + size > MAX_SIZE: flush_merged()

 

 

2. DISPATCH 관련 라벨 단독 처리

• 현재 labels에 "DISPATCH_ENTRY" 또는 "DISPATCH_EXIT"이 포함되어 있을 시, 
  병합하지 않고, 현재까지의 chunk들을 flush