Searched refs:AsDeviceMemoryBase (Results 1 – 14 of 14) sorted by relevance
175 memory_bases.push_back(argument.Buffer({}).AsDeviceMemoryBase()); in ExecuteAsyncOnStream()204 CHECK(!it->second.AsDeviceMemoryBase().is_null()); in ExecuteAsyncOnStream()207 it->second.AsDeviceMemoryBase()); in ExecuteAsyncOnStream()210 cross_program_prefetch_addr = it->second.AsDeviceMemoryBase(); in ExecuteAsyncOnStream()
43 virtual stream_executor::DeviceMemoryBase AsDeviceMemoryBase() const = 0;
47 se_mem.memory = ApiConverter::ToC(mem.AsDeviceMemoryBase()); in ToC()
99 .AsDeviceMemoryBase(); in MemoryForAllocation()183 const_cast<void*>(buffer.AsDeviceMemoryBase().opaque())); in ExecuteComputeFunction()302 registered_buffer.AsDeviceMemoryBase().size()); in CreateResultShapedBuffer()304 /*src=*/registered_buffer.AsDeviceMemoryBase().opaque(), in CreateResultShapedBuffer()317 result_buffer = buffer.AsDeviceMemoryBase(); in CreateResultShapedBuffer()
22 tensorflow::se::DeviceMemoryBase MaybeOwningDeviceMemory::AsDeviceMemoryBase() in AsDeviceMemoryBase() function in xla::MaybeOwningDeviceMemory
52 tensorflow::se::DeviceMemoryBase AsDeviceMemoryBase() const;
348 buffer_tree.element({}).AsDeviceMemoryBase(); in WriteRootTupleIndexTable()355 elements.push_back(buffer_tree.element({i}).AsDeviceMemoryBase()); in WriteRootTupleIndexTable()
75 shaped_buffer.set_buffer(index_buffer.second.AsDeviceMemoryBase(), in ToShapedBuffer()
156 elements.push_back(mem->Buffer(element_index).AsDeviceMemoryBase()); in FixTupleTableAsync()159 mem->Buffer(index).AsDeviceMemoryBase(); in FixTupleTableAsync()232 se::DeviceMemory<int8>(mutable_input_mem->AsDeviceMemoryBase()), in UpdateDynamicInputs()
64 out_it->second = in_it->second.AsDeviceMemoryBase(); in ExecuteAsyncOnStream()
275 out_it->second = it->second.AsDeviceMemoryBase(); in MaybeOwningShapeTreeToShapedBuffer()
320 .AsDeviceMemoryBase(); in BufferForAllocation()
235 execution_input->Buffer(index).AsDeviceMemoryBase(); in UpdateDynamicInputs()
179 return buffer->AsDeviceMemoryBase(); in ToShapedBuffer()