tlc-pack / relax

Apache License 2.0
193 stars 58 forks source link

[Op][Layout] Add layout_transform operator in Relax. #403

Closed psrivas2 closed 1 year ago

psrivas2 commented 1 year ago

Adds layout_transform operator in Relax as part of Relax layout planning.

layout_transform takes an input tensor x and an attribute index_map of tir::IndexMap type. It also has an optional pad_value attribute. It transforms x as per the index_map attribute and returns the transformed tensor.

This operator for now allows layout transformations that introduce implicit padding. For example, transforming a tensor of shape (10,) using the lambda i: (i//4, i%4). The output shape will be (3, 4) with two elements padded. The optional pad_value is used to pad if specified, otherwise the compiler is free to choose any value to pad.

psrivas2 commented 1 year ago

cc @masahi @sunggg @YuchenJin

psrivas2 commented 1 year ago

Thanks @MasterJH5574 @masahi @tqchen and @YuchenJin for review!