Support FFI (main) (#357)

Co-authored-by: Lu Teng <[email protected]>

Author: mayuyuace

third_party/openxla.patch

@@ -1,5 +1,6 @@
 cc_binary(
     name = "pjrt_plugin_xpu.so",
+    linkopts = ["-Wl,-rpath,$$ORIGIN/../intel_extension_for_openxla/service/gpu"],
     linkshared = True,
     visibility = ["//visibility:public"],
     deps = [

xla/BUILD

@@ -21,6 +21,24 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "sycl_custom_call",
+    srcs = [
+        "sycl_custom_call.cc",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//xla/service:onednn_util",
+        "//xla/service/gpu:sycl_onednn",
+        "@xla//xla/ffi",
+        "@xla//xla/ffi:ffi_api",
+        "@xla//xla/stream_executor",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+    ],
+    alwayslink = 1,
+)
+
 xetla_library(
     name = "onednn_matmul_utils",
     srcs = ["onednn_matmul_utils.cc"],
@@ -29,27 +47,7 @@ xetla_library(
         ":scratch_allocator",
         "//xla/service:onednn_util",
         "//xla/service/gpu/xetla/gemm:gemm_kernel",
-        "//xla/stream_executor/sycl:sycl_executor",
-        "@com_google_absl//absl/algorithm:container",
-        "@com_google_absl//absl/types:span",
-        "@onednn_gpu//:onednn_gpu",
-        "@tsl//tsl/framework:numeric_types",
-        "@tsl//tsl/platform:statusor",
-        "@tsl//tsl/platform:types",
-        "@xetla//:xetla_header",
-        "@xla//xla:shape_util",
-        "@xla//xla:status_macros",
-        "@xla//xla:statusor",
-        "@xla//xla:types",
-        "@xla//xla:util",
-        "@xla//xla:xla_data_proto_cc",
-        "@xla//xla/hlo/ir:hlo",
-        "@xla//xla/mlir_hlo",
-        "@xla//xla/mlir_hlo:lhlo_gpu",
-        "@xla//xla/service/gpu:backend_configs_cc",
-        "@xla//xla/service/gpu:ir_emission_utils",
         "@xla//xla/service/gpu:matmul_utils",
-        "@xla//xla/stream_executor:stream_executor_headers",
     ],
 )
 
@@ -58,7 +56,7 @@ cc_library(
     srcs = ["gemm_impl_picker.cc",],
     hdrs = ["gemm_impl_picker.h"],
     deps = [
-        ":onednn_matmul_utils",
+        ":sycl_onednn",
         "//xla/stream_executor/sycl:hw_info",
         "@com_google_absl//absl/algorithm:container",
         "@tsl//tsl/platform:errors",
@@ -179,31 +177,50 @@ xpu_library(
     ],
 )
 
+cc_import(
+    name = "sycl_onednn",
+    hdrs = [
+        "sycl_onednn.h",
+        "onednn_gpu_conv_runner.h",
+        "onednn_matmul_utils.h",
+    ],
+    shared_library = ":sycl_onednn.so",
+    visibility = ["//visibility:public"],
+    deps = [
+        ":scratch_allocator",
+        "@xla//xla/service/gpu:gpu_conv_runner",
+        "@xla//xla/service/gpu:thunk",
+        "@xla//xla/service/gpu:matmul_utils",
+    ],
+)
+
+cc_binary(
+    name = "sycl_onednn.so",
+    srcs = [
+        "sycl_onednn.cc",
+        "sycl_onednn.h",
+    ],
+    linkshared = True,
+    deps = [
+        ":onednn_gpu_conv_runner",
+        ":onednn_matmul_utils",
+    ],
+)
+
 cc_library(
     name = "onednn_gpu_conv_runner",
-    srcs = ["onednn_gpu_conv_runner.cc"],
-    hdrs = ["onednn_gpu_conv_runner.h"],
+    srcs = [
+        "onednn_gpu_conv_runner.cc",
+        "onednn_gpu_conv_runner.h",
+    ],
     deps = [
         ":scratch_allocator",
         "//xla/service:onednn_util",
-        "@com_google_absl//absl/strings",
-        "@tsl//tsl/framework:numeric_types",
-        "@xla//xla:shape_util",
-        "@xla//xla:status",
-        "@xla//xla:status_macros",
-        "@xla//xla:statusor",
-        "@xla//xla:types",
-        "@xla//xla:util",
-        "@xla//xla:xla_data_proto_cc",
-        "@xla//xla/hlo/ir:hlo",
-        "@xla//xla/service/gpu:backend_configs_cc",
-        "@xla//xla/service/gpu:cublas_cudnn",
+        "@xla//xla/ffi",
+        "@xla//xla/ffi:ffi_api",
         "@xla//xla/service/gpu:gpu_conv_runner",
         "@xla//xla/service/gpu:stream_executor_util",
         "@xla//xla/service/gpu:thunk",
-        "@xla//xla/stream_executor",
-        "@xla//xla/stream_executor/gpu:gpu_stream",
-        "@xla//xla/stream_executor/gpu:gpu_types_header",
     ],
 )
 
@@ -300,4 +317,4 @@ cc_library(
         "@xla//xla/service:hlo_pass",
         "@xla//xla/service:pattern_matcher",
     ],
-)
+)

xla/service/gpu/BUILD

@@ -55,26 +55,6 @@ bool IsXetlaSupport(const GemmConfig& config) {
   return xetla_support;
 }
 
-absl::StatusOr<se::gpu::BlasLt::Epilogue> AsBlasLtEpilogue(
-    GemmBackendConfig_Epilogue epilogue) {
-  switch (epilogue) {
-    case GemmBackendConfig::DEFAULT:
-      return se::gpu::BlasLt::Epilogue::kDefault;
-    case GemmBackendConfig::RELU:
-      return se::gpu::BlasLt::Epilogue::kReLU;
-    case GemmBackendConfig::GELU:
-      return se::gpu::BlasLt::Epilogue::kGELU;
-    case GemmBackendConfig::BIAS:
-      return se::gpu::BlasLt::Epilogue::kBias;
-    case GemmBackendConfig::BIAS_RELU:
-      return se::gpu::BlasLt::Epilogue::kBiasThenReLU;
-    case GemmBackendConfig::BIAS_GELU:
-      return se::gpu::BlasLt::Epilogue::kBiasThenGELU;
-    default:
-      return absl::InternalError("Unsupported Epilogue.");
-  }
-}
-
 absl::StatusOr<absl::Duration> GetExecuteTime(
     const HloInstruction* gemm, const AutotuneConfig& autotune_config) {
   se::DeviceMemoryAllocator* allocator = autotune_config.GetAllocator();
@@ -122,7 +102,7 @@ absl::StatusOr<absl::Duration> GetExecuteTime(
             autotune_config, rng_state));
   }
 
-  TF_ASSIGN_OR_RETURN(auto epilogue, AsBlasLtEpilogue(gemm_config.epilogue()));
+  TF_ASSIGN_OR_RETURN(auto epilogue, SYCLGemm::AsSYCLEpilogue(gemm_config.epilogue()));
   se::OwningScratchAllocator<> scratch_allocator(
       stream->parent()->device_ordinal(), autotune_config.GetAllocator());
 
@@ -264,4 +244,4 @@ absl::StatusOr<bool> GemmAlgorithmPicker::Run(
 }
 
 }  // namespace gpu
-}  // namespace xla
+}  // namespace xla