Generates a mask representing the effective area of data and padded area of data using iota and dynamic_size. For example, given a dimension of 7 elements and 5 effective elements: iota = [0, 1, 2, 3, 4, 5, 6] broadcast_dynamic_size = [5, 5, 5, 5, 5, 5, 5] mask = lt(iota, broadcast_dynamic_size) = [t, t, t, t, t, f, f] Once the mask is generated, the input data is then padded using the mask and
| 151 | // mask and pad value. |
| 152 | // |
| 153 | HloInstruction* PadWithScalar(HloInstruction* inst, int64 dim, |
| 154 | HloInstruction* dynamic_size, |
| 155 | HloInstruction* padding_scalar) { |
| 156 | const Shape mask_shape = |
| 157 | ShapeUtil::ChangeElementType(inst->shape(), xla::S32); |
| 158 | const Shape pred_shape = |
| 159 | ShapeUtil::ChangeElementType(inst->shape(), xla::PRED); |
| 160 | HloComputation* computation = inst->parent(); |
| 161 | HloInstruction* iota = |
| 162 | computation->AddInstruction(HloInstruction::CreateIota(mask_shape, dim)); |
| 163 | |
| 164 | HloInstruction* broadcasted_effective_size = computation->AddInstruction( |
| 165 | HloInstruction::CreateBroadcast(mask_shape, dynamic_size, {})); |
| 166 | HloInstruction* pred = |
| 167 | computation->AddInstruction(HloInstruction::CreateCompare( |
| 168 | pred_shape, iota, broadcasted_effective_size, |
| 169 | ComparisonDirection::kLt)); |
| 170 | |
| 171 | HloInstruction* broadcasted_identity_value = computation->AddInstruction( |
| 172 | HloInstruction::CreateBroadcast(inst->shape(), padding_scalar, {})); |
| 173 | HloInstruction* padded = computation->AddInstruction( |
| 174 | HloInstruction::CreateTernary(inst->shape(), HloOpcode::kSelect, pred, |
| 175 | inst, broadcasted_identity_value)); |
| 176 | return padded; |
| 177 | } |
| 178 | |
| 179 | // In a reshape if a dynamic dimension is splitted into multiple output |
| 180 | // dimensions, we need to rewrite the input of the reshape. |
no test coverage detected