From 9cde4e3dc589ab783757b863a747477bc6fafc91 Mon Sep 17 00:00:00 2001
From: Kunwar Grover <groverkss@gmail.com>
Date: Fri, 22 Dec 2023 12:37:15 +0530
Subject: [PATCH] [VectorExt] Add folders for to_simt and to_simd (#15997)

This patch adds folders for forward operand of one operation to result
of another operation if it sees them in a chain.
---
 .../Dialect/VectorExt/IR/VectorExtOps.td      |  2 +
 .../lib/Dialect/VectorExt/IR/VectorExtOps.cpp | 16 ++++++
 .../Dialect/iree_vector_ext/canonicalize.mlir | 49 +++++++++++++++++++
 3 files changed, 67 insertions(+)
 create mode 100644 llvm-external-projects/iree-dialects/test/Dialect/iree_vector_ext/canonicalize.mlir

diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/VectorExt/IR/VectorExtOps.td b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/VectorExt/IR/VectorExtOps.td
index c77adb1c755c..8dff0d7e7a0a 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/VectorExt/IR/VectorExtOps.td
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/VectorExt/IR/VectorExtOps.td
@@ -60,6 +60,7 @@ def IREEVectorExt_ToSIMDOp : IREEVectorExt_PureOp<"to_simd",
   );
   let extraClassDeclaration = [{}];
   let assemblyFormat = "$input attr-dict `:` type($input) `->` type($output)";
+  let hasFolder = 1;
 }
 
 def IREEVectorExt_ToSIMTOp : IREEVectorExt_PureOp<"to_simt",
@@ -78,6 +79,7 @@ def IREEVectorExt_ToSIMTOp : IREEVectorExt_PureOp<"to_simt",
   );
   let extraClassDeclaration = [{}];
   let assemblyFormat = "$input attr-dict `:` type($input) `->` type($output)";
+  let hasFolder = 1;
 }
 
 #endif  // IREE_DIALECT_VECTOREXT_OPS
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/VectorExt/IR/VectorExtOps.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/VectorExt/IR/VectorExtOps.cpp
index 8b7b2be83ab7..49c77b0ba469 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/VectorExt/IR/VectorExtOps.cpp
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/VectorExt/IR/VectorExtOps.cpp
@@ -36,6 +36,22 @@ LogicalResult LayoutConflictResolutionOp::verify() {
   return failure();
 }
 
+// to_simd -> to_simt
+OpFoldResult ToSIMDOp::fold(FoldAdaptor) {
+  if (auto simtOp = getOperand().getDefiningOp<ToSIMTOp>()) {
+    return simtOp.getOperand();
+  }
+  return {};
+}
+
+// to_simt -> to_simd
+OpFoldResult ToSIMTOp::fold(FoldAdaptor) {
+  if (auto simdOp = getOperand().getDefiningOp<ToSIMDOp>()) {
+    return simdOp.getOperand();
+  }
+  return {};
+}
+
 // clang-format off
 #define GET_OP_CLASSES
 #include "iree-dialects/Dialect/VectorExt/IR/VectorExtOps.cpp.inc" // IWYU pragma: keep
diff --git a/llvm-external-projects/iree-dialects/test/Dialect/iree_vector_ext/canonicalize.mlir b/llvm-external-projects/iree-dialects/test/Dialect/iree_vector_ext/canonicalize.mlir
new file mode 100644
index 000000000000..f053c931aca8
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/test/Dialect/iree_vector_ext/canonicalize.mlir
@@ -0,0 +1,49 @@
+// RUN: iree-dialects-opt --canonicalize --split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @to_simt_to_simd_fold
+// CHECK-SAME: (%[[SIMD:.*]]: vector<64x64xf32>) -> vector<64x64xf32>
+func.func @to_simt_to_simd_fold(%simd: vector<64x64xf32>) -> vector<64x64xf32> {
+  // Both to_simt and to_simd should be dce-ed after folding.
+  // CHECK-NOT: iree_vector_ext.to_simt
+  %simt = iree_vector_ext.to_simt %simd : vector<64x64xf32> -> vector<4x4x4xf32>
+  // CHECK-NOT: iree_vector_ext.to_simd
+  %simd_out = iree_vector_ext.to_simd %simt : vector<4x4x4xf32> -> vector<64x64xf32> 
+  // CHECK: return %[[SIMD]]
+  func.return %simd_out : vector<64x64xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @to_simd_to_simt_fold
+// CHECK-SAME: (%[[SIMT:.*]]: vector<4x4x4xf32>) -> vector<4x4x4xf32>
+func.func @to_simd_to_simt_fold(%simt: vector<4x4x4xf32>) -> vector<4x4x4xf32> {
+  // Both to_simt and to_simd should be dce-ed after folding.
+  // CHECK-NOT: iree_vector_ext.to_simt
+  %simd = iree_vector_ext.to_simd %simt : vector<4x4x4xf32> -> vector<64x64xf32>
+  // CHECK-NOT: iree_vector_ext.to_simd
+  %simt_out = iree_vector_ext.to_simt %simd : vector<64x64xf32> -> vector<4x4x4xf32> 
+  // CHECK: return %[[SIMT]]
+  func.return %simt_out : vector<4x4x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @to_simd_to_simt_multi_use
+// CHECK-SAME: (%[[SIMT:.*]]: vector<4x4x4xf32>)
+func.func @to_simd_to_simt_multi_use(%simt: vector<4x4x4xf32>) -> (vector<4x4x4xf16>, vector<64x64xf32>) {
+  // The to_simd operation should not be dce-ed after folding because it is returned.
+  // CHECK: %[[SIMD:.*]] = iree_vector_ext.to_simd %[[SIMT]] : vector<4x4x4xf32> -> vector<64x64xf32>
+  %simd = iree_vector_ext.to_simd %simt : vector<4x4x4xf32> -> vector<64x64xf32>
+  // The to_simt operation should be dce-ed after folding.
+  // CHECK-NOT: iree_vector_ext.to_simt
+  %simt_out = iree_vector_ext.to_simt %simd : vector<64x64xf32> -> vector<4x4x4xf32> 
+
+  // Check if the folding happened correctly.
+  // CHECK: %[[TRUNCED:.*]] = arith.truncf %[[SIMT]]
+  %trunced = arith.truncf %simt_out : vector<4x4x4xf32> to vector<4x4x4xf16>
+
+  // CHECK: return %[[TRUNCED]], %[[SIMD]]
+  func.return %trunced, %simd : vector<4x4x4xf16>, vector<64x64xf32>
+}
+
+// -----