From 92edcff8ec369449a7d8f4eef796aadbeb3b728f Mon Sep 17 00:00:00 2001
From: technowhizz <7688823+technowhizz@users.noreply.github.com>
Date: Wed, 27 Mar 2024 13:37:08 +0000
Subject: [PATCH] Add support for specfying vfio id's

---
 roles/iommu/README.md         | 10 ++++------
 roles/iommu/handlers/main.yml | 12 ++++++++++++
 roles/iommu/tasks/main.yml    | 37 +++++++++++++++++++++++++++++++++--
 3 files changed, 51 insertions(+), 8 deletions(-)
 create mode 100644 roles/iommu/handlers/main.yml

diff --git a/roles/iommu/README.md b/roles/iommu/README.md
index cb9882b..13c9fad 100644
--- a/roles/iommu/README.md
+++ b/roles/iommu/README.md
@@ -4,6 +4,10 @@
 
 - `iommu_pt`: Default true, Configure passthrough mode, which doesn't require DMA translation.
 
+## Host Vars
+
+- `iommu_vfio_pci_ids`: Can optionally be set with the pci id of the device to pass-through.
+
 ## Example playbook
 
 ```
@@ -13,10 +17,4 @@
   tasks:
     - import_role:
       name: stackhpc.linux.iommu
-  handlers:
-    - name: reboot
-      fail:
-        msg: "Please reboot your hypervisor and re-run your host configure to continue"
-      become: true
-
 ```
diff --git a/roles/iommu/handlers/main.yml b/roles/iommu/handlers/main.yml
new file mode 100644
index 0000000..579d882
--- /dev/null
+++ b/roles/iommu/handlers/main.yml
@@ -0,0 +1,12 @@
+---
+- name: Regenerate initramfs
+  ansible.builtin.shell: |-
+    #!/bin/bash
+    set -eux
+    dracut -v -f /boot/initramfs-$(uname -r).img $(uname -r)
+  become: true
+  changed_when: true
+- name: Reboot
+  ansible.builtin.reboot:
+    msg: "Rebooting the hypervisor"
+  become: true
diff --git a/roles/iommu/tasks/main.yml b/roles/iommu/tasks/main.yml
index 075a964..ba8075c 100644
--- a/roles/iommu/tasks/main.yml
+++ b/roles/iommu/tasks/main.yml
@@ -1,13 +1,45 @@
 ---
+- name: Template dracut config for vfio
+  ansible.builtin.blockinfile:
+    path: /etc/dracut.conf.d/gpu-vfio.conf
+    block: |
+      add_drivers+="vfio vfio_iommu_type1 vfio_pci vfio_virqfd"
+    owner: root
+    group: root
+    mode: "0660"
+    create: true
+  become: true
+  when: iommu_vfio_pci_ids is defined
+  notify:
+    - Regenerate initramfs
+    - reboot
+
+- name: Add vfio to modules-load.d
+  ansible.builtin.blockinfile:
+    path: /etc/modules-load.d/vfio.conf
+    block: |
+      vfio
+      vfio_iommu_type1
+      vfio_pci
+      vfio_virqfd
+    owner: root
+    group: root
+    mode: "0664"
+    create: true
+  become: true
+  when: iommu_vfio_pci_ids is defined
+  notify: reboot
+
 - name: Add iommu to kernel command line (Intel)
   ansible.builtin.include_role:
     name: stackhpc.linux.grubcmdline
   vars:
-    kernel_cmdline: # noqa: var-naming[no-role-prefix]
-      - intel_iommu=on
+    kernel_cmdline: "{{ ['intel_iommu=on'] + (['vfio-pci.ids=' + iommu_vfio_pci_ids] if iommu_vfio_pci_ids is defined else []) }}"  # noqa: var-naming[no-role-prefix]
     kernel_cmdline_remove: # noqa: var-naming[no-role-prefix]
       - ^intel_iommu=
+      - ^vfio-pci\.ids=
   when: ansible_facts.processor | select('search', 'Intel') | list | length > 0
+  notify: reboot
 
 - name: Set iommu=pt
   ansible.builtin.include_role:
@@ -18,3 +50,4 @@
     kernel_cmdline_remove: # noqa: var-naming[no-role-prefix]
       - ^iommu=
   when: iommu_pt
+  notify: reboot