[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

@@ -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 = []

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):