gvisor.dev/gvisor@v0.0.0-20240520182842-f9d4d51c7e0f/pkg/abi/tpu/tpu.go (about)

     1  // Copyright 2023 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package tpu defines constants used to interact with TPUs. The constants are
    16  // derived from those listed in  https://github.com/tensorflow/tpu/blob/master/tools/driver/drivers/char/tpu_common
    17  package tpu
    18  
    19  const (
    20  	// SizeOfTPUV4InterruptList is the total number of valid
    21  	// (BAR Index, Register Offset) pairs.
    22  	SizeOfTPUV4InterruptList = uint64(45)
    23  
    24  	// NumberOfTPUV4PageTables is the number of gasket page tables used by the
    25  	// TPU V4 driver.
    26  	NumberOfTPUV4PageTables = uint64(1)
    27  
    28  	// TPUV4DeviceID is the PCI device ID of TPU V4 hardware.
    29  	TPUV4DeviceID = 0x005E
    30  
    31  	// SizeOfTPUV4liteInterruptList is the total number of valid
    32  	// (BAR Index, Register Offset) pairs.
    33  	SizeOfTPUV4liteInterruptList = uint64(37)
    34  
    35  	// NumberOfTPUV4litePageTables is the number of gasket page tables used by the
    36  	// TPU V4 driver
    37  	NumberOfTPUV4litePageTables = uint64(1)
    38  
    39  	// TPUV4liteDeviceID is the PCI device ID of TPU V4lite hardware.
    40  	TPUV4liteDeviceID = 0x0056
    41  
    42  	// TPUV5eDeviceID is the PCI device ID of TPU V5e hardware.
    43  	TPUV5eDeviceID = 0x0063
    44  )
    45  
    46  // TPUV4InterruptsMap maps BAR indices to valid register offsets.
    47  var (
    48  	TPUV4InterruptsMap = map[uint64]map[uint64]struct{}{
    49  		2: map[uint64]struct{}{
    50  			0x15b0008: struct{}{},
    51  			0x15b0000: struct{}{},
    52  			0x16b0008: struct{}{},
    53  			0x16b0000: struct{}{},
    54  			0x17b0008: struct{}{},
    55  			0x17b0000: struct{}{},
    56  			0x18b0008: struct{}{},
    57  			0x18b0000: struct{}{},
    58  			0x19b0020: struct{}{},
    59  			0x19b0000: struct{}{},
    60  			0x19b0008: struct{}{},
    61  			0x19b0010: struct{}{},
    62  			0x19b0018: struct{}{},
    63  			0x1ab0020: struct{}{},
    64  			0x1ab0000: struct{}{},
    65  			0x1ab0008: struct{}{},
    66  			0x1ab0010: struct{}{},
    67  			0x1ab0018: struct{}{},
    68  			0x4720000: struct{}{},
    69  			0x1bb0000: struct{}{},
    70  			0x1bb0008: struct{}{},
    71  			0x1bb0010: struct{}{},
    72  			0x1bb0018: struct{}{},
    73  			0x90000:   struct{}{},
    74  			0xb0000:   struct{}{},
    75  			0xd0000:   struct{}{},
    76  			0xf0000:   struct{}{},
    77  			0x110000:  struct{}{},
    78  			0x130000:  struct{}{},
    79  			0x150000:  struct{}{},
    80  			0x170000:  struct{}{},
    81  			0x190000:  struct{}{},
    82  			0x1b0000:  struct{}{},
    83  			0x1d0000:  struct{}{},
    84  			0x1f0000:  struct{}{},
    85  			0x210000:  struct{}{},
    86  			0x230000:  struct{}{},
    87  			0x250000:  struct{}{},
    88  			0x270000:  struct{}{},
    89  			0x290000:  struct{}{},
    90  			0x2b0000:  struct{}{},
    91  			0x2d0000:  struct{}{},
    92  			0x2f0000:  struct{}{},
    93  			0x310000:  struct{}{},
    94  			0x4720018: struct{}{},
    95  		},
    96  	}
    97  
    98  	// TPUV4liteInterruptsMap maps BAR indices to valid register offsets.
    99  	TPUV4liteInterruptsMap = map[uint64]map[uint64]struct{}{
   100  		2: map[uint64]struct{}{
   101  			0x19b0020: struct{}{},
   102  			0x19b0000: struct{}{},
   103  			0x19b0008: struct{}{},
   104  			0x19b0010: struct{}{},
   105  			0x19b0018: struct{}{},
   106  			0x1ab0020: struct{}{},
   107  			0x1ab0000: struct{}{},
   108  			0x1ab0008: struct{}{},
   109  			0x1ab0010: struct{}{},
   110  			0x1ab0018: struct{}{},
   111  			0x4720000: struct{}{},
   112  			0x1bb0000: struct{}{},
   113  			0x1bb0008: struct{}{},
   114  			0x1bb0010: struct{}{},
   115  			0x1bb0018: struct{}{},
   116  			0x90000:   struct{}{},
   117  			0xb0000:   struct{}{},
   118  			0xd0000:   struct{}{},
   119  			0xf0000:   struct{}{},
   120  			0x110000:  struct{}{},
   121  			0x130000:  struct{}{},
   122  			0x150000:  struct{}{},
   123  			0x170000:  struct{}{},
   124  			0x190000:  struct{}{},
   125  			0x1b0000:  struct{}{},
   126  			0x1d0000:  struct{}{},
   127  			0x1f0000:  struct{}{},
   128  			0x210000:  struct{}{},
   129  			0x230000:  struct{}{},
   130  			0x250000:  struct{}{},
   131  			0x270000:  struct{}{},
   132  			0x290000:  struct{}{},
   133  			0x2b0000:  struct{}{},
   134  			0x2d0000:  struct{}{},
   135  			0x2f0000:  struct{}{},
   136  			0x310000:  struct{}{},
   137  			0x4720018: struct{}{},
   138  		},
   139  	}
   140  )