| 893 | } |
| 894 | |
| 895 | bool BufferAssigner::MaybeAssignBuffer(BufferAllocation* allocation, |
| 896 | const HloBuffer& hlo_buffer, |
| 897 | BufferAssignment* assignment) { |
| 898 | CHECK(!assignment->HasAllocation(hlo_buffer)) |
| 899 | << "buffer " << hlo_buffer << " already has an allocation assigned."; |
| 900 | |
| 901 | VLOG(4) << "Trying to assign " << hlo_buffer << " size " |
| 902 | << assignment->HloBufferSize(hlo_buffer) |
| 903 | << " to allocation: " << *allocation; |
| 904 | |
| 905 | if (hlo_buffer.color() != allocation->color()) { |
| 906 | VLOG(4) << "Can't assign: buffer has color " << hlo_buffer.color() |
| 907 | << " and allocation has color " << allocation->color() << "."; |
| 908 | return false; |
| 909 | } |
| 910 | |
| 911 | if (assignment->HloBufferSize(hlo_buffer) > allocation->size()) { |
| 912 | VLOG(4) << "Can't assign: buffer is larger than allocation (" |
| 913 | << assignment->HloBufferSize(hlo_buffer) << " > " |
| 914 | << allocation->size() << ")"; |
| 915 | return false; |
| 916 | } |
| 917 | |
| 918 | if (allocation->is_readonly()) { |
| 919 | VLOG(4) << "Can't assign: allocation is readonly"; |
| 920 | return false; |
| 921 | } |
| 922 | |
| 923 | if (!must_not_live_out_.empty()) { |
| 924 | if (allocation->maybe_live_out()) { |
| 925 | // If a buffer maybe live out, the allocation cannot contain any node from |
| 926 | // the "must_not_live_out_" set. |
| 927 | for (const HloValue* value : hlo_buffer.values()) { |
| 928 | if (must_not_live_out_.count(value->instruction()->opcode()) > 0) { |
| 929 | VLOG(4) << "Can't assign: " << value->instruction()->ToString() |
| 930 | << " cannot live out of the module"; |
| 931 | return false; |
| 932 | } |
| 933 | } |
| 934 | } |
| 935 | // The above check is not enough -- There could be the case where an |
| 936 | // allocation can be not live out and contains an instruction with opcode |
| 937 | // from the "must_not_live_out_" set, but assigning a live out buffer to |
| 938 | // that allocation makes the allocation live out and also contains |
| 939 | // instruction from the "must_not_live_out_" set. |
| 940 | if (assignment->alias_analysis().BufferLivesOut(hlo_buffer)) { |
| 941 | for (const auto& buffer_offset_size : allocation->assigned_buffers()) { |
| 942 | if (must_not_live_out_.count( |
| 943 | buffer_offset_size.first->instruction()->opcode()) > 0) { |
| 944 | VLOG(4) << "Can't assign: " << buffer_offset_size.first->instruction() |
| 945 | << " cannot live out of the module"; |
| 946 | return false; |
| 947 | } |
| 948 | } |
| 949 | } |
| 950 | } |
| 951 | |
| 952 | if (!allocation->is_reusable()) { |
nothing calls this directly
no test coverage detected