[flang-commits] [flang] [flang] Handle vector operands in logical op codegen (PR #225407)

[via flang-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-fir-hlfir

Author: Mattéo Rizza Murgier (matteo-rm)

<details>
<summary>Changes</summary>

FIR logical ops take `AnyLogicalOrIntegerLike` operands; this includes vector types. This PR makes the `normalizeLogicalToI1` and `extendI1ToType` utils type-agnostic so that they don't assert when passing a non-scalar type.

Fixes e.g. the crash:
```mlir
; fir-opt --fir-to-llvm-ir="target=aarch64" crash.mlir
func.func @<!-- -->neqv_vector(%a: vector<4xi32>, %b: vector<4xi32>) -> vector<4xi32> {
  %0 = fir.neqv %a, %b : vector<4xi32>
  return %0 : vector<4xi32>
}
```

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


2 Files Affected:

- (modified) flang/lib/Optimizer/CodeGen/CodeGen.cpp (+9-8) 
- (modified) flang/test/Fir/logical-convert.fir (+47) 


``````````diff
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index 8285de795118e..327f782ce3539 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -61,6 +61,7 @@
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Target/LLVMIR/Import.h"
@@ -4572,26 +4573,26 @@ struct NegcOpConversion : public fir::FIROpConversion<fir::NegcOp> {
   }
 };
 
-/// Normalize a logical value to i1 by comparing with zero.
+/// Normalize a logical value, or a vector thereof, to i1 by comparing with
+/// zero.
 static mlir::Value
 normalizeLogicalToI1(mlir::ConversionPatternRewriter &rewriter,
                      mlir::Location loc, mlir::Value value) {
   mlir::Type ty = value.getType();
-  auto i1Ty = mlir::IntegerType::get(rewriter.getContext(), 1);
-  if (ty == i1Ty)
+  if (mlir::getElementTypeOrSelf(ty).isSignlessInteger(1))
     return value;
-  mlir::Value zero = fir::genConstantIndex(loc, ty, rewriter, 0);
+  mlir::Value zero = mlir::LLVM::ConstantOp::create(rewriter, loc, ty,
+                                                    rewriter.getZeroAttr(ty));
   return mlir::LLVM::ICmpOp::create(rewriter, loc,
                                     mlir::LLVM::ICmpPredicate::ne, value, zero);
 }
 
-/// Extend an i1 value to the given integer type. Returns the value unchanged
-/// if it is already the target type.
+/// Extend an i1 value, or a vector thereof, to the given integer type. Returns
+/// the value unchanged if it is already the target type.
 static mlir::Value extendI1ToType(mlir::ConversionPatternRewriter &rewriter,
                                   mlir::Location loc, mlir::Value i1Val,
                                   mlir::Type toTy) {
-  auto i1Ty = mlir::IntegerType::get(rewriter.getContext(), 1);
-  if (toTy == i1Ty)
+  if (toTy == i1Val.getType())
     return i1Val;
   return mlir::LLVM::ZExtOp::create(rewriter, loc, toTy, i1Val);
 }
diff --git a/flang/test/Fir/logical-convert.fir b/flang/test/Fir/logical-convert.fir
index 0ebda8d5c29e7..c8ab48fcd6970 100644
--- a/flang/test/Fir/logical-convert.fir
+++ b/flang/test/Fir/logical-convert.fir
@@ -564,6 +564,53 @@ func.func @test_neqv_l4(%arg0: !fir.logical<4>, %arg1: !fir.logical<4>) -> !fir.
   return %0 : !fir.logical<4>
 }
 // -----
+// CHECK-LABEL: @test_eqv_vector_i32
+// CHECK: [[ZERO0:%[0-9]*]] = llvm.mlir.constant(dense<0> : vector<4xi32>) : vector<4xi32>
+// CHECK: [[A:%[0-9]*]] = llvm.icmp "ne" %arg0, [[ZERO0]] : vector<4xi32>
+// CHECK: [[ZERO1:%[0-9]*]] = llvm.mlir.constant(dense<0> : vector<4xi32>) : vector<4xi32>
+// CHECK: [[B:%[0-9]*]] = llvm.icmp "ne" %arg1, [[ZERO1]] : vector<4xi32>
+// CHECK: [[RES:%[0-9]*]] = llvm.icmp "eq" [[A]], [[B]] : vector<4xi1>
+// CHECK: [[EXT:%[0-9]*]] = llvm.zext [[RES]] : vector<4xi1> to vector<4xi32>
+// CHECK: llvm.return [[EXT]] : vector<4xi32>
+func.func @test_eqv_vector_i32(%arg0: vector<4xi32>, %arg1: vector<4xi32>) -> vector<4xi32> {
+  %0 = fir.eqv %arg0, %arg1 : vector<4xi32>
+  return %0 : vector<4xi32>
+}
+// -----
+// CHECK-LABEL: @test_neqv_vector_i32
+// CHECK: [[ZERO0:%[0-9]*]] = llvm.mlir.constant(dense<0> : vector<4xi32>) : vector<4xi32>
+// CHECK: [[A:%[0-9]*]] = llvm.icmp "ne" %arg0, [[ZERO0]] : vector<4xi32>
+// CHECK: [[ZERO1:%[0-9]*]] = llvm.mlir.constant(dense<0> : vector<4xi32>) : vector<4xi32>
+// CHECK: [[B:%[0-9]*]] = llvm.icmp "ne" %arg1, [[ZERO1]] : vector<4xi32>
+// CHECK: [[RES:%[0-9]*]] = llvm.icmp "ne" [[A]], [[B]] : vector<4xi1>
+// CHECK: [[EXT:%[0-9]*]] = llvm.zext [[RES]] : vector<4xi1> to vector<4xi32>
+// CHECK: llvm.return [[EXT]] : vector<4xi32>
+func.func @test_neqv_vector_i32(%arg0: vector<4xi32>, %arg1: vector<4xi32>) -> vector<4xi32> {
+  %0 = fir.neqv %arg0, %arg1 : vector<4xi32>
+  return %0 : vector<4xi32>
+}
+// -----
+// CHECK-LABEL: @test_logical_or_vector_i32
+// CHECK: [[ZERO0:%[0-9]*]] = llvm.mlir.constant(dense<0> : vector<4xi32>) : vector<4xi32>
+// CHECK: [[A:%[0-9]*]] = llvm.icmp "ne" %arg0, [[ZERO0]] : vector<4xi32>
+// CHECK: [[ZERO1:%[0-9]*]] = llvm.mlir.constant(dense<0> : vector<4xi32>) : vector<4xi32>
+// CHECK: [[B:%[0-9]*]] = llvm.icmp "ne" %arg1, [[ZERO1]] : vector<4xi32>
+// CHECK: [[RES:%[0-9]*]] = llvm.or [[A]], [[B]] : vector<4xi1>
+// CHECK: [[EXT:%[0-9]*]] = llvm.zext [[RES]] : vector<4xi1> to vector<4xi32>
+// CHECK: llvm.return [[EXT]] : vector<4xi32>
+func.func @test_logical_or_vector_i32(%arg0: vector<4xi32>, %arg1: vector<4xi32>) -> vector<4xi32> {
+  %0 = fir.logical_or %arg0, %arg1 : vector<4xi32>
+  return %0 : vector<4xi32>
+}
+// -----
+// CHECK-LABEL: @test_logical_and_vector_i1
+// CHECK: [[RES:%[0-9]*]] = llvm.and %arg0, %arg1 : vector<4xi1>
+// CHECK: llvm.return [[RES]] : vector<4xi1>
+func.func @test_logical_and_vector_i1(%arg0: vector<4xi1>, %arg1: vector<4xi1>) -> vector<4xi1> {
+  %0 = fir.logical_and %arg0, %arg1 : vector<4xi1>
+  return %0 : vector<4xi1>
+}
+// -----
 // Test on i1 type (from comparisons)
 // CHECK-LABEL: @test_logical_and_i1
 // CHECK: [[RES:%[0-9]*]] = llvm.and %arg0, %arg1 : i1

``````````

</details>


https://github.com/llvm/llvm-project/pull/225407