Skip to content

Support For Roofline Analysis Of Pallas Kernels #1805

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
+488 −3
Open

Support For Roofline Analysis Of Pallas Kernels #1805

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions xprof/convert/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -884,6 +884,7 @@ cc_library(
":xplane_to_tf_functions",
":xprof_thread_pool_executor",
"@com_google_absl//absl/cleanup",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
Expand Down
110 changes: 110 additions & 0 deletions xprof/convert/xplane_to_op_stats.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ limitations under the License.
#include <ostream>
#include <string>
#include <vector>
#include <queue>

#include "absl/cleanup/cleanup.h"
#include "absl/container/flat_hash_set.h"
Expand Down Expand Up @@ -66,6 +67,7 @@ limitations under the License.
#include "xprof/utils/kernel_stats_utils.h"
#include "xprof/utils/op_utils.h"
#include "xprof/utils/xprof_gpu_cost_analysis_types.h"
#include "absl/container/flat_hash_map.h"

namespace tensorflow {
namespace profiler {
Expand All @@ -77,6 +79,7 @@ using ::tsl::profiler::kGpuPlanePrefix;
using ::tsl::profiler::kTpuPlanePrefix;
using tsl::profiler::Timespan;
using ::tsl::profiler::XPlaneBuilder;
using tsl::profiler::kXlaOpLineName;

std::string Hostname(const XSpace& space) {
if (space.hostnames().empty()) return "localhost";
Expand Down Expand Up @@ -380,6 +383,113 @@ OpStats ConvertXSpaceToOpStats(const XSpace& space,
.has_value()) {
op_metrics_db =
ConvertTpuDeviceTraceXPlaneToOpMetricsDb(*device_plane);
XPlaneVisitor visitorSecond =
tsl::profiler::CreateTfXPlaneVisitor(device_plane);
std::queue<XEventVisitor> custom_call_blocks;
visitorSecond.ForEachLine([&](const XLineVisitor& line) {
if (line.Name() == "XLA TraceMe") {
line.ForEachEvent([&](const XEventVisitor& event) {
if (absl::StartsWith(event.Name(), "__block_")){
custom_call_blocks.push(event);
}
});
}
});
absl::flat_hash_map<std::string,
absl::flat_hash_map<std::string, uint>>
custom_call_to_block_count;

XPlaneVisitor xlaEvents =
tsl::profiler::CreateTfXPlaneVisitor(device_plane);
xlaEvents.ForEachLine([&](const XLineVisitor& line) {
if (line.Name() == kXlaOpLineName) {
line.ForEachEvent([&](const XEventVisitor& event) {
tsl::profiler::Timespan custom_call_timespan =
GetDeviceEventTimespan(event);
bool custom_call = false;
event.Metadata().ForEachStat([&]
(const XStatVisitor& stat) {
if (stat.Type().has_value()) {
switch (static_cast<StatType>(*stat.Type())) {
case StatType::kHloCategory:
custom_call =
(stat.StrOrRefValue() == "custom-call");
break;
default:
break;
}
}
});
if (custom_call){
while (!custom_call_blocks.empty()){
tsl::profiler::Timespan ccall_blck_timespan =
GetDeviceEventTimespan(custom_call_blocks.front());
if ((custom_call_timespan.begin_ps() <=
ccall_blck_timespan.begin_ps()) &&
(ccall_blck_timespan.end_ps() <=
custom_call_timespan.end_ps())
){
custom_call_to_block_count[event.DisplayName()]
[std::string(custom_call_blocks.front().
Name())] += 1;
custom_call_blocks.pop();
}else{
break;
}
}
}
});
}
});
for (OpMetrics& op_metrics :
*op_metrics_db.mutable_metrics_db()) {
const HloInstructionWrapper* instr_wrapper =
GetHloInstruction(hlo_module_map,
op_metrics.hlo_module_id(), op_metrics.name());
if (instr_wrapper != nullptr) {
if (instr_wrapper->Category() == "custom-call"){
uint64 total_flops = 0;
uint64 total_bytes_accessed = 0;
bool has_block_costs =
custom_call_to_block_count.contains(op_metrics.name());
if (has_block_costs){
for (auto&[block_name, occurrence] :
custom_call_to_block_count[op_metrics.name()]){
auto block_cost_pair = instr_wrapper->
GetCustomCallBlockCosts(block_name);
if (block_cost_pair.has_value()){
OpMetrics* child_metric =
op_metrics.mutable_children()->add_metrics_db();
child_metric->set_name(block_name);
child_metric->set_occurrences(occurrence);
child_metric->set_flops(
block_cost_pair.value().first);
child_metric->set_model_flops(
block_cost_pair.value().first);
child_metric->set_bytes_accessed(
block_cost_pair.value().second);
total_flops +=
(occurrence*block_cost_pair.value().first);
total_bytes_accessed +=
(occurrence*block_cost_pair.value().second);
}else{
LOG(WARNING) << "No Costs Found for : " << block_name;
}
}
if (instr_wrapper->FusedChildren().empty()){
LOG(INFO) << "Custom - Call Name: "
<< op_metrics.name() << " Total Flops: "
<< total_flops
<< " Total Bytes Accessed: " <<
total_bytes_accessed;
op_metrics.set_flops(total_flops);
op_metrics.set_bytes_accessed(total_bytes_accessed);
op_metrics.set_model_flops(total_flops);
}
}
}
}
}
UpdateOpMetricsDbFromHloModuleMap(op_metrics_db, hlo_module_map);
}
}
Expand Down
24 changes: 21 additions & 3 deletions xprof/utils/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -398,23 +398,41 @@ cc_library(
srcs = ["hlo_module_map.cc"],
hdrs = ["hlo_module_map.h"],
deps = [
":backend_configs_proto_cc",
":hlo_cost_analysis_wrapper",
":hlo_module_utils",
":hlo_proto_map",
":hlo_proto_to_module",
":performance_info_wrapper",
"@com_google_absl//absl/base:no_destructor",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@tsl//tsl/platform:path",
"@tsl//tsl/profiler/lib:traceme_encode",
"@tsl//tsl/profiler/protobuf:xplane_proto_cc",
"//third_party/llvm/llvm-project/mlir:ArithDialect",
"//third_party/llvm/llvm-project/mlir:DataLayoutInterfaces",
"//third_party/llvm/llvm-project/mlir:FuncDialect",
"//third_party/llvm/llvm-project/mlir:IR",
"//third_party/llvm/llvm-project/mlir:LLVMDialect",
"//third_party/llvm/llvm-project/mlir:MathDialect",
"//third_party/llvm/llvm-project/mlir:MemRefDialect",
"//third_party/llvm/llvm-project/mlir:Pass",
"//third_party/llvm/llvm-project/mlir:SCFDialect",
"//third_party/llvm/llvm-project/mlir:Support",
"//third_party/llvm/llvm-project/mlir:VectorDialect",
"//third_party/protobuf/json",
"//third_party/protobuf/util:json_util",
# "//third_party/py/jax/jaxlib/mosaic:tpu_dialect",
"@xla//xla/hlo/ir:hlo",
"@xla//xla/mlir_hlo",
"@xla//xla/pjrt:mlir_to_hlo",
"@xla//xla/service:hlo_cost_analysis",
"@xla//xla/service:hlo_proto_cc",
"@xla//xla/tsl/profiler/convert:xla_op_utils",
"@tsl//tsl/platform:path",
"@tsl//tsl/profiler/lib:traceme_encode",
"@tsl//tsl/profiler/protobuf:xplane_proto_cc",
],
)

Expand Down
Loading