[REFACTOR][Rename] MLC_LLM_SOURCE_DIR and TVM_SOURCE_DIR source directory env (#2378)

tqchen · web-flow · commit 6dd6c8962e4d · 2024-05-22T12:59:53.000-04:00
* [REFACTOR] Rename use MLC_LLM_SOURCE_DIR for source directory

This PR updates to use MLC_LLM_SOURCE_DIR to specify the
directory of mlc llm source directory.

The reason for this update is that the term XXX_HOME was usually
meant to be used in different scenarios in ML frameworks.

For example, both torch and huggingface have TORCH_HOME and HF_HOME
pointing to their local cache directory.

The variable MLC_LLM_SOURCE_DIR is aligned with cmake naming convention
(CMAKE_SOURCE_DIR).

We will have followup PR to udpate MLC_CACHE_DIR to MLC_LLM_HOME, following
the existing practices.

* Update env to point to TVM_SOURCE_DIR
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -54,15 +54,15 @@ set(USE_PROFILER OFF)
 set(USE_GTEST OFF)
 set(USE_LIBBACKTRACE OFF)
 set(BUILD_DUMMY_LIBTVM ON)
-if (NOT DEFINED TVM_HOME)
-  if(DEFINED ENV{TVM_HOME})
-    set(TVM_HOME "$ENV{TVM_HOME}")
+if (NOT DEFINED TVM_SOURCE_DIR)
+  if(DEFINED ENV{TVM_SOURCE_DIR})
+    set(TVM_SOURCE_DIR "$ENV{TVM_SOURCE_DIR}")
   else()
-    set(TVM_HOME 3rdparty/tvm)
-  endif(DEFINED ENV{TVM_HOME})
-endif (NOT DEFINED TVM_HOME)
-message(STATUS "TVM_HOME: ${TVM_HOME}")
-add_subdirectory(${TVM_HOME} tvm EXCLUDE_FROM_ALL)
+    set(TVM_SOURCE_DIR 3rdparty/tvm)
+  endif(DEFINED ENV{TVM_SOURCE_DIR})
+endif (NOT DEFINED TVM_SOURCE_DIR)
+message(STATUS "TVM_SOURCE_DIR: ${TVM_SOURCE_DIR}")
+add_subdirectory(${TVM_SOURCE_DIR} tvm EXCLUDE_FROM_ALL)
 
 set(MLC_LLM_RUNTIME_LINKER_LIB "")
 set(TOKENZIER_CPP_PATH 3rdparty/tokenizers-cpp)
@@ -74,10 +74,10 @@ add_library(mlc_llm_objs OBJECT ${MLC_LLM_SRCS})
 
 set(
   MLC_LLM_INCLUDES
-  ${TVM_HOME}/include
-  ${TVM_HOME}/3rdparty/dlpack/include
-  ${TVM_HOME}/3rdparty/dmlc-core/include
-  ${TVM_HOME}/3rdparty/picojson
+  ${TVM_SOURCE_DIR}/include
+  ${TVM_SOURCE_DIR}/3rdparty/dlpack/include
+  ${TVM_SOURCE_DIR}/3rdparty/dmlc-core/include
+  ${TVM_SOURCE_DIR}/3rdparty/picojson
 )
 
 set(MLC_LLM_COMPILE_DEFS ${MLC_LLM_COMPILE_DEFS} DMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>)
@@ -100,7 +100,7 @@ target_link_libraries(mlc_llm PRIVATE tokenizers_cpp)
 
 find_library(
   FLASH_ATTN_LIBRARY flash_attn
-  HINTS ${TVM_HOME}/*/3rdparty/libflash_attn/src
+  HINTS ${TVM_SOURCE_DIR}/*/3rdparty/libflash_attn/src
 )
 
 if (FLASH_ATTN_LIBRARY STREQUAL "FLASH_ATTN_LIBRARY-NOTFOUND")
diff --git a/android/mlc4j/CMakeLists.txt b/android/mlc4j/CMakeLists.txt
@@ -10,10 +10,10 @@ set(MLC_LLM_BINARY_DIR mlc_llm)
 set(MLC_LLM_COMPILE_DEFS TVM_LOG_CUSTOMIZE=1)
 add_subdirectory(${MLC_LLM_DIR} ${MLC_LLM_BINARY_DIR} EXCLUDE_FROM_ALL)
 
-if (NOT DEFINED TVM_HOME)
-  set(TVM_HOME ${MLC_LLM_DIR}/3rdparty/tvm)
-endif (NOT DEFINED TVM_HOME)
-message(STATUS "TVM_HOME: ${TVM_HOME}")
+if (NOT DEFINED TVM_SOURCE_DIR)
+  set(TVM_SOURCE_DIR ${MLC_LLM_DIR}/3rdparty/tvm)
+endif (NOT DEFINED TVM_SOURCE_DIR)
+message(STATUS "TVM_SOURCE_DIR: ${TVM_SOURCE_DIR}")
 
 find_package(Java REQUIRED)
 find_package(JNI REQUIRED)
@@ -25,29 +25,29 @@ include(UseJava)
 
 
 file(GLOB_RECURSE javasources
-    ${TVM_HOME}/jvm/core/src/main/java/org/apache/tvm/*.java
+    ${TVM_SOURCE_DIR}/jvm/core/src/main/java/org/apache/tvm/*.java
     ${ANDROID_DIR}/src/java/*.java
 )
 set(JNI_HEADER ${CMAKE_BINARY_DIR}/jni_header)
 add_jar(tvm4j_core ${javasources} GENERATE_NATIVE_HEADERS tvm4jheaders DESTINATION ${JNI_HEADER})
 
 add_custom_command(
-  TARGET tvm4j_core POST_BUILD 
+  TARGET tvm4j_core POST_BUILD
   COMMAND ${CMAKE_COMMAND} -E copy ${JNI_HEADER}/org_apache_tvm_LibInfo.h ${JNI_HEADER}/org_apache_tvm_native_c_api.h
 )
 
 add_library(model_android STATIC IMPORTED)
 set_target_properties(model_android PROPERTIES IMPORTED_LOCATION ${ANDROID_BIN_DIR}/lib/libmodel_android.a)
 
-add_library(tvm4j_runtime_packed SHARED ${TVM_HOME}/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc)
+add_library(tvm4j_runtime_packed SHARED ${TVM_SOURCE_DIR}/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc)
 
 target_include_directories(tvm4j_runtime_packed PUBLIC
   ${JNI_INCLUDE_DIRS}
   ${JNI_HEADER}
   ${ANDROID_DIR}/src/cpp
-  ${TVM_HOME}/3rdparty/dlpack/include
-  ${TVM_HOME}/3rdparty/dmlc-core/include
-  ${TVM_HOME}/include
+  ${TVM_SOURCE_DIR}/3rdparty/dlpack/include
+  ${TVM_SOURCE_DIR}/3rdparty/dmlc-core/include
+  ${TVM_SOURCE_DIR}/include
 )
 
 target_link_libraries(tvm4j_runtime_packed
diff --git a/android/mlc4j/prepare_libs.py b/android/mlc4j/prepare_libs.py
@@ -81,10 +81,10 @@ def main(mlc_llm_home: Path):
     logger.info('Entering "%s" for MLC LLM and tvm4j build.', os.path.abspath(build_path))
     os.chdir(build_path)
     # Generate config.cmake if TVM Home is set.
-    if "TVM_HOME" in os.environ:
-        logger.info('Set TVM_HOME to "%s"', os.environ["TVM_HOME"])
+    if "TVM_SOURCE_DIR" in os.environ:
+        logger.info('Set TVM_SOURCE_DIR to "%s"', os.environ["TVM_SOURCE_DIR"])
         with open("config.cmake", "w", encoding="utf-8") as file:
-            print("set(TVM_HOME ${%s})" % os.environ["TVM_HOME"], file=file)
+            print("set(TVM_SOURCE_DIR ${%s})" % os.environ["TVM_SOURCE_DIR"], file=file)
 
     # - Run cmake, build and install
     run_cmake(mlc_llm_home / "android" / "mlc4j")
@@ -98,11 +98,11 @@ def main(mlc_llm_home: Path):
     parser.add_argument(
         "--mlc-llm-home",
         type=Path,
-        default=os.environ.get("MLC_LLM_HOME", None),
+        default=os.environ.get("MLC_LLM_SOURCE_DIR", None),
         help="The path to MLC LLM source",
     )
     parsed = parser.parse_args()
     if parsed.mlc_llm_home is None:
         parsed.mlc_llm_home = Path(os.path.abspath(os.path.curdir)).parent.parent
-    os.environ["MLC_LLM_HOME"] = str(parsed.mlc_llm_home)
+    os.environ["MLC_LLM_SOURCE_DIR"] = str(parsed.mlc_llm_home)
     main(parsed.mlc_llm_home)
diff --git a/ci/task/test_model_compile.sh b/ci/task/test_model_compile.sh
@@ -20,10 +20,10 @@ elif [[ ${GPU} == metal ]]; then
 elif [[ ${GPU} == wasm* ]]; then
 	TARGET=wasm
 	pip install --pre -U -f https://mlc.ai/wheels mlc-ai-nightly
-	export TVM_HOME=$(dirname $(python -c 'import tvm; print(tvm.__file__)'))
-	export MLC_LLM_HOME=$(pwd)
-	cd $TVM_HOME/web/ && make -j${NUM_THREADS} && cd -
-	cd $MLC_LLM_HOME/web/ && make -j${NUM_THREADS} && cd -
+	export TVM_SOURCE_DIR=$(dirname $(python -c 'import tvm; print(tvm.__file__)'))
+	export MLC_LLM_SOURCE_DIR=$(pwd)
+	cd $TVM_SOURCE_DIR/web/ && make -j${NUM_THREADS} && cd -
+	cd $MLC_LLM_SOURCE_DIR/web/ && make -j${NUM_THREADS} && cd -
 elif [[ ${GPU} == ios ]]; then
 	TARGET=ios
 	pip install --pre -U --force-reinstal -f https://mlc.ai/wheels mlc-ai-nightly
diff --git a/cmake/gen_cmake_config.py b/cmake/gen_cmake_config.py
@@ -6,12 +6,12 @@
     tvm_home = ""  # pylint: disable=invalid-name
 
     tvm_home = input(
-        "Enter TVM_HOME in absolute path. If not specified, 3rdparty/tvm will be used by default: "
+        "Enter TVM_SOURCE_DIR in absolute path. If not specified, 3rdparty/tvm will be used by default: "
     )
     if len(tvm_home) == 0:
         tvm_home = "3rdparty/tvm"  # pylint: disable=invalid-name
 
-    cmake_config_str = f"set(TVM_HOME {tvm_home})\n"
+    cmake_config_str = f"set(TVM_SOURCE_DIR {tvm_home})\n"
     cmake_config_str += "set(CMAKE_BUILD_TYPE RelWithDebInfo)\n"
     backends = [
         Backend("CUDA", "USE_CUDA", "Use CUDA? (y/n): "),
diff --git a/docs/compilation/package_libraries_and_weights.rst b/docs/compilation/package_libraries_and_weights.rst
@@ -26,7 +26,7 @@ After cloning, the basic usage of ``mlc_llm package`` is as the following.
 
 .. code:: bash
 
-    export MLC_LLM_HOME=/path/to/mlc-llm
+    export MLC_LLM_SOURCE_DIR=/path/to/mlc-llm
     cd /path/to/app  # The app root directory which contains "mlc-package-config.json".
                      # E.g., "ios/MLCChat" or "android/MLCChat"
     mlc_llm package
@@ -200,8 +200,8 @@ Command ``mlc_llm package`` can optionally take the arguments below:
 
 ``--mlc-llm-home``
     The path to MLC LLM source code (cloned from https://github.com/mlc-ai/mlc-llm).
-    By default, it is the ``$MLC_LLM_HOME`` environment variable.
-    If neither ``$MLC_LLM_HOME`` or ``--mlc-llm-home`` is specified, error will be reported.
+    By default, it is the ``$MLC_LLM_SOURCE_DIR`` environment variable.
+    If neither ``$MLC_LLM_SOURCE_DIR`` or ``--mlc-llm-home`` is specified, error will be reported.
 
 ``--output`` / ``-o``
     The output directory of ``mlc_llm package`` command.
diff --git a/docs/deploy/android.rst b/docs/deploy/android.rst
@@ -48,7 +48,7 @@ Please ensure that the JDK versions for Android Studio and JAVA_HOME are the sam
 
 **TVM Unity runtime** is placed under `3rdparty/tvm <https://github.com/mlc-ai/mlc-llm/tree/main/3rdparty>`__ in MLC LLM, so there is no need to install anything extra. Set up the following environment variable:
 
-- ``export TVM_HOME=/path/to/mlc-llm/3rdparty/tvm``.
+- ``export TVM_SOURCE_DIR=/path/to/mlc-llm/3rdparty/tvm``.
 
 (Optional) **TVM Unity compiler** Python package (:ref:`install <tvm-unity-prebuilt-package>` or :ref:`build from source <tvm-unity-build-from-source>`). It is *NOT* required if models are prebuilt, but to compile PyTorch models from HuggingFace in the following section, the compiler is a must-dependency.
 
@@ -63,7 +63,7 @@ Check if **environment variable** are properly set as the last check. One way to
   export ANDROID_NDK=...  # Android NDK toolchain
   export TVM_NDK_CC=...   # Android NDK clang
   export JAVA_HOME=...    # Java
-  export TVM_HOME=...     # TVM Unity runtime
+  export TVM_SOURCE_DIR=...     # TVM Unity runtime
 
 
 Build Android App from Source
@@ -105,7 +105,7 @@ We have a one-line command to build and prepare all the model libraries:
 .. code:: bash
 
    cd /path/to/MLCChat  # e.g., "android/MLCChat"
-   export MLC_LLM_HOME=/path/to/mlc-llm  # e.g., "../.."
+   export MLC_LLM_SOURCE_DIR=/path/to/mlc-llm  # e.g., "../.."
    mlc_llm package
 
 This command mainly executes the following two steps:
diff --git a/docs/deploy/ios.rst b/docs/deploy/ios.rst
@@ -67,7 +67,7 @@ We have a one-line command to build and prepare all the model libraries:
 .. code:: bash
 
    cd /path/to/MLCChat  # e.g., "ios/MLCChat"
-   export MLC_LLM_HOME=/path/to/mlc-llm  # e.g., "../.."
+   export MLC_LLM_SOURCE_DIR=/path/to/mlc-llm  # e.g., "../.."
    mlc_llm package
 
 This command mainly executes the following two steps:
diff --git a/docs/install/emcc.rst b/docs/install/emcc.rst
@@ -21,20 +21,20 @@ Validate that emcc is accessible in shell
 
     emcc --version
 
-Step 2: Set TVM_HOME and MLC_LLM_HOME
+Step 2: Set TVM_SOURCE_DIR and MLC_LLM_SOURCE_DIR
 -------------------------------------
 
 We need to set a path to a tvm source in order to build tvm runtime.
 Note that you do not need to build tvm unity from the source. The source here is only used to build the web runtime component.
 Set environment variable in your shell startup profile in to point to ``3rdparty/tvm`` (if preferred, you could also
 point to your own TVM address if you installed TVM from source).
 
-Besides, we also need to set ``MLC_LLM_HOME`` so that we can locate ``mlc_wasm_runtime.bc`` when compiling a model library wasm.
+Besides, we also need to set ``MLC_LLM_SOURCE_DIR`` so that we can locate ``mlc_wasm_runtime.bc`` when compiling a model library wasm.
 
 .. code:: bash
 
-    export TVM_HOME=/path/to/3rdparty/tvm
-    export MLC_LLM_HOME=/path/to/mlc-llm
+    export TVM_SOURCE_DIR=/path/to/3rdparty/tvm
+    export MLC_LLM_SOURCE_DIR=/path/to/mlc-llm
 
 
 Step 3: Prepare Wasm Runtime
@@ -43,25 +43,25 @@ Step 3: Prepare Wasm Runtime
 First, we need to obtain a copy of the mlc-llm source code for the setup script
 
 .. code:: bash
-    
+
     git clone https://github.com/mlc-ai/mlc-llm.git --recursive
     cd mlc-llm
 
 Now we can prepare wasm runtime using the script in mlc-llm repo
 
 .. code:: bash
-    
+
     ./web/prep_emcc_deps.sh
 
 We can then validate the outcome
 
 .. code:: bash
 
-    >>> echo ${TVM_HOME}
+    >>> echo ${TVM_SOURCE_DIR}
 
     /path/set/in/step2
 
-    >>> ls -l ${TVM_HOME}/web/dist/wasm/*.bc
+    >>> ls -l ${TVM_SOURCE_DIR}/web/dist/wasm/*.bc
 
     tvmjs_support.bc
     wasm_runtime.bc
diff --git a/docs/install/mlc_llm.rst b/docs/install/mlc_llm.rst
@@ -216,8 +216,8 @@ There are two ways to do so:
 
        .. code-tab :: bash Install via environment variable
 
-          export MLC_LLM_HOME=/path-to-mlc-llm
-          export PYTHONPATH=$MLC_LLM_HOME/python:$PYTHONPATH
+          export MLC_LLM_SOURCE_DIR=/path-to-mlc-llm
+          export PYTHONPATH=$MLC_LLM_SOURCE_DIR/python:$PYTHONPATH
           alias mlc_llm="python -m mlc_llm"
 
        .. code-tab :: bash Install via pip local project
diff --git a/ios/prepare_libs.sh b/ios/prepare_libs.sh
@@ -7,7 +7,7 @@ function help {
     echo -e "  -h,  --help                          Prints this help\n"
 }
 
-MLC_LLM_HOME="${MLC_LLM_HOME:-..}"
+MLC_LLM_SOURCE_DIR="${MLC_LLM_SOURCE_DIR:-..}"
 is_simulator="false"
 arch="arm64"
 
@@ -54,7 +54,7 @@ fi
 
 mkdir -p build/ && cd build/
 
-cmake $MLC_LLM_HOME\
+cmake $MLC_LLM_SOURCE_DIR\
   -DCMAKE_BUILD_TYPE=$type\
   -DCMAKE_SYSTEM_NAME=iOS\
   -DCMAKE_SYSTEM_VERSION=14.0\
@@ -72,5 +72,5 @@ cmake --build . --config release --target mlc_llm_static -j
 cmake --build . --target install --config release -j
 cd ..
 
-rm -rf $MLC_LLM_HOME/ios/MLCSwift/tvm_home
-ln -s $MLC_LLM_HOME/3rdparty/tvm $MLC_LLM_HOME/ios/MLCSwift/tvm_home
+rm -rf $MLC_LLM_SOURCE_DIR/ios/MLCSwift/tvm_home
+ln -s $MLC_LLM_SOURCE_DIR/3rdparty/tvm $MLC_LLM_SOURCE_DIR/ios/MLCSwift/tvm_home
diff --git a/python/mlc_llm/cli/package.py b/python/mlc_llm/cli/package.py
@@ -24,7 +24,7 @@ def _parse_package_config(path: Union[str, Path]) -> Path:
         return path
 
     def _parse_mlc_llm_home(path: str) -> Path:
-        os.environ["MLC_LLM_HOME"] = path
+        os.environ["MLC_LLM_SOURCE_DIR"] = path
         return Path(path)
 
     def _parse_output(path: Union[str, Path]) -> Path:
@@ -42,8 +42,8 @@ def _parse_output(path: Union[str, Path]) -> Path:
     parser.add_argument(
         "--mlc-llm-home",
         type=_parse_mlc_llm_home,
-        default=os.environ.get("MLC_LLM_HOME", None),
-        help=HELP["mlc_llm_home"] + " (default: the $MLC_LLM_HOME environment variable)",
+        default=os.environ.get("MLC_LLM_SOURCE_DIR", None),
+        help=HELP["mlc_llm_home"] + " (default: the $MLC_LLM_SOURCE_DIR environment variable)",
     )
     parser.add_argument(
         "--output",
@@ -58,7 +58,7 @@ def _parse_output(path: Union[str, Path]) -> Path:
             "MLC LLM home is not specified. "
             "Please obtain a copy of MLC LLM source code by "
             "cloning https://github.com/mlc-ai/mlc-llm, and set environment variable "
-            '"MLC_LLM_HOME=path/to/mlc-llm"'
+            '"MLC_LLM_SOURCE_DIR=path/to/mlc-llm"'
         )
     package(
         package_config_path=parsed.package_config,
diff --git a/python/mlc_llm/support/auto_target.py b/python/mlc_llm/support/auto_target.py
@@ -220,15 +220,15 @@ def build(mod: IRModule, args: "CompileArgs", pipeline=None):
         # Try to locate `mlc_wasm_runtime.bc`
         bc_path = None
         bc_candidates = ["web/dist/wasm/mlc_wasm_runtime.bc"]
-        if os.environ.get("MLC_LLM_HOME", None):
-            mlc_source_home_dir = os.environ["MLC_LLM_HOME"]
+        if os.environ.get("MLC_LLM_SOURCE_DIR", None):
+            mlc_source_home_dir = os.environ["MLC_LLM_SOURCE_DIR"]
             bc_candidates.append(
                 os.path.join(mlc_source_home_dir, "web", "dist", "wasm", "mlc_wasm_runtime.bc")
             )
         error_info = (
             "Cannot find library: mlc_wasm_runtime.bc\n"
             + "Make sure you have run `./web/prep_emcc_deps.sh` and "
-            + "`export MLC_LLM_HOME=/path/to/mlc-llm` so that we can locate the file. "
+            + "`export MLC_LLM_SOURCE_DIR=/path/to/mlc-llm` so that we can locate the file. "
             + "We tried to look at candidate paths:\n"
         )
         for candidate in bc_candidates:
diff --git a/rust/README.md b/rust/README.md
@@ -9,7 +9,7 @@ To set up the MLC-LLM Rust package, please follow these steps:
 
 **Step 2:** Define the environment variables for TVM and MLC-LLM by running the following commands in your terminal:
 ```bash
-export TVM_HOME=/path/to/tvm
+export TVM_SOURCE_DIR=/path/to/tvm
 export MLC_HOME=/path/to/mlc-llm
 ```
 
diff --git a/web/Makefile b/web/Makefile
@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-TVM_ROOT=$(TVM_HOME)
+TVM_ROOT=$(TVM_SOURCE_DIR)
 MLC_LLM_ROOT=$(shell cd ..; pwd)
 
 INCLUDE_FLAGS = -I$(TVM_ROOT) -I$(TVM_ROOT)/include\
diff --git a/web/prep_emcc_deps.sh b/web/prep_emcc_deps.sh
