tensorflow/compiler/mlir/tensorflow/transforms/tpu_space_to_depth_pass.cc
// Handles padding before convolution for space to depth transform.
Handles padding before convolution for space to depth transform.
Handles padding before convolution for space to depth transform .
LogicalResult HandlePad(TF::PadOp op, int32_t kernel_size, int32_t block_size) {
  auto ranked_type = op.input().getType().dyn_cast<RankedTensorType>();
  if (!ranked_type) return failure();
  auto pad_input_shape = ranked_type.getShape();
  Location loc = op.getLoc();
  OpBuilder builder(op);
  builder.setInsertionPoint(op);
  auto padding_type = RankedTensorType::get({4, 2}, builder.getIntegerType(32));
  int32_t pad_total = kernel_size - 1;
  int32_t pad_beg = (pad_total / 2 + 1) / block_size;
  int32_t pad_end = (pad_total / 2) / block_size;
  SmallVector<int32_t, 8> values = {0,       0,       pad_beg, pad_end,
                                    pad_beg, pad_end, 0,       0};
  auto paddings = DenseIntElementsAttr::get(padding_type, values);
  op.setOperand(1, builder.create<TF::ConstOp>(loc, paddings));
  auto input = op.getOperand(0);
  SmallVector<int64_t, 4> transform_shape = {
      pad_input_shape[0], pad_input_shape[1] / block_size,
      pad_input_shape[2] / block_size,
      pad_input_shape[3] * block_size * block_size};
  if (auto cast_op = dyn_cast_or_null<TF::CastOp>(input.getDefiningOp()))
    if (failed(HandleCast(cast_op, transform_shape))) return failure();
  auto transform_result_type =
      RankedTensorType::get(transform_shape, getElementTypeOrSelf(input));
  input.setType(transform_result_type);
  op.setOperand(0, input);
  return success();
}
LogicalResult HandlePad ( TF :: PadOp op , int32_t kernel_size , int32_t block_size ) { auto ranked_type = op . input ( ) . getType ( ) . dyn_cast < RankedTensorType > ( ) ; if ( ! ranked_type ) return failure ( ) ; auto pad_input_shape = ranked_type . getShape ( ) ; Location loc = op . getLoc ( ) ; OpBuilder builder ( op ) ; builder . setInsertionPoint ( op ) ; auto padding_type = RankedTensorType :: get ( { 4 , 2 } , builder . getIntegerType ( 32 ) ) ; int32_t pad_total = kernel_size - 1 ; int32_t pad_beg = ( pad_total / 2 + 1 ) / block_size ; int32_t pad_end = ( pad_total / 2 ) / block_size ; SmallVector < int32_t , 8 > values = { 0 , 0 , pad_beg , pad_end , pad_beg , pad_end , 0 , 0 } ; auto paddings = DenseIntElementsAttr :: get ( padding_type , values ) ; op . setOperand ( 1 , builder . create < TF :: ConstOp > ( loc , paddings ) ) ; auto input = op . getOperand ( 0 ) ; SmallVector < int64_t , 4 > transform_shape = { pad_input_shape [ 0 ] , pad_input_shape [ 1 ] / block_size , pad_input_shape [ 2 ] / block_size , pad_input_shape [ 3 ] * block_size * block_size } ; if ( auto cast_op = dyn_cast_or_null < TF :: CastOp > ( input . getDefiningOp ( ) ) ) if ( failed ( HandleCast ( cast_op , transform_shape ) ) ) return failure ( ) ; auto transform_result_type = RankedTensorType :: get ( transform_shape , getElementTypeOrSelf ( input ) ) ; input . setType ( transform_result_type ) ; op . setOperand ( 0 , input ) ; return success ( ) ; }
Handles padding before convolution for space to depth transform.
Handles padding before convolution for space to depth transform .
// Calculate paddings. // Update pad_op paddings. // Set input type. // Input of the pad op could be a cast op.