Fix dot product result type setting

lib/Conversion/TorchToMhlo/Linear.cpp

@@ -71,8 +71,8 @@ Value getPermutedTensor(PatternRewriter &rewriter, Operation *op, Value input,
   return result.getResult();
 }
 
-void castContractingDim(PatternRewriter &rewriter, Operation *op, Value &lhs,
-                        Value &rhs, int64_t lhsContractingDim,
+RankedTensorType castContractingDim(PatternRewriter &rewriter, Operation *op, Value &lhs,
+                        Value &rhs, int64_t lhsResultDim, int64_t rhsResultDim, int64_t lhsContractingDim,
                         int64_t rhsContractingDim) {
   auto lhsTy = lhs.getType().dyn_cast<RankedTensorType>();
   auto rhsTy = rhs.getType().dyn_cast<RankedTensorType>();
@@ -98,6 +98,16 @@ void castContractingDim(PatternRewriter &rewriter, Operation *op, Value &lhs,
       rhs = rewriter.create<tensor::CastOp>(op->getLoc(), newRankTy, rhs);
     }
   }
+  SmallVector<int64_t> outShape;
+  outShape.append(lhsShape.begin(), lhsShape.end());
+  for (int k = 0; k < lhsShape.size(); ++ k) {
+    if (outShape[k] == ShapedType::kDynamicSize && rhsShape[k] > 0) {
+        outShape[k] = rhsShape[k];
+    }
+  }
+  outShape[lhsResultDim] = lhsShape[lhsResultDim];
+  outShape[rhsResultDim] = rhsShape[rhsResultDim];
+  return RankedTensorType::get(outShape, lhsTy.getElementType());
 }
 
 void getBmmBroadcast(PatternRewriter &rewriter, Operation *op, Value &inpLhs,
@@ -212,6 +222,9 @@ class ConvertAtenMatmulBaseOp : public ConvertAtenOp<AtenOpT> {
                       options.dimSizeIndexBits);
     }
     auto batchDims = llvm::to_vector<4>(llvm::seq<int64_t>(0, nBatchDims));
+
+    auto lhsResultDim = nBatchDims;
+    auto rhsResultDim = nBatchDims + 1;
     auto lhsContractingDim = nBatchDims + 1;
     auto rhsContractingDim = nBatchDims;
     if (lhsRank == 1)
@@ -224,17 +237,18 @@ class ConvertAtenMatmulBaseOp : public ConvertAtenOp<AtenOpT> {
             /*rhsBatchingDimensions=*/batchDims,
             /*lhsContractingDimensions=*/{lhsContractingDim},
             /*rhsContractingDimensions=*/{rhsContractingDim});
-    auto resultTy = ConvertAtenOp<AtenOpT>::getTypeConverter()
-                        ->convertType(op.getType())
-                        .template cast<RankedTensorType>();
-
-    castContractingDim(rewriter, op, lhs, rhs, lhsContractingDim,
+    auto outTy = castContractingDim(rewriter, op, lhs, rhs, 
+		    lhsResultDim, rhsResultDim, lhsContractingDim,
                        rhsContractingDim);
     output = rewriter
-                 .create<mhlo::DotGeneralOp>(op->getLoc(), resultTy, lhs, rhs,
+                 .create<mhlo::DotGeneralOp>(op->getLoc(), outTy, lhs, rhs,
                                              dotDimensionNumbers, nullptr)
                  .getResult();
-
+    auto resultTy = ConvertAtenOp<AtenOpT>::getTypeConverter()
+                        ->convertType(op.getType())
+                        .template cast<RankedTensorType>();
+    output = rewriter
+                 .create<tensor::CastOp>(op->getLoc(), resultTy, output).getResult();
     return success();
   }
 
@@ -386,10 +400,14 @@ class ConvertAtenLinearOp : public ConvertAtenMatmulBaseOp<AtenOpT> {
     auto resultRank = std::max(lhsTy.getRank(), rhsTy.getRank());
     auto nBatchDims = resultRank - 2;
     auto batchDims = llvm::to_vector<4>(llvm::seq<int64_t>(0, nBatchDims));
+
+    auto lhsResultDim = nBatchDims;
+    auto rhsResultDim = nBatchDims + 1;
     auto lhsContractingDim = nBatchDims + 1;
     auto rhsContractingDim = nBatchDims;
 
-    castContractingDim(rewriter, op, lhs, rhs, lhsContractingDim,
+    auto outTy = castContractingDim(rewriter, op, lhs, rhs,
+		   lhsResultDim, rhsResultDim, lhsContractingDim,
                        rhsContractingDim);
     mhlo::DotDimensionNumbersAttr dotDimensionNumbers =
         mhlo::DotDimensionNumbersAttr::get(
@@ -398,24 +416,22 @@ class ConvertAtenLinearOp : public ConvertAtenMatmulBaseOp<AtenOpT> {
             /*rhsBatchingDimensions=*/batchDims,
             /*lhsContractingDimensions=*/{lhsContractingDim},
             /*rhsContractingDimensions=*/{rhsContractingDim});
-
-    auto resultTy =
-        ConvertAtenOp<AtenOpT>::getTypeConverter()->convertType(op.getType());
-
     Value matmulOutput = rewriter.create<mhlo::DotGeneralOp>(
-        op->getLoc(), resultTy, lhs, rhs, dotDimensionNumbers, nullptr);
+        op->getLoc(), outTy, lhs, rhs, dotDimensionNumbers, nullptr);
 
     Value matmulPlusBias = matmulOutput;
     if (!biasTy.template isa<Torch::NoneType>()) {
       // Bias addition broadcasts to the matmul output shape.
       matmulPlusBias =
           rewriter
-              .create<chlo::BroadcastAddOp>(op->getLoc(), resultTy,
+              .create<chlo::BroadcastAddOp>(op->getLoc(), outTy,
                                             matmulOutput, bias, nullptr)
               .getResult();
     }
 
-    rewriter.replaceOpWithNewOp<mhlo::ConvertOp>(op, resultTy, matmulPlusBias);
+    auto resultTy =
+        ConvertAtenOp<AtenOpT>::getTypeConverter()->convertType(op.getType());
+    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, resultTy, matmulPlusBias);
     return success();
   }
 };