diff --git a/include/triton/Dialect/CMakeLists.txt b/include/triton/Dialect/CMakeLists.txt index 8eea29534..819e542c8 100644 --- a/include/triton/Dialect/CMakeLists.txt +++ b/include/triton/Dialect/CMakeLists.txt @@ -1 +1,2 @@ -add_subdirectory(Triton) \ No newline at end of file +add_subdirectory(Triton) +# add_subdirectory(TritonGPU) diff --git a/include/triton/Dialect/TritonGPU/CMakeLists.txt b/include/triton/Dialect/TritonGPU/CMakeLists.txt new file mode 100644 index 000000000..f33061b2d --- /dev/null +++ b/include/triton/Dialect/TritonGPU/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/include/triton/Dialect/TritonGPU/IR/Dialect.h b/include/triton/Dialect/TritonGPU/IR/Dialect.h new file mode 100644 index 000000000..bd08e1195 --- /dev/null +++ b/include/triton/Dialect/TritonGPU/IR/Dialect.h @@ -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_ diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td new file mode 100644 index 000000000..38e2afbe4 --- /dev/null +++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td @@ -0,0 +1,65 @@ +#ifndef TRITONGPU_ATTRDEFS +#define TRITONGPU_ATTRDEFS + +include "TritonGPUDialect.td" + +class TritonGPU_Attr traits = []> + : AttrDef; + +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 diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td new file mode 100644 index 000000000..bc2adfd94 --- /dev/null +++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td @@ -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 diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td new file mode 100644 index 000000000..934d01f5b --- /dev/null +++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td @@ -0,0 +1,27 @@ +#ifndef TRITONGPU_OPS +#define TRITONGPU_OPS + +include "TritonGPUDialect.td" +// include "" + +def TTG_Op traits = []> : + Op; + +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