Some progress on TritonGPU
This commit is contained in:
@@ -1 +1,2 @@
|
|||||||
add_subdirectory(Triton)
|
add_subdirectory(Triton)
|
||||||
|
# add_subdirectory(TritonGPU)
|
||||||
|
1
include/triton/Dialect/TritonGPU/CMakeLists.txt
Normal file
1
include/triton/Dialect/TritonGPU/CMakeLists.txt
Normal file
@@ -0,0 +1 @@
|
|||||||
|
add_subdirectory(IR)
|
5
include/triton/Dialect/TritonGPU/IR/Dialect.h
Normal file
5
include/triton/Dialect/TritonGPU/IR/Dialect.h
Normal file
@@ -0,0 +1,5 @@
|
|||||||
|
#ifndef TRITON_DIALECT_TRITONGPU_IR_DIALECT_H_
|
||||||
|
#define TRITON_DIALECT_TRITONGPU_IR_DIALECT_H_
|
||||||
|
|
||||||
|
|
||||||
|
#endif // TRITON_DIALECT_TRITONGPU_IR_DIALECT_H_
|
65
include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td
Normal file
65
include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td
Normal file
@@ -0,0 +1,65 @@
|
|||||||
|
#ifndef TRITONGPU_ATTRDEFS
|
||||||
|
#define TRITONGPU_ATTRDEFS
|
||||||
|
|
||||||
|
include "TritonGPUDialect.td"
|
||||||
|
|
||||||
|
class TritonGPU_Attr<string name, list<Trait> traits = []>
|
||||||
|
: AttrDef<TritonGPU_Dialect, name, traits>;
|
||||||
|
|
||||||
|
def TritonGPUSharedEncodingAttr : TritonGPU_Attr<"TritonGPUSharedEncoding"> {
|
||||||
|
let mnemonic = "shared (memory) encoding";
|
||||||
|
|
||||||
|
let description = [{
|
||||||
|
Example:
|
||||||
|
|
||||||
|
```mlir
|
||||||
|
#SMEM = #triton_gpu.encoding<{
|
||||||
|
vec = 8,
|
||||||
|
perPhase = 8,
|
||||||
|
maxPhase = 1
|
||||||
|
}>
|
||||||
|
```
|
||||||
|
}];
|
||||||
|
|
||||||
|
let parameters = (
|
||||||
|
ins
|
||||||
|
// swizzle info
|
||||||
|
"unsigned":$vec, "unsigned":$perPhase, "unsigned":$maxPhase
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
def TritonGPUCoalescedEncodingAttr : TritonGPU_Attr<"TritonGPUCoalescedEncoding"> {
|
||||||
|
let mnemonic = "coalesced encoding";
|
||||||
|
|
||||||
|
let description = [{}];
|
||||||
|
|
||||||
|
let parameters = (
|
||||||
|
ins
|
||||||
|
ArrayRefParameter<"unsigned">:$nanoTileSize,
|
||||||
|
ArrayRefParameter<"unsigned">:$microTileSize,
|
||||||
|
ArrayRefParameter<"unsigned">:$shapePerCTA,
|
||||||
|
ArrayRefParameter<"unsigned">:$reptitions
|
||||||
|
);
|
||||||
|
|
||||||
|
// let genVerifyDecl = 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
def TritonGPUMmaEncodingAttr : TritonGPU_Attr<"TritonGPUMmaEncoding"> {
|
||||||
|
let mnemonic = "mma encoding";
|
||||||
|
|
||||||
|
let description = [{}];
|
||||||
|
|
||||||
|
let parameters = (
|
||||||
|
ins
|
||||||
|
ArrayRefParameter<"unsigned">:$fragmentPerWarp,
|
||||||
|
ArrayRefParameter<"unsigned">:$shapePerWarp,
|
||||||
|
ArrayRefParameter<"unsigned">:$warpPerTile,
|
||||||
|
ArrayRefParameter<"unsigned">:$shapePerTile,
|
||||||
|
ArrayRefParameter<"unsigned">:$reptitions,
|
||||||
|
ArrayRefParameter<"unsigned">:$contigPerThread
|
||||||
|
);
|
||||||
|
|
||||||
|
// let genVerifyDecl = 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
#endif
|
20
include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td
Normal file
20
include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td
Normal file
@@ -0,0 +1,20 @@
|
|||||||
|
#ifndef TRITONGPU_DIALECT
|
||||||
|
#define TRITONGPU_DIALECT
|
||||||
|
|
||||||
|
include "mlir/IR/OpBase.td"
|
||||||
|
|
||||||
|
def TritonGPU_Dialect : Dialect {
|
||||||
|
let name = "triton_gpu";
|
||||||
|
|
||||||
|
let cppNamespace = "::mlir::triton::gpu";
|
||||||
|
|
||||||
|
let description = [{
|
||||||
|
Triton GPU Dialect.
|
||||||
|
}];
|
||||||
|
|
||||||
|
let dependentDialects = [
|
||||||
|
"triton::TritonDialect"
|
||||||
|
];
|
||||||
|
}
|
||||||
|
|
||||||
|
#endif
|
27
include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
Normal file
27
include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
Normal file
@@ -0,0 +1,27 @@
|
|||||||
|
#ifndef TRITONGPU_OPS
|
||||||
|
#define TRITONGPU_OPS
|
||||||
|
|
||||||
|
include "TritonGPUDialect.td"
|
||||||
|
// include ""
|
||||||
|
|
||||||
|
def TTG_Op<string mnemonic, list<Trait> traits = []> :
|
||||||
|
Op<TritonGPU_Dialect, mnemonic, traits>;
|
||||||
|
|
||||||
|
def TTG_ConvertLayoutOp : TTG_Op<"convert_layout",
|
||||||
|
[NoSideEffect, SameOperandsAndResultType]> {
|
||||||
|
let summary = "convert layout";
|
||||||
|
|
||||||
|
let arguments = (ins TT_Tensor:$src);
|
||||||
|
|
||||||
|
let results = (outs TT_Tensor:$result);
|
||||||
|
}
|
||||||
|
|
||||||
|
def TTG_AsyncWaitOp : TTG_Op<"async_wait"> {
|
||||||
|
let summary = "async wait";
|
||||||
|
|
||||||
|
let arguments = (ins I32Attr:$num);
|
||||||
|
}
|
||||||
|
|
||||||
|
// def TTG_CopyAsyncOp : TTG_Op<"copy_async"> {}
|
||||||
|
|
||||||
|
#endif
|
Reference in New Issue
Block a user