tvm::topi::cuda Namespace Reference

Functions
tvm::te::Tensor	dense_cuda (const Target &target, const tvm::te::Tensor &data, const tvm::te::Tensor &weight, const tvm::te::Tensor &bias, const DataType &out_dtype)
	Implementation of dense for CUDA backend. More...
Schedule	schedule_dense (const Target &target, const Array< Tensor > &outs)
	Create a CUDA schedule for dense. More...
Schedule	schedule_injective_from_existing (Schedule sch, const Tensor &out)
	Updates an existing schedule for the given injective ops. More...
Schedule	schedule_injective (const Target &target, const Array< Tensor > &outs)
	Create a CUDA schedule for the given output tensors. More...
Schedule	schedule_pool (const Target &target, const Array< Tensor > &outs)
	Create a CUDA schedule for pool. More...
Schedule	schedule_global_pool (const Target &target, const Array< Tensor > &outs)
	Create a CUDA schedule for global_pool. More...
Schedule	ScheduleReduce (const Target &target, Operation op, Schedule sch, bool is_idx_reduce=false)
	Schedule a given reduce operation. More...
void	TraverseBeforeReduce (Schedule s, Operation op)
	Recursively traverse operator inputs, setting injective inputs to be computed inline. More...
void	TraverseAfterReduce (const Target &target, Schedule s, Operation op)
	Schedule a reduce op, then invoke TraverseBeforeReduce on each of the op's inputs. More...
Schedule	schedule_reduce (const Target &target, Array< Tensor > outs)
	Create a CUDA schedule for a reduce operation. More...
Schedule	schedule_softmax (const Target &target, const Array< Tensor > &outs)
	Create a CUDA schedule for the given softmax output tensors. More...

Function Documentation

dense_cuda()

tvm::te::Tensor tvm::topi::cuda::dense_cuda ( const Target &  target, const tvm::te::Tensor &  data, const tvm::te::Tensor &  weight, const tvm::te::Tensor &  bias, const DataType &  out_dtype  )

inline

Implementation of dense for CUDA backend.

Parameters

target	The target device
data	Tensor with shape [batch, in_dim]
weight	Tensor with shape [out_dim, in_dim]
bias	Tensor with shape [out_dim]. Optional; to omit bias, pass Tensor()
out_dtype	Output data type. Used for mixed precision.

Returns

Tensor with shape [batch, out_dim]

schedule_dense()

Schedule tvm::topi::cuda::schedule_dense ( const Target &  target, const Array< Tensor > &  outs  )

inline

Create a CUDA schedule for dense.

Parameters

target	The target to generate a schedule for.
outs	The output tensors.

Returns

A schedule for the given ops.

schedule_global_pool()

Schedule tvm::topi::cuda::schedule_global_pool ( const Target &  target, const Array< Tensor > &  outs  )

inline

Create a CUDA schedule for global_pool.

Parameters

target	The target to generate a schedule for.
outs	The output tensors.

Returns

A schedule for the given ops.

schedule_injective()

Schedule tvm::topi::cuda::schedule_injective ( const Target &  target, const Array< Tensor > &  outs  )

inline

Create a CUDA schedule for the given output tensors.

Parameters

target	The target to generate a schedule for.
outs	The output tensors.

Returns

A schedule for the given ops.

schedule_injective_from_existing()

Schedule tvm::topi::cuda::schedule_injective_from_existing ( Schedule  sch, const Tensor &  out  )

inline

Updates an existing schedule for the given injective ops.

Parameters

sch	The schedule to update.
out	The tensor representing the injective op.

Returns

The updated schedule.

schedule_pool()

Schedule tvm::topi::cuda::schedule_pool ( const Target &  target, const Array< Tensor > &  outs  )

inline

Create a CUDA schedule for pool.

Parameters

target	The target to generate a schedule for.
outs	The output tensors.

Returns

A schedule for the given ops.

schedule_reduce()

Schedule tvm::topi::cuda::schedule_reduce	(	const Target &	target,
		Array< Tensor >	outs
	)

Create a CUDA schedule for a reduce operation.

Parameters

target	The target to generate a schedule for.
outs	The output tensors.

Returns

A schedule for the given ops.

schedule_softmax()

Schedule tvm::topi::cuda::schedule_softmax ( const Target &  target, const Array< Tensor > &  outs  )

inline

Create a CUDA schedule for the given softmax output tensors.

Parameters

target	The target to generate a schedule for.
outs	The output tensors.

Returns

A schedule for the given ops.

ScheduleReduce()

Schedule tvm::topi::cuda::ScheduleReduce	(	const Target &	target,
		Operation	op,
		Schedule	sch,
		bool	is_idx_reduce = false
	)

Schedule a given reduce operation.

Parameters

target	The target to generate a schedule for.
op	The operation representing the injective operation.
sch	The schedule to apply this scheduling to
is_idx_reduce	Pass true to schedule a reduce op that returns an index, such as argmax or argmin.

Returns

The schedule given by sch

TraverseAfterReduce()

void tvm::topi::cuda::TraverseAfterReduce	(	const Target &	target,
		Schedule	s,
		Operation	op
	)

Schedule a reduce op, then invoke TraverseBeforeReduce on each of the op's inputs.

Parameters

target	The target to generate a schedule for.
s	The schedule we are building
op	The reduce op

TraverseBeforeReduce()

void tvm::topi::cuda::TraverseBeforeReduce	(	Schedule	s,
		Operation	op
	)

Recursively traverse operator inputs, setting injective inputs to be computed inline.

Parameters

s	The schedule we are building
op	The current op in the traversal