| 900 | } |
| 901 | |
| 902 | ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { |
| 903 | if (debug_options_.xla_hlo_graph_sharding_color()) { |
| 904 | if (!instr->has_sharding()) { |
| 905 | return kDashedBorder; |
| 906 | } |
| 907 | auto it = sharding_colors_.find(instr->sharding()); |
| 908 | if (it != sharding_colors_.end()) { |
| 909 | return it->second; |
| 910 | } |
| 911 | ColorScheme color = static_cast<ColorScheme>( |
| 912 | kBlue + (next_shard_color_++ % (kDashedBorder - kBlue))); |
| 913 | sharding_colors_.emplace(instr->sharding(), color); |
| 914 | return color; |
| 915 | } |
| 916 | |
| 917 | // Choose different weights of orange for small vs large parameters. This |
| 918 | // distinction is often important, especially in fusion nodes. |
| 919 | auto parameter_color = IsSmall(instr) ? kOrange : kDarkOrange; |
| 920 | |
| 921 | // Special case: If this instruction has a parameter merged into it, paint it |
| 922 | // the same color as a parameter. Unless the merged-in parameter is a |
| 923 | // parameter to a fusion node that is bound to a constant -- these aren't |
| 924 | // "real" parameters from the user's perspective. |
| 925 | if (absl::c_any_of(instr->operands(), [&](const HloInstruction* operand) { |
| 926 | return operand->opcode() == HloOpcode::kParameter && |
| 927 | ShouldMergeIntoUsers(operand) && |
| 928 | TryGetFusionParameterConstant(operand) == nullptr; |
| 929 | })) { |
| 930 | return parameter_color; |
| 931 | } |
| 932 | |
| 933 | // Pick different colors or shapes for instructions which are particularly |
| 934 | // expensive (eg, dot) and those which are unusual in some way or unique |
| 935 | // (eg, parameter). |
| 936 | switch (instr->opcode()) { |
| 937 | case HloOpcode::kAbs: |
| 938 | case HloOpcode::kAdd: |
| 939 | case HloOpcode::kAnd: |
| 940 | case HloOpcode::kAtan2: |
| 941 | case HloOpcode::kBitcastConvert: |
| 942 | case HloOpcode::kCeil: |
| 943 | case HloOpcode::kClamp: |
| 944 | case HloOpcode::kClz: |
| 945 | case HloOpcode::kCompare: |
| 946 | case HloOpcode::kComplex: |
| 947 | case HloOpcode::kConvert: |
| 948 | case HloOpcode::kCos: |
| 949 | case HloOpcode::kDivide: |
| 950 | case HloOpcode::kExp: |
| 951 | case HloOpcode::kExpm1: |
| 952 | case HloOpcode::kFloor: |
| 953 | case HloOpcode::kImag: |
| 954 | case HloOpcode::kIota: |
| 955 | case HloOpcode::kIsFinite: |
| 956 | case HloOpcode::kLog: |
| 957 | case HloOpcode::kLog1p: |
| 958 | case HloOpcode::kMaximum: |
| 959 | case HloOpcode::kMinimum: |
nothing calls this directly
no test coverage detected