[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.

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

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

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>
+}
+
+// -----