class habana::custom_op::HabanaCustomOpDescriptor

Overview

The descriptor contains all the necessary information to define custom TPC op run on PyTorch HPU:

#define REGISTER_CUSTOM_OP_ATTRIBUTES(                                    \
    schema_name, guid, input_desc, output_desc, param_func)               \
{                                                                         \
    habana::custom_op::NodeDesc node_desc{guid, schema_name, param_func}; \
    habana::custom_op::HabanaCustomOpDescriptor op_desc{                  \
        node_desc, inputs_desc, outputs_desc};                            \
    habana::custom_op::registerKernel(op_desc);                           \
}

#define HPU_PARAMS_STUB(structname) \
size = sizeof(structname);    \
auto params = std::make_shared<structname>()

/**
* @brief Descriptor for custom op containing all necessary information to
* define user HPU TPC kernel.

* The user is responsible to define all node TPC info within NodeDesc, all
* inputs/outputs info within vectors of InputDesc/OutputDesc. User needs
* to register descriptor with the macro REGISTER_CUSTOM_OP_ATTRIBUTES. User
* will retrieve his descriptor from registry, and call execute with inputs.
*/
class HabanaCustomOpDescriptor {
public:
HabanaCustomOpDescriptor(
    NodeDesc node_desc,
    const std::vector<InputDesc>& inputs,
    const std::vector<OutputDesc>& outputs)
    : node_desc_(node_desc), inputs_(inputs), outputs_(outputs) {}
HabanaCustomOpDescriptor() {}

/**
* @brief Actual call by user C++ to op
*
* @param inputs all values by order to op execution
* @return std::vector<at::Tensor> vector of op results. If single output
* vector with size()==1.
*/
std::vector<at::Tensor> execute(const std::vector<c10::IValue>& inputs);

/**
* @brief Get the Custom Op Descriptor object
*
* @param op schema registration name which used in
* REGISTER_CUSTOM_OP_ATTRIBUTES
* @return const HabanaCustomOpDescriptor
*/
static const HabanaCustomOpDescriptor getCustomOpDescriptor(std::string op);

std::string getSchemaName() const;
std::string getGuid() const;
unsigned getInputsSize() const;
unsigned getOutputsSize() const;
const std::vector<InputDesc>& getInputs() const;
const std::vector<OutputDesc>& getOutputs() const;
bool hasUserParamsFunc() const;
const allocate_user_params_func& getUserParamsAllocFunc() const;
bool hasOutputShapeFunc(unsigned index) const;
const compute_output_shape_function& getOutputShapeFunc(unsigned index) const;

Detailed Documentation

The descriptor contains all the necessary information to define custom TPC Op running on PyTorch HPU.

Main Methods For User

#define REGISTER_CUSTOM_OP_ATTRIBUTES(                                    \
    schema_name, guid, input_desc, output_desc, param_func)               \
{                                                                         \
    habana::custom_op::NodeDesc node_desc{guid, schema_name, param_func}; \
    habana::custom_op::HabanaCustomOpDescriptor op_desc{                  \
        node_desc, inputs_desc, outputs_desc};                            \
    habana::custom_op::registerKernel(op_desc);                           \
}

• Use this macro to construct HabanaCustomOpDescriptor and register as an HPU op:

static const HabanaCustomOpDescriptor getCustomOpDescriptor(std::string op);

• Getting HabanaCustomOpDescriptor from registry used with the macro REGISTER_CUSTOM_OP_ATTRIBUTES:

std::vector<at::Tensor> execute(const std::vector<c10::IValue>& inputs);

• The actual execution call by the user with the current Op instance inputs:

#define HPU_PARAMS_STUB(structname) \
size = sizeof(structname);    \
auto params = std::make_shared<structname>()

Note

Use this macro to allocate user param in callback function if needed.