The problem specification for timeloop defines the tensor workload (e.g., a CNN layer, an MLP layer, or a GEMM kernel) for the tool to model. It starts with the root key problem and uses two second-level keywords: shape and instance to specify the two fundamental properties of a tensor workload.
# problem.yaml
problem:
shape:
...
instance:
...
A problem shape defines the operation space and data spaces of a tensor workload. The way a Timeloop user specifies the shape of a tensor workload follows closely with the Einstein summation notation (einsum). First, the specification requires a name
for the problem shape. Second, the specification requires a list of dimensions
that defines the compute (operation) space of the tensor workload. Finally, the specification requires the projection relation between the compute space and the data space (input and output tensors) using data-spaces
key words. For example, the following specification defines the shape of a 1D convolution tensor workload:
problem:
shape:
name: Conv1D
dimensions: [ R, P ]
data-spaces:
- name: Weights
projection:
- [ [R] ]
- name: Inputs
projection:
- [ [P], [R] ]
- name: Outputs
projection:
- [ [P] ]
read-write: True
instance:
...
Note that the output tensor is defined as read-write: True
, while the other input tensors are read-only with the read-write
keyword unspecified.
With this problem shape specification, Timeloop is able to model tensor workloads that can be represnted with an einsum. In practice, because many tensor workloads share the same shape (e.g, Conv2D layers), Timeloop supports reading an existing problem shape defined in another file, for example the ones under the problem-shapes
directory.
problem:
shape: cnn-layer
...
Optionally, in the problem shape specification, the user can also specify coefficients
to define the projection relation in data-spaces
in a sum of product (SOP) form. These optional coefficients are useful when the user needs to specify a complex projection relationship, such as in convolutions with strides and dilations. For example, the following problem shape specification defines a 2D convolution with strides and dilations:
shape:
name: "CNN-Layer"
dimensions: [ R, S, P, Q, C, K, N ]
coefficients:
- name: Wstride
default: 1
- name: Hstride
default: 1
- name: Wdilation
default: 1
- name: Hdilation
default: 1
data-spaces:
- name: Weights
projection:
- [ [C] ]
- [ [K] ]
- [ [R] ]
- [ [S] ]
- name: Inputs
projection:
- [ [N] ]
- [ [C] ]
# SOP form: R*Wdilation + P*Wstride
- [ [R, Wdilation], [P, Wstride] ]
# SOP form: S*Hdilation + Q*Hstride
- [ [S, Hdilation], [Q, Hstride] ]
- name: Outputs
projection:
- [ [N] ]
- [ [K] ]
- [ [Q] ]
- [ [P] ]
read-write: True
In summary, the keys under shapes
are:
name
: a stringdimension
: a list of unique, single-character dimension namesdata-spaces
: a list of tensors, each defined with name
, projection
as a list of lists with dimension variables, and read-write
.With the problem shape defined, the problem instance specification determines the actual bound of each tensor dimension, and the constant coeffient (e.g, strides) used in the problem shape specification. For example, for the 1D convolution example above, an problem instance can be:
problem:
shape:
...
instance:
R: 3
P: 16
Using the provided problem shape that defines 2D convolution with strides and dilations (problem-shapes/cnn-layer.yaml
), an user can specify a CNN convolution layer (Conv2D), such as AlexNet layer1 with batch size 1, as the following:
problem:
shape: cnn-layer
instance:
C: 3
Hdilation: 1
Hstride: 4
K: 96
N: 1
P: 55
Q: 55
R: 11
S: 11
Wdilation: 1
Wstride: 4