| 943 | } |
| 944 | |
| 945 | bool FindFusedBatchNormEx(const RemapperContext& ctx, int node_index, |
| 946 | FusedBatchNormEx* matched) { |
| 947 | // Root of the pattern must be a Relu. |
| 948 | // TODO(ezhulenev): Forward control dependencies. |
| 949 | const auto* node_view = ctx.graph_view.GetNode(node_index); |
| 950 | const auto* node_def = node_view->node(); |
| 951 | // TODO(lyandy): Forward controls for patterns with control dependencies. |
| 952 | if (!IsRelu(*node_def) || HasControlFaninOrFanout(*node_view)) return false; |
| 953 | |
| 954 | // Returns true iff the node is a compatible FusedBatchNorm node. |
| 955 | const auto valid_batch_norm = |
| 956 | [&](const utils::MutableNodeView& fused_batch_norm) -> bool { |
| 957 | const auto* fused_batch_norm_node_def = fused_batch_norm.node(); |
| 958 | if (!IsFusedBatchNorm(*fused_batch_norm_node_def)) return false; |
| 959 | |
| 960 | if (DisableMKL()) return false; |
| 961 | |
| 962 | DataType t_dtype = GetDataTypeFromAttr(*fused_batch_norm_node_def, "T"); |
| 963 | if (t_dtype != DT_FLOAT && t_dtype != DT_BFLOAT16) return false; |
| 964 | |
| 965 | // Get the FusedBatchNorm training mode. |
| 966 | bool is_training; |
| 967 | if (!GetNodeAttr(*fused_batch_norm_node_def, kIsTraining, &is_training) |
| 968 | .ok()) |
| 969 | return false; |
| 970 | // In training mode we rely on cuDNN for computing FusedBatchNorm with side |
| 971 | // inputs and activation, and it has its own limitations. In inference mode |
| 972 | // we have a custom CUDA kernel that doesn't not have these constraints. |
| 973 | if (is_training && NodeIsOnGpu(fused_batch_norm_node_def)) { |
| 974 | // cuDNN only supports NHWC data layout. |
| 975 | string data_format; |
| 976 | if (!GetNodeAttr(*fused_batch_norm_node_def, kDataFormat, &data_format) |
| 977 | .ok()) |
| 978 | return false; |
| 979 | if (data_format != "NHWC") return false; |
| 980 | |
| 981 | // Data type must be DT_HALF. |
| 982 | if (t_dtype != DT_HALF) return false; |
| 983 | |
| 984 | // Channel dimension must be a multiple of 4. |
| 985 | const auto& props = ctx.graph_properties.GetInputProperties( |
| 986 | fused_batch_norm_node_def->name()); |
| 987 | |
| 988 | const bool valid_channel_dim = !props.empty() && |
| 989 | props[0].shape().dim_size() == 4 && |
| 990 | props[0].shape().dim(3).size() % 4 == 0; |
| 991 | if (!valid_channel_dim) return false; |
| 992 | |
| 993 | // cuDNN must support CUDNN_BATCHNORM_SPATIAL_PERSISTENT mode. |
| 994 | if (!BatchnormSpatialPersistentEnabled()) return false; |
| 995 | } |
| 996 | |
| 997 | // FusedBatchNormV2 and V3 have an extra type parameter. |
| 998 | if ((fused_batch_norm_node_def->op() != "FusedBatchNorm") && |
| 999 | !HasDataType(fused_batch_norm_node_def, DT_FLOAT, "U")) |
| 1000 | return false; |
| 1001 | |
| 1002 | // Check that only one node consumes the 0-th output of a FusedBatchNorm. |
no test coverage detected