Google TPU (Tensor Processing Unit)

Google TPU 아키텍처: MXU/HBM/ICI 토폴로지, Gasket/ACCEL 드라이버, PCIe/커스텀 인터커넥트, Edge TPU/Coral, XLA/JAX 컴파일 파이프라인, TPU Pod 확장

전제 조건: PCI / PCIe 서브시스템디바이스 드라이버 문서를 먼저 읽으세요. TPU는 PCIe 엔드포인트로 연결되며, 커널 드라이버 프레임워크의 이해가 필요합니다.
일상 비유: TPU는 행렬 연산 전용 공장과 비슷합니다. 범용 CPU가 다양한 작업을 처리하는 것과 달리, TPU의 MXU(Matrix Multiply Unit)는 오직 텐서 곱셈에만 특화되어 극도의 효율을 달성합니다.

TPU (Tensor Processing Unit)는 Google이 설계한 ASIC 기반 도메인 특화 가속기로, 행렬 연산(텐서 연산)에 최적화되어 머신러닝 학습(Training)과 추론(Inference) 워크로드를 처리합니다. Cloud TPU는 Google 데이터센터에서 PCIe 또는 자체 호스트 인터페이스로 연결되며, Edge TPU는 PCIe M.2 또는 USB 폼 팩터로 임베디드 환경에서 추론 가속을 제공합니다.

커널 관점: TPU는 PCIe 디바이스로서 BAR 매핑, DMA 전송, MSI-X 인터럽트, IOMMU 연동 등 표준 PCIe 메커니즘을 활용합니다. Linux 커널에서는 drivers/staging/gasket/ (Gasket 프레임워크)와 drivers/accel/ (ACCEL 서브시스템, 6.2+)을 통해 가속기 드라이버를 지원합니다.

TPU 세대별 하드웨어 아키텍처

세대연도용도연산 성능HBM인터커넥트공정
TPU v12015추론 전용92 TOPS (INT8)— (DRAM 8 GB)PCIe 3.0 x1628nm
TPU v22017학습 + 추론45 TFLOPS (bf16)HBM2 16 GBICI (4링크)16nm
TPU v32018학습 + 추론123 TFLOPS (bf16)HBM2 32 GBICI (6링크)16nm (수냉)
TPU v42021학습 + 추론275 TFLOPS (bf16)HBM2e 32 GBICI (6링크, 3D Torus)7nm
TPU v5e2023효율 최적화197 TFLOPS (bf16)HBM2e 16 GBICIN/A
TPU v5p2023대규모 학습459 TFLOPS (bf16)HBM2e 95 GBICIN/A
TPU v6e (Trillium)2024차세대 효율918 TFLOPS (bf16)HBM 32 GBICIN/A
Edge TPU2019추론 전용4 TOPS (INT8)PCIe / USBN/A

MXU (Matrix Multiply Unit) — 시스톨릭 어레이

TPU의 핵심 연산 유닛은 MXU (Matrix Multiply Unit)으로, 시스톨릭 어레이(systolic array) 구조를 사용합니다. TPU v2/v3는 칩당 2개의 TensorCore를 가지며, 각 TensorCore에 128×128 시스톨릭 어레이 MXU가 탑재됩니다.

128×128 시스톨릭 어레이 (MXU) — 데이터 흐름 w[0] w[1] w[2] ··· w[127] Weight (preloaded, 상→하) x[0] x[1] x[2] ··· x[127] Input (좌→우) PE PE PE ··· ··· PE PE PE PE ··· ··· PE PE PE PE ··· ··· PE PE PE PE ··· ··· PE y[0] y[1] y[2] ··· y[127] Output (부분합 누적, 상→하) partial sum 각 PE (Processing Element) = 1× MAC (Multiply-Accumulate) / 사이클 accumulator += weight × input → 결과를 아래 PE로 전달 (weight stationary 방식) 128 × 128 = 16,384 PE → 1 사이클에 16,384 MAC → bf16 기준 32,768 FLOPS/사이클
항목설명
PE 격자128 × 12816,384개 Processing Element
MAC/사이클16,384각 PE가 1 MAC 연산 수행
FLOPS/사이클 (bf16)32,768MAC = multiply + add → ×2
TPU v4 피크~275 TFLOPS (bf16)~1.05 GHz × 32,768 × 4 MXU × 2 (bf16 packing)
커널 관점: 호스트 CPU는 MXU를 직접 프로그래밍하지 않습니다. XLA/PJRT 컴파일러가 HLO(High Level Operations)를 TPU 명령어로 변환하고, 커널 드라이버는 DMA로 명령어 버퍼와 데이터를 전송하고, 완료 인터럽트를 수신하며, 메모리 매핑을 관리하는 역할을 합니다.
bf16 (bfloat16): Google이 TPU용으로 설계한 16비트 부동소수점 형식. IEEE fp32와 동일한 8비트 지수부(exponent)를 가져 표현 범위(dynamic range)가 fp32와 같으면서, 7비트 가수부(mantissa)로 메모리 대역폭을 절반으로 줄입니다. 딥러닝 학습에서 fp32와 거의 동등한 수렴 특성을 보입니다.

TPU 메모리 아키텍처

TPU v4 메모리 계층 구조 TensorCore (×2 per chip) MXU 0 128×128 bf16 MXU 1 128×128 bf16 Vector Unit (SIMD ALU) VMEM (16 MB) 벡터 레지스터 파일 SMEM 스칼라 레지스터 파일 CMEM — 공통 메모리 Cross-Core Shared Memory HBM2e — 메인 메모리 32 GB / 1.6 TB/s 대역폭 ICI — Inter-Chip Interconnect 6 링크 / 다른 TPU 칩 직접 연결 커널 드라이버 관점 • HBM은 BAR를 통해 MMIO로 매핑되거나 DMA를 통해 접근 • VMEM/SMEM은 TPU 내부 SRAM — 드라이버가 직접 접근하지 않음 • 호스트 ↔ HBM 데이터 전송은 DMA 엔진이 담당

ICI (Inter-Chip Interconnect)

TPU v2 이후 칩 간 통신을 위한 ICI (Inter-Chip Interconnect)를 제공합니다. ICI는 PCIe를 사용하지 않는 Google 자체 설계 고속 직렬 링크로, 다수의 TPU 칩을 직접 연결하여 Pod 단위의 대규모 분산 학습을 가능하게 합니다.

세대ICI 링크 수토폴로지링크당 대역폭Pod 규모
TPU v242D Torus (16×16)~496 Gbps256 칩 (64 TF/pod)
TPU v362D Torus (32×16)~656 Gbps1,024 칩
TPU v463D Torus (4×4×4 ~ 큐브)~2.4 Tbps (총)4,096 칩
TPU v5p다수3D Torus향상8,960 칩
TPU v4 Pod — 3D Torus ICI 토폴로지 (4×4 상면도) z ↑ x y Chip 0,0 Chip 1,0 Chip 2,0 Chip 3,0 Chip 0,1 Chip 1,1 Chip 2,1 Chip 3,1 Chip 0,2 Chip 1,2 Chip 2,2 Chip 3,2 Chip 0,3 Chip 1,3 Chip 2,3 Chip 3,3 3D Torus 특성 (4,096칩 = 4×4×4 큐브 블록 × 토러스) ICI 직접 링크 토러스 랩어라운드 각 칩 6개 이웃 (±x, ±y, ±z) · AllReduce bisection BW 극대화 · 단일 링크 장애 시 대체 경로 커널 관점: ICI는 칩 내부 HW 관리, 호스트 드라이버는 PCIe 연결만 관리

TPU의 PCIe 인터페이스

TPU는 호스트 시스템과 PCIe 버스를 통해 연결됩니다. Edge TPU (Coral)는 표준 PCIe 엔드포인트로, Cloud TPU는 커스텀 호스트 인터페이스 칩을 통해 PCIe로 연결됩니다.

디바이스PCIe 규격Vendor IDDevice IDBAR 구성
Edge TPU (Apex)PCIe Gen2 x11ac1 (Google)089aBAR0: 레지스터, BAR2: 펌웨어/데이터
Cloud TPU v4PCIe Gen3 x161ae0 (Google)(비공개)다수 BAR: CSR, DMA, HBM 윈도우
# Edge TPU (Coral PCIe Accelerator) 확인
lspci -nn | grep -i google
# 03:00.0 System peripheral [0880]: Global Unichip Corp. [1ac1:089a]

# 상세 정보
lspci -vvv -s 03:00.0
#   Region 0: Memory at f7000000 (64-bit, prefetchable) [size=16M]
#   Region 2: Memory at f8000000 (64-bit, prefetchable) [size=2M]
#   Capabilities: [80] MSI-X: Enable+ Count=128 Masked-

# Edge TPU BAR 매핑 구조
# BAR0 (16 MB): Apex CSR (Control/Status Registers)
#   - 인터럽트 상태/마스크 레지스터
#   - DMA 디스크립터 링 베이스/크기
#   - 펌웨어 상태 레지스터
# BAR2 (2 MB): 명령어 버퍼 / 데이터 교환 영역

Gasket 프레임워크 (커널 TPU 드라이버)

Google은 TPU 계열 가속기를 위한 커널 드라이버 프레임워크로 Gasket (Google ASIC Software, Kernel Extensions, and Tools)을 개발했습니다. 이 프레임워크는 drivers/staging/gasket/에 위치하며, PCIe 가속기 디바이스의 공통 기능을 추상화합니다.

소스 파일역할
gasket_core.cPCI probe/remove, char device 등록, BAR 매핑, ioctl 디스패치(Dispatch)
gasket_ioctl.c표준 ioctl 처리 (페이지 테이블 매핑, 이벤트 fd, 디바이스 리셋)
gasket_page_table.c2단계 디바이스 페이지 테이블 관리 (가상주소 → DMA 주소)
gasket_interrupt.cMSI-X 인터럽트 설정 및 eventfd 연동
gasket_sysfs.csysfs 속성 노출 (펌웨어 버전, 디바이스 상태 등)
apex_driver.cEdge TPU (Apex) 전용 드라이버 — Gasket 위에 구현
/* Gasket 프레임워크 핵심 구조체 (drivers/staging/gasket/) */

/* gasket_driver_desc — 디바이스별 드라이버 기술자 */
struct gasket_driver_desc {
    const char *name;               /* 드라이버 이름 */
    const char *driver_version;      /* 드라이버 버전 문자열 */

    /* PCI BAR 기술 */
    struct gasket_bar_desc bar_descriptions[GASKET_NUM_BARS];
    int num_page_tables;             /* 디바이스 페이지 테이블 수 */
    struct gasket_page_table_config page_table_configs[...];

    /* 인터럽트 기술 */
    int num_interrupts;
    struct gasket_interrupt_desc interrupts[...];

    /* 콜백 */
    int (*device_open_cb)(struct gasket_dev *dev);
    int (*device_close_cb)(struct gasket_dev *dev);
    int (*device_reset_cb)(struct gasket_dev *dev);
    enum gasket_status (*device_status_cb)(struct gasket_dev *dev);
    long (*ioctl_handler_cb)(...);  /* 디바이스별 커스텀 ioctl */
};

/* gasket_dev — 디바이스 인스턴스 (probe 시 생성) */
struct gasket_dev {
    struct pci_dev *pci_dev;
    struct device *dev;
    struct cdev cdev;               /* char device */
    int dev_idx;                     /* /dev/apex_N 인덱스 */

    /* BAR 매핑 */
    struct gasket_bar_data bar_data[GASKET_NUM_BARS];

    /* 페이지 테이블 (디바이스 IOVA → 호스트 물리 주소) */
    struct gasket_page_table *page_table[...];

    /* 인터럽트 */
    struct gasket_interrupt_data *interrupt_data;

    const struct gasket_driver_desc *driver_desc;
    struct mutex mutex;
    int ownership_owned;             /* 독점 소유권 */
};

Edge TPU (Apex) 드라이버 상세

Edge TPU 드라이버(apex_driver.c)는 Gasket 프레임워크 위에 구현된 구체적인 디바이스 드라이버로, Google Coral PCIe/M.2 가속기를 지원합니다.

/* Edge TPU (Apex) 드라이버 — Gasket 기반 구현 */
/* drivers/staging/gasket/apex_driver.c */

#define APEX_DRIVER_NAME     "apex"
#define APEX_DRIVER_VERSION  "1.0"

/* PCI ID 테이블 */
static const struct pci_device_id apex_pci_ids[] = {
    { PCI_DEVICE(0x1ac1, 0x089a) },  /* Edge TPU */
    { 0, }
};

/* Apex BAR 기술 */
enum {
    APEX_BAR_INDEX  = 0,  /* BAR0: CSR 레지스터 (16 MB) */
    APEX_BAR2_INDEX = 2,  /* BAR2: 펌웨어/데이터 (2 MB) */
};

/* 주요 CSR 오프셋 */
#define APEX_BAR2_REG_SCU_BASE    0x1A300  /* SCU (System Control) */
#define APEX_BAR2_REG_KERNEL_HIB  0x1A400  /* Host Interface Block */

/* Apex 드라이버 기술자 — Gasket 프레임워크에 등록 */
static const struct gasket_driver_desc apex_desc = {
    .name            = APEX_DRIVER_NAME,
    .driver_version  = APEX_DRIVER_VERSION,
    .major           = 120,
    .minor           = 0,
    .module          = THIS_MODULE,

    .bar_descriptions = {
        /* BAR0: 레지스터 공간 */
        { .size = 0x1000000, .permissions = 0660 },
        /* BAR2: 데이터 공간 */
        { .size = 0x200000,  .permissions = 0660 },
    },

    .num_page_tables = 1,
    .page_table_configs = { {
        .id = 0,
        .mode = GASKET_PAGE_TABLE_MODE_NORMAL,
        .total_entries = 8192,     /* 4K × 8192 = 32 MB 매핑 가능 */
        .base_reg = 0x1A488,       /* 페이지 테이블 베이스 레지스터 */
    } },

    .num_interrupts  = 1,
    .interrupts      = { {
        .index  = 0,
        .reg    = 0x1A318,          /* 인터럽트 상태 CSR */
        .packing = GASKET_IRQ_UNPACKED,
    } },

    .device_reset_cb  = apex_reset,
    .device_status_cb = apex_get_status,
};

/* Apex probe — Gasket 프레임워크가 호출 */
static int apex_pci_probe(struct pci_dev *pci_dev,
                          const struct pci_device_id *id)
{
    int ret;
    /* Gasket 프레임워크에 디바이스 등록 */
    ret = gasket_pci_add_device(pci_dev, &apex_desc);
    if (ret)
        return ret;
    /* /dev/apex_N 디바이스 노드 생성 */
    /* BAR 매핑, MSI-X 설정, 펌웨어 로드 등 Gasket이 처리 */
    return 0;
}

Gasket 디바이스 페이지 테이블

TPU는 자체 디바이스 페이지 테이블을 가지며, 호스트 사용자 공간(User Space) 가상 주소(Virtual Address)를 디바이스가 접근 가능한 DMA 주소로 변환합니다. Gasket 프레임워크가 이 페이지 테이블의 생명주기를 관리합니다.

Gasket 2단계 페이지 테이블 구조 사용자 공간 VA Gasket ioctl 디바이스 페이지 테이블 Level 0 (Directory) Entry 0 L1 DMA 주소 Entry 1 L1 DMA 주소 Entry 2 L1 DMA 주소 ... Level 1 (Page Table) PTE 0 PTE 1 PTE 2 ... PTE 0 PTE 1 PTE 2 ... PTE 0 PTE 1 PTE 2 ... 각 PTE: 호스트 물리 페이지 DMA 주소 TPU DMA 엔진이 이 페이지 테이블을 참조하여 호스트 메모리에 직접 접근 (scatter-gather DMA) 사용자 VA → get_user_pages_fast() → dma_map_page() → 디바이스 PT 엔트리 기록
/* 페이지 테이블 매핑 ioctl */
struct gasket_page_table_ioctl {
    u64 page_table_index;   /* 페이지 테이블 ID */
    u64 size;               /* 매핑 크기 (바이트) */
    u64 host_address;       /* 사용자 공간 가상 주소 */
    u64 device_address;     /* 디바이스 측 주소 (IOVA) */
};

/* Gasket 페이지 테이블 매핑 흐름:
 *
 * 1. 사용자 → ioctl(GASKET_IOCTL_MAP_BUFFER, &pt_ioctl)
 * 2. gasket_page_table_map():
 *    a. get_user_pages_fast() — 사용자 페이지 핀
 *    b. dma_map_page() — 각 페이지의 DMA 주소 획득
 *    c. 디바이스 페이지 테이블 엔트리에 DMA 주소 기록
 *    d. BAR의 페이지 테이블 베이스 레지스터 갱신
 * 3. TPU DMA 엔진이 디바이스 페이지 테이블을 참조하여
 *    호스트 메모리에 scatter-gather DMA 수행
 */

TPU DMA 데이터 전송 흐름

호스트 ↔ TPU 데이터 전송 흐름 Host CPU PCIe 채널 TPU Chip User Buffer (pinned memory) 모델 가중치 DMA → HBM (모델 가중치 로드) Input Data (pinned memory) 입력 텐서 DMA → HBM (입력 텐서) TensorCore (MXU 텐서 연산) Output Buffer (결과 수신) ← 출력 텐서 DMA HBM (출력 텐서) ← MSI-X IRQ (DMA 완료 인터럽트) 단계 1. 사용자 공간: 입력 데이터 버퍼 준비, ioctl로 매핑 요청 2. 커널 드라이버: get_user_pages_fast()로 페이지 핀, DMA 매핑 설정 3. 커널 드라이버: 명령어 큐(Doorbell)에 DMA 전송 명령 기록 4. TPU DMA 엔진: 호스트 메모리 → HBM으로 데이터 전송 5. TPU TensorCore: MXU에서 텐서 연산 실행 6. TPU DMA 엔진: HBM → 호스트 메모리로 결과 전송 7. TPU: MSI-X 인터럽트로 호스트에 완료 알림 8. 커널 드라이버: eventfd를 통해 사용자 공간에 통지

TPU 인터럽트 처리 (MSI-X / eventfd)

/* Gasket 인터럽트 처리 구조 */

/* 1. MSI-X 벡터 할당 (probe 시) */
static int gasket_interrupt_init(struct gasket_dev *dev)
{
    int ret;
    /* MSI-X 벡터 할당 */
    ret = pci_alloc_irq_vectors(dev->pci_dev,
                               dev->num_interrupts,
                               dev->num_interrupts,
                               PCI_IRQ_MSIX);
    if (ret < 0)
        return ret;

    /* 각 벡터에 ISR 등록 */
    for (int i = 0; i < dev->num_interrupts; i++) {
        int irq = pci_irq_vector(dev->pci_dev, i);
        request_irq(irq, gasket_interrupt_handler,
                    0, dev->driver_desc->name, dev);
    }
    return 0;
}

/* 2. ISR — eventfd를 통해 사용자 공간에 통지 */
static irqreturn_t gasket_interrupt_handler(int irq, void *ctx)
{
    struct gasket_dev *dev = ctx;
    struct eventfd_ctx *eventfd;

    /* 인터럽트 상태 레지스터 읽기 및 클리어 */
    u64 status = readq(dev->bar_data[0].virt_base + IRQ_STATUS_REG);
    writeq(status, dev->bar_data[0].virt_base + IRQ_CLEAR_REG);

    /* 등록된 eventfd에 시그널 전달 */
    eventfd = dev->interrupt_data->eventfds[0];
    if (eventfd)
        eventfd_signal(eventfd, 1);

    return IRQ_HANDLED;
}

/* 3. 사용자 공간에서 eventfd 등록 (ioctl) */
/*
 * 사용자 프로그램:
 *   int efd = eventfd(0, EFD_NONBLOCK);
 *   struct gasket_interrupt_eventfd ie = {
 *       .interrupt = 0,        // 인터럽트 인덱스
 *       .event_fd  = efd,      // eventfd 파일 디스크립터
 *   };
 *   ioctl(tpu_fd, GASKET_IOCTL_SET_EVENTFD, &ie);
 *
 *   // 비동기 완료 대기
 *   struct pollfd pfd = { .fd = efd, .events = POLLIN };
 *   poll(&pfd, 1, timeout_ms);
 */

TPU 펌웨어 로딩

/*
 * Edge TPU 펌웨어 로딩 과정
 *
 * 1. 드라이버 probe 시 request_firmware() 호출
 *    - 펌웨어 파일: /lib/firmware/google/apex_firmware.bin
 *    - 또는 유저가 sysfs를 통해 커스텀 펌웨어 로드
 *
 * 2. 펌웨어를 BAR2 영역에 DMA 전송
 *    - BAR2에 매핑된 디바이스 메모리에 직접 기록
 *    - 또는 DMA coherent 버퍼를 통해 전송
 *
 * 3. CSR에 펌웨어 실행 명령 기록
 *    - SCU 레지스터에 리셋 해제 + 실행 비트 설정
 *
 * 4. 펌웨어 준비 상태 폴링
 *    - HIB (Host Interface Block) 상태 레지스터 확인
 *    - 타임아웃 내에 READY 상태 전환 확인
 */

static int apex_reset(struct gasket_dev *gasket_dev)
{
    struct apex_dev *apex = gasket_dev_get_drvdata(gasket_dev);

    /* 1. 칩 소프트 리셋 */
    writeq(SCU_RESET_VALUE,
           gasket_dev->bar_data[2].virt_base + APEX_BAR2_REG_SCU_BASE);

    /* 2. 리셋 완료 대기 */
    usleep_range(100, 200);

    /* 3. 펌웨어 로드 (이미 호스트 메모리에 로드된 경우) */
    memcpy_toio(gasket_dev->bar_data[2].virt_base + FW_LOAD_OFFSET,
                apex->fw_data, apex->fw_size);

    /* 4. 펌웨어 실행 시작 */
    writeq(HIB_START_VALUE,
           gasket_dev->bar_data[2].virt_base + APEX_BAR2_REG_KERNEL_HIB);

    /* 5. READY 상태 폴링 */
    return readq_poll_timeout(
        gasket_dev->bar_data[2].virt_base + HIB_STATUS_REG,
        status, status == FW_STATUS_READY,
        100,          /* 100 µs 간격 */
        5000000);    /* 5초 타임아웃 */
}

ACCEL 서브시스템 (Linux 6.2+)

Linux 6.2부터 도입된 ACCEL (Accelerator) 서브시스템은 ML/AI 가속기를 위한 표준 커널 프레임워크입니다. DRM 코어 인프라를 재활용(Recycling)하면서 /dev/accel/accelN 별도 디바이스 노드를 제공하여 GPU와 가속기를 명확히 분리합니다.

구분DRM (/dev/dri/)ACCEL (/dev/accel/)
용도GPU 렌더링/디스플레이ML/AI 연산 가속
KMS필요 (디스플레이 출력)불필요 (연산 전용)
GEMGPU 버퍼 관리가속기 메모리 관리(Memory Management)
스케줄러(Scheduler)drm_sched (GPU 작업)디바이스별 자체 스케줄러
주요 드라이버i915, amdgpu, nouveauhabanalabs (Gaudi), qaic, ivpu
디바이스 노드/dev/dri/cardN, /dev/dri/renderDN/dev/accel/accelN
/* ACCEL 드라이버 등록 예시 (DRM 기반) */
#include <drm/drm_accel.h>
#include <drm/drm_drv.h>

static const struct drm_driver my_accel_driver = {
    .driver_features = DRIVER_COMPUTE_ACCEL,  /* 핵심: ACCEL 플래그 */
    .name            = "my_tpu",
    .desc            = "My TPU Accelerator",
    .date            = "20240101",
    .major           = 1,
    .minor           = 0,

    /* GEM 오브젝트 콜백 */
    .gem_create_object = my_gem_create_object,

    /* ioctl 테이블 (디바이스별 커스텀 명령) */
    .ioctls           = my_accel_ioctls,
    .num_ioctls       = ARRAY_SIZE(my_accel_ioctls),

    /* 파일 오퍼레이션 */
    .fops = &my_accel_fops,
};

/*
 * DRIVER_COMPUTE_ACCEL 플래그의 효과:
 * - /dev/accel/accelN 디바이스 노드 자동 생성 (/dev/dri/ 대신)
 * - KMS 관련 ioctl 비활성화 (모드 설정 불필요)
 * - DRM 코어의 메모리 관리/ioctl 인프라는 그대로 활용
 * - sysfs: /sys/class/accel/accel0/
 */

/* PCI probe에서 ACCEL 디바이스 초기화 */
static int my_tpu_probe(struct pci_dev *pdev,
                        const struct pci_device_id *id)
{
    struct drm_device *drm;
    int ret;

    /* DRM/ACCEL 디바이스 할당 */
    drm = drm_dev_alloc(&my_accel_driver, &pdev->dev);
    if (IS_ERR(drm))
        return PTR_ERR(drm);

    /* PCI 리소스 설정 */
    ret = pcim_enable_device(pdev);
    pci_set_master(pdev);
    dma_set_mask_and_coherent(&pdev->dev, DMA_BIT_MASK(64));

    /* 디바이스 등록 → /dev/accel/accel0 생성 */
    ret = drm_dev_register(drm, 0);
    return ret;
}
ACCEL vs Gasket: Gasket은 Google의 자체 프레임워크(staging)로 Edge TPU에 사용되고, ACCEL은 Linux 메인라인 표준 프레임워크(6.2+)로 DRM 인프라를 재활용합니다. 새로운 가속기 드라이버는 ACCEL 서브시스템을 사용하는 것이 권장됩니다.

TPU 소프트웨어 스택 (커널-사용자 공간 인터페이스)

User Space TensorFlow JAX PyTorch/XLA PJRT libtpu.so XLA Compiler (HLO → TPU 명령어) ioctl() / mmap() / eventfd Kernel Space Gasket / ACCEL Framework PCI Subsystem DMA Engine MSI-X IRQ IOMMU BAR MMIO TPU Hardware (MXU + HBM + DMA + ICI)

TPU 주요 ioctl 인터페이스

ioctl기능설명
GASKET_IOCTL_RESET디바이스 리셋TPU 칩 소프트 리셋 + 펌웨어 재로드
GASKET_IOCTL_MAP_BUFFER버퍼 매핑사용자 버퍼를 디바이스 페이지 테이블에 등록
GASKET_IOCTL_UNMAP_BUFFER버퍼 해제디바이스 페이지 테이블 엔트리 제거 + 페이지 언핀
GASKET_IOCTL_SET_EVENTFD이벤트 등록MSI-X 인터럽트에 eventfd 연결
GASKET_IOCTL_CLEAR_EVENTFD이벤트 해제eventfd 연결 해제
GASKET_IOCTL_NUMBER_PAGE_TABLES페이지 테이블 수지원하는 디바이스 페이지 테이블 수 조회
GASKET_IOCTL_PAGE_TABLE_SIZE테이블 크기특정 페이지 테이블의 최대 엔트리 수 조회
GASKET_IOCTL_SIMPLE_PAGE_TABLE_SIZE단순 테이블 크기단순(1단계) 페이지 테이블 크기 조회

TPU sysfs 인터페이스

# Edge TPU sysfs 경로
ls /sys/class/apex/apex_0/

# 디바이스 속성
cat /sys/class/apex/apex_0/device/vendor     # 0x1ac1 (Google)
cat /sys/class/apex/apex_0/device/device     # 0x089a (Edge TPU)
cat /sys/class/apex/apex_0/temp              # 칩 온도 (밀리도)
cat /sys/class/apex/apex_0/status            # 디바이스 상태

# ACCEL 서브시스템 (Linux 6.2+)
ls /sys/class/accel/accel0/
cat /sys/class/accel/accel0/device/vendor
cat /sys/class/accel/accel0/device/device

# PCIe 링크 상태 확인
lspci -vvv -s 03:00.0 | grep -E "LnkSta|LnkCap"
# LnkCap: Port #0, Speed 5GT/s, Width x1
# LnkSta: Speed 5GT/s, Width x1

# DMA 상태 확인
cat /sys/kernel/debug/dma-buf/bufinfo        # DMA-BUF 사용 현황
cat /proc/iomem | grep -i apex               # BAR 매핑 주소

Edge TPU 드라이버 설치 및 확인

# Edge TPU (Coral) PCIe 드라이버 — Gasket 모듈 로드

# 1. 커널 모듈 확인/로드
modprobe gasket       # Gasket 코어 프레임워크
modprobe apex         # Edge TPU (Apex) 드라이버
lsmod | grep apex
# apex                    28672  0
# gasket                  77824  1 apex

# 2. 디바이스 노드 확인
ls -la /dev/apex_*
# crw-rw---- 1 root apex 120, 0 ... /dev/apex_0

# 3. udev 규칙 (권한 설정)
# /etc/udev/rules.d/65-apex.rules
# SUBSYSTEM=="apex", MODE="0660", GROUP="apex"

# 4. dmesg 로그 확인
dmesg | grep -i apex
# [  2.345] apex 0000:03:00.0: Apex device found
# [  2.346] apex 0000:03:00.0: enabling device (0000 -> 0002)
# [  2.347] apex 0000:03:00.0: successfully loaded firmware

# 5. 추론 테스트 (TensorFlow Lite + Edge TPU Runtime)
# Python에서:
# import tflite_runtime.interpreter as tflite
# interpreter = tflite.Interpreter(
#     model_path='model_edgetpu.tflite',
#     experimental_delegates=[
#         tflite.load_delegate('libedgetpu.so.1')])

TPU 관련 Kconfig 옵션

옵션설명
CONFIG_STAGINGStaging 드라이버 활성화 (Gasket 포함)
CONFIG_STAGING_GASKET_FRAMEWORKGasket 코어 프레임워크
CONFIG_STAGING_APEX_DRIVEREdge TPU (Apex) 드라이버
CONFIG_DRM_ACCELACCEL 서브시스템 (Linux 6.2+, DRM 기반)
CONFIG_DRMDRM 코어 (ACCEL 의존성)
CONFIG_PCI_MSIMSI/MSI-X (TPU 인터럽트)
CONFIG_IOMMU_SUPPORTIOMMU (DMA 보안 격리)
CONFIG_DMA_SHARED_BUFFERDMA-BUF (버퍼 공유)
CONFIG_FW_LOADER펌웨어 로딩 인프라

TPU vs GPU: 커널 드라이버 관점 비교

특성GPU (DRM/KMS)TPU (Gasket/ACCEL)
디스플레이 출력KMS로 제어 (CRTC, Encoder, Connector)없음 (연산 전용)
메모리 관리GEM/TTM (GPU VRAM)디바이스 페이지 테이블 (HBM)
작업 제출커맨드 링 버퍼 (GPU 명령어)DMA 디스크립터 + Doorbell CSR
스케줄링drm_sched (GPU 컨텍스트 전환)하드웨어 자체 스케줄링
사용자 APIOpenGL/Vulkan → DRM ioctlXLA/PJRT → Gasket ioctl
칩 간 통신NVLink/xGMI (제한적)ICI (대규모 3D Torus)
데이터 타입fp32/fp16/int (다목적)bf16/int8 (ML 특화)
프로그래밍범용 셰이더시스톨릭 어레이 고정 파이프라인(Pipeline)

TPU 관련 커널 소스 구조

경로설명
drivers/staging/gasket/Gasket 프레임워크 (Edge TPU 지원)
drivers/staging/gasket/gasket_core.cPCI probe, BAR 매핑, char device 관리
drivers/staging/gasket/gasket_page_table.c디바이스 페이지 테이블 (DMA 주소 변환)
drivers/staging/gasket/gasket_interrupt.cMSI-X 설정, eventfd 통지
drivers/staging/gasket/apex_driver.cEdge TPU (Apex) 드라이버
drivers/accel/ACCEL 서브시스템 (Linux 6.2+)
drivers/accel/habanalabs/Intel Gaudi 가속기 드라이버
drivers/accel/qaic/Qualcomm Cloud AI 100 드라이버
drivers/accel/ivpu/Intel VPU (NPU) 드라이버
include/drm/drm_accel.hACCEL 서브시스템 API 헤더

TPU Pod 토폴로지 — 3D Torus

TPU v4부터 칩 간 연결에 3D Torus ICI 토폴로지를 채용합니다. TPU v4 Pod는 4×4×4 = 64개 칩을 하나의 큐브(Cube)로 구성하며, 각 축의 양 끝이 연결되어(wrap-around) 토러스 구조를 형성합니다. 이 구조 덕분에 임의의 두 칩 사이 최대 홉(hop) 수가 절반으로 줄어들어, 대규모 분산 학습에서 집합 통신(Collective Communication) 지연이 크게 감소합니다.

TPU v5p는 이 토러스를 더 큰 규모로 확장하여 최대 8,960개 칩을 하나의 Pod로 연결할 수 있습니다. Pod 슬라이스(Pod Slice)는 전체 Pod 중 사용자에게 할당되는 연속적인 칩 부분 집합을 의미하며, Google Cloud에서 TPU를 프로비저닝할 때 슬라이스 단위로 요청합니다. 전체 Pod(Full Pod)는 하나의 물리적 토러스 네트워크를 공유하는 모든 칩의 집합입니다.

4×4 Torus 슬라이스 (하나의 XY 평면, Z축 생략) Chip [0,0] Chip [0,1] Chip [0,2] Chip [0,3] wrap wrap Chip [1,0] Chip [1,1] Chip [1,2] Chip [1,3] Chip [2,0] Chip [2,1] Chip [2,2] Chip [2,3] Chip [3,0] Chip [3,1] Chip [3,2] Chip [3,3] wrap (Y축) ICI 직접 연결 Wrap-around 연결 (토러스) 실제 3D: Z축 4층 적층 (여기서는 XY 1개 면만 표시) TPU v4 Pod: 4×4×4 = 64칩/큐브 — 최대 4,096칩/Pod

Allreduce 최적화: Torus 토폴로지에서는 각 축 방향으로 독립적인 링(Ring)을 구성할 수 있습니다. 3D Torus의 경우 X, Y, Z 세 축에 대해 병렬로 reduce-scatter와 all-gather를 수행하여, 단순 Ring Allreduce 대비 통신량을 ~3배 줄일 수 있습니다. XLA 런타임은 토폴로지 정보를 기반으로 최적의 집합 통신 패턴을 자동 선택합니다.

개념설명
Pod Slice전체 Pod 중 사용자에게 할당되는 연속 칩 부분집합 (예: v4-32 = 32칩 슬라이스)
Full Pod하나의 물리적 3D Torus를 공유하는 모든 칩 (TPU v4 = 최대 4,096칩)
ICI 대역폭TPU v4 기준 링크당 ~24 GB/s, 칩당 6링크 = 144 GB/s 양방향
Torus 최대 홉N×N×N Torus에서 최대 홉 = 3×⌊N/2⌋ (Mesh 대비 절반)

XLA/JAX 컴파일 파이프라인

TPU에서 실행되는 모든 연산은 XLA (Accelerated Linear Algebra) 컴파일러를 통해 TPU 네이티브 명령어로 변환됩니다. JAX, TensorFlow, PyTorch/XLA 등 프론트엔드 프레임워크가 생성한 연산 그래프는 HLO(High Level Operations) IR로 표현되며, XLA가 이를 최적화하고 TPU 하드웨어에 맞게 낮춥니다(lowering).

XLA/JAX 컴파일 파이프라인 JAX / TF Python 코드 HLO IR 연산 그래프 XLA Optimizer 융합/레이아웃/타일링 TPU Lowering MXU 명령어 생성 TPU Executable DMA 전송 → 실행 PJRT (Pretty-much Just a Runtime) 프레임워크 ↔ 백엔드 추상화 계층 — 디바이스 관리, 버퍼 할당, 실행 디스패치 SPMD (Single Program Multiple Data) 파티셔닝 GSPMD 파티셔너 텐서 분할 전략 결정 (sharding) Collective Ops 삽입 all-reduce, all-gather, reduce-scatter 칩별 실행 동일 프로그램, 다른 데이터 파티션 혼합 정밀도 (Mixed Precision) BF16 행렬곱 FP32 누적/갱신 FP8 (v5p+)

PJRT (Pretty-much Just a Runtime)는 프레임워크(JAX, TensorFlow, PyTorch/XLA)와 하드웨어 백엔드(TPU, GPU, CPU) 사이의 표준화된 런타임 인터페이스입니다. PJRT를 통해 프레임워크는 디바이스 종류에 관계없이 동일한 API로 버퍼 할당, 컴파일, 실행을 요청할 수 있으며, 커널 드라이버와의 상호작용(DMA 전송, 인터럽트 대기 등)은 PJRT 백엔드가 처리합니다.

SPMD 파티셔닝: XLA의 GSPMD(Generalized SPMD) 파티셔너는 사용자가 지정한 샤딩 어노테이션(jax.sharding)을 기반으로 텐서를 여러 칩에 분할합니다. 분할된 각 칩은 동일한 프로그램을 실행하되 서로 다른 데이터 파티션을 처리하며, 필요 시 ICI를 통해 집합 연산(all-reduce, all-gather)을 수행합니다.

BF16/FP8 혼합 정밀도: XLA는 행렬곱 연산에 BF16을 사용하고 누적(accumulation) 및 가중치 갱신에는 FP32를 사용하는 혼합 정밀도 전략을 자동 적용합니다. TPU v5p 이후에서는 FP8(E4M3/E5M2) 데이터 타입도 지원하여 MXU 처리량을 추가로 향상시킵니다.

TPU VM 아키텍처

Google Cloud의 TPU 접근 방식은 두 가지입니다: 레거시 TPU NodeTPU VM. TPU VM은 TPU v2-alpha 이후 기본 모드로, 사용자가 TPU 칩에 직접 연결된 호스트 머신에 SSH로 접근할 수 있습니다.

TPU VM vs 레거시 TPU Node 아키텍처 TPU VM (현재 기본 모드) 사용자 SSH 직접 접근 TPU VM Host Linux (커스텀 커널 모듈 가능) PCIe / 커스텀 TPU Chips MXU + HBM ICI ··· 레거시 TPU Node (비권장) 사용자 VM GCE 인스턴스 gRPC TPU Server (숨겨진) gRPC proxy → 드라이버 TPU Chips MXU + HBM TPU VM의 장점 1. 낮은 지연: gRPC 프록시 제거 → 호스트-TPU 간 직접 통신 2. 직접 디바이스 접근: /dev/accel*, sysfs, perf 프로파일링 가능 3. 커스텀 커널 모듈: 사용자 정의 드라이버/모듈 로드 가능 (Gasket 커스터마이징 등) 4. Multislice: 여러 Pod 슬라이스를 DCN(Data Center Network)으로 연결하여 초대규모 학습

TPU Node (레거시) 모드에서는 사용자의 GCE VM과 TPU 호스트가 분리되어 있어, gRPC를 통해 원격으로 TPU에 접근했습니다. 이 방식은 네트워크 지연이 추가되고, TPU 호스트의 파일 시스템이나 디바이스 노드에 직접 접근할 수 없는 제약이 있었습니다.

TPU VM 모드에서는 사용자가 TPU 칩에 물리적으로 연결된 호스트 머신에 직접 SSH 접속합니다. 이를 통해 /dev/accel* 디바이스 노드에 직접 접근하고, sysfs를 통해 TPU 상태를 모니터링하며, 필요 시 커스텀 커널 모듈을 로드할 수 있습니다. JAX/TensorFlow 라이브러리를 호스트에서 직접 실행하므로 gRPC 오버헤드가 제거됩니다.

Multislice 학습: 단일 Pod 슬라이스의 칩 수를 초과하는 대규모 모델 학습 시, 여러 Pod 슬라이스를 Google의 DCN(Data Center Network)으로 연결하는 Multislice 구성을 사용합니다. 각 슬라이스 내부는 ICI로 고속 통신하고, 슬라이스 간에는 DCN을 통해 집합 연산을 수행합니다. XLA의 GSPMD 파티셔너가 ICI/DCN 경계를 인식하여 통신 패턴을 최적화합니다.

# TPU VM 접속 및 상태 확인
$ gcloud compute tpus tpu-vm ssh my-tpu-node --zone=us-central1-a

# TPU 디바이스 노드 확인
$ ls /dev/accel*
/dev/accel0  /dev/accel1  /dev/accel2  /dev/accel3

# TPU 칩 상태 확인 (sysfs)
$ cat /sys/class/accel/accel0/device/status
alive

# Multislice 환경 변수 예시
$ export TPU_WORKER_ID=0
$ export TPU_WORKER_HOSTNAMES="10.0.0.1,10.0.0.2,10.0.0.3,10.0.0.4"

전력 및 냉각 관리

TPU의 전력 소비와 냉각 방식은 세대별로 크게 달라집니다. 특히 TPU v3부터 수냉(Liquid Cooling)이 도입되어 고밀도 배치에서도 안정적인 열 관리를 보장합니다.

세대TDP (칩당)냉각 방식비고
TPU v1~40 W공냉 (Air)추론 전용, PCIe 카드 폼 팩터
TPU v2~250 W공냉 (Air)학습 + 추론, 커스텀 보드
TPU v3~450 W수냉 (Liquid)최초 수냉 도입 — 전용 냉각수 순환 시스템
TPU v4~275 W수냉 (Liquid)7nm 공정으로 전력 효율 개선
TPU v5p~400 W (추정)수냉 (Liquid)성능 대비 전력 효율 최적화
Edge TPU~2 W공냉/수동임베디드, USB/PCIe M.2 폼 팩터

데이터센터 수냉 시스템: TPU v3 이후 칩은 직접 수냉 방식을 사용합니다. 각 TPU 보드에 냉각수 유로(cold plate)가 부착되어 칩에서 발생하는 열을 냉각수로 직접 전달합니다. 데이터센터 수준에서는 냉각 타워(cooling tower)와 열교환기를 통해 냉각수 온도를 관리하며, 이를 통해 공냉 대비 PUE(Power Usage Effectiveness)를 크게 개선합니다.

Edge TPU 열 관리 — 커널 연동: Edge TPU(Coral)는 임베디드 환경에서 동작하므로 리눅스 커널의 thermal 프레임워크와 연동됩니다. Apex 드라이버는 sysfs를 통해 칩 온도와 쓰로틀링(throttling) 상태를 노출하며, 설정 온도 초과 시 자동으로 클록 주파수를 낮추거나 추론 요청을 거부합니다.

# Edge TPU 온도 확인 (Coral USB Accelerator)
$ cat /sys/class/apex/apex_0/temp
62

# Edge TPU 쓰로틀링 상태 확인
$ cat /sys/class/apex/apex_0/trip_point0_temp
85000
$ cat /sys/class/apex/apex_0/trip_point1_temp
100000

# 커널 thermal zone 연동 확인
$ ls /sys/class/thermal/thermal_zone*/type
# apex_thermal 등의 이름으로 등록됨

# Edge TPU 전력 상태 확인 (PCIe M.2 모듈)
$ lspci -vv -s 03:00.0 | grep -i power
    Capabilities: [80] Power Management version 3
    Status: D0 PME-Enable+ DSel=0 DScale=0 PME-
Edge TPU 열 보호: trip_point0(일반적으로 85도C)에 도달하면 추론 속도가 감소하며, trip_point1(100도C)에 도달하면 디바이스가 정지합니다. 임베디드 제품 설계 시 방열판 크기와 주변 온도를 충분히 고려하여 열 설계 전력(TDP) 마진을 확보해야 합니다.

참고자료