[llvm-branch-commits] [clang] [HLSL][Matrix] Add Matrix splat support for booleans (PR #175809)

[via llvm-branch-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-hlsl

Author: Deric C. (Icohedron)

<details>
<summary>Changes</summary>

Fixes #<!-- -->175808

This PR adds support for boolean matrix splats by adding tests and fixing a bug in `CodeGenFunction::EmitToMemory` when the type of a boolean matrix already matches the type expected of a load/store. 

This PR also addresses the todo comment in `clang/lib/Sema/SemaExpr.cpp` regarding support for boolean matrix splats.

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


3 Files Affected:

- (modified) clang/lib/CodeGen/CGExpr.cpp (+4) 
- (modified) clang/lib/Sema/SemaExpr.cpp (+20-8) 
- (modified) clang/test/CodeGenHLSL/BasicFeatures/MatrixSplat.hlsl (+49) 


``````````diff
diff --git a/clang/lib/CodeGen/CGExpr.cpp b/clang/lib/CodeGen/CGExpr.cpp
index a8e53e77b2e8c..30308b5b109a6 100644
--- a/clang/lib/CodeGen/CGExpr.cpp
+++ b/clang/lib/CodeGen/CGExpr.cpp
@@ -2218,6 +2218,10 @@ llvm::Value *CodeGenFunction::EmitToMemory(llvm::Value *Value, QualType Ty) {
 
   if (Ty->isExtVectorBoolType() || Ty->isConstantMatrixBoolType()) {
     llvm::Type *StoreTy = convertTypeForLoadStore(Ty, Value->getType());
+
+    if (Value->getType() == StoreTy)
+      return Value;
+
     if (StoreTy->isVectorTy() && StoreTy->getScalarSizeInBits() >
                                      Value->getType()->getScalarSizeInBits())
       return Builder.CreateZExt(Value, StoreTy);
diff --git a/clang/lib/Sema/SemaExpr.cpp b/clang/lib/Sema/SemaExpr.cpp
index 4d787a60eba3b..5e849ee69379d 100644
--- a/clang/lib/Sema/SemaExpr.cpp
+++ b/clang/lib/Sema/SemaExpr.cpp
@@ -7899,14 +7899,26 @@ ExprResult Sema::prepareMatrixSplat(QualType MatrixTy, Expr *SplattedExpr) {
   assert(DestElemTy->isFloatingType() ||
          DestElemTy->isIntegralOrEnumerationType());
 
-  // TODO: Add support for boolean matrix once exposed
-  // https://github.com/llvm/llvm-project/issues/170920
-  ExprResult CastExprRes = SplattedExpr;
-  CastKind CK = PrepareScalarCast(CastExprRes, DestElemTy);
-  if (CastExprRes.isInvalid())
-    return ExprError();
-  SplattedExpr = CastExprRes.get();
-
+  CastKind CK;
+  if (SplattedExpr->getType()->isBooleanType()) {
+    // As with vectors, we want `true` to become -1 when splatting, and we
+    // need a two-step cast if the destination element type is floating.
+    if (DestElemTy->isFloatingType()) {
+      // Cast boolean to signed integral, then to floating.
+      ExprResult CastExprRes = ImpCastExprToType(SplattedExpr, Context.IntTy,
+                                                 CK_BooleanToSignedIntegral);
+      SplattedExpr = CastExprRes.get();
+      CK = CK_IntegralToFloating;
+    } else {
+      CK = CK_BooleanToSignedIntegral;
+    }
+  } else {
+    ExprResult CastExprRes = SplattedExpr;
+    CK = PrepareScalarCast(CastExprRes, DestElemTy);
+    if (CastExprRes.isInvalid())
+      return ExprError();
+    SplattedExpr = CastExprRes.get();
+  }
   return ImpCastExprToType(SplattedExpr, DestElemTy, CK);
 }
 
diff --git a/clang/test/CodeGenHLSL/BasicFeatures/MatrixSplat.hlsl b/clang/test/CodeGenHLSL/BasicFeatures/MatrixSplat.hlsl
index 802c418f1dad5..618bdbdfe4102 100644
--- a/clang/test/CodeGenHLSL/BasicFeatures/MatrixSplat.hlsl
+++ b/clang/test/CodeGenHLSL/BasicFeatures/MatrixSplat.hlsl
@@ -23,6 +23,17 @@ void ConstantFloatSplat() {
     float2x2 M = 3.25;
 }
 
+// CHECK-LABEL: define hidden void @_Z17ConstantBoolSplatv(
+// CHECK-SAME: ) #[[ATTR0]] {
+// CHECK-NEXT:  [[ENTRY:.*:]]
+// CHECK-NEXT:    [[M:%.*]] = alloca [9 x i32], align 4
+// CHECK-NEXT:    store <9 x i32> splat (i32 1), ptr [[M]], align 4
+// CHECK-NEXT:    ret void
+//
+void ConstantBoolSplat() {
+    bool3x3 M = true;
+}
+
 // CHECK-LABEL: define hidden void @_Z12DynamicSplatf(
 // CHECK-SAME: float noundef nofpclass(nan inf) [[VALUE:%.*]]) #[[ATTR0]] {
 // CHECK-NEXT:  [[ENTRY:.*:]]
@@ -39,6 +50,25 @@ void DynamicSplat(float Value) {
     float3x3 M = Value;
 }
 
+// CHECK-LABEL: define hidden void @_Z16DynamicBoolSplatb(
+// CHECK-SAME: i1 noundef [[VALUE:%.*]]) #[[ATTR0]] {
+// CHECK-NEXT:  [[ENTRY:.*:]]
+// CHECK-NEXT:    [[VALUE_ADDR:%.*]] = alloca i32, align 4
+// CHECK-NEXT:    [[M:%.*]] = alloca [16 x i32], align 4
+// CHECK-NEXT:    [[STOREDV:%.*]] = zext i1 [[VALUE]] to i32
+// CHECK-NEXT:    store i32 [[STOREDV]], ptr [[VALUE_ADDR]], align 4
+// CHECK-NEXT:    [[TMP0:%.*]] = load i32, ptr [[VALUE_ADDR]], align 4
+// CHECK-NEXT:    [[LOADEDV:%.*]] = trunc i32 [[TMP0]] to i1
+// CHECK-NEXT:    [[SPLAT_SPLATINSERT:%.*]] = insertelement <16 x i1> poison, i1 [[LOADEDV]], i64 0
+// CHECK-NEXT:    [[SPLAT_SPLAT:%.*]] = shufflevector <16 x i1> [[SPLAT_SPLATINSERT]], <16 x i1> poison, <16 x i32> zeroinitializer
+// CHECK-NEXT:    [[TMP1:%.*]] = zext <16 x i1> [[SPLAT_SPLAT]] to <16 x i32>
+// CHECK-NEXT:    store <16 x i32> [[TMP1]], ptr [[M]], align 4
+// CHECK-NEXT:    ret void
+//
+void DynamicBoolSplat(bool Value) {
+    bool4x4 M = Value;
+}
+
 // CHECK-LABEL: define hidden void @_Z13CastThenSplatDv4_f(
 // CHECK-SAME: <4 x float> noundef nofpclass(nan inf) [[VALUE:%.*]]) #[[ATTR0]] {
 // CHECK-NEXT:  [[ENTRY:.*:]]
@@ -55,3 +85,22 @@ void DynamicSplat(float Value) {
 void CastThenSplat(float4 Value) {
     float3x3 M = (float) Value;
 }
+
+// CHECK-LABEL: define hidden void @_Z17BoolCastThenSplatDv3_i(
+// CHECK-SAME: <3 x i32> noundef [[VALUE:%.*]]) #[[ATTR0]] {
+// CHECK-NEXT:  [[ENTRY:.*:]]
+// CHECK-NEXT:    [[VALUE_ADDR:%.*]] = alloca <3 x i32>, align 16
+// CHECK-NEXT:    [[M:%.*]] = alloca [4 x i32], align 4
+// CHECK-NEXT:    store <3 x i32> [[VALUE]], ptr [[VALUE_ADDR]], align 16
+// CHECK-NEXT:    [[TMP0:%.*]] = load <3 x i32>, ptr [[VALUE_ADDR]], align 16
+// CHECK-NEXT:    [[TOBOOL:%.*]] = icmp ne <3 x i32> [[TMP0]], zeroinitializer
+// CHECK-NEXT:    [[CAST_VTRUNC:%.*]] = extractelement <3 x i1> [[TOBOOL]], i32 0
+// CHECK-NEXT:    [[SPLAT_SPLATINSERT:%.*]] = insertelement <4 x i1> poison, i1 [[CAST_VTRUNC]], i64 0
+// CHECK-NEXT:    [[SPLAT_SPLAT:%.*]] = shufflevector <4 x i1> [[SPLAT_SPLATINSERT]], <4 x i1> poison, <4 x i32> zeroinitializer
+// CHECK-NEXT:    [[TMP1:%.*]] = zext <4 x i1> [[SPLAT_SPLAT]] to <4 x i32>
+// CHECK-NEXT:    store <4 x i32> [[TMP1]], ptr [[M]], align 4
+// CHECK-NEXT:    ret void
+//
+void BoolCastThenSplat(int3 Value) {
+    bool2x2 M = (bool) Value;
+}

``````````

</details>


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