[AMD] Added a canonicalizer to ConcatOp (#7273)

Added a canonicalization pattern to `concatOp`. The pattern removes a
`concatOp` if it has a single input operand.
This scenario can potentially happen as a result of ops refinement. A
corresponding lit-test is included

Author: ravil-mobile

test/TritonGPU/amd/amd-canonicalize-extract-slice.mlir

@@ -22,3 +22,16 @@ module attributes {"ttg.compute-capability" = 0 : i32, "ttg.num-ctas" = 1 : i32,
     tt.return %2 : tensor<32x64xf32, #blocked>
   }
 }
+
+// -----
+
+#blocked = #ttg.blocked<{sizePerThread = [1, 8], threadsPerWarp = [8, 8], warpsPerCTA = [4, 1], order = [1, 0]}>
+module attributes {"ttg.compute-capability" = 0 : i32, "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 64 : i32} {
+  tt.func @canonicalize_singleton_concat(%arg0: tensor<128x128xf32, #blocked>) -> tensor<128x128xf32, #blocked> {
+    // CHECK-LABEL: tt.func @canonicalize_singleton_concat
+
+    %1 = amdgpu.concat %arg0: tensor<128x128xf32, #blocked> -> tensor<128x128xf32, #blocked>
+    // CHECK: tt.return %arg0 : tensor<128x128xf32, #blocked>
+    tt.return %1 : tensor<128x128xf32, #blocked>
+  }
+}

third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOps.td

@@ -234,6 +234,7 @@ def ConcatOp : TT_AMDGPU_Op<"concat", [Pure]> {
   }];
 
   let hasVerifier = 1;
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//

third_party/amd/lib/Dialect/TritonAMDGPU/IR/Dialect.cpp

@@ -446,4 +446,23 @@ LogicalResult ConcatOp::verify() {
 
   return success();
 }
+
+// This pattern removes a concatOp if it has a single input operand.
+// This scenario can potentially happen as a result of ops refinement.
+mlir::LogicalResult foldConcatOpFromSingleSource(amdgpu::ConcatOp op,
+                                                 PatternRewriter &rewriter) {
+  auto sources = op.getSources();
+  if (sources.size() == 1) {
+    auto source = sources.front();
+    auto result = op.getResult();
+    result.replaceAllUsesWith(source);
+    return success();
+  }
+  return failure();
+}
+
+void ConcatOp::getCanonicalizationPatterns(mlir::RewritePatternSet &patterns,
+                                           mlir::MLIRContext *context) {
+  patterns.add(foldConcatOpFromSingleSource);
+}
 } // namespace mlir::triton::amdgpu