DataLayout Structure

Development pre-RFC
#1

Given https://github.com/dmlc/tvm/pull/1103, we are starting to have a general sense of data layout definition. I want to follow up on this since I think this is a design decision that we need to make. This RFC is used to discuss this problem

In short, I think it may be helpful to introduce a formal DataLayout to TVM’s Node system (possibly in buffer.h ), so that TOPI and upstream compiler can benefit from this. First of all, I like the stying convention NCHW, and NCHW16c convention introduced by @yzhliu as in https://github.com/dmlc/nnvm/pull/447 which is compatible with common deep network notation and extends to folded data layout.

Besides the data layout string format, it might be helpful to design a useful internal data structure to handle layout.

Formally a data layout is a function mapping(possibly bijective) from the existing index (i0, i1, …, im) to the final stored index (j0, j1, … jm).

Abstract Interface of Layout

We essentially need a bijective mapping.

struct DataLayoutNode {
   // Final shape of the underlying array, given the shape of the normal layout
   virtual Array<Expr> ForwardShape(Array<Expr> shape);
   // Final index of the underlying array, given the normal layout.
   virtual Array<Expr> ForwardIndex(Array<Expr> index);
   // Given store index, recover the original representation space index.
   virtual Array<Expr> BackwardIndex(Array<Expr> store_index);
};

Options for concrete data types

There are two possible options on how we can use data structure to store this.

Option 1: Reuse the IterVarRelation in the schedule.

struct DataLayoutNode {
   // The original axis
   Array<IterVar> orig_axis;
   // The final axis in the store order
   Array<IterVar> store_axis;
   // The relation of final store axis to orig_axis, so far it can only be split.
   Array<IterVarRelation> relations;
};

Option 2: Direct Record the mapping

Quite general, but need rules to check the bijective-ness.

struct DataLayoutNode {
   // The original axis, with symbolic shape
   Array<IterVar> orig_axis;
   // The shape of the stored array
   Array<Expr> shape;
   // expression of each location, on how original location can be mapped
   // to the store location, example 
   // [i0 / 16, i1, i0 % 16]
   Array<Expr> store_axis;
};

Canonicalization

Since layout is more of a canonical type information, it would be great we will get a single instance of canonical layout when-ever we create the same value. One possible approach is to take the abstract interface backed by a simple convention, before we extrapolate.

#2

In option 1,

// The relation of final store axis to orig_axis, so far it can only be split.

Do you mean the orig_axis=NCHW and the store_axis=NCHW16c ? I think it can also be orig_axis=NCHW16c and the store_axis=NCHW, which implies a fuse relation.

and it can also be NCHW16c <-> NCHW8c, can IterVarRelation describe it? I think that will be a fuse+split.

Minor question, ‘axis’ or ‘dim’?

And would you like to introduce the name convention to tvm as well? I think the answer is no, so far we can keep layout name convention in nnvm level.

#3

orig_axis=NCHW and the store_axis=NCHW16c . You are right that NCHW16c <-> NCHW8 is fuse+ split if we really need such mapping. We are using axis to be consistent with current convention.

Since the layout convention might happen in TOPI, I think it is not too bad to trickle up to support string parsing of such layout in tvm, if we make DataLayout a canonical Node

#4

I start to make an implementation. Would you mind explain a bit more about this approach? @tqchen