MIGraphX: a compute graph optimizer

This page records my notes about exploring MIGraphX, henceforth MGX.

Overview §

The hierarchy of the model abstraction in MGX is shown as below:

flowchart LR

A[program] --> B[modules]

B --> C[instructions]

C --> D[op]

When reading a ONNX (or other formats) model, MGX will try to parse and represent it with the abstraction. MGX will first try to finalize all operations in a program (i.e., all of its modules and all instructions in the modules), then invoke the compute function defined in each operation to evaluate the model.

Parse and compile an ONNX model §

Compiler §

Pass §

Evaluate and run a compiled program §

After creating a program, users can now call the eval function to

Quick ref: common data structures §

Type erasure §

Many data structures in MGX follows the type erasure design pattern to hide the type information and expose the unified interfaces to outside. Usually, the xxx struct contains the xxx_impl unique pointer with the real data member stored.

Common code snippets §

template <class T, class F, F f> // NOLINT
using manage_ptr = std::unique_ptr<T, manage_deleter<F, f>>;
#define MIGRAPHX_MANAGE_PTR(T, F) \
    migraphx::manage_ptr<std::remove_pointer_t<T>, decltype(&F), &F>

Macro MIGRAPHX_MANAGE_PTR will return a std::unique_ptr with the destructor set.

Shapes and arguments §

struct MIGRAPHX_EXPORT shape
{
    shape(type_t t, std::vector<std::size_t> l);
    shape(type_t t, std::vector<std::size_t> l, std::vector<std::size_t> s);
    // for dynamic tensor shape
    struct MIGRAPHX_EXPORT dynamic_dimension
    {
        std::size_t min = 0;
        std::size_t max = 0;
        std::set<std::size_t> optimals{};
    };
 
    std::size_t elements() const;
    std::size_t bytes() const;
};

shape is described as dimension sizes and optional stride sizes. In additional to fixed shape tensors, MGX also supports shapes with dynamic dimensions (with min and max values alongside the dimension).

struct MIGRAPHX_EXPORT argument : raw_data<argument>
{
    template <class T>
    argument(shape s, T* d)
        : m_shape(std::move(s))
    {
        assign_buffer([d] { return reinterpret_cast<char*>(d); });
    }
 
    template <class T>
    argument(shape s, std::shared_ptr<T> d)
        : m_shape(std::move(s))
    {
        assign_buffer([d] { return reinterpret_cast<char*>(d.get()); });
    }
    
    private:
    void assign_buffer(std::function<char*()> d);
    struct data_t
    {
        std::function<char*()> get = nullptr;
        std::vector<data_t> sub = {};
        data_t share() const;
        static data_t from_args(const std::vector<argument>& args);
    };
    argument(const shape& s, const data_t& d);
    shape m_shape;
    data_t m_data{};
};

argument is the structure to place the substantial data described by shape.

Instructions and op §

using instruction_ref = std::list<instruction>::iterator;
// members in an instruction
struct MIGRAPHX_EXPORT instruction
{
    operation op;
    shape result{};
    std::vector<instruction_ref> output;
    std::vector<instruction_ref> arguments;
    std::vector<module_ref> module_args;
    literal lit;
    bool normalized       = false;
    std::size_t target_id = 0;
};

Modules §

struct MIGRAPHX_EXPORT module
{
    std::unique_ptr<module_impl> impl;
};
 
struct module_impl
{
    // A list is used to keep references to an instruction stable
    std::list<instruction> instructions;
    std::unordered_set<instruction*> instruction_set;
    std::string name;
    uint32_t nparams = 0;
    bool bypass      = false;
};

Programs §

struct MIGRAPHX_EXPORT program
{
    std::unique_ptr<program_impl> impl;
};
 
struct program_impl
{
    // A map is used to keep references to modules of the program
    std::unordered_map<std::string, module> modules;
    std::vector<context> contexts;
    std::vector<target> targets;
};

Target and context §

struct context
{
    context(std::size_t device_id = 0, std::size_t n = value_of(MIGRAPHX_NSTREAMS{}, 1))
        : current_device(std::make_shared<hip_device>(device_id, n)),
          begin_event(create_event()),
          finish_event(create_event())
    {
    }
    
    private:
    // TODO: Make this a vector to support multiple devices
    std::shared_ptr<hip_device> current_device;
    std::vector<shared<hip_event_ptr>> events;
    bool exhaustive_tune = false;
    bool measure_perf    = false;
    // for event perf timing
    shared<hip_event_ptr> start_event = nullptr;
    shared<hip_event_ptr> stop_event  = nullptr;
    // for stream syncronization
    shared<hip_event_ptr> begin_event  = nullptr;
    shared<hip_event_ptr> finish_event = nullptr;
    problem_cache pc{};
};

context is an abstraction for managing HIP device (GPU, stream and event).

struct MIGRAPHX_GPU_EXPORT target
{
    std::string name() const;
    std::vector<pass> get_passes(migraphx::context& gctx, const compile_options& options) const;
    migraphx::context get_context() const;
    argument copy_to(const argument& arg) const;
    argument copy_from(const argument& arg) const;
    argument allocate(const shape& s) const;
};

target rather defines the possible “actions” (both compile time and execution time) taken on the device.

Ranges §

template <class Iterator>
struct iterator_range
{
    Iterator start;
    Iterator last;
 
    Iterator begin() const { return start; }
 
    Iterator end() const { return last; }
};
 
template <class Iterator, MIGRAPHX_REQUIRES(not std::is_integral<Iterator>{})>
iterator_range<Iterator> range(Iterator start, Iterator last)
{
    return {start, last};
}
 
inline iterator_range<iota_iterator> range(std::ptrdiff_t start, std::ptrdiff_t last)
{
    return {{start, {}}, {last, {}}};
}
inline iterator_range<iota_iterator> range(std::ptrdiff_t last) { return range(0, last); }

Use case: for (auto i : range(N)) {}, which will iterative i from 0 to N - 1.

Useful links §