[Triton] Proton 프로파일러에서 불필요한 lock 추가 제거

PR 링크: triton-lang/triton#9257 상태: Merged | 변경: +129 / -121

들어가며

Triton의 Proton 프로파일러는 GPU 커널 실행 중 메트릭을 수집한다. 프로파일링은 실제 워크로드와 동시에 실행되므로 오버헤드를 최소화해야 한다. 이 PR은 Data 클래스의 lock 전략을 최적화하여 불필요한 shared_mutex 획득을 줄인다.

핵심 코드 분석

Before: 모든 phase 접근에 mutex 사용

class Data : public ScopeInterface {
  std::size_t currentPhase{0};
  std::size_t completePhase{kNoCompletePhase};

  size_t getCurrentPhase() const {
    std::shared_lock<std::shared_mutex> lock(mutex);
    return currentPhase;
  }

  // PhaseStore: 항상 lock 후 접근
  void *getOrCreatePtr(size_t phase) override {
    auto slot = getOrCreateSlot(phase);
    std::unique_lock<std::shared_mutex> slotLock(slot->mutex);
    if (!slot->value) {
      slot->value = std::make_unique<T>();
    }
    return slot->value.get();
  }
};

After: atomic + 조건부 lock

class Data : public ScopeInterface {
  std::atomic<std::size_t> currentPhase{0};  // atomic으로 변경
  std::size_t completeUpToPhase{kNoCompletePhase};

  struct PhaseInfo {
    size_t current{0};
    size_t completeUpTo{kNoCompletePhase};
    bool isComplete(size_t phase) const {
      return completeUpTo != kNoCompletePhase && completeUpTo >= phase;
    }
  };

  // 현재 phase에만 lock 필요
  [[nodiscard]] std::unique_lock<std::shared_mutex>
  lockIfCurrentPhase(size_t phase) {
    std::unique_lock<std::shared_mutex> lock(mutex, std::defer_lock);
    const auto currentPhaseValue = currentPhase.load(std::memory_order_relaxed);
    if (phase == currentPhaseValue) {
      lock.lock();
    }
    // 과거 phase는 application thread가 접근하지 않으므로 lock 불필요
    return lock;
  }
};

// PhaseStore: getPtr/createPtr 분리
void *createPtr(size_t phase) override {
  std::shared_ptr<Slot> slot;
  {
    std::unique_lock<std::shared_mutex> lock(phasesMutex);
    auto &entry = phases[phase];
    if (!entry)
      entry = std::make_shared<Slot>();
    slot = entry;
  }
  {
    std::unique_lock<std::shared_mutex> slotLock(slot->mutex);
    if (!slot->value)
      slot->value = std::make_unique<T>();
    return slot->value.get();
  }
}

void *getPtr(size_t phase) override {
  return getSlot(phase)->value.get();  // lock 없이 직접 접근
}

왜 이게 좋은가

1. currentPhase를 atomic으로: 단순 읽기에 mutex를 획득하지 않아 contention이 줄어든다.
2. 조건부 locking: 현재 phase에만 lock을 걸고, 이미 완료된 phase는 lock-free로 접근한다.
3. PhaseStore getPtr/createPtr 분리: 이미 생성된 phase의 데이터를 읽을 때 불필요한 생성 로직과 lock을 건너뛴다.
4. API 정리: addScopeMetrics/addEntryMetrics를 addMetrics로 통합하고, PhaseInfo 구조체로 phase 상태를 원자적으로 조회한다.

정리

이 PR은 Proton 프로파일러의 Data 클래스에서 lock 전략을 최적화했다. currentPhase를 std::atomic으로 변경하고, 완료된 phase는 lock-free 접근을 허용하며, PhaseStore의 read/write 경로를 분리하여 프로파일링 오버헤드를 줄였다.

참고 자료

• triton-lang/triton#9257

---

이 글은 AI를 활용하여 PR의 핵심 변경사항을 분석하고 정리한 것입니다. 실제 코드의 맥락은 원본 PR을 참고해 주세요.