| 904 | } |
| 905 | |
| 906 | Status IrEmitterUnnested::EmitScatter( |
| 907 | Thunk* thunk, HloInstruction* scatter, |
| 908 | const llvm_ir::ElementGenerator& scatter_indices_gen, |
| 909 | const llvm_ir::ElementGenerator& updates_gen) { |
| 910 | const HloInstruction* operand = scatter->operand(0); |
| 911 | const HloInstruction* scatter_indices = scatter->operand(1); |
| 912 | const HloInstruction* updates = scatter->operand(2); |
| 913 | const ScatterDimensionNumbers& dim_numbers = |
| 914 | scatter->scatter_dimension_numbers(); |
| 915 | CHECK(ShapeUtil::Equal(scatter->shape(), operand->shape())); |
| 916 | |
| 917 | auto loop_body_emitter = [&](const IrArray::Index& index) -> Status { |
| 918 | std::vector<llvm::Value*> raw_window_multidim; |
| 919 | std::vector<llvm::Value*> input_scatter_multidim; |
| 920 | std::vector<int64> raw_window_bounds; |
| 921 | |
| 922 | // Partition the index into window indices and scatter indices. |
| 923 | for (int64 i = 0, e = index.size(); i != e; ++i) { |
| 924 | // For window indices also remember the window size, this comes in handy |
| 925 | // later. |
| 926 | if (absl::c_binary_search(dim_numbers.update_window_dims(), i)) { |
| 927 | raw_window_multidim.push_back(index[i]); |
| 928 | raw_window_bounds.push_back(updates->shape().dimensions(i)); |
| 929 | } else { |
| 930 | input_scatter_multidim.push_back(index[i]); |
| 931 | } |
| 932 | } |
| 933 | DCHECK_EQ(raw_window_multidim.size(), |
| 934 | dim_numbers.update_window_dims_size()); |
| 935 | |
| 936 | // Apply inserted_window_dims to the window dimensions. |
| 937 | int64 raw_window_multidim_idx = 0; |
| 938 | std::vector<llvm::Value*> input_window_multidim; |
| 939 | std::vector<int64> input_window_bounds; |
| 940 | for (int64 i = 0, e = operand->shape().rank(); i != e; ++i) { |
| 941 | if (absl::c_binary_search(dim_numbers.inserted_window_dims(), i)) { |
| 942 | input_window_bounds.push_back(1); // Trivial dimension. |
| 943 | input_window_multidim.push_back(index.GetConstantWithIndexType(0)); |
| 944 | } else { |
| 945 | input_window_bounds.push_back( |
| 946 | raw_window_bounds[raw_window_multidim_idx]); |
| 947 | input_window_multidim.push_back( |
| 948 | raw_window_multidim[raw_window_multidim_idx]); |
| 949 | ++raw_window_multidim_idx; |
| 950 | } |
| 951 | } |
| 952 | DCHECK_EQ(input_window_multidim.size(), operand->shape().rank()); |
| 953 | |
| 954 | // Insert a 1 dimension at the end if index_vector_dim requests one. |
| 955 | Shape scatter_indices_shape = scatter_indices->shape(); |
| 956 | if (dim_numbers.index_vector_dim() == scatter_indices_shape.rank()) { |
| 957 | scatter_indices_shape.add_dimensions(1); |
| 958 | scatter_indices_shape.mutable_layout()->add_minor_to_major( |
| 959 | dim_numbers.index_vector_dim()); |
| 960 | } |
| 961 | |
| 962 | // Now load the indices corresponding to the current window from |
| 963 | // scatter_indices. |
nothing calls this directly
no test coverage detected