[FRONTEND][BACKEND] Added trans instruction; made flash attention bwd pass work (#943)

2022-12-03 09:58:24 -08:00
parent 4d64589b22
commit 8edfe813a5
12 changed files with 310 additions and 143 deletions
--- a/include/triton/Dialect/Triton/IR/TritonOps.td
+++ b/include/triton/Dialect/Triton/IR/TritonOps.td
@@ -295,6 +295,18 @@ def TT_CatOp : TT_Op<"cat", [NoSideEffect,
    let assemblyFormat = "$lhs `,` $rhs attr-dict `:` functional-type(operands, results)";
 }

+def TT_TransOp : TT_Op<"trans", [NoSideEffect,
+                               SameOperandsAndResultElementType]> {
+
+    let summary = "transpose a tensor";
+
+    let arguments = (ins TT_Tensor:$src);
+
+    let results = (outs TT_Tensor:$result);
+
+    let assemblyFormat = "$src attr-dict `:` functional-type(operands, results)";
+}
+
 //
 // SPMD Ops
 //