wehub-resource-sync
82 lines
4.0 KiB
CMake
82 lines
4.0 KiB
CMake
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
if(USE_CUDA AND USE_CUTLASS)
|
|
tvm_file_glob(GLOB CUTLASS_CONTRIB_SRC
|
|
src/relax/backend/contrib/cutlass/*.cc
|
|
)
|
|
list(APPEND COMPILER_SRCS ${CUTLASS_CONTRIB_SRC})
|
|
|
|
set(FPA_INTB_GEMM_TVM_BINDING ON)
|
|
set(FPA_INTB_GEMM_TVM_HOME ${PROJECT_SOURCE_DIR})
|
|
|
|
### Build cutlass runtime objects for fpA_intB_gemm using its cutlass submodule
|
|
add_subdirectory(${PROJECT_SOURCE_DIR}/3rdparty/cutlass_fpA_intB_gemm)
|
|
target_include_directories(fpA_intB_gemm PRIVATE
|
|
${PROJECT_SOURCE_DIR}/3rdparty/cutlass_fpA_intB_gemm
|
|
${PROJECT_SOURCE_DIR}/3rdparty/cutlass_fpA_intB_gemm/cutlass/include
|
|
)
|
|
target_link_libraries(fpA_intB_gemm_tvm PRIVATE tvm_ffi_header)
|
|
|
|
set(CUTLASS_FPA_INTB_RUNTIME_SRCS "")
|
|
list(APPEND CUTLASS_FPA_INTB_RUNTIME_SRCS src/runtime/extra/contrib/cutlass/weight_preprocess.cc)
|
|
add_library(fpA_intB_cutlass_objs OBJECT ${CUTLASS_FPA_INTB_RUNTIME_SRCS})
|
|
target_link_libraries(fpA_intB_cutlass_objs PRIVATE tvm_runtime_extra_defs)
|
|
target_include_directories(fpA_intB_cutlass_objs PRIVATE
|
|
${PROJECT_SOURCE_DIR}/3rdparty/cutlass_fpA_intB_gemm
|
|
${PROJECT_SOURCE_DIR}/3rdparty/cutlass_fpA_intB_gemm/cutlass/include
|
|
)
|
|
target_link_libraries(tvm_runtime_extra PRIVATE fpA_intB_cutlass_objs)
|
|
|
|
### Build cutlass runtime objects for flash attention
|
|
add_subdirectory(${PROJECT_SOURCE_DIR}/3rdparty/libflash_attn)
|
|
target_include_directories(flash_attn PRIVATE
|
|
${PROJECT_SOURCE_DIR}/3rdparty/libflash_attn
|
|
${PROJECT_SOURCE_DIR}/3rdparty/libflash_attn/cutlass/include
|
|
)
|
|
|
|
### Build cutlass runtime objects using TVM's 3rdparty/cutlass submodule
|
|
set(CUTLASS_DIR ${PROJECT_SOURCE_DIR}/3rdparty/cutlass)
|
|
set(TVM_CUTLASS_RUNTIME_SRCS "")
|
|
|
|
if(CMAKE_CUDA_ARCHITECTURES MATCHES "90a")
|
|
list(APPEND TVM_CUTLASS_RUNTIME_SRCS src/runtime/extra/contrib/cutlass/fp16_group_gemm_sm90.cu)
|
|
list(APPEND TVM_CUTLASS_RUNTIME_SRCS src/runtime/extra/contrib/cutlass/fp8_group_gemm_sm90.cu)
|
|
list(APPEND TVM_CUTLASS_RUNTIME_SRCS src/runtime/extra/contrib/cutlass/fp8_gemm.cu)
|
|
list(APPEND TVM_CUTLASS_RUNTIME_SRCS src/runtime/extra/contrib/cutlass/fp8_groupwise_scaled_gemm_sm90.cu)
|
|
endif()
|
|
if(CMAKE_CUDA_ARCHITECTURES MATCHES "100a")
|
|
list(APPEND TVM_CUTLASS_RUNTIME_SRCS src/runtime/extra/contrib/cutlass/fp16_group_gemm_sm100.cu)
|
|
list(APPEND TVM_CUTLASS_RUNTIME_SRCS src/runtime/extra/contrib/cutlass/fp8_groupwise_scaled_gemm_sm100.cu)
|
|
list(APPEND TVM_CUTLASS_RUNTIME_SRCS src/runtime/extra/contrib/cutlass/fp8_groupwise_scaled_group_gemm_sm100.cu)
|
|
endif()
|
|
if(TVM_CUTLASS_RUNTIME_SRCS)
|
|
add_library(tvm_cutlass_objs OBJECT ${TVM_CUTLASS_RUNTIME_SRCS})
|
|
target_compile_options(tvm_cutlass_objs PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-lineinfo --expt-relaxed-constexpr>)
|
|
target_include_directories(tvm_cutlass_objs PRIVATE
|
|
${CUTLASS_DIR}/include
|
|
${PROJECT_SOURCE_DIR}/3rdparty/cutlass_fpA_intB_gemm/cutlass_extensions/include
|
|
)
|
|
target_link_libraries(tvm_cutlass_objs PRIVATE tvm_runtime_extra_defs)
|
|
# Note: enable this to get more detailed logs for cutlass kernels
|
|
# target_compile_definitions(tvm_cutlass_objs PRIVATE CUTLASS_DEBUG_TRACE_LEVEL=2)
|
|
target_link_libraries(tvm_runtime_extra PRIVATE tvm_cutlass_objs)
|
|
endif()
|
|
|
|
message(STATUS "Build with CUTLASS")
|
|
endif()
|