From 41f478f0419781f6968de0661ce570916c9de863 Mon Sep 17 00:00:00 2001
From: Kiran Chandramohan <kiran.chandramohan@arm.com>
Date: Tue, 18 Jul 2023 09:09:58 +0000
Subject: [PATCH] [Flang][HLFIR] Relax size check for dot_product intrinsic

If the size of one of the operand arrays is not known at compile
time, do not issue a size mismatch error sinc they could match at
runtime.

Fixes the compilation error in polyhedron/induct2.

Reviewed By: tblah, vzakhari

Differential Revision: https://reviews.llvm.org/D155302
---
 flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp |  4 +++-
 flang/test/Lower/HLFIR/dot_product.f90    | 12 ++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
index 51b44d7..37a1d9a 100644
--- a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
+++ b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
@@ -733,7 +733,9 @@ mlir::LogicalResult hlfir::DotProductOp::verify() {
   int64_t lhsSize = lhsShape[0];
   int64_t rhsSize = rhsShape[0];
 
-  if (lhsSize != rhsSize)
+  constexpr int64_t unknownExtent = fir::SequenceType::getUnknownExtent();
+  if ((lhsSize != unknownExtent) && (rhsSize != unknownExtent) &&
+      (lhsSize != rhsSize))
     return emitOpError("both arrays must have the same size");
 
   if (mlir::isa<fir::LogicalType>(lhsEleTy) !=
diff --git a/flang/test/Lower/HLFIR/dot_product.f90 b/flang/test/Lower/HLFIR/dot_product.f90
index 47e623a..890dc4a 100644
--- a/flang/test/Lower/HLFIR/dot_product.f90
+++ b/flang/test/Lower/HLFIR/dot_product.f90
@@ -70,3 +70,15 @@ endsubroutine
 ! CHECK-NEXT:    hlfir.assign %[[PROD]] to %[[RES_VAR]]#0 : i32, !fir.ref<i32>
 ! CHECK-NEXT:    return
 ! CHECK-NEXT:   }
+
+! CHECK-LABEL: func.func @_QPdot_product5
+! CHECK:    %[[LHS:.*]]:2 = hlfir.declare %{{.*}} {uniq_name = "_QFdot_product5Elhs"} : (!fir.box<!fir.array<?xi32>>) -> (!fir.box<!fir.array<?xi32>>, !fir.box<!fir.array<?xi32>>)
+! CHECK:    %[[C3:.*]] = arith.constant 3 : index
+! CHECK:    %[[RHS_SHAPE:.*]] = fir.shape %[[C3]] : (index) -> !fir.shape<1>
+! CHECK:    %[[RHS:.*]]:2 = hlfir.declare %{{.*}}(%[[RHS_SHAPE]]) {uniq_name = "_QFdot_product5Erhs"} : (!fir.ref<!fir.array<3xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<3xi32>>, !fir.ref<!fir.array<3xi32>>)
+! CHECK:    {{.*}} = hlfir.dot_product %[[LHS]]#0 %[[RHS]]#0 {fastmath = #arith.fastmath<contract>} : (!fir.box<!fir.array<?xi32>>, !fir.ref<!fir.array<3xi32>>) -> i32
+subroutine dot_product5(lhs, rhs, res)
+  integer :: lhs(:), rhs(3)
+  integer :: res
+  res = dot_product(lhs, rhs)
+endsubroutine
-- 
2.7.4