[BACKEND] Fold transpose(splat_const) (#5259)

Add folding for a transpose of a splat constant. --------- Co-authored-by: peterbell10 <[email protected]>
triton-lang · Nov 26, 2024 · deee78f · deee78f
1 parent 22e212b
commit deee78f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/lib/Dialect/Triton/IR/Ops.cpp b/lib/Dialect/Triton/IR/Ops.cpp
@@ -199,6 +199,11 @@ OpFoldResult TransOp::fold(FoldAdaptor adaptor) {
     return getResult();
   }
 
+  // Eliminate splat constant transpose ops.
+  if (auto attr =
+          llvm::dyn_cast_if_present<SplatElementsAttr>(adaptor.getSrc()))
+    return attr.reshape(getType());
+
   return {};
 }
 

diff --git a/test/Triton/canonicalize.mlir b/test/Triton/canonicalize.mlir
@@ -173,3 +173,14 @@ tt.func @fold_broadcast_constant_pattern(%cst : f32) -> tensor<8x2xf32> {
     // CHECK-NEXT: tt.return %[[cst]] : tensor<8x2xf32>
     tt.return %bst_out : tensor<8x2xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @fold_transpose_constant
+tt.func @fold_transpose_constant() -> tensor<128x16xf32> {
+    // CHECK: %[[cst:.*]] = arith.constant dense<1.000000e+00> : tensor<128x16xf32>
+    %cst = arith.constant dense<1.0> : tensor<16x128xf32>
+    %r = tt.trans %cst {order = array<i32: 1, 0>} : tensor<16x128xf32> -> tensor<128x16xf32>
+    // CHECK-NEXT: tt.return %[[cst]] : tensor<128x16xf32>
+    tt.return %r : tensor<128x16xf32>
+}