[llvm] [InstCombine] Handle fixed-width results in get_active_lane_mask fold (PR #185317)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: tudinhh (tudinhh)

<details>
<summary>Changes</summary>

The optimization introduced in #<!-- -->183329 incorrectly assumed that any extraction from a scalable active lane mask used a scalable index. When the result of a `llvm.vector.extract` is a fixed-width vector, the index should not be multiplied by vscale.

This PR adds a check to ensure the index is only scaled by VScaleMin when the return type of the extraction is a scalable vector, not fixed-width.

Fixes #<!-- -->185271

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


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp (+3-1) 
- (modified) llvm/test/Transforms/InstCombine/get_active_lane_mask.ll (+13) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index 8185f009bc9d0..83a215eaba6fa 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3929,7 +3929,9 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
                        m_Value(), m_ConstantInt(ALMUpperBound)))) {
       const auto &Attrs = II->getFunction()->getAttributes().getFnAttrs();
       unsigned VScaleMin = Attrs.getVScaleRangeMin();
-      if (ExtractIdx * VScaleMin >= ALMUpperBound->getZExtValue())
+      unsigned ScaleFactor =
+          cast<VectorType>(ReturnType)->isScalableTy() ? VScaleMin : 1;
+      if (ExtractIdx * ScaleFactor >= ALMUpperBound->getZExtValue())
         return replaceInstUsesWith(CI,
                                    ConstantVector::getNullValue(ReturnType));
     }
diff --git a/llvm/test/Transforms/InstCombine/get_active_lane_mask.ll b/llvm/test/Transforms/InstCombine/get_active_lane_mask.ll
index 530cd9efa5b9c..44ce3f73fab67 100644
--- a/llvm/test/Transforms/InstCombine/get_active_lane_mask.ll
+++ b/llvm/test/Transforms/InstCombine/get_active_lane_mask.ll
@@ -76,3 +76,16 @@ define <4 x i1> @ext_has_active_lanes() {
   %ext = tail call <4 x i1> @llvm.vector.extract.v4i1.nxv16i1(<vscale x 16 x i1> %wide.alm, i64 4)
   ret <4 x i1> %ext
 }
+
+define <4 x i1> @extract_fixed_from_active_lane_mask() vscale_range(4, 4) {
+; CHECK-LABEL: define <4 x i1> @extract_fixed_from_active_lane_mask(
+; CHECK-SAME: ) #[[ATTR1:[0-9]+]] {
+; CHECK-NEXT:    [[MASK:%.*]] = call <vscale x 4 x i1> @llvm.get.active.lane.mask.nxv4i1.i32(i32 0, i32 15)
+; CHECK-NEXT:    [[EXT:%.*]] = call <4 x i1> @llvm.vector.extract.v4i1.nxv4i1(<vscale x 4 x i1> [[MASK]], i64 4)
+; CHECK-NEXT:    ret <4 x i1> [[EXT]]
+;
+  %mask = call <vscale x 4 x i1> @llvm.get.active.lane.mask.nxv4i1.i32(i32 0, i32 15)
+  %ext = call <4 x i1> @llvm.vector.extract.v4i1.nxv4i1(<vscale x 4 x i1> %mask, i64 4)
+  ret <4 x i1> %ext
+}
+

``````````

</details>


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