tensorflow/compiler/mlir/tensorflow/transforms/tpu_space_to_depth_pass.cc
// Applies space to depth transform for the first convolution on TPU device.
Applies space to depth transform for the first convolution on TPU device.
Applies space to depth transform for the first convolution on TPU device .
void HandleFirstConvolution(TF::Conv2DOp conv2d, int64_t block_size) {
  auto input_tensor_type =
      conv2d.input().getType().dyn_cast<RankedTensorType>();
  auto filter_tensor_type =
      conv2d.filter().getType().dyn_cast<RankedTensorType>();
  if (!input_tensor_type || !filter_tensor_type) return;
  auto filter_shape = filter_tensor_type.getShape();
  SmallVector<int32_t, 4> old_filter_shape(filter_shape.begin(),
                                           filter_shape.end());
  auto conv2d_input = conv2d.input();
  if (auto block_arg = conv2d_input.dyn_cast<mlir::BlockArgument>()) {
    HandleFuncOp(block_arg.getOwner()->getParentOp());
  }
  if (auto pad_op = dyn_cast_or_null<TF::PadOp>(conv2d_input.getDefiningOp())) {
    if (failed(HandlePad(pad_op, filter_shape[0], block_size))) return;
    auto pad_input = pad_op.input();
    if (auto block_arg = pad_input.dyn_cast<mlir::BlockArgument>()) {
      HandleFuncOp(block_arg.getOwner()->getParentOp());
    }
  }
  HandleConv2DInput(conv2d, block_size);
  HandleConv2DStride(conv2d);
  HandleConv2DFilter(conv2d, block_size);
  filter_shape = conv2d.filter().getType().cast<RankedTensorType>().getShape();
  SmallVector<int32_t, 4> new_filter_shape(filter_shape.begin(),
                                           filter_shape.end());
  if (!conv2d_input.getDefiningOp()) return;
  for (Operation* user : conv2d_input.getDefiningOp()->getUsers()) {
    if (auto backprop = dyn_cast<TF::Conv2DBackpropFilterOp>(user)) {
      HandleConv2DBackPropFilter(backprop, old_filter_shape, new_filter_shape,
                                 block_size);
    }
  }
}
void HandleFirstConvolution ( TF :: Conv2DOp conv2d , int64_t block_size ) { auto input_tensor_type = conv2d . input ( ) . getType ( ) . dyn_cast < RankedTensorType > ( ) ; auto filter_tensor_type = conv2d . filter ( ) . getType ( ) . dyn_cast < RankedTensorType > ( ) ; if ( ! input_tensor_type || ! filter_tensor_type ) return ; auto filter_shape = filter_tensor_type . getShape ( ) ; SmallVector < int32_t , 4 > old_filter_shape ( filter_shape . begin ( ) , filter_shape . end ( ) ) ; auto conv2d_input = conv2d . input ( ) ; if ( auto block_arg = conv2d_input . dyn_cast < mlir :: BlockArgument > ( ) ) { HandleFuncOp ( block_arg . getOwner ( ) -> getParentOp ( ) ) ; } if ( auto pad_op = dyn_cast_or_null < TF :: PadOp > ( conv2d_input . getDefiningOp ( ) ) ) { if ( failed ( HandlePad ( pad_op , filter_shape [ 0 ] , block_size ) ) ) return ; auto pad_input = pad_op . input ( ) ; if ( auto block_arg = pad_input . dyn_cast < mlir :: BlockArgument > ( ) ) { HandleFuncOp ( block_arg . getOwner ( ) -> getParentOp ( ) ) ; } } HandleConv2DInput ( conv2d , block_size ) ; HandleConv2DStride ( conv2d ) ; HandleConv2DFilter ( conv2d , block_size ) ; filter_shape = conv2d . filter ( ) . getType ( ) . cast < RankedTensorType > ( ) . getShape ( ) ; SmallVector < int32_t , 4 > new_filter_shape ( filter_shape . begin ( ) , filter_shape . end ( ) ) ; if ( ! conv2d_input . getDefiningOp ( ) ) return ; for ( Operation * user : conv2d_input . getDefiningOp ( ) -> getUsers ( ) ) { if ( auto backprop = dyn_cast < TF :: Conv2DBackpropFilterOp > ( user ) ) { HandleConv2DBackPropFilter ( backprop , old_filter_shape , new_filter_shape , block_size ) ; } } }
Applies space to depth transform for the first convolution on TPU device.
Applies space to depth transform for the first convolution on TPU device .
// Check if input and filter type are RankedTensorType. // Book keeping filter shape for padding and backprop filter rewrite. // Handles input. // Change on device function type/shape. // Rewrite pad_op before Convolutioin. // Change on device function type/shape. // Handle Conv2D input, stride and filter. // Book keeping new filter shape for backprop filter rewrite. // Filter shape is defined in HandleConv2DFilter, thus it is RankedTensorType. // Rewrite Conv2DBackPropFilter that is the user of first convolution's input.