Changes for lightweight wheels (#56)

* Update build_rocm.sh * Update build_rocm.sh * Update build_rocm.sh * Update build_rocm.sh * Update build_rocm.sh * Update build_rocm.sh * Update build_rocm.sh * Update build_rocm.sh * Update build_rocm.sh * Update build_rocm.sh * Update build_wheel.sh * Update build_rocm.sh * Update build_rocm.sh * Update build_wheel.sh * Update build_rocm.sh * Update build_wheel.sh * Update build_wheel.sh * Update build_wheel.sh to work on OVERRIDE_PACKAGE_VERSION * Update build_wheel.sh * Update the whl name with lw * Update build_common.sh to use BUILD_LIGHTWEIGHT flag * Update build_rocm.sh to ensure that only libmagma.so is included with BUILD_LIGHTWEIGHT enabled * Update build_common.sh to correct how basename command was applied to individual elements in DEPS_LIST * Update build_common.sh * Update build_common.sh * Don't be quiet * Update build_common.sh * Update build_common.sh * Update build_common.sh * Update build_common.sh * Update build_rocm.sh * Update build_rocm.sh * Update build_rocm.sh * Update build_common.sh * Update build_common.sh * Update build_common.sh * Update build_wheel.sh * Update build_common.sh * remove unwanted changes * Update build_rocm.sh * Update build_wheel.sh * Update build_rocm.sh * Update build_common.sh * Update build_wheel.sh * Update build_wheel.sh * Update build_wheel.sh * Update build_common.sh * Update build_common.sh * Update build_rocm.sh * Update quoting * Add back "/" into src list * Add TODO comment for libaotriton_v2.so Will be needed when cherry picking changes to rocm6.3 builder branch --------- Co-authored-by: Pruthvi Madugundu <[email protected]> Co-authored-by: Jithun Nair <[email protected]>
ROCm · Sep 17, 2024 · 3946b4b · 3946b4b
1 parent 08ad11b
commit 3946b4b
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 36 deletions.
diff --git a/manywheel/build_common.sh b/manywheel/build_common.sh
@@ -59,6 +59,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
 

diff --git a/manywheel/build_rocm.sh b/manywheel/build_rocm.sh
@@ -210,6 +210,14 @@ 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
@@ -231,51 +239,56 @@ 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[*]}
 )
+if [[ "$BUILD_LIGHTWEIGHT" != "1" ]]; then
+    DEPS_SONAME+=(${OS_SO_FILES[*]})
+fi
 
-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 [[ $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+=(${MIOPEN_SHARE_FILES[@]/#/$MIOPEN_SHARE_SRC/})
-    DEPS_AUX_DSTLIST+=(${MIOPEN_SHARE_FILES[@]/#/$MIOPEN_SHARE_DST/})
+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
 
-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_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))
+
+        DEPS_AUX_SRCLIST+=(${MIOPEN_SHARE_FILES[@]/#/$MIOPEN_SHARE_SRC/})
+        DEPS_AUX_DSTLIST+=(${MIOPEN_SHARE_FILES[@]/#/$MIOPEN_SHARE_DST/})
+    fi
+
+    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
-    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
 # No triton dependency for now on 3.12 since we don't have binaries for it
 # and torch.compile doesn't work.