49 lines
1.3 KiB
C++
49 lines
1.3 KiB
C++
#ifndef TRITON_DIALECT_TRITONGPU_IR_DIALECT_H_
|
|
#define TRITON_DIALECT_TRITONGPU_IR_DIALECT_H_
|
|
|
|
#include "mlir/Dialect/GPU/GPUDialect.h"
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/IR/BuiltinOps.h"
|
|
#include "mlir/IR/Dialect.h"
|
|
|
|
// TritonGPU depends on Triton
|
|
#include "triton/Dialect/Triton/IR/Dialect.h"
|
|
|
|
#include "triton/Dialect/TritonGPU/IR/Dialect.h.inc"
|
|
#include "triton/Dialect/TritonGPU/IR/Traits.h"
|
|
|
|
#define GET_ATTRDEF_CLASSES
|
|
#include "triton/Dialect/Triton/IR/AttrInterfaces.h.inc"
|
|
#include "triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.h.inc"
|
|
|
|
#define GET_OP_CLASSES
|
|
#include "triton/Dialect/TritonGPU/IR/Ops.h.inc"
|
|
|
|
namespace mlir {
|
|
namespace triton {
|
|
namespace gpu {
|
|
|
|
unsigned getElemsPerThread(Type type);
|
|
|
|
SmallVector<unsigned> getThreadsPerWarp(const Attribute &layout);
|
|
|
|
SmallVector<unsigned> getWarpsPerCTA(const Attribute &layout);
|
|
|
|
SmallVector<unsigned> getSizePerThread(const Attribute &layout);
|
|
|
|
SmallVector<unsigned> getContigPerThread(Attribute layout);
|
|
|
|
SmallVector<unsigned> getThreadsPerCTA(const Attribute &layout);
|
|
|
|
SmallVector<unsigned> getShapePerCTA(const Attribute &layout);
|
|
|
|
SmallVector<unsigned> getOrder(const Attribute &layout);
|
|
|
|
bool isaDistributedLayout(const Attribute &layout);
|
|
|
|
} // namespace gpu
|
|
} // namespace triton
|
|
} // namespace mlir
|
|
|
|
#endif // TRITON_DIALECT_TRITONGPU_IR_DIALECT_H_
|