[pull] master from tensorflow:master

third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD

@@ -79,6 +79,7 @@ cc_library(
         "@xla//xla/python:__pkg__",
         "//learning/brain/tfrc/executor/stream_executor:__pkg__",
         "//net/grpc/internal/src/core/ext/xprof_profiler:__pkg__",
+        "//third_party/py/torch_tpu/pjrt:__pkg__",
     ]),
     deps = [
         ":profiler_interface",
@@ -135,6 +136,7 @@ cc_library(
         "@xla//xla/tsl/profiler:internal",
         "@xla//xla/tsl/profiler:xla_profiler_backends",
         "//net/grpc/internal/src/core/ext/xprof_profiler:__pkg__",
+        "//third_party/py/torch_tpu/pjrt:__pkg__",
     ]),
     deps = [
         "//tsl/profiler/protobuf:xplane_proto_cc",

third_party/xla/xla/backends/autotuner/BUILD

@@ -35,6 +35,7 @@ cc_library(
     hdrs = ["autotuner.h"],
     deps = [
         ":autotuner_cache_interface",
+        ":backends_proto_cc",
         ":codegen_backend",
         ":profiler",
         "//xla:autotune_results_proto_cc",

third_party/xla/xla/backends/autotuner/autotuner.cc

@@ -44,6 +44,7 @@ limitations under the License.
 #include "google/protobuf/text_format.h"
 #include "xla/autotuning.pb.h"
 #include "xla/backends/autotuner/autotuner_cache_interface.h"
+#include "xla/backends/autotuner/backends.pb.h"
 #include "xla/backends/autotuner/codegen_backend.h"
 #include "xla/backends/autotuner/profiler.h"
 #include "xla/hlo/ir/hlo_instruction.h"

third_party/xla/xla/backends/gpu/codegen/kernels/BUILD

@@ -409,15 +409,12 @@ cc_library(
     name = "ptx_custom_kernel",
     srcs = ["ptx_custom_kernel.cc"],
     hdrs = ["ptx_custom_kernel.h"],
-    tags = [
-        "cuda-only",
-        "gpu",
-    ],
     visibility = [":friends"],
     deps = [
         ":custom_kernel",
         "//xla/stream_executor:kernel",
         "//xla/stream_executor:kernel_args",
+        "//xla/stream_executor:kernel_args_packing_spec",
         "//xla/stream_executor:kernel_spec",
         "//xla/stream_executor:launch_dim",
         "@com_google_absl//absl/status:statusor",

third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.cc

@@ -16,15 +16,18 @@ limitations under the License.
 #include "xla/backends/gpu/codegen/kernels/ptx_custom_kernel.h"
 
 #include <cstddef>
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "xla/backends/gpu/codegen/kernels/custom_kernel.h"
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/kernel_args.h"
+#include "xla/stream_executor/kernel_args_packing_spec.h"
 #include "xla/stream_executor/kernel_spec.h"
 #include "xla/stream_executor/launch_dim.h"
 
@@ -54,7 +57,7 @@ absl::StatusOr<CustomKernel> GetPtxCustomKernel(std::string kernel_name,
           ptx, kernel_name, /*arity=*/num_args, KernelArgsPacking);
   return CustomKernel(std::move(kernel_name), kernel_spec, block_dim,
                       thread_dim, shared_memory_bytes);
-};
+}
 
 absl::StatusOr<CustomKernel> GetPtxCustomKernel(
     std::string kernel_name, absl::string_view ptx, int num_args,
@@ -65,17 +68,33 @@ absl::StatusOr<CustomKernel> GetPtxCustomKernel(
           ptx, kernel_name, /*arity=*/num_args, KernelArgsPacking);
   return CustomKernel(std::move(kernel_name), kernel_spec, block_dim,
                       thread_dim, cluster_dim, shared_memory_bytes);
-};
+}
 
 absl::StatusOr<CustomKernel> GetOwnedPtxCustomKernel(
     std::string kernel_name, std::string ptx, int num_args,
     se::BlockDim block_dim, se::ThreadDim thread_dim,
     size_t shared_memory_bytes) {
   se::KernelLoaderSpec kernel_spec =
       se::KernelLoaderSpec::CreateOwningCudaPtxInMemorySpec(
-          ptx, kernel_name, /*arity=*/num_args, KernelArgsPacking);
+          std::move(ptx), kernel_name, /*arity=*/num_args, KernelArgsPacking);
   return CustomKernel(std::move(kernel_name), kernel_spec, block_dim,
                       thread_dim, shared_memory_bytes);
-};
+}
+
+absl::StatusOr<CustomKernel> CreateOwnedCubinCustomKernel(
+    std::string kernel_name, std::vector<uint8_t> cubin, int num_args,
+    se::BlockDim block_dim, se::ThreadDim thread_dim,
+    size_t shared_memory_bytes) {
+  se::KernelArgsPackingSpec packing_spec;
+  for (int i = 0; i < num_args; i++) {
+    packing_spec.AddAddressArgument(i);
+  }
+
+  se::KernelLoaderSpec kernel_spec =
+      se::KernelLoaderSpec::CreateOwningCudaCubinInMemorySpec(
+          std::move(cubin), kernel_name, /*arity=*/num_args, packing_spec);
+  return CustomKernel(std::move(kernel_name), std::move(kernel_spec), block_dim,
+                      thread_dim, shared_memory_bytes);
+}
 
 }  // namespace xla::gpu::kernel

third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.h

@@ -43,5 +43,10 @@ absl::StatusOr<CustomKernel> GetOwnedPtxCustomKernel(
     se::BlockDim block_dim, se::ThreadDim thread_dim,
     size_t shared_memory_bytes = 0);
 
+absl::StatusOr<CustomKernel> CreateOwnedCubinCustomKernel(
+    std::string kernel_name, std::vector<uint8_t> cubin, int num_args,
+    se::BlockDim block_dim, se::ThreadDim thread_dim,
+    size_t shared_memory_bytes);
+
 }  // namespace xla::gpu::kernel
 #endif  // XLA_BACKENDS_GPU_CODEGEN_KERNELS_PTX_CUSTOM_KERNEL_H_

third_party/xla/xla/backends/gpu/codegen/llvm/BUILD

@@ -29,6 +29,9 @@ cc_library(
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla/backends/gpu/codegen:fusion_emitter",
+        "//xla/backends/gpu/codegen/kernels:custom_kernel",
+        "//xla/backends/gpu/codegen/kernels:ptx_custom_kernel",
+        "//xla/backends/gpu/runtime:custom_kernel_thunk",
         "//xla/backends/gpu/runtime:kernel_thunk",
         "//xla/backends/gpu/runtime:thunk",
         "//xla/backends/gpu/runtime:thunk_id",
@@ -122,3 +125,15 @@ cc_library(
         "@tsl//tsl/platform:logging",
     ],
 )
+
+cc_library(
+    name = "llvm_ir_compiler",
+    hdrs = ["llvm_ir_compiler.h"],
+    deps = [
+        "//xla:xla_proto_cc",
+        "//xla/stream_executor:device_description",
+        "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/status:statusor",
+        "@llvm-project//llvm:Core",
+    ],
+)