diff --git a/third_party/tsl/third_party/spirv_llvm_translator/BUILD b/third_party/tsl/third_party/spirv_llvm_translator/BUILD new file mode 100644 index 00000000000000..e69de29bb2d1d6 diff --git a/third_party/tsl/third_party/spirv_llvm_translator/spirv_llvm_translator.BUILD b/third_party/tsl/third_party/spirv_llvm_translator/spirv_llvm_translator.BUILD new file mode 100644 index 00000000000000..557e2e8f50edd2 --- /dev/null +++ b/third_party/tsl/third_party/spirv_llvm_translator/spirv_llvm_translator.BUILD @@ -0,0 +1,34 @@ +cc_library( + name = "spirv_llvm_translator", + srcs = glob([ + "lib/SPIRV/libSPIRV/*.cpp", + "lib/SPIRV/libSPIRV/*.hpp", + "lib/SPIRV/libSPIRV/*.h", + "lib/SPIRV/Mangler/*.cpp", + "lib/SPIRV/Mangler/*.h", + "lib/SPIRV/*.cpp", + "lib/SPIRV/*.hpp", + "lib/SPIRV/*.h", + ]), + hdrs = glob(["include/*"]), + includes = [ + "include/", + "lib/SPIRV/", + "lib/SPIRV/Mangler/", + "lib/SPIRV/libSPIRV/", + ], + visibility = ["//visibility:public"], + deps = [ + "@llvm-project//llvm:Analysis", + "@llvm-project//llvm:BitWriter", + "@llvm-project//llvm:CodeGen", + "@llvm-project//llvm:Core", + "@llvm-project//llvm:Demangle", + "@llvm-project//llvm:IRReader", + "@llvm-project//llvm:Linker", + "@llvm-project//llvm:Passes", + "@llvm-project//llvm:Support", + "@llvm-project//llvm:TransformUtils", + "@spirv_headers//:spirv_cpp_headers", + ], +) diff --git a/third_party/tsl/third_party/spirv_llvm_translator/spirv_llvm_translator.patch b/third_party/tsl/third_party/spirv_llvm_translator/spirv_llvm_translator.patch new file mode 100644 index 00000000000000..0d5dfe9bd753dd --- /dev/null +++ b/third_party/tsl/third_party/spirv_llvm_translator/spirv_llvm_translator.patch @@ -0,0 +1,20 @@ +diff --git a/lib/SPIRV/SPIRVInternal.h b/lib/SPIRV/SPIRVInternal.h +index a828add8..924e13b4 100644 +--- a/lib/SPIRV/SPIRVInternal.h ++++ b/lib/SPIRV/SPIRVInternal.h +@@ -179,11 +179,12 @@ typedef SPIRVMap IntBoolOpMap; + "-v512:512:512-v1024:1024:1024" + + enum SPIRAddressSpace { +- SPIRAS_Private, ++ SPIRAS_Generic, + SPIRAS_Global, +- SPIRAS_Constant, ++ SPIRAS_Internal, + SPIRAS_Local, +- SPIRAS_Generic, ++ SPIRAS_Constant, ++ SPIRAS_Private, + SPIRAS_GlobalDevice, + SPIRAS_GlobalHost, + SPIRAS_Input, diff --git a/third_party/tsl/workspace2.bzl b/third_party/tsl/workspace2.bzl index d5004e732eece7..8e8f9df51d1c68 100644 --- a/third_party/tsl/workspace2.bzl +++ b/third_party/tsl/workspace2.bzl @@ -605,6 +605,22 @@ def _tf_repositories(): urls = tf_mirror_urls("https://github.com/google/glog/archive/refs/tags/v0.4.0.tar.gz"), ) + tf_http_archive( + name = "spirv_headers", + sha256 = "1a248f4199f4a30b2e7304ed4d62f731765bab12d9fa4c5abb189a5e2d57f1f5", + strip_prefix = "SPIRV-Headers-1c6bb2743599e6eb6f37b2969acc0aef812e32e3", + urls = tf_mirror_urls("https://github.com/KhronosGroup/SPIRV-Headers/archive/1c6bb2743599e6eb6f37b2969acc0aef812e32e3.tar.gz"), + ) + + tf_http_archive( + name = "spirv_llvm_translator", + sha256 = "8d33cb29a480457152ce0166bb823fddfba89dea6ce4e7d4eccef7509951f73d", + strip_prefix = "SPIRV-LLVM-Translator-0aab1249acbfb8d206c24988f0974638c290e931", + build_file = "//third_party/spirv_llvm_translator:spirv_llvm_translator.BUILD", + patch_file = ["//third_party/spirv_llvm_translator:spirv_llvm_translator.patch"], + urls = tf_mirror_urls("https://github.com/KhronosGroup/SPIRV-LLVM-Translator/archive/0aab1249acbfb8d206c24988f0974638c290e931.tar.gz"), + ) + def workspace(): # Check the bazel version before executing any repository rules, in case # those rules rely on the version we require here. diff --git a/xla/service/gpu/llvm_gpu_backend/BUILD b/xla/service/gpu/llvm_gpu_backend/BUILD index 2e20618056e35f..e968518cf437dd 100644 --- a/xla/service/gpu/llvm_gpu_backend/BUILD +++ b/xla/service/gpu/llvm_gpu_backend/BUILD @@ -2,6 +2,10 @@ load( "@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured", ) +load( + "@local_config_sycl//sycl:build_defs.bzl", + "if_sycl_is_configured", +) load("//xla:xla.bzl", "xla_cc_test") load("//xla/tsl:tsl.bzl", "internal_visibility") @@ -80,6 +84,8 @@ cc_library( "@local_config_rocm//rocm:rocm_headers", "@llvm-project//llvm:AMDGPUCodeGen", "@llvm-project//llvm:AMDGPUAsmParser", + ]) + if_sycl_is_configured([ + "@spirv_llvm_translator//:spirv_llvm_translator", ]), ) diff --git a/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc b/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc index c84dd8f6092bed..752fdc61ec4b1e 100644 --- a/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc +++ b/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc @@ -101,6 +101,11 @@ limitations under the License. #include "third_party/gpus/cuda/include/cuda.h" #endif +#if TENSORFLOW_USE_SYCL +#include "LLVMSPIRVLib.h" +#include "LLVMSPIRVOpts.h" +#endif // TENSORFLOW_USE_SYCL + namespace xla { namespace gpu { namespace { @@ -396,7 +401,9 @@ absl::Status LinkAndOptimizeModule( llvm::CGSCCAnalysisManager cgam; llvm::ModuleAnalysisManager mam; - fam.registerPass([&] { return target_machine->getTargetIRAnalysis(); }); + if (target_machine) { + fam.registerPass([&] { return target_machine->getTargetIRAnalysis(); }); + } llvm::PipelineTuningOptions pto; pto.SLPVectorization = true; @@ -1075,5 +1082,92 @@ absl::StatusOr> CompileToHsaco( } // namespace amdgpu +namespace { +std::unique_ptr SPIRGetTargetMachine( + llvm::Triple target_triple, se::GpuComputeCapability gpu_version, + const DebugOptions& debug_options) { + return nullptr; +} + +absl::Status SPIRTargetModuleLinker( + llvm::Module* module, se::GpuComputeCapability gpu_version, + const DebugOptions& debug_options, + const std::string& device_bitcode_dir_path) { + return OkStatus(); +} + +absl::StatusOr EmitModuleToSpir( + llvm::Module* module, se::GpuComputeCapability gpu_version, + const DebugOptions& debug_options) { +#if TENSORFLOW_USE_SYCL + SPIRV::TranslatorOpts::ExtensionsStatusMap ExtensionsStatus; + SPIRV::TranslatorOpts opts(SPIRV::VersionNumber::MaximumVersion, + ExtensionsStatus); + opts.enableAllExtensions(); // enable all SPIR-V extension first + + std::ostringstream oss; + std::string err; + bool success = llvm::writeSpirv(module, opts, oss, err); + if (!success) { + return xla::Internal("Fails to convert LLVM as SPIR-V: %s", err); + } + return oss.str(); +#else + return absl::UnimplementedError("Not implemented for SYCL"); +#endif +} + +void SPIRBackendInit(const DebugOptions& debug_options) { + FeedLLVMWithFlags({"-slp-vectorize-hor=false"}); + + FeedLLVMWithFlags({ + "-slp-min-reg-size=64", + "-slp-max-reg-size=64", + }); + + llvm_ir::InitializeLLVMCommandLineOptions( + debug_options.xla_backend_extra_options()); + + llvm::PassRegistry* registry = llvm::PassRegistry::getPassRegistry(); + InitializePasses(registry); +} +} // namespace + +namespace spir { +absl::StatusOr> CompileToSpir( + llvm::Module* module, se::GpuComputeCapability gpu_version, + const DebugOptions& debug_options) { + std::string libdevice_dir_path; + static absl::once_flag backend_init_flag; + absl::call_once(backend_init_flag, SPIRBackendInit, debug_options); + + std::string spir; + { + XLA_SCOPED_LOGGING_TIMER("Compile module " + module->getName().str()); + + // If the module has no functions or globals, there's nothing to compile. + if (module->empty() && module->global_empty()) { + VLOG(2) << "Module '" << module->getName().str() + << "' is empty. Skipping compilation."; + return std::vector(); + } + + llvm::Triple default_target_triple("spir64-unknown-unknown"); + std::unique_ptr target_machine = + SPIRGetTargetMachine(default_target_triple, gpu_version, debug_options); + + TF_RETURN_IF_ERROR(LinkAndOptimizeModule( + module, gpu_version, debug_options, libdevice_dir_path, + SPIRTargetModuleLinker, default_target_triple, target_machine.get(), + kDefaultInlineThreshold)); + + // Lower optimized LLVM module to SPIR. + TF_ASSIGN_OR_RETURN(spir, + EmitModuleToSpir(module, gpu_version, debug_options)); + } + return std::vector(spir.begin(), spir.end()); +} +} // namespace spir + } // namespace gpu } // namespace xla diff --git a/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.h b/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.h index 3d67bf043e6444..f5ddd25647085e 100644 --- a/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.h +++ b/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.h @@ -73,6 +73,13 @@ absl::StatusOr> CompileToHsaco( const std::string& module_config_cache_key); } // namespace amdgpu +namespace spir { +// Compiles the argument module and returns it. +absl::StatusOr> CompileToSpir( + llvm::Module* module, se::GpuComputeCapability gpu_version, + const DebugOptions& debug_options); +} // namespace spir + } // namespace gpu } // namespace xla