PyTorch 내부 구조 11 - Operator Schema, Dispatch Key, Meta 함수
custom op를 제대로 등록하려면 구현 이전에 schema와 dispatch 구조를 먼저 분명히 해야 한다
왜 schema가 먼저인가
operator를 만든다는 것은 함수 하나를 만드는 일이 아니다. 입력/출력 규약, aliasing 성질, shape 추론 가능성, backend별 등록 지점을 함께 정의하는 일이다.
특히 torch.compile이나 graph capture를 고려하면 meta 수준 shape/dtype 추론도 중요해진다.
dispatch key 감각
같은 operator라도 CPU, CUDA, Autograd, Meta 등 여러 구현이 달릴 수 있다. dispatch key는 그 구현 선택 경로를 조직한다. 이 감각이 있어야 custom op가 왜 특정 backend에서만 동작하는지 이해할 수 있다.
다음 글에서는 backward 구현 패턴을 본다. forward를 붙인 다음에는 gradient semantics를 어디에 둘지 결정해야 한다.