[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.
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 c77adb1..8dff0d7 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 @@
);
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 @@
);
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 8b7b2be..49c77b0 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 @@
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 0000000..f053c93
--- /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>
+}
+
+// -----
