Vengineerの妄想(準備期間)

人生は短いけど、長いです。人生を楽しみましょう!

Google Tensor G2 の Edge TPU のデバイスドライバからわかったこと!(その2)

はじめに

vengineer.hatenablog.com

の続きとして、

Google Tensor G2 の Edge TPU のデバイスドライバからわかったこと!(その2)

その1の後半で取り上げて、GXPデバイスドライバソースコードも眺めてみたら、もう少しわかったことがありました。

GXPソースコード

Google Tensor G2 の Edge TPU (janeiro) の下記のmailbox の部分にあった、EXT_DSP_MAILBOXES がありました。

#define EDGETPU_NUM_VII_MAILBOXES 7
#define EDGETPU_NUM_P2P_MAILBOXES 0
#define EDGETPU_NUM_EXT_DSP_MAILBOXES 4
#define EDGETPU_NUM_EXT_AOC_MAILBOXES 1
#define EDGETPU_NUM_EXT_MAILBOXES (EDGETPU_NUM_EXT_DSP_MAILBOXES + EDGETPU_NUM_EXT_AOC_MAILBOXES)
#define EDGETPU_NUM_MAILBOXES (EDGETPU_NUM_VII_MAILBOXES + EDGETPU_NUM_EXT_MAILBOXES + 1)

これに対して、GXP の gxp.h の中に下記のようなコードを見つけました。

struct gxp_tpu_mbx_queue_ioctl {
    __u32 tpu_fd; /* TPU virtual device group fd */
    /*
    * Bitfield indicating which virtual cores to allocate and map the
    * buffers for.
    * To map for virtual core X, set bit X in this field, i.e. `1 << X`.
    *
    * This field is not used by the unmap IOCTL, which always unmaps the
    * buffers for all cores it had been mapped for.
    */
    __u32 virtual_core_list;
    /*
    * The user address of an edgetpu_mailbox_attr struct, containing
    * cmd/rsp queue size, mailbox priority and other relevant info.
    * This structure is defined in edgetpu.h in the TPU driver.
    */
    __u64 attr_ptr;
};
/*
 * Map TPU-DSP mailbox cmd/rsp queue buffers.
 *
 * The client must have allocated a virtual device.
 */
#define GXP_MAP_TPU_MBX_QUEUE \
   _IOW(GXP_IOCTL_BASE, 13, struct gxp_tpu_mbx_queue_ioctl)
/*
 * Un-map TPU-DSP mailbox cmd/rsp queue buffers previously mapped by
 * GXP_MAP_TPU_MBX_QUEUE.
 *
 * Only the @tpu_fd field will be used. Other fields will be fetched
 * from the kernel's internal records. It is recommended to use the argument
 * that was passed in GXP_MAP_TPU_MBX_QUEUE to un-map the buffers.
 *
 * The client must have allocated a virtual device.
 */
#define GXP_UNMAP_TPU_MBX_QUEUE \
   _IOW(GXP_IOCTL_BASE, 14, struct gxp_tpu_mbx_queue_ioctl)

また、gxp-platform.c にも、TPUというコードがあります。

#if (IS_ENABLED(CONFIG_GXP_TEST) || IS_ENABLED(CONFIG_ANDROID)) && !IS_ENABLED(CONFIG_GXP_GEM5)
#include <soc/google/tpu-ext.h>
#endif

edgetpu_ext_driver_cmd コマンドを使っています。

static int gxp_map_tpu_mbx_queue(struct gxp_client *client,
                 struct gxp_tpu_mbx_queue_ioctl __user *argp)
{
#if (IS_ENABLED(CONFIG_GXP_TEST) || IS_ENABLED(CONFIG_ANDROID)) && !IS_ENABLED(CONFIG_GXP_GEM5)
    struct gxp_dev *gxp = client->gxp;
    struct edgetpu_ext_mailbox_info *mbx_info;
    struct gxp_tpu_mbx_queue_ioctl ibuf;
    struct edgetpu_ext_client_info gxp_tpu_info;
    u32 phys_core_list = 0;
    u32 virtual_core_list;
    u32 core_count;
    int ret = 0;
    if (!gxp->tpu_dev.mbx_paddr) {
        dev_err(gxp->dev, "%s: TPU is not available for interop\n",
            __func__);
        return -EINVAL;
    }
    if (copy_from_user(&ibuf, argp, sizeof(ibuf)))
        return -EFAULT;
    down_write(&client->semaphore);
    if (!check_client_has_available_vd(client, "GXP_MAP_TPU_MBX_QUEUE")) {
        ret = -ENODEV;
        goto out_unlock_client_semaphore;
    }
    down_read(&gxp->vd_semaphore);
    virtual_core_list = ibuf.virtual_core_list;
    core_count = hweight_long(virtual_core_list);
    phys_core_list = gxp_vd_virt_core_list_to_phys_core_list(
        client->vd, virtual_core_list);
    if (!phys_core_list) {
        dev_err(gxp->dev, "%s: invalid virtual core list 0x%x\n",
            __func__, virtual_core_list);
        ret = -EINVAL;
        goto out;
    }
    mbx_info =
        kmalloc(sizeof(struct edgetpu_ext_mailbox_info) + core_count *
            sizeof(struct edgetpu_ext_mailbox_descriptor),
            GFP_KERNEL);
    if (!mbx_info) {
        ret = -ENOMEM;
        goto out;
    }
    if (client->tpu_file) {
        dev_err(gxp->dev, "Mappings already exist for TPU mailboxes");
        ret = -EBUSY;
        goto out_free;
    }
    gxp_tpu_info.tpu_fd = ibuf.tpu_fd;
    gxp_tpu_info.mbox_map = phys_core_list;
    gxp_tpu_info.attr = (struct edgetpu_mailbox_attr __user *)ibuf.attr_ptr;
    ret = edgetpu_ext_driver_cmd(gxp->tpu_dev.dev,
                     EDGETPU_EXTERNAL_CLIENT_TYPE_DSP,
                     ALLOCATE_EXTERNAL_MAILBOX, &gxp_tpu_info,
                     mbx_info);
    if (ret) {
        dev_err(gxp->dev, "Failed to allocate ext TPU mailboxes %d",
            ret);
        goto out_free;
    }
    /*
    * If someone is attacking us through this interface -
    * it's possible that ibuf.tpu_fd here is already a different file from
    * the one passed to edgetpu_ext_driver_cmd() (if the runtime closes the
    * FD and opens another file exactly between the TPU driver call above
    * and the fget below).
    * But the worst consequence of this attack is we fget() ourselves (GXP
    * FD), which only leads to memory leak (because the file object has a
    * reference to itself). The race is also hard to hit so we don't insist
    * on preventing it.
    */
    client->tpu_file = fget(ibuf.tpu_fd);
    if (!client->tpu_file) {
        edgetpu_ext_driver_cmd(gxp->tpu_dev.dev,
                       EDGETPU_EXTERNAL_CLIENT_TYPE_DSP,
                       FREE_EXTERNAL_MAILBOX, &gxp_tpu_info,
                       NULL);
        ret = -EINVAL;
        goto out_free;
    }
    /* Align queue size to page size for iommu map. */
    mbx_info->cmdq_size = ALIGN(mbx_info->cmdq_size, PAGE_SIZE);
    mbx_info->respq_size = ALIGN(mbx_info->respq_size, PAGE_SIZE);
    ret = gxp_dma_map_tpu_buffer(gxp, client->vd, virtual_core_list,
                     phys_core_list, mbx_info);
    if (ret) {
        dev_err(gxp->dev, "Failed to map TPU mailbox buffer %d", ret);
        fput(client->tpu_file);
        client->tpu_file = NULL;
        edgetpu_ext_driver_cmd(gxp->tpu_dev.dev,
                       EDGETPU_EXTERNAL_CLIENT_TYPE_DSP,
                       FREE_EXTERNAL_MAILBOX, &gxp_tpu_info,
                       NULL);
        goto out_free;
    }
    client->mbx_desc.phys_core_list = phys_core_list;
    client->mbx_desc.virt_core_list = virtual_core_list;
    client->mbx_desc.cmdq_size = mbx_info->cmdq_size;
    client->mbx_desc.respq_size = mbx_info->respq_size;
out_free:
    kfree(mbx_info);
out:
    up_read(&gxp->vd_semaphore);
out_unlock_client_semaphore:
    up_write(&client->semaphore);
    return ret;
#else
    return -ENODEV;
#endif
}
static int gxp_unmap_tpu_mbx_queue(struct gxp_client *client,
                   struct gxp_tpu_mbx_queue_ioctl __user *argp)
{
#if (IS_ENABLED(CONFIG_GXP_TEST) || IS_ENABLED(CONFIG_ANDROID)) && !IS_ENABLED(CONFIG_GXP_GEM5)
    struct gxp_dev *gxp = client->gxp;
    struct gxp_tpu_mbx_queue_ioctl ibuf;
    struct edgetpu_ext_client_info gxp_tpu_info;
    int ret = 0;
    if (copy_from_user(&ibuf, argp, sizeof(ibuf)))
        return -EFAULT;
    down_write(&client->semaphore);
    if (!client->vd) {
        dev_err(gxp->dev,
            "GXP_UNMAP_TPU_MBX_QUEUE requires the client allocate a VIRTUAL_DEVICE\n");
        ret = -ENODEV;
        goto out;
    }
    if (!client->tpu_file) {
        dev_err(gxp->dev, "No mappings exist for TPU mailboxes");
        ret = -EINVAL;
        goto out;
    }
    gxp_dma_unmap_tpu_buffer(gxp, client->vd, client->mbx_desc);
    gxp_tpu_info.tpu_fd = ibuf.tpu_fd;
    edgetpu_ext_driver_cmd(gxp->tpu_dev.dev,
                   EDGETPU_EXTERNAL_CLIENT_TYPE_DSP,
                   FREE_EXTERNAL_MAILBOX, &gxp_tpu_info, NULL);
    fput(client->tpu_file);
    client->tpu_file = NULL;
out:
    up_write(&client->semaphore);
    return ret;
#else
    return -ENODEV;
#endif
}

static int gxp_platform_probe(struct platform_device *pdev) の中で TPU を使っています。tpu-device がデバイス名のようです。tpu-device から reg から tpu mailbox のレジスタのベースアドレスを見つけています。

 tpu_found = true;
    /* Get TPU device from device tree */
    np = of_parse_phandle(dev->of_node, "tpu-device", 0);
    if (IS_ERR_OR_NULL(np)) {
        dev_warn(dev, "No tpu-device in device tree\n");
        tpu_found = false;
    }
    tpu_pdev = of_find_device_by_node(np);
    if (!tpu_pdev) {
        dev_err(dev, "TPU device not found\n");
        tpu_found = false;
    }
    /* get tpu mailbox register base */
    ret = of_property_read_u64_index(np, "reg", 0, &base_addr);
    of_node_put(np);
    if (ret) {
        dev_warn(dev, "Unable to get tpu-device base address\n");
        tpu_found = false;
    }
    /* get gxp-tpu mailbox register offset */
    ret = of_property_read_u64(dev->of_node, "gxp-tpu-mbx-offset",
                   &offset);
    if (ret) {
        dev_warn(dev, "Unable to get tpu-device mailbox offset\n");
        tpu_found = false;
    }
    if (tpu_found) {
        gxp->tpu_dev.dev = &tpu_pdev->dev;
        get_device(gxp->tpu_dev.dev);
        gxp->tpu_dev.mbx_paddr = base_addr + offset;
    } else {
        dev_warn(dev, "TPU will not be available for interop\n");
        gxp->tpu_dev.mbx_paddr = 0;
    }
    ret = gxp_dma_init(gxp);
    if (ret) {
        dev_err(dev, "Failed to initialize GXP DMA interface\n");
        goto err_put_tpu_dev;
    }

Makefile にもありました

# Access TPU driver's exported symbols.
KBUILD_EXTRA_SYMBOLS += ../google-modules/edgetpu/janeiro/drivers/edgetpu/Module.symvers

ZEBU

gxp-config.h に、ZEBU なる文字がありました。これ、Synopsys の ZEBU のことでしょうか?

#if defined(CONFIG_GXP_ZEBU) || defined(CONFIG_GXP_IP_ZEBU)
#define GXP_TIME_DELAY_FACTOR 20
#else
#define GXP_TIME_DELAY_FACTOR 1
#endif

Makefile にもありました。

# If building via make directly, specify target platform by adding
#     "GXP_PLATFORM=<target>"
# With one of the following values:
#     - CLOUDRIPPER
#     - ZEBU
#     - IP_ZEBU
# Defaults to building for CLOUDRIPPER if not otherwise specified.
GXP_PLATFORM ?= CLOUDRIPPER
GXP_CHIP ?= AMALTHEA

CLOUDRIPPER は、どうやら、Google Tensor G2 開発用ボードの codename のようです。G1 開発用ボードは、Slider のようです。

おわりに

Google Tensor G2の Edge TPU と DSP が繋がっていることが分かりました。