From 5c34e2686df56b4d7feff0b75ea454bb238bde14 Mon Sep 17 00:00:00 2001
From: Joel Smith <joelsmith@tenstorrent.com>
Date: Fri, 13 Dec 2024 17:10:00 +0000
Subject: [PATCH] update script

---
 scripts/iommu_detect.sh | 65 +++++++++++++++++++++--------------------
 1 file changed, 33 insertions(+), 32 deletions(-)

diff --git a/scripts/iommu_detect.sh b/scripts/iommu_detect.sh
index c1525b3a..2ff94f72 100755
--- a/scripts/iommu_detect.sh
+++ b/scripts/iommu_detect.sh
@@ -6,7 +6,6 @@ GS_PID="faca"
 WH_PID="401e"
 BH_PID="b140"
 
-# Find all Tenstorrent PCI devices with full domain info
 tt_devices=$(lspci -D -d ${TT_VID}: | cut -d' ' -f1)
 
 if [ -z "$tt_devices" ]; then
@@ -17,43 +16,45 @@ fi
 found_gs_wh=false
 found_bh=false
 
+# First identify devices
 for dev in $tt_devices; do
     device_id=$(lspci -D -n -s "$dev" | cut -d' ' -f3 | cut -d: -f2)
-    echo "Checking device $dev (ID: $device_id):"
-    
     case $device_id in
-        $GS_PID|$WH_PID)
-            found_gs_wh=true
-            if [ -f "/sys/bus/pci/devices/${dev}/iommu_group/type" ]; then
-                iommu_type=$(cat "/sys/bus/pci/devices/${dev}/iommu_group/type")
-                if [[ "$iommu_type" == *"DMA"* ]]; then
-                    echo "  WARNING: Grayskull/Wormhole device with IOMMU enabled (type: $iommu_type) - this configuration is not supported"
-                else
-                    echo "  Grayskull/Wormhole device detected - hugepages required"
-                fi
-            else
-                echo "  Grayskull/Wormhole device detected - hugepages required"
-            fi
-            ;;
-        $BH_PID)
-            found_bh=true
-            if [ -f "/sys/bus/pci/devices/${dev}/iommu_group/type" ]; then
-                iommu_type=$(cat "/sys/bus/pci/devices/${dev}/iommu_group/type")
-                if [[ "$iommu_type" == *"DMA"* ]]; then
-                    echo "  Blackhole device with IOMMU enabled (type: $iommu_type) - hugepages optional"
-                else
-                    echo "  Blackhole device with IOMMU in passthrough mode (type: $iommu_type) - hugepages required"
-                fi
-            else
-                echo "  Blackhole device with no IOMMU configuration - hugepages required"
-            fi
-            ;;
-        *)
-            echo "  Unknown device ID: $device_id"
-            ;;
+        $GS_PID|$WH_PID) found_gs_wh=true ;;
+        $BH_PID) found_bh=true ;;
     esac
 done
 
+# Check and output for each device
+for dev in $tt_devices; do
+    device_id=$(lspci -D -n -s "$dev" | cut -d' ' -f3 | cut -d: -f2)
+    echo "Checking device $dev (ID: $device_id):"
+    
+    # Check IOMMU status for this device
+    iommu_enabled=false
+    iommu_type="none"
+    if [ -f "/sys/bus/pci/devices/${dev}/iommu_group/type" ]; then
+        iommu_type=$(cat "/sys/bus/pci/devices/${dev}/iommu_group/type")
+        [[ "$iommu_type" == *"DMA"* ]] && iommu_enabled=true
+    fi
+
+    if [[ "$device_id" == "$GS_PID" || "$device_id" == "$WH_PID" ]]; then
+        if [ "$iommu_enabled" = true ]; then
+            echo "  WARNING: Grayskull/Wormhole device with IOMMU enabled (type: $iommu_type) - this configuration is not supported"
+        else
+            echo "  Grayskull/Wormhole device detected - hugepages required"
+        fi
+    elif [[ "$device_id" == "$BH_PID" ]]; then
+        if [ "$iommu_enabled" = true ]; then
+            echo "  Blackhole device with IOMMU enabled (type: $iommu_type) - hugepages optional"
+        else
+            echo "  Blackhole device with no IOMMU/passthrough (type: $iommu_type) - hugepages required"
+        fi
+    else
+        echo "  Unknown device ID: $device_id"
+    fi
+done
+
 echo -e "\nSummary:"
 if [ "$found_gs_wh" = true ]; then
     echo "- System has Grayskull/Wormhole devices - hugepages required"