| 1158 | } |
| 1159 | |
| 1160 | Status IrEmitter::HandleConvolution(HloInstruction* convolution) { |
| 1161 | auto lhs = convolution->operand(0); |
| 1162 | auto rhs = convolution->operand(1); |
| 1163 | TF_RETURN_IF_ERROR(ElementTypesSameAndSupported( |
| 1164 | /*instruction=*/*convolution, /*operands=*/{lhs, rhs}, |
| 1165 | /*supported_types=*/{F16, F32, F64, C64, C128})); |
| 1166 | |
| 1167 | // TODO(tonywy): Add PotentiallyImplementedAsMKLConvolution to support |
| 1168 | // different data layouts. |
| 1169 | if (PotentiallyImplementedAsEigenConvolution(*convolution, |
| 1170 | target_machine_features_)) { |
| 1171 | const Shape& lhs_shape = lhs->shape(); |
| 1172 | const Shape& rhs_shape = rhs->shape(); |
| 1173 | const Shape& convolution_shape = convolution->shape(); |
| 1174 | // The input, kernel and output agree with respect to layout. |
| 1175 | if (LayoutUtil::IsMonotonicWithDim0Major(lhs_shape.layout()) && |
| 1176 | LayoutUtil::IsMonotonicWithDim0Major(rhs_shape.layout()) && |
| 1177 | LayoutUtil::IsMonotonicWithDim0Major(convolution_shape.layout())) { |
| 1178 | // We lower 1D convolutions into calls to the same Eigen function as 2D |
| 1179 | // convolutions, except that we pretend that the 1D convolution is really |
| 1180 | // a 2D convolution with the missing dimension set to 1. We also adjust |
| 1181 | // the padding, dilation parameters as needed. |
| 1182 | bool one_dim_convolution = lhs_shape.dimensions_size() == 3; |
| 1183 | llvm::Value* lhs_address = GetEmittedValueFor(lhs); |
| 1184 | llvm::Value* rhs_address = GetEmittedValueFor(rhs); |
| 1185 | TF_RETURN_IF_ERROR(EmitTargetAddressForOp(convolution)); |
| 1186 | |
| 1187 | const ConvolutionDimensionNumbers& dnums = |
| 1188 | convolution->convolution_dimension_numbers(); |
| 1189 | |
| 1190 | // Input tensor. |
| 1191 | const Shape& input_shape = convolution->operand(0)->shape(); |
| 1192 | int64 input_batch = input_shape.dimensions(dnums.input_batch_dimension()); |
| 1193 | int64 input_rows = |
| 1194 | input_shape.dimensions(dnums.input_spatial_dimensions(0)); |
| 1195 | int64 input_cols = |
| 1196 | one_dim_convolution |
| 1197 | ? 1 |
| 1198 | : input_shape.dimensions(dnums.input_spatial_dimensions(1)); |
| 1199 | int64 input_channels = |
| 1200 | input_shape.dimensions(dnums.input_feature_dimension()); |
| 1201 | |
| 1202 | // Kernel tensor. |
| 1203 | const Shape& kernel_shape = convolution->operand(1)->shape(); |
| 1204 | int64 kernel_rows = |
| 1205 | kernel_shape.dimensions(dnums.kernel_spatial_dimensions(0)); |
| 1206 | int64 kernel_cols = |
| 1207 | one_dim_convolution |
| 1208 | ? 1 |
| 1209 | : kernel_shape.dimensions(dnums.kernel_spatial_dimensions(1)); |
| 1210 | int64 kernel_channels = |
| 1211 | kernel_shape.dimensions(dnums.kernel_input_feature_dimension()); |
| 1212 | int64 kernel_filters = |
| 1213 | kernel_shape.dimensions(dnums.kernel_output_feature_dimension()); |
| 1214 | |
| 1215 | // Output tensor. |
| 1216 | const Shape& convolution_shape = convolution->shape(); |
| 1217 | int64 output_rows = |
nothing calls this directly
no test coverage detected