PR #11424: [XLA:GPU] Add SPIRV-LLVM-Translator and translation pass

Imported from GitHub PR #11424 It is a sub PR of #9042 to add spirv-llvm-translator and translation pass Copybara import of the project: -- 6d8ce76 by Sheng, Yang <[email protected]>: Add SPIRV-LLVM-Translator and translation pass update comments and SPIRV-LLVM-Translator commit Merging this change closes #11424 COPYBARA_INTEGRATE_REVIEW=#11424 from Intel-tensorflow:yang/llvm-spirv 6d8ce76 PiperOrigin-RevId: 661272459
openxla · Aug 9, 2024 · e8a7075 · e8a7075
1 parent 321c62b
commit e8a7075
Show file tree

Hide file tree

Showing 9 changed files with 150 additions and 0 deletions.
diff --git a/opensource_only.files b/opensource_only.files
@@ -34,6 +34,7 @@ third_party/py/python_init_toolchains.bzl:
 third_party/py/python_repo.bzl:
 third_party/python_runtime/BUILD:
 third_party/repo.bzl:
+third_party/spirv_llvm_translator/spirv_llvm_translator.BUILD:
 third_party/stablehlo/BUILD:
 tools/toolchains/BUILD:
 tools/toolchains/clang6/BUILD:

diff --git a/third_party/spirv_llvm_translator/BUILD b/third_party/spirv_llvm_translator/BUILD
@@ -0,0 +1,7 @@
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+# spirv_llvm_translator license placeholder
diff --git a/third_party/spirv_llvm_translator/spirv_llvm_translator.BUILD b/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/spirv_llvm_translator/spirv_llvm_translator.patch b/third_party/spirv_llvm_translator/spirv_llvm_translator.patch
@@ -0,0 +1,25 @@
+diff --git a/lib/SPIRV/SPIRVInternal.h b/lib/SPIRV/SPIRVInternal.h
+index a828add8..924e13b4 100644
+
+Spir backend uses different addrspace representations link with nvptx backend link.
+We reorder the enum value here so that we can make XLA LLVM codegen simple(avoiding
+changing addrspace based on device backend everywhere)
+
+--- a/lib/SPIRV/SPIRVInternal.h
++++ b/lib/SPIRV/SPIRVInternal.h
+@@ -179,11 +179,12 @@ typedef SPIRVMap<Op, Op, IntBoolOpMapId> 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/opensource_only.files b/third_party/tsl/opensource_only.files
@@ -93,6 +93,7 @@ third_party/remote_config/remote_platform_configure.bzl:
 third_party/repo.bzl:
 third_party/six.BUILD:
 third_party/snappy.BUILD:
+third_party/spirv_llvm_translator/spirv_llvm_translator.BUILD:
 third_party/systemlibs/BUILD.tpl:
 third_party/systemlibs/BUILD:
 third_party/systemlibs/absl_py.BUILD:

diff --git a/third_party/tsl/third_party/spirv_llvm_translator/BUILD b/third_party/tsl/third_party/spirv_llvm_translator/BUILD
@@ -0,0 +1,7 @@
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+# spirv_llvm_translator license placeholder
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
@@ -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
@@ -0,0 +1,25 @@
+diff --git a/lib/SPIRV/SPIRVInternal.h b/lib/SPIRV/SPIRVInternal.h
+index a828add8..924e13b4 100644
+
+Spir backend uses different addrspace representations link with nvptx backend link.
+We reorder the enum value here so that we can make XLA LLVM codegen simple(avoiding
+changing addrspace based on device backend everywhere)
+
+--- a/lib/SPIRV/SPIRVInternal.h
++++ b/lib/SPIRV/SPIRVInternal.h
+@@ -179,11 +179,12 @@ typedef SPIRVMap<Op, Op, IntBoolOpMapId> 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
@@ -591,6 +591,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 = "11d835c60297b26532c05c3f3b581ba7a2787b5ae7399e94f72c392169216f11",
+        strip_prefix = "SPIRV-Headers-b73e168ca5e123dcf3dea8a34b19a5130f421ae1",
+        urls = tf_mirror_urls("https://github.com/KhronosGroup/SPIRV-Headers/archive/b73e168ca5e123dcf3dea8a34b19a5130f421ae1.tar.gz"),
+    )
+
+    tf_http_archive(
+        name = "spirv_llvm_translator",
+        sha256 = "d499769f4fd1e0ce9d4dbd3622ee7e3e641b5623dcdf811521e3e7c0bdb1e6c2",
+        strip_prefix = "SPIRV-LLVM-Translator-dad1f0eaab8047a4f73c50ed5f3d1694b78aae97",
+        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/dad1f0eaab8047a4f73c50ed5f3d1694b78aae97.tar.gz"),
+    )
+
 # buildifier: disable=unnamed-macro
 def workspace():
     # Check the bazel version before executing any repository rules, in case