[mlir][vector] Implement InferTypeOpInterface on vector.to_elements

[matthias-springer]

Just for convenience. This auto-generates an additional builder that infers the result type.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

Just for convenience. This auto-generates an additional builder that infers the result type.

---

Full diff: https://github.com/llvm/llvm-project/pull/153172.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Vector/IR/VectorOps.td (+1-1)
• (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+10)

diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index b3b8afdd8b4c1..30c1d97ba58f1 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -754,7 +754,7 @@ def Vector_FMAOp :
 }
 
 def Vector_ToElementsOp : Vector_Op<"to_elements", [
-    Pure,
+    InferTypeOpAdaptor, Pure,
     ShapedTypeMatchesElementCountAndTypes<"source", "elements">]> {
   let summary = "operation that decomposes a vector into all its scalar elements";
   let description = [{
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index cb4783d26a114..74e48b59b6460 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -2402,6 +2402,16 @@ LogicalResult ToElementsOp::fold(FoldAdaptor adaptor,
   return foldToElementsFromElements(*this, results);
 }
 
+LogicalResult
+ToElementsOp::inferReturnTypes(MLIRContext *ctx, std::optional<Location> loc,
+                               ToElementsOp::Adaptor adaptor,
+                               SmallVectorImpl<Type> &inferredReturnTypes) {
+  auto vecType = cast<VectorType>(adaptor.getSource().getType());
+  Type elType = vecType.getElementType();
+  inferredReturnTypes.append(vecType.getNumElements(), elType);
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // FromElementsOp
 //===----------------------------------------------------------------------===//

[llvmbot]

@llvm/pr-subscribers-mlir-vector

Author: Matthias Springer (matthias-springer)

Changes

Just for convenience. This auto-generates an additional builder that infers the result type.

---

Full diff: https://github.com/llvm/llvm-project/pull/153172.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Vector/IR/VectorOps.td (+1-1)
• (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+10)

diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index b3b8afdd8b4c1..30c1d97ba58f1 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -754,7 +754,7 @@ def Vector_FMAOp :
 }
 
 def Vector_ToElementsOp : Vector_Op<"to_elements", [
-    Pure,
+    InferTypeOpAdaptor, Pure,
     ShapedTypeMatchesElementCountAndTypes<"source", "elements">]> {
   let summary = "operation that decomposes a vector into all its scalar elements";
   let description = [{
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index cb4783d26a114..74e48b59b6460 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -2402,6 +2402,16 @@ LogicalResult ToElementsOp::fold(FoldAdaptor adaptor,
   return foldToElementsFromElements(*this, results);
 }
 
+LogicalResult
+ToElementsOp::inferReturnTypes(MLIRContext *ctx, std::optional<Location> loc,
+                               ToElementsOp::Adaptor adaptor,
+                               SmallVectorImpl<Type> &inferredReturnTypes) {
+  auto vecType = cast<VectorType>(adaptor.getSource().getType());
+  Type elType = vecType.getElementType();
+  inferredReturnTypes.append(vecType.getNumElements(), elType);
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // FromElementsOp
 //===----------------------------------------------------------------------===//