kernels

package
v0.0.0-...-3233e8b Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Dec 8, 2023 License: BSD-3-Clause Imports: 13 Imported by: 0

Documentation

Index

Constants

View Source 
const (
	TpuCompilationCacheServiceExternal_GetTpuProgram_FullMethodName = "/tensorflow.tpu.TpuCompilationCacheServiceExternal/GetTpuProgram"
)

Variables

View Source 
var (
	CompilationCacheFetchTarget_name = map[int32]string{
		0: "INVALID",
		1: "MAIN",
		2: "SHARDING",
		3: "UNSHARDING",
	}
	CompilationCacheFetchTarget_value = map[string]int32{
		"INVALID":    0,
		"MAIN":       1,
		"SHARDING":   2,
		"UNSHARDING": 3,
	}
)

Enum value maps for CompilationCacheFetchTarget.

View Source 
var (
	TPUHostTransferProto_TransferDirection_name = map[int32]string{
		0: "NONE",
		1: "DEVICE_TO_HOST",
		2: "HOST_TO_DEVICE",
	}
	TPUHostTransferProto_TransferDirection_value = map[string]int32{
		"NONE":           0,
		"DEVICE_TO_HOST": 1,
		"HOST_TO_DEVICE": 2,
	}
)

Enum value maps for TPUHostTransferProto_TransferDirection.

View Source 
var File_tensorflow_core_tpu_kernels_tpu_compilation_cache_common_proto protoreflect.FileDescriptor
View Source 
var File_tensorflow_core_tpu_kernels_tpu_compilation_cache_proto protoreflect.FileDescriptor
View Source 
var File_tensorflow_core_tpu_kernels_tpu_compile_proto protoreflect.FileDescriptor
View Source 
var File_tensorflow_core_tpu_kernels_tpu_executable_info_proto protoreflect.FileDescriptor
View Source 
var TpuCompilationCacheServiceExternal_ServiceDesc = grpc.ServiceDesc{
	ServiceName: "tensorflow.tpu.TpuCompilationCacheServiceExternal",
	HandlerType: (*TpuCompilationCacheServiceExternalServer)(nil),
	Methods: []grpc.MethodDesc{
		{
			MethodName: "GetTpuProgram",
			Handler:    _TpuCompilationCacheServiceExternal_GetTpuProgram_Handler,
		},
	},
	Streams:  []grpc.StreamDesc{},
	Metadata: "tensorflow/core/tpu/kernels/tpu_compilation_cache.proto",
}

TpuCompilationCacheServiceExternal_ServiceDesc is the grpc.ServiceDesc for TpuCompilationCacheServiceExternal service. It's only intended for direct use with grpc.RegisterService, and not to be introspected or modified (even as a copy)

Functions

func RegisterTpuCompilationCacheServiceExternalServer 

func RegisterTpuCompilationCacheServiceExternalServer(s grpc.ServiceRegistrar, srv TpuCompilationCacheServiceExternalServer)

Types

type CompilationCacheFetchTarget 

type CompilationCacheFetchTarget int32

Target type for compilation cache fetch operation. 

const (
	CompilationCacheFetchTarget_INVALID    CompilationCacheFetchTarget = 0
	CompilationCacheFetchTarget_MAIN       CompilationCacheFetchTarget = 1
	CompilationCacheFetchTarget_SHARDING   CompilationCacheFetchTarget = 2
	CompilationCacheFetchTarget_UNSHARDING CompilationCacheFetchTarget = 3
)

func (CompilationCacheFetchTarget) Descriptor 

func (CompilationCacheFetchTarget) Descriptor() protoreflect.EnumDescriptor

func (CompilationCacheFetchTarget) Enum 

func (x CompilationCacheFetchTarget) Enum() *CompilationCacheFetchTarget

func (CompilationCacheFetchTarget) EnumDescriptor deprecated 

func (CompilationCacheFetchTarget) EnumDescriptor() ([]byte, []int)

Deprecated: Use CompilationCacheFetchTarget.Descriptor instead.

func (CompilationCacheFetchTarget) Number 

func (x CompilationCacheFetchTarget) Number() protoreflect.EnumNumber

func (CompilationCacheFetchTarget) String 

func (x CompilationCacheFetchTarget) String() string

func (CompilationCacheFetchTarget) Type 

func (CompilationCacheFetchTarget) Type() protoreflect.EnumType

type GetTpuProgramRequest 

type GetTpuProgramRequest struct {

	// Types that are assignable to KeyOneof:
	//
	//	*GetTpuProgramRequest_Key
	//	*GetTpuProgramRequest_UidAndIndex
	KeyOneof    isGetTpuProgramRequest_KeyOneof `protobuf_oneof:"key_oneof"`
	FetchTarget CompilationCacheFetchTarget     `` /* 143-byte string literal not displayed */
	// contains filtered or unexported fields
}

func (*GetTpuProgramRequest) Descriptor deprecated 

func (*GetTpuProgramRequest) Descriptor() ([]byte, []int)

Deprecated: Use GetTpuProgramRequest.ProtoReflect.Descriptor instead.

func (*GetTpuProgramRequest) GetFetchTarget 

func (x *GetTpuProgramRequest) GetFetchTarget() CompilationCacheFetchTarget

func (*GetTpuProgramRequest) GetKey 

func (x *GetTpuProgramRequest) GetKey() string

func (*GetTpuProgramRequest) GetKeyOneof 

func (m *GetTpuProgramRequest) GetKeyOneof() isGetTpuProgramRequest_KeyOneof

func (*GetTpuProgramRequest) GetUidAndIndex 

func (x *GetTpuProgramRequest) GetUidAndIndex() *TpuCompilationUidAndIndex

func (*GetTpuProgramRequest) ProtoMessage 

func (*GetTpuProgramRequest) ProtoMessage()

func (*GetTpuProgramRequest) ProtoReflect 

func (x *GetTpuProgramRequest) ProtoReflect() protoreflect.Message

func (*GetTpuProgramRequest) Reset 

func (x *GetTpuProgramRequest) Reset()

func (*GetTpuProgramRequest) String 

func (x *GetTpuProgramRequest) String() string

type GetTpuProgramRequest_Key 

type GetTpuProgramRequest_Key struct {
	Key string `protobuf:"bytes,1,opt,name=key,proto3,oneof"`
}

type GetTpuProgramRequest_UidAndIndex 

type GetTpuProgramRequest_UidAndIndex struct {
	UidAndIndex *TpuCompilationUidAndIndex `protobuf:"bytes,2,opt,name=uid_and_index,json=uidAndIndex,proto3,oneof"`
}

type GetTpuProgramResponseExternal 

type GetTpuProgramResponseExternal struct {
	Proto               *GetTpuProgramResponseExternal_Blob `protobuf:"bytes,1,opt,name=proto,proto3" json:"proto,omitempty"`
	HostComputeMetadata *tf2xla.HostComputeMetadata         `protobuf:"bytes,2,opt,name=host_compute_metadata,json=hostComputeMetadata,proto3" json:"host_compute_metadata,omitempty"`
	MayModifyVariables  bool                                `protobuf:"varint,3,opt,name=may_modify_variables,json=mayModifyVariables,proto3" json:"may_modify_variables,omitempty"`
	CompilerMetadata    *GetTpuProgramResponseExternal_Blob `protobuf:"bytes,4,opt,name=compiler_metadata,json=compilerMetadata,proto3" json:"compiler_metadata,omitempty"`
	// Whether the program is empty, which could be true for sharding/unsharding
	// entries.
	IsEmpty bool `protobuf:"varint,5,opt,name=is_empty,json=isEmpty,proto3" json:"is_empty,omitempty"`
	// contains filtered or unexported fields
}

Response for GetTpuProgram RPC.

func (*GetTpuProgramResponseExternal) Descriptor deprecated 

func (*GetTpuProgramResponseExternal) Descriptor() ([]byte, []int)

Deprecated: Use GetTpuProgramResponseExternal.ProtoReflect.Descriptor instead.

func (*GetTpuProgramResponseExternal) GetCompilerMetadata 

func (x *GetTpuProgramResponseExternal) GetCompilerMetadata() *GetTpuProgramResponseExternal_Blob

func (*GetTpuProgramResponseExternal) GetHostComputeMetadata 

func (x *GetTpuProgramResponseExternal) GetHostComputeMetadata() *tf2xla.HostComputeMetadata

func (*GetTpuProgramResponseExternal) GetIsEmpty 

func (x *GetTpuProgramResponseExternal) GetIsEmpty() bool

func (*GetTpuProgramResponseExternal) GetMayModifyVariables 

func (x *GetTpuProgramResponseExternal) GetMayModifyVariables() bool

func (*GetTpuProgramResponseExternal) GetProto 

func (x *GetTpuProgramResponseExternal) GetProto() *GetTpuProgramResponseExternal_Blob

func (*GetTpuProgramResponseExternal) ProtoMessage 

func (*GetTpuProgramResponseExternal) ProtoMessage()

func (*GetTpuProgramResponseExternal) ProtoReflect 

func (x *GetTpuProgramResponseExternal) ProtoReflect() protoreflect.Message

func (*GetTpuProgramResponseExternal) Reset 

func (x *GetTpuProgramResponseExternal) Reset()

func (*GetTpuProgramResponseExternal) String 

func (x *GetTpuProgramResponseExternal) String() string

type GetTpuProgramResponseExternal_Blob 

type GetTpuProgramResponseExternal_Blob struct {
	Data []byte `protobuf:"bytes,1,opt,name=data,proto3" json:"data,omitempty"`
	// contains filtered or unexported fields
}

func (*GetTpuProgramResponseExternal_Blob) Descriptor deprecated 

func (*GetTpuProgramResponseExternal_Blob) Descriptor() ([]byte, []int)

Deprecated: Use GetTpuProgramResponseExternal_Blob.ProtoReflect.Descriptor instead.

func (*GetTpuProgramResponseExternal_Blob) GetData 

func (x *GetTpuProgramResponseExternal_Blob) GetData() []byte

func (*GetTpuProgramResponseExternal_Blob) ProtoMessage 

func (*GetTpuProgramResponseExternal_Blob) ProtoMessage()

func (*GetTpuProgramResponseExternal_Blob) ProtoReflect 

func (x *GetTpuProgramResponseExternal_Blob) ProtoReflect() protoreflect.Message

func (*GetTpuProgramResponseExternal_Blob) Reset 

func (x *GetTpuProgramResponseExternal_Blob) Reset()

func (*GetTpuProgramResponseExternal_Blob) String 

func (x *GetTpuProgramResponseExternal_Blob) String() string

type TPUExecutableInfoProto 

type TPUExecutableInfoProto struct {

	// The shapes of the inputs and outputs.
	InputShapes []*data.ShapeProto `protobuf:"bytes,2,rep,name=input_shapes,json=inputShapes,proto3" json:"input_shapes,omitempty"`
	OutputShape *data.ShapeProto   `protobuf:"bytes,3,opt,name=output_shape,json=outputShape,proto3" json:"output_shape,omitempty"`
	// Dynamic output indices indicate which outputs have dynamic dimensions.
	DynamicOutputIndices []*TPUExecutableInfoProto_ShapeIndex `protobuf:"bytes,11,rep,name=dynamic_output_indices,json=dynamicOutputIndices,proto3" json:"dynamic_output_indices,omitempty"`
	// For each resource variable output, what was the index of the corresponding
	// input and was it updated? The indices are sorted by input order.
	VariableIndices []*TPUExecutableInfoProto_UpdateIndexPair `protobuf:"bytes,10,rep,name=variable_indices,json=variableIndices,proto3" json:"variable_indices,omitempty"`
	// The shapes of the outputs when represented as Tensors. These may not
	// match the output_shape values because we may flatten tensors to avoid
	// excess padding.
	OutputTensorShapes []*framework.TensorShapeProto `protobuf:"bytes,8,rep,name=output_tensor_shapes,json=outputTensorShapes,proto3" json:"output_tensor_shapes,omitempty"`
	// Optional session module for passing XLA computations between TPUCompileOp
	// and TPUExecuteOp. This is needed to support the
	// --xla_dump_hlo_snapshots flag.
	SessionModule *service.HloSnapshot `protobuf:"bytes,5,opt,name=session_module,json=sessionModule,proto3" json:"session_module,omitempty"`
	// The physical device ids assigned to the replicated cores.
	DeviceAssignment *data.DeviceAssignmentProto `protobuf:"bytes,6,opt,name=device_assignment,json=deviceAssignment,proto3" json:"device_assignment,omitempty"`
	// contains filtered or unexported fields
}

A serialization of TPUExecutable. Only includes fields necessary to load and execute a program on a worker node.

func (*TPUExecutableInfoProto) Descriptor deprecated 

func (*TPUExecutableInfoProto) Descriptor() ([]byte, []int)

Deprecated: Use TPUExecutableInfoProto.ProtoReflect.Descriptor instead.

func (*TPUExecutableInfoProto) GetDeviceAssignment 

func (x *TPUExecutableInfoProto) GetDeviceAssignment() *data.DeviceAssignmentProto

func (*TPUExecutableInfoProto) GetDynamicOutputIndices 

func (x *TPUExecutableInfoProto) GetDynamicOutputIndices() []*TPUExecutableInfoProto_ShapeIndex

func (*TPUExecutableInfoProto) GetInputShapes 

func (x *TPUExecutableInfoProto) GetInputShapes() []*data.ShapeProto

func (*TPUExecutableInfoProto) GetOutputShape 

func (x *TPUExecutableInfoProto) GetOutputShape() *data.ShapeProto

func (*TPUExecutableInfoProto) GetOutputTensorShapes 

func (x *TPUExecutableInfoProto) GetOutputTensorShapes() []*framework.TensorShapeProto

func (*TPUExecutableInfoProto) GetSessionModule 

func (x *TPUExecutableInfoProto) GetSessionModule() *service.HloSnapshot

func (*TPUExecutableInfoProto) GetVariableIndices 

func (x *TPUExecutableInfoProto) GetVariableIndices() []*TPUExecutableInfoProto_UpdateIndexPair

func (*TPUExecutableInfoProto) ProtoMessage 

func (*TPUExecutableInfoProto) ProtoMessage()

func (*TPUExecutableInfoProto) ProtoReflect 

func (x *TPUExecutableInfoProto) ProtoReflect() protoreflect.Message

func (*TPUExecutableInfoProto) Reset 

func (x *TPUExecutableInfoProto) Reset()

func (*TPUExecutableInfoProto) String 

func (x *TPUExecutableInfoProto) String() string

type TPUExecutableInfoProto_ShapeIndex 

type TPUExecutableInfoProto_ShapeIndex struct {
	Index []int32 `protobuf:"varint,1,rep,packed,name=index,proto3" json:"index,omitempty"`
	// contains filtered or unexported fields
}

func (*TPUExecutableInfoProto_ShapeIndex) Descriptor deprecated 

func (*TPUExecutableInfoProto_ShapeIndex) Descriptor() ([]byte, []int)

Deprecated: Use TPUExecutableInfoProto_ShapeIndex.ProtoReflect.Descriptor instead.

func (*TPUExecutableInfoProto_ShapeIndex) GetIndex 

func (x *TPUExecutableInfoProto_ShapeIndex) GetIndex() []int32

func (*TPUExecutableInfoProto_ShapeIndex) ProtoMessage 

func (*TPUExecutableInfoProto_ShapeIndex) ProtoMessage()

func (*TPUExecutableInfoProto_ShapeIndex) ProtoReflect 

func (x *TPUExecutableInfoProto_ShapeIndex) ProtoReflect() protoreflect.Message

func (*TPUExecutableInfoProto_ShapeIndex) Reset 

func (x *TPUExecutableInfoProto_ShapeIndex) Reset()

func (*TPUExecutableInfoProto_ShapeIndex) String 

func (x *TPUExecutableInfoProto_ShapeIndex) String() string

type TPUExecutableInfoProto_UpdateIndexPair 

type TPUExecutableInfoProto_UpdateIndexPair struct {
	Index   int32 `protobuf:"varint,1,opt,name=index,proto3" json:"index,omitempty"`
	Updated bool  `protobuf:"varint,2,opt,name=updated,proto3" json:"updated,omitempty"`
	// contains filtered or unexported fields
}

func (*TPUExecutableInfoProto_UpdateIndexPair) Descriptor deprecated 

func (*TPUExecutableInfoProto_UpdateIndexPair) Descriptor() ([]byte, []int)

Deprecated: Use TPUExecutableInfoProto_UpdateIndexPair.ProtoReflect.Descriptor instead.

func (*TPUExecutableInfoProto_UpdateIndexPair) GetIndex 

func (x *TPUExecutableInfoProto_UpdateIndexPair) GetIndex() int32

func (*TPUExecutableInfoProto_UpdateIndexPair) GetUpdated 

func (x *TPUExecutableInfoProto_UpdateIndexPair) GetUpdated() bool

func (*TPUExecutableInfoProto_UpdateIndexPair) ProtoMessage 

func (*TPUExecutableInfoProto_UpdateIndexPair) ProtoMessage()

func (*TPUExecutableInfoProto_UpdateIndexPair) ProtoReflect 

func (x *TPUExecutableInfoProto_UpdateIndexPair) ProtoReflect() protoreflect.Message

func (*TPUExecutableInfoProto_UpdateIndexPair) Reset 

func (x *TPUExecutableInfoProto_UpdateIndexPair) Reset()

func (*TPUExecutableInfoProto_UpdateIndexPair) String 

func (x *TPUExecutableInfoProto_UpdateIndexPair) String() string

type TPUHostTransferInfoProto 

type TPUHostTransferInfoProto struct {
	HostTransfers []*TPUHostTransferProto `protobuf:"bytes,1,rep,name=host_transfers,json=hostTransfers,proto3" json:"host_transfers,omitempty"`
	// contains filtered or unexported fields
}

func (*TPUHostTransferInfoProto) Descriptor deprecated 

func (*TPUHostTransferInfoProto) Descriptor() ([]byte, []int)

Deprecated: Use TPUHostTransferInfoProto.ProtoReflect.Descriptor instead.

func (*TPUHostTransferInfoProto) GetHostTransfers 

func (x *TPUHostTransferInfoProto) GetHostTransfers() []*TPUHostTransferProto

func (*TPUHostTransferInfoProto) ProtoMessage 

func (*TPUHostTransferInfoProto) ProtoMessage()

func (*TPUHostTransferInfoProto) ProtoReflect 

func (x *TPUHostTransferInfoProto) ProtoReflect() protoreflect.Message

func (*TPUHostTransferInfoProto) Reset 

func (x *TPUHostTransferInfoProto) Reset()

func (*TPUHostTransferInfoProto) String 

func (x *TPUHostTransferInfoProto) String() string

type TPUHostTransferProto 

type TPUHostTransferProto struct {

	// Channel identifier assigned by compiler and used in host commands.
	Channel int64 `protobuf:"varint,1,opt,name=channel,proto3" json:"channel,omitempty"`
	// Direction of the transfer operation.
	Direction TPUHostTransferProto_TransferDirection `` /* 127-byte string literal not displayed */
	// Channel identifier prodided by XLA client.
	Key string `protobuf:"bytes,3,opt,name=key,proto3" json:"key,omitempty"`
	// Shape of the data to be transferred (including layout).
	Shape *data.ShapeProto `protobuf:"bytes,5,opt,name=shape,proto3" json:"shape,omitempty"`
	// Address of the device buffer in HBM (byte offset).
	BufferOffset int64 `protobuf:"varint,6,opt,name=buffer_offset,json=bufferOffset,proto3" json:"buffer_offset,omitempty"`
	// Original data type for this host transfer before X64 rewrite.
	OriginalType data.PrimitiveType `protobuf:"varint,7,opt,name=original_type,json=originalType,proto3,enum=xla.PrimitiveType" json:"original_type,omitempty"`
	// If this host transfer is a splitted X64 transfer, specifies whether this
	// transfer is for lower bits.
	IsLowerBits bool `protobuf:"varint,8,opt,name=is_lower_bits,json=isLowerBits,proto3" json:"is_lower_bits,omitempty"`
	// The name of host side command handler.
	HostHandlerName string `protobuf:"bytes,9,opt,name=host_handler_name,json=hostHandlerName,proto3" json:"host_handler_name,omitempty"`
	// contains filtered or unexported fields
}

Metadata for a data transfer between device and host.

func (*TPUHostTransferProto) Descriptor deprecated 

func (*TPUHostTransferProto) Descriptor() ([]byte, []int)

Deprecated: Use TPUHostTransferProto.ProtoReflect.Descriptor instead.

func (*TPUHostTransferProto) GetBufferOffset 

func (x *TPUHostTransferProto) GetBufferOffset() int64

func (*TPUHostTransferProto) GetChannel 

func (x *TPUHostTransferProto) GetChannel() int64

func (*TPUHostTransferProto) GetDirection 

func (x *TPUHostTransferProto) GetDirection() TPUHostTransferProto_TransferDirection

func (*TPUHostTransferProto) GetHostHandlerName 

func (x *TPUHostTransferProto) GetHostHandlerName() string

func (*TPUHostTransferProto) GetIsLowerBits 

func (x *TPUHostTransferProto) GetIsLowerBits() bool

func (*TPUHostTransferProto) GetKey 

func (x *TPUHostTransferProto) GetKey() string

func (*TPUHostTransferProto) GetOriginalType 

func (x *TPUHostTransferProto) GetOriginalType() data.PrimitiveType

func (*TPUHostTransferProto) GetShape 

func (x *TPUHostTransferProto) GetShape() *data.ShapeProto

func (*TPUHostTransferProto) ProtoMessage 

func (*TPUHostTransferProto) ProtoMessage()

func (*TPUHostTransferProto) ProtoReflect 

func (x *TPUHostTransferProto) ProtoReflect() protoreflect.Message

func (*TPUHostTransferProto) Reset 

func (x *TPUHostTransferProto) Reset()

func (*TPUHostTransferProto) String 

func (x *TPUHostTransferProto) String() string

type TPUHostTransferProto_TransferDirection 

type TPUHostTransferProto_TransferDirection int32
const (
	TPUHostTransferProto_NONE           TPUHostTransferProto_TransferDirection = 0
	TPUHostTransferProto_DEVICE_TO_HOST TPUHostTransferProto_TransferDirection = 1
	TPUHostTransferProto_HOST_TO_DEVICE TPUHostTransferProto_TransferDirection = 2
)

func (TPUHostTransferProto_TransferDirection) Descriptor 

func (TPUHostTransferProto_TransferDirection) Descriptor() protoreflect.EnumDescriptor

func (TPUHostTransferProto_TransferDirection) Enum 

func (x TPUHostTransferProto_TransferDirection) Enum() *TPUHostTransferProto_TransferDirection

func (TPUHostTransferProto_TransferDirection) EnumDescriptor deprecated 

func (TPUHostTransferProto_TransferDirection) EnumDescriptor() ([]byte, []int)

Deprecated: Use TPUHostTransferProto_TransferDirection.Descriptor instead.

func (TPUHostTransferProto_TransferDirection) Number 

func (x TPUHostTransferProto_TransferDirection) Number() protoreflect.EnumNumber

func (TPUHostTransferProto_TransferDirection) String 

func (x TPUHostTransferProto_TransferDirection) String() string

func (TPUHostTransferProto_TransferDirection) Type 

func (TPUHostTransferProto_TransferDirection) Type() protoreflect.EnumType

type TpuCompilationCacheServiceExternalClient 

type TpuCompilationCacheServiceExternalClient interface {
	// This method requests the cached proto that the TPU execute op has been
	// instructed to execute.
	GetTpuProgram(ctx context.Context, in *GetTpuProgramRequest, opts ...grpc.CallOption) (*GetTpuProgramResponseExternal, error)
}

TpuCompilationCacheServiceExternalClient is the client API for TpuCompilationCacheServiceExternal service.

For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream.

func NewTpuCompilationCacheServiceExternalClient 

func NewTpuCompilationCacheServiceExternalClient(cc grpc.ClientConnInterface) TpuCompilationCacheServiceExternalClient

type TpuCompilationCacheServiceExternalServer 

type TpuCompilationCacheServiceExternalServer interface {
	// This method requests the cached proto that the TPU execute op has been
	// instructed to execute.
	GetTpuProgram(context.Context, *GetTpuProgramRequest) (*GetTpuProgramResponseExternal, error)
	// contains filtered or unexported methods
}

TpuCompilationCacheServiceExternalServer is the server API for TpuCompilationCacheServiceExternal service. All implementations must embed UnimplementedTpuCompilationCacheServiceExternalServer for forward compatibility

type TpuCompilationRequestProto 

type TpuCompilationRequestProto struct {

	// A flag reserved for using experimental version of the compilation. By
	// default the value should be false.
	UseExperimental bool `protobuf:"varint,1,opt,name=use_experimental,json=useExperimental,proto3" json:"use_experimental,omitempty"`
	// Use mlir to lower computation(s) to Hlo.
	UseMlir bool `protobuf:"varint,2,opt,name=use_mlir,json=useMlir,proto3" json:"use_mlir,omitempty"`
	// If true, returns hlo metadatas.
	ReturnHloProtos bool `protobuf:"varint,3,opt,name=return_hlo_protos,json=returnHloProtos,proto3" json:"return_hlo_protos,omitempty"`
	// If true, unloads cache on session close.
	UnloadCacheOnSessionClose bool `` /* 143-byte string literal not displayed */
	// Compilation metadata.
	Metadata *tpu.TPUCompileMetadataProto `protobuf:"bytes,5,opt,name=metadata,proto3" json:"metadata,omitempty"`
	// Computation argument shapes.
	ArgShapes []*framework.TensorShapeProto `protobuf:"bytes,6,rep,name=arg_shapes,json=argShapes,proto3" json:"arg_shapes,omitempty"`
	// Input tensor that gives const guarantee to the TF runtime.
	GuaranteedConstants []*framework.TensorProto `protobuf:"bytes,7,rep,name=guaranteed_constants,json=guaranteedConstants,proto3" json:"guaranteed_constants,omitempty"`
	// MLIR module definition.
	MlirModule string `protobuf:"bytes,8,opt,name=mlir_module,json=mlirModule,proto3" json:"mlir_module,omitempty"`
	// A set of named functions used as the input to lowering to Hlo when mlir is
	// not used.
	FdefLib *framework.FunctionDefLibrary `protobuf:"bytes,9,opt,name=fdef_lib,json=fdefLib,proto3" json:"fdef_lib,omitempty"`
	// The version of the graph definition used to lower TF function to Hlo.
	GraphDefVersion int32 `protobuf:"varint,10,opt,name=graph_def_version,json=graphDefVersion,proto3" json:"graph_def_version,omitempty"`
	// Function containing the computation to compile.
	Function *framework.NameAttrList `protobuf:"bytes,11,opt,name=function,proto3" json:"function,omitempty"`
	// contains filtered or unexported fields
}

TPU compilation request for compiling computations into XLA HLO IR and build TPU programs.

func (*TpuCompilationRequestProto) Descriptor deprecated 

func (*TpuCompilationRequestProto) Descriptor() ([]byte, []int)

Deprecated: Use TpuCompilationRequestProto.ProtoReflect.Descriptor instead.

func (*TpuCompilationRequestProto) GetArgShapes 

func (x *TpuCompilationRequestProto) GetArgShapes() []*framework.TensorShapeProto

func (*TpuCompilationRequestProto) GetFdefLib 

func (x *TpuCompilationRequestProto) GetFdefLib() *framework.FunctionDefLibrary

func (*TpuCompilationRequestProto) GetFunction 

func (x *TpuCompilationRequestProto) GetFunction() *framework.NameAttrList

func (*TpuCompilationRequestProto) GetGraphDefVersion 

func (x *TpuCompilationRequestProto) GetGraphDefVersion() int32

func (*TpuCompilationRequestProto) GetGuaranteedConstants 

func (x *TpuCompilationRequestProto) GetGuaranteedConstants() []*framework.TensorProto

func (*TpuCompilationRequestProto) GetMetadata 

func (x *TpuCompilationRequestProto) GetMetadata() *tpu.TPUCompileMetadataProto

func (*TpuCompilationRequestProto) GetMlirModule 

func (x *TpuCompilationRequestProto) GetMlirModule() string

func (*TpuCompilationRequestProto) GetReturnHloProtos 

func (x *TpuCompilationRequestProto) GetReturnHloProtos() bool

func (*TpuCompilationRequestProto) GetUnloadCacheOnSessionClose 

func (x *TpuCompilationRequestProto) GetUnloadCacheOnSessionClose() bool

func (*TpuCompilationRequestProto) GetUseExperimental 

func (x *TpuCompilationRequestProto) GetUseExperimental() bool

func (*TpuCompilationRequestProto) GetUseMlir 

func (x *TpuCompilationRequestProto) GetUseMlir() bool

func (*TpuCompilationRequestProto) ProtoMessage 

func (*TpuCompilationRequestProto) ProtoMessage()

func (*TpuCompilationRequestProto) ProtoReflect 

func (x *TpuCompilationRequestProto) ProtoReflect() protoreflect.Message

func (*TpuCompilationRequestProto) Reset 

func (x *TpuCompilationRequestProto) Reset()

func (*TpuCompilationRequestProto) String 

func (x *TpuCompilationRequestProto) String() string

type TpuCompilationUidAndIndex 

type TpuCompilationUidAndIndex struct {
	Uid        int64 `protobuf:"varint,1,opt,name=uid,proto3" json:"uid,omitempty"`
	ProtoIndex int32 `protobuf:"varint,2,opt,name=proto_index,json=protoIndex,proto3" json:"proto_index,omitempty"`
	// contains filtered or unexported fields
}

func (*TpuCompilationUidAndIndex) Descriptor deprecated 

func (*TpuCompilationUidAndIndex) Descriptor() ([]byte, []int)

Deprecated: Use TpuCompilationUidAndIndex.ProtoReflect.Descriptor instead.

func (*TpuCompilationUidAndIndex) GetProtoIndex 

func (x *TpuCompilationUidAndIndex) GetProtoIndex() int32

func (*TpuCompilationUidAndIndex) GetUid 

func (x *TpuCompilationUidAndIndex) GetUid() int64

func (*TpuCompilationUidAndIndex) ProtoMessage 

func (*TpuCompilationUidAndIndex) ProtoMessage()

func (*TpuCompilationUidAndIndex) ProtoReflect 

func (x *TpuCompilationUidAndIndex) ProtoReflect() protoreflect.Message

func (*TpuCompilationUidAndIndex) Reset 

func (x *TpuCompilationUidAndIndex) Reset()

func (*TpuCompilationUidAndIndex) String 

func (x *TpuCompilationUidAndIndex) String() string

type UnimplementedTpuCompilationCacheServiceExternalServer 

type UnimplementedTpuCompilationCacheServiceExternalServer struct {
}

UnimplementedTpuCompilationCacheServiceExternalServer must be embedded to have forward compatible implementations.

func (UnimplementedTpuCompilationCacheServiceExternalServer) GetTpuProgram 

func (UnimplementedTpuCompilationCacheServiceExternalServer) GetTpuProgram(context.Context, *GetTpuProgramRequest) (*GetTpuProgramResponseExternal, error)

type UnsafeTpuCompilationCacheServiceExternalServer 

type UnsafeTpuCompilationCacheServiceExternalServer interface {
	// contains filtered or unexported methods
}

UnsafeTpuCompilationCacheServiceExternalServer may be embedded to opt out of forward compatibility for this service. Use of this interface is not recommended, as added methods to TpuCompilationCacheServiceExternalServer will result in compilation errors.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.