21 lines
509 B
C++
21 lines
509 B
C++
|
#ifndef LAYOUT_3D_HPP
|
||
|
|
#define LAYOUT_3D_HPP
|
||
|
|
#include "kernel_operator.h"
|
||
|
|
#include "catlass/catlass.hpp"
|
||
|
|
class Layout3D {
|
||
|
|
int64_t strides[2];
|
||
|
|
public:
|
||
|
|
CATLASS_DEVICE
|
||
|
|
Layout3D() {}
|
||
|
|
CATLASS_DEVICE
|
||
|
|
Layout3D(int64_t stride0, int64_t stride1) {
|
||
|
|
strides[0] = stride0;
|
||
|
|
strides[1] = stride1;
|
||
|
|
}
|
||
|
|
CATLASS_DEVICE
|
||
|
|
int64_t operator() (int64_t dim0, int64_t dim1, int64_t dim2) {
|
||
|
|
return dim0 * strides[0] + dim1 * strides[1] + dim2;
|
||
|
|
}
|
||
|
|
};
|
||
|
|
#endif // LAYOUT_3D_HPP
|