[xla:gpu] Support creating ThunkInfo from HLO tensorflow#6224

PiperOrigin-RevId: 574728456
Intel-tensorflow · Oct 19, 2023 · 3fbde8e · 3fbde8e
1 parent 9e4a43c
commit 3fbde8e
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 13 deletions.
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
@@ -743,6 +743,7 @@ cc_library(
     deps = [
         ":buffer_allocations",
         ":gpu_executable_run_options",
+        "//xla/hlo/ir:hlo",
         "//xla/service:executable",
         "//xla/stream_executor",
         "//xla/translate/mhlo_to_hlo:location_exporter",

diff --git a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
@@ -1418,11 +1418,9 @@ Status IrEmitterUnnested::EmitCholeskyThunk(const HloInstruction* instr) {
 
   ThunkSequence thunks;
 
-  // TODO(b/304613751): We need to construct real ThunkInfo from instr for Xprof
-  // integration.
   if (operand_buffer != a_buffer) {
     thunks.push_back(std::make_unique<DeviceToDeviceCopyThunk>(
-        Thunk::ThunkInfo(nullptr),
+        Thunk::ThunkInfo::WithProfileAnnotation(instr),
         /*source_buffer=*/operand_buffer,
         /*destination_buffer=*/a_buffer,
         /*mem_size=*/ShapeUtil::ByteSizeOf(shape),
@@ -1431,7 +1429,7 @@ Status IrEmitterUnnested::EmitCholeskyThunk(const HloInstruction* instr) {
   }
 
   thunks.push_back(std::make_unique<CholeskyThunk>(
-      Thunk::ThunkInfo(nullptr), options,
+      Thunk::ThunkInfo::WithProfileAnnotation(instr), options,
       PtxOptsFromDebugOptions(ir_emitter_context_->debug_options()), a_buffer,
       workspace_buffer, info_buffer, shape.element_type(), batch_size, n));
 
@@ -1440,7 +1438,7 @@ Status IrEmitterUnnested::EmitCholeskyThunk(const HloInstruction* instr) {
     AddThunkToThunkSequence(std::move(thunks[0]));
   } else {
     AddThunkToThunkSequence(std::make_unique<SequentialThunk>(
-        Thunk::ThunkInfo(nullptr), std::move(thunks)));
+        Thunk::ThunkInfo::WithProfileAnnotation(instr), std::move(thunks)));
   }
 
   return OkStatus();

diff --git a/third_party/xla/xla/service/gpu/thunk.cc b/third_party/xla/xla/service/gpu/thunk.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <string>
 
 #include "absl/strings/str_format.h"
+#include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/translate/mhlo_to_hlo/location_exporter.h"
 
 namespace xla {
@@ -130,5 +131,13 @@ Thunk::ThunkInfo Thunk::ThunkInfo::WithProfileAnnotation(mlir::Operation* op) {
   return thunk_info;
 }
 
+Thunk::ThunkInfo Thunk::ThunkInfo::WithProfileAnnotation(
+    const HloInstruction* instr) {
+  ThunkInfo thunk_info(nullptr);
+  thunk_info.profile_annotation =
+      absl::StrFormat("Thunk:#hlo_op=%s#", instr->name());
+  return thunk_info;
+}
+
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/gpu/thunk.h b/third_party/xla/xla/service/gpu/thunk.h
@@ -27,6 +27,7 @@ limitations under the License.
 
 #include "absl/types/span.h"
 #include "mlir/IR/Operation.h"  // from @llvm-project
+#include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/service/gpu/buffer_allocations.h"
 #include "xla/service/gpu/gpu_executable_run_options.h"
 #include "xla/service/service_executable_run_options.h"
@@ -115,19 +116,20 @@ class Thunk {
 
   struct ThunkInfo {
     explicit ThunkInfo(mlir::Operation* op) : op(op) {}
-    std::optional<int64_t> profile_index;
+    static ThunkInfo WithProfileAnnotation(mlir::Operation* op);
+    static ThunkInfo WithProfileAnnotation(const HloInstruction* instr);
+
     std::string profile_annotation;
+    // TODO(b/304613751): This is only needed by the LMHLO. Remove this when
+    // LMHLO is removed from the runtime pipeline.
     mlir::Operation* op;
-
-    static ThunkInfo WithProfileAnnotation(mlir::Operation* op);
   };
 
   // The hlo_instruction argument is meant to be the instruction this thunk was
   // generated from, but Thunk never uses this argument other than to save it
   // to Thunk::hlo_instruction, so it can be null.
   Thunk(Kind kind, ThunkInfo thunk_info)
       : kind_(kind),
-        profile_index_(thunk_info.profile_index),
         profile_annotation_(thunk_info.profile_annotation),
         op_(thunk_info.op) {}
   virtual ~Thunk() = default;
@@ -177,12 +179,8 @@ class Thunk {
 
   static absl::string_view KindToString(Thunk::Kind kind);
 
- protected:
-  std::optional<int64_t> profile_index() const { return profile_index_; }
-
  private:
   Kind kind_;
-  std::optional<int64_t> profile_index_;
   std::string profile_annotation_;
   mlir::Operation* op_;
 };