[go: nahoru, domu]
Skip to content

Commit

Permalink
Remove reference to mlir::Operations in Thunks.
Browse files Browse the repository at this point in the history 
The `mlir::Operation` in `Thunk` (`Thunk::op_`) is never set to a value different from `nullptr`.

Conveniently, this also removes the need for `Thunk::ClearCompileTimeInfo`, which in turn leaves `ForAllThunks` (in `compile_module_to_llvm_ir.cc`) unused.

I'm doing this first so that I can then (a) move `ForAllThunks` to its own file, and (b) change its behavior to actually iterate over all `Thunks`: for example, `ForAllThunks` currently does not iterate over the `embedded_thunk` in `AddressComputationThunk`. This behavior will be safer to change after there's no remaining users of the function.

PiperOrigin-RevId: 634879646
  • Loading branch information
@tensorflower-gardener
tensorflower-gardener committed May 17, 2024
1 parent 43e54f6 commit d928140
Show file tree
Hide file tree
Showing 11 changed files with 46 additions and 88 deletions.
5 changes: 1 addition & 4 deletions third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -234,10 +234,7 @@ absl::StatusOr<CompileModuleResults> CompileModuleToLlvmIr(
RecordHloToLlvmDuration(end_usecs - start_usecs);
}

auto thunk_sequence = ir_emitter->ConsumeThunkSequence();
ForAllThunks([](Thunk* thunk) { thunk->ClearCompileTimeInfo(); },
thunk_sequence.get());
results.executable = std::move(thunk_sequence);
results.executable = ir_emitter->ConsumeThunkSequence();

return results;
}
Expand Down
2 changes: 0 additions & 2 deletions third_party/xla/xla/service/gpu/gpu_compiler.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -404,8 +404,6 @@ GpuThunkAotCompilationResult::LoadExecutable(
ir_emitter->EmitHloComputation(hlo_module->entry_computation()));
std::unique_ptr<ThunkSequence> thunk_sequence =
ir_emitter->ConsumeThunkSequence();
ForAllThunks([](Thunk* thunk) { thunk->ClearCompileTimeInfo(); },
thunk_sequence.get());

// Get all other fields required by GpuExecutable.
std::vector<GpuExecutable::ConstantInfo> constants =
Expand Down
2 changes: 1 addition & 1 deletion third_party/xla/xla/service/gpu/runtime/address_computation_thunk.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ AddressComputationThunk::AddressComputationThunk(
std::vector<std::optional<uint64_t>> offset_byte_sizes)
: Thunk(Kind::kAddressComputation, thunk_info),
embedded_thunk_(std::make_unique<SequentialThunk>(
ThunkInfo(thunk_info.op), std::move(*embedded_thunk))),
ThunkInfo(), std::move(*embedded_thunk))),
embedded_thunk_arguments_(std::move(arguments)),
fake_allocations_(std::move(fake_allocations)),
offset_buffer_indices_(std::move(offset_buffer_indices)),
Expand Down
63 changes: 28 additions & 35 deletions third_party/xla/xla/service/gpu/runtime/address_computation_thunk_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -130,15 +130,14 @@ TEST(AddressComputationThunkTest, SlicedGemm) {
// Creating embedded GEMM thunk.
ThunkSequence seq;
seq.emplace_back(std::make_unique<GemmThunk>(
Thunk::ThunkInfo(nullptr), config.value(), slice_lhs_fake, slice_rhs,
slice_out, slice_workspace, /*deterministic=*/true));
Thunk::ThunkInfo(), config.value(), slice_lhs_fake, slice_rhs, slice_out,
slice_workspace, /*deterministic=*/true));

// Wrapping address computation thunk around the GEMM thunk.
std::vector<BufferAllocation::Slice> lhs_offsets{slice_lhs_offset_0,
slice_lhs_offset_1};
AddressComputationThunk thunk(
Thunk::ThunkInfo(nullptr),
std::make_unique<ThunkSequence>(std::move(seq)),
Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
{slice_lhs, slice_rhs, slice_out, slice_workspace},
std::move(fake_allocations),
{lhs_offsets, std::nullopt, std::nullopt, std::nullopt},
Expand Down Expand Up @@ -284,7 +283,7 @@ TEST(AddressComputationThunkTest, SlicedNonContiguousGemm) {
// Creating embedded GEMM thunk.
ThunkSequence seq;
seq.emplace_back(std::make_unique<GemmThunk>(
Thunk::ThunkInfo(nullptr), config.value(), slice_lhs_fake, slice_rhs_fake,
Thunk::ThunkInfo(), config.value(), slice_lhs_fake, slice_rhs_fake,
slice_out, slice_workspace, /*deterministic=*/true));

// Wrapping address computation thunk around the GEMM thunk.
Expand All @@ -293,8 +292,7 @@ TEST(AddressComputationThunkTest, SlicedNonContiguousGemm) {
std::vector<BufferAllocation::Slice> rhs_offsets{slice_rhs_offset_0,
slice_rhs_offset_1};
AddressComputationThunk thunk(
Thunk::ThunkInfo(nullptr),
std::make_unique<ThunkSequence>(std::move(seq)),
Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
{slice_lhs, slice_rhs, slice_out, slice_workspace},
std::move(fake_allocations),
{lhs_offsets, rhs_offsets, std::nullopt, std::nullopt},
Expand Down Expand Up @@ -449,7 +447,7 @@ TEST(AddressComputationThunkTest, MulipleSlicedOperandsGemm) {
// Creating embedded GEMM thunk.
ThunkSequence seq;
seq.emplace_back(std::make_unique<GemmThunk>(
Thunk::ThunkInfo(nullptr), config.value(), slice_lhs_fake, slice_rhs_fake,
Thunk::ThunkInfo(), config.value(), slice_lhs_fake, slice_rhs_fake,
slice_out, slice_workspace, /*deterministic=*/true));

// Wrapping address computation thunk around the GEMM thunk.
Expand All @@ -458,8 +456,7 @@ TEST(AddressComputationThunkTest, MulipleSlicedOperandsGemm) {
std::vector<BufferAllocation::Slice> rhs_offsets{slice_rhs_offset_0,
slice_rhs_offset_1};
AddressComputationThunk thunk(
Thunk::ThunkInfo(nullptr),
std::make_unique<ThunkSequence>(std::move(seq)),
Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
{slice_lhs, slice_rhs, slice_out, slice_workspace},
std::move(fake_allocations),
{lhs_offsets, rhs_offsets, std::nullopt, std::nullopt},
Expand Down Expand Up @@ -627,17 +624,17 @@ TEST(AddressComputationThunkTest, SlicedMemcpy) {
// Creating embedded custom call thunk.
ThunkSequence seq;
seq.emplace_back(std::make_unique<CustomCallThunk>(
Thunk::ThunkInfo(nullptr), registration->handler, operands, results,
Thunk::ThunkInfo(), registration->handler, operands, results,
/*attributes=*/CustomCallThunk::AttributesMap(),
/*called_computation=*/nullptr));

// Wrapping address computation thunk around the custom call thunk.
std::vector<BufferAllocation::Slice> slice_offsets{
slice_offset_0, slice_offset_1, slice_offset_2, slice_offset_3};
AddressComputationThunk thunk(
Thunk::ThunkInfo(nullptr),
std::make_unique<ThunkSequence>(std::move(seq)), {slice_src, slice_dst},
std::move(fake_allocations), {slice_offsets, std::nullopt},
Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
{slice_src, slice_dst}, std::move(fake_allocations),
{slice_offsets, std::nullopt},
{ShapeUtil::MakeShape(PrimitiveType::S32, {8, 8, 10, 8}), std::nullopt},
// Make sure to pass a dst shape with the same rank as src shape (i.e.
// original slice result and not bitcasted one)
Expand Down Expand Up @@ -785,7 +782,7 @@ TEST(AddressComputationThunkTest, SlicedOutputMemcpy) {
// Creating embedded custom call thunk.
ThunkSequence seq;
seq.emplace_back(std::make_unique<CustomCallThunk>(
Thunk::ThunkInfo(nullptr), registration->handler, operands, results,
Thunk::ThunkInfo(), registration->handler, operands, results,
/*attributes=*/CustomCallThunk::AttributesMap(),
/*called_computation=*/nullptr));

Expand All @@ -797,9 +794,9 @@ TEST(AddressComputationThunkTest, SlicedOutputMemcpy) {
slice_dst_offset_0, slice_dst_offset_1, slice_dst_offset_2,
slice_dst_offset_3};
AddressComputationThunk thunk(
Thunk::ThunkInfo(nullptr),
std::make_unique<ThunkSequence>(std::move(seq)), {slice_src, slice_dst},
std::move(fake_allocations), {slice_src_offsets, slice_dst_offsets},
Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
{slice_src, slice_dst}, std::move(fake_allocations),
{slice_src_offsets, slice_dst_offsets},
{ShapeUtil::MakeShape(PrimitiveType::S32, {8, 8, 10, 2}),
ShapeUtil::MakeShape(PrimitiveType::S32, {2, 2, 2, 2})},
// Make sure to pass a dst shape with the same rank as src shape (i.e.
Expand Down Expand Up @@ -966,15 +963,14 @@ TEST(AddressComputationThunkTest, SlicedGemmArbitraryArgumentOrder) {
// Creating embedded GEMM thunk.
ThunkSequence seq;
seq.emplace_back(std::make_unique<GemmThunk>(
Thunk::ThunkInfo(nullptr), config.value(), slice_lhs_fake, slice_rhs_fake,
Thunk::ThunkInfo(), config.value(), slice_lhs_fake, slice_rhs_fake,
slice_out_fake, slice_workspace_fake, /*deterministic=*/true));

// Wrapping address computation thunk around the GEMM thunk.
std::vector<BufferAllocation::Slice> lhs_offsets{slice_lhs_offset_0,
slice_lhs_offset_1};
AddressComputationThunk thunk(
Thunk::ThunkInfo(nullptr),
std::make_unique<ThunkSequence>(std::move(seq)),
Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
{slice_lhs, slice_rhs, slice_out, slice_workspace},
std::move(fake_allocations),
{lhs_offsets, std::nullopt, std::nullopt, std::nullopt},
Expand Down Expand Up @@ -1115,15 +1111,14 @@ TEST(AddressComputationThunkTest, SlicedGemmArbitraryNumberOfArguments) {
// Creating embedded GEMM thunk.
ThunkSequence seq;
seq.emplace_back(std::make_unique<GemmThunk>(
Thunk::ThunkInfo(nullptr), config.value(), slice_lhs_fake, slice_rhs_fake,
Thunk::ThunkInfo(), config.value(), slice_lhs_fake, slice_rhs_fake,
slice_out_fake, slice_workspace_fake, /*deterministic=*/true));

// Wrapping address computation thunk around the GEMM thunk.
std::vector<BufferAllocation::Slice> lhs_offsets{slice_lhs_offset_0,
slice_lhs_offset_1};
AddressComputationThunk thunk(
Thunk::ThunkInfo(nullptr),
std::make_unique<ThunkSequence>(std::move(seq)),
Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
{slice_lhs, slice_rhs, slice_out, slice_workspace},
std::move(fake_allocations),
{lhs_offsets, std::nullopt, std::nullopt, std::nullopt},
Expand Down Expand Up @@ -1257,15 +1252,14 @@ TEST(AddressComputationThunkTest, SlicedTupledOperandGemm) {
// Creating embedded GEMM thunk.
ThunkSequence seq;
seq.emplace_back(std::make_unique<GemmThunk>(
Thunk::ThunkInfo(nullptr), config.value(), slice_lhs_fake, slice_rhs,
slice_out, slice_workspace, /*deterministic=*/true));
Thunk::ThunkInfo(), config.value(), slice_lhs_fake, slice_rhs, slice_out,
slice_workspace, /*deterministic=*/true));

// Wrapping address computation thunk around the GEMM thunk.
std::vector<BufferAllocation::Slice> lhs_offsets{slice_lhs_offset_0,
slice_lhs_offset_1};
AddressComputationThunk thunk(
Thunk::ThunkInfo(nullptr),
std::make_unique<ThunkSequence>(std::move(seq)),
Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
{slice_lhs, slice_rhs, slice_out, slice_workspace},
std::move(fake_allocations),
{lhs_offsets, std::nullopt, std::nullopt, std::nullopt},
Expand Down Expand Up @@ -1427,7 +1421,7 @@ TEST(AddressComputationThunkTest, SlicedMemcpyOOB) {
// Creating embedded custom call thunk.
ThunkSequence seq;
seq.emplace_back(std::make_unique<CustomCallThunk>(
Thunk::ThunkInfo(nullptr), registration->handler, operands, results,
Thunk::ThunkInfo(), registration->handler, operands, results,
/*attributes=*/CustomCallThunk::AttributesMap(),
/*called_computation=*/nullptr));

Expand All @@ -1439,9 +1433,9 @@ TEST(AddressComputationThunkTest, SlicedMemcpyOOB) {
slice_dst_offset_0, slice_dst_offset_1, slice_dst_offset_2,
slice_dst_offset_3};
AddressComputationThunk thunk(
Thunk::ThunkInfo(nullptr),
std::make_unique<ThunkSequence>(std::move(seq)), {slice_src, slice_dst},
std::move(fake_allocations), {slice_src_offsets, slice_dst_offsets},
Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
{slice_src, slice_dst}, std::move(fake_allocations),
{slice_src_offsets, slice_dst_offsets},
{ShapeUtil::MakeShape(PrimitiveType::S32, {8, 8, 10, 2}),
ShapeUtil::MakeShape(PrimitiveType::S32, {2, 2, 2, 2})},
// Make sure to pass a dst shape with the same rank as src shape (i.e.
Expand Down Expand Up @@ -1609,15 +1603,14 @@ TEST(AddressComputationThunkTest, SlicedOperandsSameBufferGemm) {
// Creating embedded GEMM thunk.
ThunkSequence seq;
seq.emplace_back(std::make_unique<GemmThunk>(
Thunk::ThunkInfo(nullptr), config.value(), slice_lhs_fake, slice_rhs_fake,
Thunk::ThunkInfo(), config.value(), slice_lhs_fake, slice_rhs_fake,
slice_out_fake, slice_workspace_fake, /*deterministic=*/true));

// Wrapping address computation thunk around the GEMM thunk.
std::vector<BufferAllocation::Slice> lhs_offsets{slice_lhs_offset_0,
slice_lhs_offset_1};
AddressComputationThunk thunk(
Thunk::ThunkInfo(nullptr),
std::make_unique<ThunkSequence>(std::move(seq)),
Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
{slice_lhs, slice_rhs, slice_out, slice_workspace},
std::move(fake_allocations),
{lhs_offsets, std::nullopt, std::nullopt, std::nullopt},
Expand Down
24 changes: 12 additions & 12 deletions third_party/xla/xla/service/gpu/runtime/command_buffer_thunk_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,7 @@ TEST(CommandBufferThunkTest, MemcpyCmd) {
commands.Emplace<MemcpyDeviceToDeviceCmd>(s0, slice_b, slice_a, byte_length);

// Construct a thunk with command sequence.
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());

se::StreamExecutorMemoryAllocator allocator(executor);
ServiceExecutableRunOptions run_options;
Expand Down Expand Up @@ -186,7 +186,7 @@ TEST(CommandBufferThunkTest, MemzeroCmd) {
commands.Emplace<MemzeroCmd>(s0, slice_a);

// Construct a thunk with command sequence.
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());

ServiceExecutableRunOptions run_options;
se::StreamExecutorMemoryAllocator allocator(executor);
Expand Down Expand Up @@ -228,7 +228,7 @@ TEST(CommandBufferThunkTest, Memset32Cmd) {
commands.Emplace<Memset32Cmd>(s0, slice_a, int32_t{84});

// Construct a thunk with command sequence.
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());

ServiceExecutableRunOptions run_options;
se::StreamExecutorMemoryAllocator allocator(executor);
Expand Down Expand Up @@ -267,7 +267,7 @@ TEST(CommandBufferThunkTest, Memset32CmdOnDifferentStreams) {
commands.Emplace<Memset32Cmd>(s1, slice1, int32_t{34});

// Construct a thunk with command sequence.
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());

ServiceExecutableRunOptions run_options;
se::StreamExecutorMemoryAllocator allocator(executor);
Expand Down Expand Up @@ -320,7 +320,7 @@ TEST(CommandBufferThunkTest, LaunchCmd) {
/*shmem_bytes=*/0);

// Construct a thunk with command sequence.
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());

ServiceExecutableRunOptions run_options;
se::StreamExecutorMemoryAllocator allocator(executor);
Expand Down Expand Up @@ -416,7 +416,7 @@ TEST(CommandBufferThunkTest, CustomAddKernelLaunchCmd) {
/*shmem_bytes=*/0);

// Construct a thunk with command sequence.
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());

ServiceExecutableRunOptions run_options;
se::StreamExecutorMemoryAllocator allocator(executor);
Expand Down Expand Up @@ -531,7 +531,7 @@ TEST(CommandBufferThunkTest, GemmCmd) {
/*deterministic=*/true);

// Construct a thunk with command sequence.
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());

ServiceExecutableRunOptions run_options;
se::StreamExecutorMemoryAllocator allocator(executor);
Expand Down Expand Up @@ -631,7 +631,7 @@ TEST(CommandBufferThunkTest, MultipleLaunchCmd) {
/*shmem_bytes=*/0);

// Construct a thunk with command sequence.
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());

ServiceExecutableRunOptions run_options;
se::StreamExecutorMemoryAllocator allocator(executor);
Expand Down Expand Up @@ -746,7 +746,7 @@ TEST(CommandBufferThunkTest, IfCmd) {
commands.Emplace<IfCmd>(s0, slice_p, std::move(then_commands));

// Construct a thunk with command sequence.
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());

ServiceExecutableRunOptions run_options;
se::StreamExecutorMemoryAllocator allocator(executor);
Expand Down Expand Up @@ -845,7 +845,7 @@ TEST(CommandBufferThunkTest, IfElseCmd) {
std::move(else_commands));

// Construct a thunk with command sequence.
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());

ServiceExecutableRunOptions run_options;
se::StreamExecutorMemoryAllocator allocator(executor);
Expand Down Expand Up @@ -934,7 +934,7 @@ TEST(CommandBufferThunkTest, CaseCmd) {
commands.Emplace<CaseCmd>(s0, slice_i, std::move(branches));

// Construct a thunk with command sequence.
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());

ServiceExecutableRunOptions run_options;
se::StreamExecutorMemoryAllocator allocator(executor);
Expand Down Expand Up @@ -1013,7 +1013,7 @@ TEST(CommandBufferThunkTest, ForCmd) {
std::move(body_commands));

// Construct a thunk with command sequence.
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());

ServiceExecutableRunOptions run_options;
se::StreamExecutorMemoryAllocator allocator(executor);
Expand Down
2 changes: 0 additions & 2 deletions third_party/xla/xla/service/gpu/runtime/copy_thunk.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,8 +45,6 @@ class DeviceToDeviceCopyThunk : public Thunk {

absl::Status ExecuteOnStream(const ExecuteParams& params) override;

void ClearCompileTimeInfo() override { Thunk::ClearCompileTimeInfo(); }

const BufferAllocation::Slice& source() const { return source_buffer_; }
const BufferAllocation::Slice& destination() const {
return destination_buffer_;
Expand Down
2 changes: 0 additions & 2 deletions third_party/xla/xla/service/gpu/runtime/kernel_thunk.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,8 +82,6 @@ class KernelThunk : public Thunk {
absl::Status Initialize(const InitializeParams& params) override;
absl::Status ExecuteOnStream(const ExecuteParams& params) override;

void ClearCompileTimeInfo() override { Thunk::ClearCompileTimeInfo(); }

const std::vector<BufferAllocation::Slice>& arguments() const {
return args_;
}
Expand Down
4 changes: 0 additions & 4 deletions third_party/xla/xla/service/gpu/runtime/memset_thunk.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,8 +37,6 @@ class MemzeroThunk : public Thunk {

absl::Status ExecuteOnStream(const ExecuteParams& params) override;

void ClearCompileTimeInfo() override { Thunk::ClearCompileTimeInfo(); }

const BufferAllocation::Slice& destination() const { return dest_; }

private:
Expand All @@ -57,8 +55,6 @@ class Memset32BitValueThunk : public Thunk {

absl::Status ExecuteOnStream(const ExecuteParams& params) override;

void ClearCompileTimeInfo() override { Thunk::ClearCompileTimeInfo(); }

const BufferAllocation::Slice& destination() const { return dest_; }
uint32_t value() const { return value_; }

Expand Down
10 changes: 1 addition & 9 deletions third_party/xla/xla/service/gpu/runtime/thunk.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@ limitations under the License.
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "mlir/IR/Operation.h" // from @llvm-project
#include "xla/executable_run_options.h"
#include "xla/ffi/execution_context.h"
#include "xla/hlo/ir/hlo_instruction.h"
Expand Down Expand Up @@ -340,16 +339,9 @@ bool IsReductionCollective(Thunk::Kind kind) {
kind == Thunk::kNcclReduceScatterStart;
}

Thunk::ThunkInfo Thunk::ThunkInfo::WithProfileAnnotation(mlir::Operation* op) {
ThunkInfo thunk_info(op);
thunk_info.profile_annotation =
mlir::mhlo::GetDebugNameFromLocation(op->getLoc());
return thunk_info;
}

Thunk::ThunkInfo Thunk::ThunkInfo::WithProfileAnnotation(
const HloInstruction* instr) {
ThunkInfo thunk_info(nullptr);
ThunkInfo thunk_info;
thunk_info.profile_annotation = instr->name();
auto gpu_backend_config = instr->backend_config<GpuBackendConfig>();
if (gpu_backend_config.ok()) {
Expand Down
Loading

0 comments on commit d928140

Please sign in to comment.