tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.cc
// Refine MLIR types based on new shape information.
Refine MLIR types based on new shape information.
Refine MLIR types based on new shape information .
Status RefineShapes(llvm::ArrayRef<TensorOrResourceShape> arg_shapes,
                    mlir::ModuleOp module) {
  auto producer_or = GetTfGraphProducerVersion(module);
  if (!producer_or.ok()) return producer_or.status();
  int64_t producer_version = producer_or.ValueOrDie();
  llvm::SmallVector<int64_t, 16> shape_backing;
  llvm::SmallVector<llvm::ArrayRef<int64_t>, 4> arg_shapes_copy;
  {
    size_t count = 0;
    for (const TensorOrResourceShape& tensor_resource_shape : arg_shapes) {
      if (tensor_resource_shape.is_resource) continue;
      count += tensor_resource_shape.shape.dims();
    }
    shape_backing.resize(count);
    arg_shapes_copy.reserve(arg_shapes.size());
    size_t offset = 0;
    for (const TensorOrResourceShape& tensor_resource_shape : arg_shapes) {
      if (tensor_resource_shape.is_resource) {
        arg_shapes_copy.push_back(llvm::ArrayRef<int64_t>());
        continue;
      }
      size_t start = offset;
      for (tensorflow::TensorShapeDim dim : tensor_resource_shape.shape) {
        shape_backing[offset] = dim.size;
        ++offset;
      }
      if (offset == start) {
        arg_shapes_copy.push_back(llvm::ArrayRef<int64_t>());
      } else {
        arg_shapes_copy.push_back(
            llvm::ArrayRef<int64_t>(&shape_backing[start], offset - start));
      }
    }
  }
  auto main_func = module.lookupSymbol<mlir::FuncOp>("main");
  mlir::StatusScopedDiagnosticHandler error_handler(module.getContext());
  mlir::LogicalResult result = mlir::TF::InferShapeForFunction(
      main_func, arg_shapes_copy, producer_version);
  if (failed(result)) {
    return error_handler.Combine(
        errors::Internal("MLIR Shape refinement failed"));
  }
  return Status::OK();
}
Status RefineShapes ( llvm :: ArrayRef < TensorOrResourceShape > arg_shapes , mlir :: ModuleOp module ) { auto producer_or = GetTfGraphProducerVersion ( module ) ; if ( ! producer_or . ok ( ) ) return producer_or . status ( ) ; int64_t producer_version = producer_or . ValueOrDie ( ) ; llvm :: SmallVector < int64_t , 16 > shape_backing ; llvm :: SmallVector < llvm :: ArrayRef < int64_t > , 4 > arg_shapes_copy ; { size_t count = 0 ; for ( const TensorOrResourceShape & tensor_resource_shape : arg_shapes ) { if ( tensor_resource_shape . is_resource ) continue ; count += tensor_resource_shape . shape . dims ( ) ; } shape_backing . resize ( count ) ; arg_shapes_copy . reserve ( arg_shapes . size ( ) ) ; size_t offset = 0 ; for ( const TensorOrResourceShape & tensor_resource_shape : arg_shapes ) { if ( tensor_resource_shape . is_resource ) { arg_shapes_copy . push_back ( llvm :: ArrayRef < int64_t > ( ) ) ; continue ; } size_t start = offset ; for ( tensorflow :: TensorShapeDim dim : tensor_resource_shape . shape ) { shape_backing [ offset ] = dim . size ; ++ offset ; } if ( offset == start ) { arg_shapes_copy . push_back ( llvm :: ArrayRef < int64_t > ( ) ) ; } else { arg_shapes_copy . push_back ( llvm :: ArrayRef < int64_t > ( & shape_backing [ start ] , offset - start ) ) ; } } } auto main_func = module . lookupSymbol < mlir :: FuncOp > ( " " ) ; mlir :: StatusScopedDiagnosticHandler error_handler ( module . getContext ( ) ) ; mlir :: LogicalResult result = mlir :: TF :: InferShapeForFunction ( main_func , arg_shapes_copy , producer_version ) ; if ( failed ( result ) ) { return error_handler . Combine ( errors :: Internal ( " " ) ) ; } return Status :: OK ( ) ; }
Refine MLIR types based on new shape information.
Refine MLIR types based on new shape information .
// Convert arg_shapes to a mlir friendly format.