tvm::topi::rocm Namespace Reference

Functions
tvm::te::Tensor	dense_rocm (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 rocm backend. More...
Schedule	schedule_dense (const Target &target, const Array< Tensor > &outs)
	Create a rocm 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 rocm schedule for the given output tensors. More...
Schedule	schedule_pool (const Target &target, const Array< Tensor > &outs)
	Create a rocm schedule for pool. More...
Schedule	schedule_global_pool (const Target &target, const Array< Tensor > &outs)
	Create a rocm schedule for global_pool. More...
Schedule	schedule_reduce (const Target &target, Array< Tensor > outs)
	Create a rocm schedule for a reduce operation. More...
Schedule	schedule_softmax (const Target &target, const Array< Tensor > &outs)
	Create a rocm schedule for the given softmax output tensors. More...

Function Documentation

dense_rocm()

tvm::te::Tensor tvm::topi::rocm::dense_rocm ( 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 rocm 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::rocm::schedule_dense ( const Target &  target, const Array< Tensor > &  outs  )

inline

Create a rocm 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::rocm::schedule_global_pool ( const Target &  target, const Array< Tensor > &  outs  )

inline

Create a rocm 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::rocm::schedule_injective ( const Target &  target, const Array< Tensor > &  outs  )

inline

Create a rocm 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::rocm::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::rocm::schedule_pool ( const Target &  target, const Array< Tensor > &  outs  )

inline

Create a rocm 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::rocm::schedule_reduce	(	const Target &	target,
		Array< Tensor >	outs
	)

Create a rocm 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::rocm::schedule_softmax ( const Target &  target, const Array< Tensor > &  outs  )

inline

Create a rocm 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.