[llvm] [llvm][InstCombine] bitcast bfloat half castpair bug (PR #79832)

[Nashe Mncube via llvm-commits]

https://github.com/nasherm updated https://github.com/llvm/llvm-project/pull/79832

>From 682363ed9c339ba08a0f524db6a14642718707e8 Mon Sep 17 00:00:00 2001
From: nasmnc01 <nashe.mncube at arm.com>
Date: Fri, 26 Jan 2024 14:30:16 +0000
Subject: [PATCH] [llvm][InstCombine] bitcast bfloat half castpair bug

Miscompilation arises due to instruction combining of cast pairs
of the type `bitcast bfloat to half` + `<FPOp> bfloat to half` or
`bitcast half to bfloat` + `<FPOp half to bfloat`. For example
`bitcast bfloat to half`+`fpext half to double` or
`bitcast bfloat to half`+`fpext bfloat to double` respectively
reduce to `fpext bfloat to double` and `fpext half to double`.
This is an incorrect conversion as it assumes the representation
of `bfloat` and `half` are equivalent due to having the same width.
As a consequence miscompilation arises.

Change-Id: Ie5b7c4b385a946325c60de5495ce3bdf087abc46
---
 llvm/lib/IR/Instructions.cpp                  |  7 ++
 .../InstCombine/bitcast-bfloat-half-mixing.ll | 87 +++++++++++++++++++
 2 files changed, 94 insertions(+)
 create mode 100644 llvm/test/Transforms/InstCombine/bitcast-bfloat-half-mixing.ll

diff --git a/llvm/lib/IR/Instructions.cpp b/llvm/lib/IR/Instructions.cpp
index 87874c3abc468..3189b547ca5fe 100644
--- a/llvm/lib/IR/Instructions.cpp
+++ b/llvm/lib/IR/Instructions.cpp
@@ -3214,6 +3214,13 @@ unsigned CastInst::isEliminableCastPair(
         return secondOp;
       return 0;
     case 6:
+      // In cast pairs bfloat and half float shouldn't be treated as equivalent
+      // if the first operation is a bitcast i.e. if we have
+      // bitcast bfloat to half + fpext half to double we shouldn't reduce to
+      // fpext bfloat to double as this isn't equal to fpext half to double.
+      // This has been generalised for all float pairs that have the same width.
+      if(SrcTy->isFloatingPointTy() && MidTy->isFloatingPointTy())
+        return 0;
       // No-op cast in first op implies secondOp as long as the SrcTy
       // is a floating point.
       if (SrcTy->isFloatingPointTy())
diff --git a/llvm/test/Transforms/InstCombine/bitcast-bfloat-half-mixing.ll b/llvm/test/Transforms/InstCombine/bitcast-bfloat-half-mixing.ll
new file mode 100644
index 0000000000000..0c4beb223536b
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/bitcast-bfloat-half-mixing.ll
@@ -0,0 +1,87 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt -passes=instcombine -S %s | FileCheck %s
+
+define double @F0(bfloat %P0) {
+; CHECK-LABEL: define double @F0(
+; CHECK-SAME: bfloat [[P0:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV0:%.*]] = bitcast bfloat [[P0]] to half
+; CHECK-NEXT:    [[TMP0:%.*]] = fpext half [[CONV0]] to double
+; CHECK-NEXT:    ret double [[TMP0]]
+;
+entry:
+  %conv0 = bitcast bfloat %P0 to half
+  %0 = fpext half %conv0 to double
+  ret double %0
+}
+
+define double @F1(half %P1) {
+; CHECK-LABEL: define double @F1(
+; CHECK-SAME: half [[P1:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV1:%.*]] = bitcast half [[P1]] to bfloat
+; CHECK-NEXT:    [[TMP0:%.*]] = fpext bfloat [[CONV1]] to double
+; CHECK-NEXT:    ret double [[TMP0]]
+;
+entry:
+  %conv1 = bitcast half %P1 to bfloat
+  %0 = fpext bfloat %conv1 to double
+  ret double %0
+}
+
+define i32 @F2(bfloat %P2) {
+; CHECK-LABEL: define i32 @F2(
+; CHECK-SAME: bfloat [[P2:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV2:%.*]] = bitcast bfloat [[P2]] to half
+; CHECK-NEXT:    [[TMP0:%.*]] = fptoui half [[CONV2]] to i32
+; CHECK-NEXT:    ret i32 [[TMP0]]
+;
+entry:
+  %conv2 = bitcast bfloat %P2 to half
+  %0 = fptoui half %conv2 to i32
+  ret i32 %0
+}
+
+define i32 @F3(half %P3) {
+; CHECK-LABEL: define i32 @F3(
+; CHECK-SAME: half [[P3:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV3:%.*]] = bitcast half [[P3]] to bfloat
+; CHECK-NEXT:    [[TMP0:%.*]] = fptoui bfloat [[CONV3]] to i32
+; CHECK-NEXT:    ret i32 [[TMP0]]
+;
+entry:
+  %conv3 = bitcast half %P3 to bfloat
+  %0 = fptoui bfloat %conv3 to i32
+  ret i32 %0
+}
+
+define i32 @F4(bfloat %P4) {
+; CHECK-LABEL: define i32 @F4(
+; CHECK-SAME: bfloat [[P4:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV4:%.*]] = bitcast bfloat [[P4]] to half
+; CHECK-NEXT:    [[TMP0:%.*]] = fptosi half [[CONV4]] to i32
+; CHECK-NEXT:    ret i32 [[TMP0]]
+;
+entry:
+  %conv4 = bitcast bfloat %P4 to half
+  %0 = fptosi half %conv4 to i32
+  ret i32 %0
+}
+
+define i32 @F5(half %P5) {
+; CHECK-LABEL: define i32 @F5(
+; CHECK-SAME: half [[P5:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV5:%.*]] = bitcast half [[P5]] to bfloat
+; CHECK-NEXT:    [[TMP0:%.*]] = fptosi bfloat [[CONV5]] to i32
+; CHECK-NEXT:    ret i32 [[TMP0]]
+;
+entry:
+  %conv5 = bitcast half %P5 to bfloat
+  %0 = fptosi bfloat %conv5 to i32
+  ret i32 %0
+}
+