Changes for lightweight wheels

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
 
@@ -117,7 +123,7 @@ if [[ -z "$PYTORCH_ROOT" ]]; then
 fi
 pushd "$PYTORCH_ROOT"
 python setup.py clean
-retry pip install -qr requirements.txt
+retry pip install -r requirements.txt
 case ${DESIRED_PYTHON} in
   cp36-cp36m)
     retry pip install -q numpy==1.11
@@ -330,11 +336,13 @@ for pkg in /$WHEELHOUSE_DIR/torch*linux*.whl /$LIBTORCH_HOUSE_DIR/libtorch*.zip;
         # copy over needed dependent .so files over and tag them with their hash
         patched=()
         for filepath in "${DEPS_LIST[@]}"; do
-            filename=$(basename $filepath)
-            destpath=$PREFIX/lib/$filename
-            if [[ "$filepath" != "$destpath" ]]; then
-                cp $filepath $destpath
-            fi
+            filename=$(basename "$filepath")  # Extract the file name
+            destpath=$PREFIX/lib/$filename    # Set the destination path
+
+            if [[ -f "$filepath" ]]; then     # Check if the file exists
+                cp "$filepath" "$destpath"    # Copy the file to the destination
+                echo "Copied $filepath to $destpath"
+            fi 
 
             # ROCm workaround for roctracer dlopens
             if [[ "$DESIRED_CUDA" == *"rocm"* ]]; then

manywheel/build_rocm.sh

@@ -210,6 +210,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
+    ROCM_SO_FILES=(
+        "libmagma.so"
+    )
+fi
+
 ROCM_SO_PATHS=()
 for lib in "${ROCM_SO_FILES[@]}"
 do
@@ -230,26 +237,31 @@ do
 done
 
 DEPS_LIST=(
-    ${ROCM_SO_PATHS[*]}
-    ${OS_SO_PATHS[*]}
+    "${ROCM_SO_PATHS[*]}"
 )
-
+if [[ "$BUILD_LIGHTWEIGHT" != "1" ]]; then
+    DEPS_LIST+=("${OS_SO_PATHS[*]}")
+fi
 DEPS_SONAME=(
-    ${ROCM_SO_FILES[*]}
-    ${OS_SO_FILES[*]}
+    "${ROCM_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_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_DSTLIST=(
-    "${ROCBLAS_LIB_FILES[@]/#/$ROCBLAS_LIB_DST/}"
-    "${HIPBLASLT_LIB_FILES[@]/#/$HIPBLASLT_LIB_DST/}"
-    "share/libdrm/amdgpu.ids"
-)
+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 [[ $ROCM_INT -ge 50500 ]]; then
     # MIOpen library files

wheel/build_wheel.sh

@@ -56,6 +56,13 @@ else
         build_number_prefix=".post$build_number"
     fi
 fi
+# Append .lw if BUILD_LIGHTWEIGHT is enabled
+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