[tvm] branch main updated: [Bugfix][CUTLASS] CUTLASS path finding (#15480)

[kp...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

kparzysz pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 2f090648d2 [Bugfix][CUTLASS] CUTLASS path finding (#15480)
2f090648d2 is described below

commit 2f090648d29988468d288517a406b96db7b7e03a
Author: Junru Shao <ju...@apache.org>
AuthorDate: Fri Aug 4 11:43:44 2023 -0700

    [Bugfix][CUTLASS] CUTLASS path finding (#15480)
    
    [Bugfix][CUTLASS] CUTLASS path finding (#15476)
    
    This PR fixes the path finding for `3rdparty/cutlass` by trying out
    different combinations. It is helpful when TVM is packaged differently.
---
 python/tvm/_ffi/libinfo.py          |  9 +++++++--
 python/tvm/contrib/cutlass/build.py | 15 ++++++++-------
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/python/tvm/_ffi/libinfo.py b/python/tvm/_ffi/libinfo.py
index 19b8b50bbf..3ec78fe278 100644
--- a/python/tvm/_ffi/libinfo.py
+++ b/python/tvm/_ffi/libinfo.py
@@ -15,8 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 """Library information."""
-import sys
 import os
+import sys
 
 
 def split_env_var(env_var, split):
@@ -169,7 +169,12 @@ def find_include_path(name=None, search_path=None, optional=False):
         source_dir = os.environ["TVM_HOME"]
     else:
         ffi_dir = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
-        source_dir = os.path.join(ffi_dir, "..", "..", "..")
+        for source_dir in ["..", "../..", "../../.."]:
+            source_dir = os.path.join(ffi_dir, source_dir)
+            if os.path.isdir(os.path.join(source_dir, "include")):
+                break
+        else:
+            raise AssertionError("Cannot find the source directory given ffi_dir: {ffi_dir}")
     third_party_dir = os.path.join(source_dir, "3rdparty")
 
     header_path = []
diff --git a/python/tvm/contrib/cutlass/build.py b/python/tvm/contrib/cutlass/build.py
index 7b5a6b97b9..cd990d0eea 100644
--- a/python/tvm/contrib/cutlass/build.py
+++ b/python/tvm/contrib/cutlass/build.py
@@ -37,13 +37,14 @@ def has_cutlass():
 
 
 def _get_cutlass_path():
-    tvm_root = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../../../../")
-    cutlass_path = os.path.join(tvm_root, "3rdparty/cutlass")
-    assert os.path.exists(cutlass_path), (
-        f"The CUTLASS root directory not found in {cutlass_path}. Currently, using CUTLASS "
-        f"requires building TVM from source."
-    )
-    return cutlass_path
+    invalid_paths = []
+    for rel in ["../../../../", "../../../", "../../"]:
+        tvm_root = os.path.join(os.path.dirname(os.path.realpath(__file__)), rel)
+        cutlass_path = os.path.join(tvm_root, "3rdparty/cutlass")
+        if os.path.exists(cutlass_path):
+            return cutlass_path
+        invalid_paths.append(cutlass_path)
+    raise AssertionError(f"The CUTLASS root directory not found in: {invalid_paths}")
 
 
 def _get_cutlass_compile_options(sm, threads, use_fast_math=False):