[DAPHNE-daphne-eu#637] Missing type inference of eigen()

This commit adds type inference to mlir::daphne::EigenOp Closes daphne-eu#637
ldirry · Nov 6, 2023 · e50c5e7 · e50c5e7
1 parent b41fa58
commit e50c5e7
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/src/ir/daphneir/DaphneInferTypesOpInterface.cpp b/src/ir/daphneir/DaphneInferTypesOpInterface.cpp
@@ -192,6 +192,11 @@ std::vector<Type> daphne::RandMatrixOp::inferTypes() {
     return {daphne::MatrixType::get(getContext(), elTy)};
 }
 
+std::vector<Type> daphne::EigenOp::inferTypes() {
+    auto evMatType = getArg().getType().dyn_cast<daphne::MatrixType>();
+    return  {evMatType.withSameElementType(), evMatType};
+}
+
 std::vector<Type> daphne::GroupJoinOp::inferTypes() {
     daphne::FrameType lhsFt = getLhs().getType().dyn_cast<daphne::FrameType>();
     daphne::FrameType rhsFt = getRhs().getType().dyn_cast<daphne::FrameType>();

diff --git a/src/ir/daphneir/DaphneOps.td b/src/ir/daphneir/DaphneOps.td
@@ -703,7 +703,8 @@ def Daphne_OrderOp : Daphne_Op<"order", [
 // Matrix decompositions & co
 // ****************************************************************************
 
-def Daphne_EigenOp : Daphne_Op<"eigenCal", [TypeFromFirstArg, DeclareOpInterfaceMethods<InferShapeOpInterface>]> {
+def Daphne_EigenOp : Daphne_Op<"eigenCal", [TypeFromFirstArg, DeclareOpInterfaceMethods<InferTypesOpInterface>,
+        DeclareOpInterfaceMethods<InferShapeOpInterface>]> {
     let arguments = (ins MatrixOf<[FloatScalar]>:$arg);
     let results = (outs MatrixOf<[FloatScalar]>:$eigenValues, MatrixOf<[FloatScalar]>:$eigenVectors);
 }