mini_jit::ir

struct Dimension
#include <Dimension.h>

The Dimension struct represents a dimension in a tensor operation. It contains information about the type of dimension (M, N, K), execution type (Prim, Seq, Shared), size, and strides for the input and output tensors.

Public Functions

inline Dimension(dim_t type, exec_t exec_type, int64_t size, int64_t stride_in0, int64_t stride_in1, int64_t stride_out)

Construct a new Dimension object.

Parameters:

  • type – Type of the dimension (M, N, K).

  • exec_type – Execution type (Prim, Seq, Shared, …).

  • size – Size of the dimension.

  • stride_in0 – Stride in the first input tensor.

  • stride_in1 – Stride in the second input tensor.

  • stride_out – Stride in the output tensor.

Public Members

dim_t type = dim_t::m

Type of the dimension (M, N, K)

exec_t exec_type = exec_t::undefined

Execution type (Prim, Seq, Shared, …)

int64_t size = 0

Dimension size.

int64_t stride_in0 = 0

Stride in the first input tensor.

int64_t stride_in1 = 0

Stride in the second input tensor.

int64_t stride_out = 0

Stride in the output tensor.

class IRConverter
#include <IRConverter.h>

The IRConverter class provides methods to convert between configuration parameters and IR Dimension representation.

Public Functions

IRConverter() = delete

Deleted constructor to prevent instantiation of the static IRConverter class.

Public Static Functions

static void convertConfigToDimensions(std::span<const dim_t> i_dim_types, std::span<const exec_t> i_exec_types, std::span<const int64_t> i_dim_sizes, std::span<const int64_t> i_strides_in0, std::span<const int64_t> i_strides_in1, std::span<const int64_t> i_strides_out, std::vector<Dimension> &o_dimensions)

Convert configuration parameters to a vector of Dimension objects.

Parameters:

  • i_dim_types – A span of dimension types (M, N, K).

  • i_exec_types – A span of execution types (Prim, Seq, Shared).

  • i_dim_sizes – A span of dimension sizes.

  • i_strides_in0 – A span of strides for the first input tensor.

  • i_strides_in1 – A span of strides for the second input tensor.

  • i_strides_out – A span of strides for the output tensor.

  • o_dimensions – A vector to store the converted Dimension objects.

static void convertDimensionsToConfig(const std::vector<Dimension> &i_dimensions, std::vector<dim_t> &o_dim_types, std::vector<exec_t> &o_exec_types, std::vector<int64_t> &o_dim_sizes, std::vector<int64_t> &o_strides_in0, std::vector<int64_t> &o_strides_in1, std::vector<int64_t> &o_strides_out)

Convert a vector of Dimension objects to configuration parameters.

Parameters:

  • i_dimensions – A vector of Dimension objects to be converted.

  • o_dim_types – A vector to store the dimension types.

  • o_exec_types – A vector to store the execution types.

  • o_dim_sizes – A vector to store the dimension sizes.

  • o_strides_in0 – A vector to store the strides for the first input tensor.

  • o_strides_in1 – A vector to store the strides for the second input tensor.

  • o_strides_out – A vector to store the strides for the output tensor.

class Optimizer
#include <Optimizer.h>

The Optimizer class provides methods to optimize tensor operations by adjusting dimensions, splitting large dimensions, and creating shared loops.

It is recommened to only call the optimize method, which will internally call the other methods in the correct order.

Public Functions

Optimizer() = delete

Deleted constructor to prevent instantiation of the static Optimizer class.

Public Static Functions

static void optimize(std::vector<Dimension> &dimensions, int64_t thread_target, int64_t max_kernel_size, int64_t min_kernel_size)

Optimize the dimensions of a tensor operation.

Parameters:

  • dimensions – A vector of dimensions to be optimized.

  • thread_target – The target number of threads for optimization.

  • max_kernel_size – The maximum size of a kernel dimension

  • min_kernel_size – The minimum size of a kernel dimension

static void optimize(std::vector<dim_t> &dim_types, std::vector<exec_t> &exec_types, std::vector<int64_t> &dim_sizes, std::vector<int64_t> &strides_in0, std::vector<int64_t> &strides_in1, std::vector<int64_t> &strides_out, int64_t thread_target, int64_t max_kernel_size, int64_t min_kernel_size)

Optimize the dimensions of a tensor operation.

Parameters:

  • dim_types – A vector of dimension types (M, N, K).

  • exec_types – A vector of execution types (Prim, Seq, Shared).

  • dim_sizes – A vector of dimension sizes.

  • strides_in0 – A vector of strides for the first input tensor.

  • strides_in1 – A vector of strides for the second input tensor.

  • strides_out – A vector of strides for the output tensor.

  • thread_target – The target number of threads for optimization.

  • max_kernel_size – The maximum size of a kernel dimension

  • min_kernel_size – The minimum size of a kernel dimension

static void identifyPrimitives(std::vector<Dimension> &dimensions)

Identify primitive dimensions in the tensor operation and adjust their order.

Parameters:

dimensions – A vector of dimensions to be processed.

static void splitDimensions(std::vector<Dimension> &dimensions, int64_t max_kernel_size, int64_t min_kernel_size)

Split large dimensions into smaller ones.

Parameters:

  • dimensions – A vector of dimensions to be processed.

  • max_kernel_size – The maximum size allowed for a kernel dimension.

  • min_kernel_size – The minimum size allowed for a kernel dimension.

static void fuseDimensions(std::vector<Dimension> &dimensions, int64_t min_kernel_size)

Fuse small dimensions into larger dimensions.

Parameters:

  • dimensions – A vector of dimensions to be processed.

  • min_kernel_size – The minimum size for a kernel dimension to be considered for fusion.

static void createSharedLoops(std::vector<Dimension> &dimensions, int64_t thread_target)

Turn sequential dimensions into shared dimensions.

Parameters:

  • dimensions – A vector of dimensions to be processed.

  • thread_target – The target number of threads for optimization.