From f77ae02b6e188d5d3a94d5eb6224521f0185945c Mon Sep 17 00:00:00 2001
From: Ethan Wee <158101733+ethanwee1@users.noreply.github.com>
Date: Wed, 9 Oct 2024 10:20:51 -0700
Subject: [PATCH 1/4] Update build_common.sh with lightweight capability

---
 manywheel/build_common.sh | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/manywheel/build_common.sh b/manywheel/build_common.sh
index 9f89bed95..477d1abec 100644
--- a/manywheel/build_common.sh
+++ b/manywheel/build_common.sh
@@ -69,6 +69,12 @@ fi
 if [[ -z "$build_number" ]]; then
     build_number=1
 fi
+if [[ "$BUILD_LIGHTWEIGHT" == "1" ]]; then
+    build_version="${build_version}.lw"
+fi
+
+echo "Final build_version: $build_version"
+
 export PYTORCH_BUILD_VERSION=$build_version
 export PYTORCH_BUILD_NUMBER=$build_number
 

From c0017792dfcce62a97c4d2dfd29ab2d0328be1e3 Mon Sep 17 00:00:00 2001
From: Ethan Wee <158101733+ethanwee1@users.noreply.github.com>
Date: Wed, 9 Oct 2024 10:26:52 -0700
Subject: [PATCH 2/4] Update build_rocm.sh with lightweight wheel capability

---
 manywheel/build_rocm.sh | 78 ++++++++++++++++++++++++-----------------
 1 file changed, 45 insertions(+), 33 deletions(-)

diff --git a/manywheel/build_rocm.sh b/manywheel/build_rocm.sh
index 86ee9feec..1f7a85d44 100755
--- a/manywheel/build_rocm.sh
+++ b/manywheel/build_rocm.sh
@@ -232,6 +232,13 @@ OTHER_FILES=$(ls $HIPBLASLT_LIB_SRC | grep -v gfx)
 HIPBLASLT_LIB_FILES=($ARCH_SPECIFIC_FILES $OTHER_FILES)
 
 # ROCm library files
+# Overwrite ROCM_SO_FILES to contain only libmagma.so if BUILD_LIGHTWEIGHT is enabled
+if [[ "$BUILD_LIGHTWEIGHT" == "1" ]]; then
+    #TODO: add libaotriton_v2.so when using preinstalled aotriton
+    ROCM_SO_FILES=(
+        "libmagma.so"
+    )
+fi
 ROCM_SO_PATHS=()
 for lib in "${ROCM_SO_FILES[@]}"
 do
@@ -253,49 +260,54 @@ done
 
 DEPS_LIST=(
     ${ROCM_SO_PATHS[*]}
-    ${OS_SO_PATHS[*]}
 )
-
+if [[ "$BUILD_LIGHTWEIGHT" != "1" ]]; then
+    DEPS_LIST+=(${OS_SO_PATHS[*]})
+fi
 DEPS_SONAME=(
     ${ROCM_SO_FILES[*]}
-    ${OS_SO_FILES[*]}
-)
-
-DEPS_AUX_SRCLIST=(
-    "${ROCBLAS_LIB_FILES[@]/#/$ROCBLAS_LIB_SRC/}"
-    "${HIPBLASLT_LIB_FILES[@]/#/$HIPBLASLT_LIB_SRC/}"
-    "/opt/amdgpu/share/libdrm/amdgpu.ids"
-)
-
-DEPS_AUX_DSTLIST=(
-    "${ROCBLAS_LIB_FILES[@]/#/$ROCBLAS_LIB_DST/}"
-    "${HIPBLASLT_LIB_FILES[@]/#/$HIPBLASLT_LIB_DST/}"
-    "share/libdrm/amdgpu.ids"
 )
+if [[ "$BUILD_LIGHTWEIGHT" != "1" ]]; then
+    DEPS_SONAME+=(${OS_SO_FILES[*]})
+fi
 
-if [[ $ROCM_INT -ge 50500 ]]; then
-    # MIOpen library files
-    MIOPEN_SHARE_SRC=$ROCM_HOME/share/miopen/db
-    MIOPEN_SHARE_DST=share/miopen/db
-    MIOPEN_SHARE_FILES=($(ls $MIOPEN_SHARE_SRC | grep -E $ARCH))
+DEPS_AUX_SRCLIST=()
+if [[ "$BUILD_LIGHTWEIGHT" != "1" ]]; then
+    DEPS_AUX_SRCLIST+=("${ROCBLAS_LIB_FILES[@]/#/$ROCBLAS_LIB_SRC/}")
+    DEPS_AUX_SRCLIST+=("${HIPBLASLT_LIB_FILES[@]/#/$HIPBLASLT_LIB_SRC/}")
+    DEPS_AUX_SRCLIST+=("/opt/amdgpu/share/libdrm/amdgpu.ids")
+fi
 
-    DEPS_AUX_SRCLIST+=(${MIOPEN_SHARE_FILES[@]/#/$MIOPEN_SHARE_SRC/})
-    DEPS_AUX_DSTLIST+=(${MIOPEN_SHARE_FILES[@]/#/$MIOPEN_SHARE_DST/})
+DEPS_AUX_DSTLIST=()
+if [[ "$BUILD_LIGHTWEIGHT" != "1" ]]; then
+    DEPS_AUX_DSTLIST+=("${ROCBLAS_LIB_FILES[@]/#/$ROCBLAS_LIB_DST/}")
+    DEPS_AUX_DSTLIST+=("${HIPBLASLT_LIB_FILES[@]/#/$HIPBLASLT_LIB_DST/}")
+    DEPS_AUX_DSTLIST+=("share/libdrm/amdgpu.ids")
 fi
+if [[ "$BUILD_LIGHTWEIGHT" != "1" ]]; then
+    if [[ $ROCM_INT -ge 50500 ]]; then
+        # MIOpen library files
+        MIOPEN_SHARE_SRC=$ROCM_HOME/share/miopen/db
+        MIOPEN_SHARE_DST=share/miopen/db
+        MIOPEN_SHARE_FILES=($(ls $MIOPEN_SHARE_SRC | grep -E $ARCH))
 
-if [[ $ROCM_INT -ge 50600 ]]; then
-    # RCCL library files
-    if [[ $ROCM_INT -ge 50700 ]]; then
-        RCCL_SHARE_SRC=$ROCM_HOME/share/rccl/msccl-algorithms
-        RCCL_SHARE_DST=share/rccl/msccl-algorithms
-    else
-        RCCL_SHARE_SRC=$ROCM_HOME/lib/msccl-algorithms
-        RCCL_SHARE_DST=lib/msccl-algorithms
+        DEPS_AUX_SRCLIST+=(${MIOPEN_SHARE_FILES[@]/#/$MIOPEN_SHARE_SRC/})
+        DEPS_AUX_DSTLIST+=(${MIOPEN_SHARE_FILES[@]/#/$MIOPEN_SHARE_DST/})
     fi
-    RCCL_SHARE_FILES=($(ls $RCCL_SHARE_SRC))
 
-    DEPS_AUX_SRCLIST+=(${RCCL_SHARE_FILES[@]/#/$RCCL_SHARE_SRC/})
-    DEPS_AUX_DSTLIST+=(${RCCL_SHARE_FILES[@]/#/$RCCL_SHARE_DST/})
+    if [[ $ROCM_INT -ge 50600 ]]; then
+        # RCCL library files
+        if [[ $ROCM_INT -ge 50700 ]]; then
+            RCCL_SHARE_SRC=$ROCM_HOME/share/rccl/msccl-algorithms
+            RCCL_SHARE_DST=share/rccl/msccl-algorithms
+        else
+            RCCL_SHARE_SRC=$ROCM_HOME/lib/msccl-algorithms
+            RCCL_SHARE_DST=lib/msccl-algorithms
+        fi
+        RCCL_SHARE_FILES=($(ls $RCCL_SHARE_SRC))
+
+        DEPS_AUX_SRCLIST+=(${RCCL_SHARE_FILES[@]/#/$RCCL_SHARE_SRC/})
+        DEPS_AUX_DSTLIST+=(${RCCL_SHARE_FILES[@]/#/$RCCL_SHARE_DST/})
 fi
 
 # Add triton install dependency

From efd7868992ec9900db3287ef678784088656f6ba Mon Sep 17 00:00:00 2001
From: Jithun Nair <37884920+jithunnair-amd@users.noreply.github.com>
Date: Thu, 17 Oct 2024 13:26:00 -0500
Subject: [PATCH 3/4] Add libaotriton_v2.so to list of libraries to bundle

---
 manywheel/build_rocm.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/manywheel/build_rocm.sh b/manywheel/build_rocm.sh
index 1f7a85d44..5a05a48ce 100755
--- a/manywheel/build_rocm.sh
+++ b/manywheel/build_rocm.sh
@@ -234,9 +234,9 @@ HIPBLASLT_LIB_FILES=($ARCH_SPECIFIC_FILES $OTHER_FILES)
 # ROCm library files
 # Overwrite ROCM_SO_FILES to contain only libmagma.so if BUILD_LIGHTWEIGHT is enabled
 if [[ "$BUILD_LIGHTWEIGHT" == "1" ]]; then
-    #TODO: add libaotriton_v2.so when using preinstalled aotriton
     ROCM_SO_FILES=(
         "libmagma.so"
+        "libaotriton_v2.so"
     )
 fi
 ROCM_SO_PATHS=()

From ab2c0f4eab7a73528b279b342fbd84ca276040b3 Mon Sep 17 00:00:00 2001
From: Ethan Wee <158101733+ethanwee1@users.noreply.github.com>
Date: Fri, 18 Oct 2024 12:23:55 -0700
Subject: [PATCH 4/4] Build error

---
 manywheel/build_rocm.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/manywheel/build_rocm.sh b/manywheel/build_rocm.sh
index 5a05a48ce..d2924b729 100755
--- a/manywheel/build_rocm.sh
+++ b/manywheel/build_rocm.sh
@@ -308,6 +308,7 @@ if [[ "$BUILD_LIGHTWEIGHT" != "1" ]]; then
 
         DEPS_AUX_SRCLIST+=(${RCCL_SHARE_FILES[@]/#/$RCCL_SHARE_SRC/})
         DEPS_AUX_DSTLIST+=(${RCCL_SHARE_FILES[@]/#/$RCCL_SHARE_DST/})
+    fi
 fi
 
 # Add triton install dependency