Expand WaveActiveAllEqual with vector and matrix arguments. (#5428)

For vector and matrix inputs to WaveActiveAllEqual, the
OpGroupNonuniformAllEqual instruction must be applied on each element of
the vector or matrix, and then have the results combined into a vector
or matrix of bools. We make that change in this PR.

Fixes #5426

Author: s-perron

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -8540,7 +8540,7 @@ SpirvEmitter::processIntrinsicCallExpr(const CallExpr *callExpr) {
     retVal = processWaveVote(callExpr, spv::Op::OpGroupNonUniformBallot);
     break;
   case hlsl::IntrinsicOp::IOP_WaveActiveAllEqual:
-    retVal = processWaveVote(callExpr, spv::Op::OpGroupNonUniformAllEqual);
+    retVal = processWaveActiveAllEqual(callExpr);
     break;
   case hlsl::IntrinsicOp::IOP_WaveActiveCountBits:
     retVal = processWaveCountBits(callExpr, spv::GroupOperation::Reduce);
@@ -9384,6 +9384,74 @@ SpirvEmitter::processWaveQuadWideShuffle(const CallExpr *callExpr,
       opcode, retType, spv::Scope::Subgroup, value, target, srcLoc);
 }
 
+SpirvInstruction *
+SpirvEmitter::processWaveActiveAllEqual(const CallExpr *callExpr) {
+  assert(callExpr->getNumArgs() == 1);
+  featureManager.requestTargetEnv(SPV_ENV_VULKAN_1_1, "Wave Operation",
+                                  callExpr->getExprLoc());
+  SpirvInstruction *arg = doExpr(callExpr->getArg(0));
+  const QualType retType = callExpr->getCallReturnType(astContext);
+
+  if (isScalarType(retType))
+    return processWaveActiveAllEqualScalar(arg, callExpr->getExprLoc());
+
+  if (isVectorType(retType))
+    return processWaveActiveAllEqualVector(arg, callExpr->getExprLoc());
+
+  assert(isMxNMatrix(retType));
+  return processWaveActiveAllEqualMatrix(arg, retType, callExpr->getExprLoc());
+}
+
+SpirvInstruction *
+SpirvEmitter::processWaveActiveAllEqualScalar(SpirvInstruction *arg,
+                                              clang::SourceLocation srcLoc) {
+  return spvBuilder.createGroupNonUniformUnaryOp(
+      srcLoc, spv::Op::OpGroupNonUniformAllEqual, astContext.BoolTy,
+      spv::Scope::Subgroup, arg);
+}
+
+SpirvInstruction *
+SpirvEmitter::processWaveActiveAllEqualVector(SpirvInstruction *arg,
+                                              clang::SourceLocation srcLoc) {
+  uint32_t vectorSize = 0;
+  QualType elementType;
+  isVectorType(arg->getAstResultType(), &elementType, &vectorSize);
+  assert(vectorSize >= 2 && "Vector size in spir-v must be at least 2");
+
+  llvm::SmallVector<SpirvInstruction *, 4> elements;
+  for (uint32_t i = 0; i < vectorSize; ++i) {
+    SpirvInstruction *element =
+        spvBuilder.createCompositeExtract(elementType, arg, {i}, srcLoc);
+    elements.push_back(processWaveActiveAllEqualScalar(element, srcLoc));
+  }
+
+  QualType booleanVectortype =
+      astContext.getExtVectorType(astContext.BoolTy, vectorSize);
+  return spvBuilder.createCompositeConstruct(booleanVectortype, elements,
+                                             srcLoc);
+}
+
+SpirvInstruction *
+SpirvEmitter::processWaveActiveAllEqualMatrix(SpirvInstruction *arg,
+                                              QualType booleanMatrixType,
+                                              clang::SourceLocation srcLoc) {
+  uint32_t numberOfRows = 0;
+  uint32_t numberOfColumns = 0;
+  QualType elementType;
+  isMxNMatrix(arg->getAstResultType(), &elementType, &numberOfRows,
+              &numberOfColumns);
+  assert(numberOfRows >= 2 && "Vector size in spir-v must be at least 2");
+
+  QualType rowType = astContext.getExtVectorType(elementType, numberOfColumns);
+  llvm::SmallVector<SpirvInstruction *, 4> rows;
+  for (uint32_t i = 0; i < numberOfRows; ++i) {
+    SpirvInstruction *row =
+        spvBuilder.createCompositeExtract(rowType, arg, {i}, srcLoc);
+    rows.push_back(processWaveActiveAllEqualVector(row, srcLoc));
+  }
+  return spvBuilder.createCompositeConstruct(booleanMatrixType, rows, srcLoc);
+}
+
 SpirvInstruction *SpirvEmitter::processIntrinsicModf(const CallExpr *callExpr) {
   // Signature is: ret modf(x, ip)
   // [in]    x: the input floating-point value.

tools/clang/lib/SPIRV/SpirvEmitter.h

@@ -608,6 +608,35 @@ class SpirvEmitter : public ASTConsumer {
   SpirvInstruction *processWaveQuadWideShuffle(const CallExpr *,
                                                hlsl::IntrinsicOp op);
 
+  /// Generates the Spir-V instructions needed to implement the given call to
+  /// WaveActiveAllEqual. Returns a pointer to the instruction that produces the
+  /// final result.
+  SpirvInstruction *processWaveActiveAllEqual(const CallExpr *);
+
+  /// Generates the Spir-V instructions needed to implement WaveActiveAllEqual
+  /// with the scalar input `arg`. Returns a pointer to the instruction that
+  /// produces the final result. srcLoc should be the source location of the
+  /// original call.
+  SpirvInstruction *
+  processWaveActiveAllEqualScalar(SpirvInstruction *arg,
+                                  clang::SourceLocation srcLoc);
+
+  /// Generates the Spir-V instructions needed to implement WaveActiveAllEqual
+  /// with the vector input `arg`. Returns a pointer to the instruction that
+  /// produces the final result. srcLoc should be the source location of the
+  /// original call.
+  SpirvInstruction *
+  processWaveActiveAllEqualVector(SpirvInstruction *arg,
+                                  clang::SourceLocation srcLoc);
+
+  /// Generates the Spir-V instructions needed to implement WaveActiveAllEqual
+  /// with the matrix input `arg`. Returns a pointer to the instruction that
+  /// produces the final result. srcLoc should be the source location of the
+  /// original call.
+  SpirvInstruction *
+  processWaveActiveAllEqualMatrix(SpirvInstruction *arg, QualType,
+                                  clang::SourceLocation srcLoc);
+
   /// Processes the NonUniformResourceIndex intrinsic function.
   SpirvInstruction *processIntrinsicNonUniformResourceIndex(const CallExpr *);
 

tools/clang/test/CodeGenSPIRV/sm6.wave-active-all-equal-matrix.hlsl

@@ -0,0 +1,53 @@
+// RUN: %dxc -T cs_6_0 -HV 2018 -E main -fspv-target-env=vulkan1.1
+
+struct S {
+    float2x2 val;
+    bool res;
+};
+
+RWStructuredBuffer<S> values;
+
+// CHECK: OpCapability GroupNonUniformVote
+
+[numthreads(32, 1, 1)]
+void main(uint3 id: SV_DispatchThreadID) {
+
+// Each element of the matrix must be extracted, and be passed to OpGroupNonUniformAllEqual. 
+// CHECK: [[ld:%\w+]] = OpLoad %mat2v2float %34
+
+// Process the first row.
+// CHECK: [[row_0:%\w+]] = OpCompositeExtract %v2float [[ld]] 0
+// CHECK: [[element_0_0:%\w+]] = OpCompositeExtract %float [[row_0]] 0
+// CHECK: [[res_0_0:%\w+]] = OpGroupNonUniformAllEqual %bool %uint_3 [[element_0_0]]
+// CHECK: [[element_0_1:%\w+]] = OpCompositeExtract %float [[row_0]] 1
+// CHECK: [[res_0_1:%\w+]] = OpGroupNonUniformAllEqual %bool %uint_3 [[element_0_1]]
+
+// Combine the results in a row for the results matrix.
+// CHECK: [[res_0:%\w+]] = OpCompositeConstruct %v2bool [[res_0_0]] [[res_0_1]]
+
+// Process the second row.
+// CHECK: [[row_1:%\w+]] = OpCompositeExtract %v2float [[ld]] 1
+// CHECK: [[element_1_0:%\w+]] = OpCompositeExtract %float [[row_1]] 0
+// CHECK: [[res_1_0:%\w+]] = OpGroupNonUniformAllEqual %bool %uint_3 [[element_1_0]]
+// CHECK: [[element_1_1:%\w+]] = OpCompositeExtract %float [[row_1]] 1
+// CHECK: [[res_1_1:%\w+]] = OpGroupNonUniformAllEqual %bool %uint_3 [[element_1_1]]
+
+
+// Combine the results in a row for the results matrix.
+// CHECK: [[res_1:%\w+]] = OpCompositeConstruct %v2bool [[res_1_0]] [[res_1_1]]
+
+// Combind the results for each row in a "matrix" for the final result.
+// CHECK: [[res_matrix:%\w+]] = OpCompositeConstruct %_arr_v2bool_uint_2 [[res_0]] [[res_1]]
+
+// Apply the `all` to the entire matrix.
+// CHECK: [[res_vec0:%\w+]] = OpCompositeExtract %v2bool [[res_matrix]] 0
+// CHECK: [[all0:%\w+]] = OpAll %bool [[res_vec0]]
+// CHECK: [[res_vec1:%\w+]] = OpCompositeExtract %v2bool %53 1
+// CHECK: [[all1:%\w+]] = OpAll %bool [[res_vec1]]
+// CHECK: [[all_vec:%\w+]] = OpCompositeConstruct %v2bool [[all0]] [[all1]]
+// CHECK: [[all:%\w+]] = OpAll %bool [[all_vec]]
+
+// Do the select.
+// CHECK: OpSelect %uint [[all]] %uint_1 %uint_0
+    values[id.x].res = all(WaveActiveAllEqual(values[id.x].val));
+}

tools/clang/test/CodeGenSPIRV/sm6.wave-active-all-equal-matrix1x1.hlsl

@@ -0,0 +1,21 @@
+// RUN: %dxc -T cs_6_0 -HV 2018 -E main -fspv-target-env=vulkan1.1
+
+struct S {
+    float1x1 val;
+    bool res;
+};
+
+RWStructuredBuffer<S> values;
+
+// CHECK: OpCapability GroupNonUniformVote
+
+[numthreads(32, 1, 1)]
+void main(uint3 id: SV_DispatchThreadID) {
+
+// For a 1x1 matrix, the spirv type should become a scalar because Spir-V cannot have a 1x1 matrix.
+
+// CHECK: [[ld:%\w+]] = OpLoad %float %32
+// CHECK: [[res:%\w+]] = OpGroupNonUniformAllEqual %bool %uint_3 [[ld]]
+// CHECK: OpSelect %uint [[res]] %uint_1 %uint_0
+    values[id.x].res = all(WaveActiveAllEqual(values[id.x].val));
+}

tools/clang/test/CodeGenSPIRV/sm6.wave-active-all-equal-scalar.hlsl

@@ -0,0 +1,17 @@
+// RUN: %dxc -T cs_6_0 -HV 2018 -E main -fspv-target-env=vulkan1.1
+
+struct S {
+    uint val;
+    bool res;
+};
+
+RWStructuredBuffer<S> values;
+
+// CHECK: OpCapability GroupNonUniformVote
+
+[numthreads(32, 1, 1)]
+void main(uint3 id: SV_DispatchThreadID) {
+// CHECK: [[eq:%\w+]] = OpGroupNonUniformAllEqual %bool %uint_3 {{%\w+}}
+// CHECK: OpSelect %uint [[eq]] %uint_1 %uint_0
+    values[id.x].res = WaveActiveAllEqual(values[id.x].val);
+}

tools/clang/test/CodeGenSPIRV/sm6.wave-active-all-equal-vector.hlsl

@@ -0,0 +1,31 @@
+// RUN: %dxc -T cs_6_0 -HV 2018 -E main -fspv-target-env=vulkan1.1
+
+struct S {
+    float4 val;
+    bool res;
+};
+
+RWStructuredBuffer<S> values;
+
+// CHECK: OpCapability GroupNonUniformVote
+
+[numthreads(32, 1, 1)]
+void main(uint3 id: SV_DispatchThreadID) {
+
+// Each element of the vector must be extracted, and be passed to OpGroupNonUniformAllEqual. 
+// CHECK: [[ld:%\w+]] = OpLoad %v4float
+// CHECK: [[element0:%\w+]] = OpCompositeExtract %float [[ld]] 0
+// CHECK: [[res0:%\w+]] = OpGroupNonUniformAllEqual %bool %uint_3 [[element0]]
+// CHECK: [[element1:%\w+]] = OpCompositeExtract %float [[ld]] 1
+// CHECK: [[res1:%\w+]] = OpGroupNonUniformAllEqual %bool %uint_3 [[element1]]
+// CHECK: [[element2:%\w+]] = OpCompositeExtract %float [[ld]] 2
+// CHECK: [[res2:%\w+]] = OpGroupNonUniformAllEqual %bool %uint_3 [[element2]]
+// CHECK: [[element3:%\w+]] = OpCompositeExtract %float [[ld]] 3
+// CHECK: [[res3:%\w+]] = OpGroupNonUniformAllEqual %bool %uint_3 [[element3]]
+
+// Then the results must be combined into a boolean vector.
+// CHECK: [[vec_res:%\w+]] = OpCompositeConstruct %v4bool [[res0]] [[res1]] [[res2]] [[res3]]
+// CHECK: [[res:%\w+]] = OpAll %bool [[vec_res]]
+// CHECK: OpSelect %uint [[res]] %uint_1 %uint_0 
+    values[id.x].res = all(WaveActiveAllEqual(values[id.x].val));
+}

tools/clang/test/CodeGenSPIRV/sm6.wave-active-all-equal.hlsl

@@ -0,0 +1,19 @@
+// RUN: %dxc -T cs_6_0 -E main -HV 2018 -fspv-target-env=vulkan1.0
+
+struct S {
+    float4 val1;
+    uint val2;
+    bool res;
+};
+
+RWStructuredBuffer<S> values;
+
+[numthreads(32, 1, 1)]
+void main(uint3 id: SV_DispatchThreadID) {
+    uint x = id.x;
+    values[x].res = WaveActiveAllEqual(values[x].val1) && WaveActiveAllEqual(values[x].val2);
+}
+
+// CHECK: sm6.wave-active-all-equal.vulkan1.0.hlsl:14:21: error: Vulkan 1.1 is required for Wave Operation but not permitted to use
+// CHECK-NEXT: values[x].res = WaveActiveAllEqual(values[x].val1) && WaveActiveAllEqual(values[x].val2);
+// CHECK: note: please specify your target environment via command line option -fspv-target-env=

tools/clang/test/CodeGenSPIRV/sm6.wave-active-all-equal.vulkan1.0.hlsl

@@ -17,8 +17,15 @@ void main(uint3 id: SV_DispatchThreadID) {
     uint x = id.x;
 // CHECK:         [[ptr:%\d+]] = OpAccessChain %_ptr_StorageBuffer_v4float %values %int_0 {{%\d+}} %int_0
 // CHECK-NEXT: [[f32val:%\d+]] = OpLoad %v4float [[ptr]]
-// TODO: The front end will return bool4 for the first call, which acutally should be bool.
-// XXXXX-NEXT:        {{%\d+}} = OpGroupNonUniformAllEqual %bool %uint_3 [[f32val]]
+// CHECK: [[element0:%\w+]] = OpCompositeExtract %float [[f32val]] 0
+// CHECK: [[res0:%\w+]] = OpGroupNonUniformAllEqual %bool %uint_3 [[element0]]
+// CHECK: [[element1:%\w+]] = OpCompositeExtract %float [[f32val]] 1
+// CHECK: [[res1:%\w+]] = OpGroupNonUniformAllEqual %bool %uint_3 [[element1]]
+// CHECK: [[element2:%\w+]] = OpCompositeExtract %float [[f32val]] 2
+// CHECK: [[res2:%\w+]] = OpGroupNonUniformAllEqual %bool %uint_3 [[element2]]
+// CHECK: [[element3:%\w+]] = OpCompositeExtract %float [[f32val]] 3
+// CHECK: [[res3:%\w+]] = OpGroupNonUniformAllEqual %bool %uint_3 [[element3]]
+// CHECK: [[vec_res:%\w+]] = OpCompositeConstruct %v4bool [[res0]] [[res1]] [[res2]] [[res3]]
 
 // CHECK:         [[ptr:%\d+]] = OpAccessChain %_ptr_StorageBuffer_uint %values %int_0 {{%\d+}} %int_1
 // CHECK-NEXT: [[u32val:%\d+]] = OpLoad %uint [[ptr]]

tools/clang/test/CodeGenSPIRV/sm6.wave-active-all-equal.vulkan1.2.hlsl

@@ -1587,11 +1587,20 @@ TEST_F(FileTest, SM6WaveActiveBallot) {
   runFileTest("sm6.wave-active-ballot.hlsl");
 }
 
-// Shader model 6.0 wave reduction
-// TODO(5410): Still unclear what should happen with WaveActiveAllEqual with a vector parameter.
-// For now, the generated SPIR-V is invalid.
-TEST_F(FileTest, SM6WaveActiveAllEqual) {
-  runFileTest("sm6.wave-active-all-equal.hlsl", Expect::Success, /*runValidation=*/ false);
+TEST_F(FileTest, SM6WaveActiveAllEqualScalar) {
+  runFileTest("sm6.wave-active-all-equal-scalar.hlsl");
+}
+TEST_F(FileTest, SM6WaveActiveAllEqualVector) {
+  runFileTest("sm6.wave-active-all-equal-vector.hlsl");
+}
+TEST_F(FileTest, SM6WaveActiveAllEqualMatrix) {
+  runFileTest("sm6.wave-active-all-equal-matrix.hlsl");
+}
+TEST_F(FileTest, SM6WaveActiveAllEqualMatrix1x1) {
+  runFileTest("sm6.wave-active-all-equal-matrix1x1.hlsl");
+}
+TEST_F(FileTest, SM6WaveActiveAllEqualVulkan1_0) {
+  runFileTest("sm6.wave-active-all-equal.vulkan1.0.hlsl", Expect::Failure);
 }
 TEST_F(FileTest, SM6WaveActiveSum) {
   runFileTest("sm6.wave-active-sum.hlsl");
@@ -2950,9 +2959,7 @@ TEST_F(FileTest, CompatibilityWithVk1p1) {
   runFileTest("sm6.quad-read-across-x.vulkan1.2.hlsl");
   runFileTest("sm6.quad-read-across-y.vulkan1.2.hlsl");
   runFileTest("sm6.quad-read-lane-at.vulkan1.2.hlsl");
-  // TODO(5410): Still unclear what should happen with WaveActiveAllEqual with a vector parameter.
-  // For now, the generated SPIR-V is invalid.
-  runFileTest("sm6.wave-active-all-equal.vulkan1.2.hlsl", Expect::Success, /*runValidation=*/false);
+  runFileTest("sm6.wave-active-all-equal.vulkan1.2.hlsl");
   runFileTest("sm6.wave-active-all-true.vulkan1.2.hlsl");
   runFileTest("sm6.wave-active-any-true.vulkan1.2.hlsl");
   runFileTest("sm6.wave-active-ballot.vulkan1.2.hlsl");