MLX
 
Loading...
Searching...
No Matches
copy.h
Go to the documentation of this file.
1// Copyright © 2023-2024 Apple Inc.
2
3#pragma once
4
6#include "mlx/stream.h"
7
8namespace mlx::core {
9
10// Generic copy inplace
12 const array& in,
13 array& out,
14 const Shape& data_shape,
15 const Strides& i_strides,
16 const Strides& o_strides,
17 int64_t i_offset,
18 int64_t o_offset,
19 CopyType ctype,
20 const Stream& s,
21 const std::optional<array>& dynamic_i_offset = std::nullopt,
22 const std::optional<array>& dynamic_o_offset = std::nullopt);
23
24void copy_gpu(const array& src, array& out, CopyType ctype, const Stream& s);
25void copy_gpu(const array& src, array& out, CopyType ctype);
26
28 const array& in,
29 array& out,
30 CopyType ctype,
31 const Stream& s);
32
34 const array& in,
35 array& out,
36 const Strides& i_strides,
37 int64_t i_offset,
38 CopyType ctype,
39 const Stream& s);
40
41// Fill the output with the scalar val
42void fill_gpu(const array& val, array& out, const Stream& s);
43
44} // namespace mlx::core
Definition array.h:24
Definition allocator.h:7
void copy_gpu_inplace(const array &in, array &out, const Shape &data_shape, const Strides &i_strides, const Strides &o_strides, int64_t i_offset, int64_t o_offset, CopyType ctype, const Stream &s, const std::optional< array > &dynamic_i_offset=std::nullopt, const std::optional< array > &dynamic_o_offset=std::nullopt)
std::vector< ShapeElem > Shape
Definition array.h:21
std::vector< int64_t > Strides
Definition array.h:22
CopyType
Definition copy.h:10
void copy_gpu(const array &src, array &out, CopyType ctype, const Stream &s)
void fill_gpu(const array &val, array &out, const Stream &s)
Definition stream.h:9