Handle the implicit cast to/from boolean type

Author: jkwak-work

source/slang/hlsl.meta.slang

@@ -6256,7 +6256,9 @@ bool all(vector<T,N> x)
             };
         }
     case wgsl:
-        __intrinsic_asm "all";
+        if (__isBool<T>())
+            __intrinsic_asm "all";
+        __intrinsic_asm "all(vec$N0<bool>($0))";
     default:
         bool result = true;
         for(int i = 0; i < N; ++i)

source/slang/slang-emit-wgsl.cpp

@@ -1317,28 +1317,31 @@ bool WGSLSourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOu
         {
             // WGSL doesn't have operator overloadings for `&&` and `||` when the operands are
             // non-scalar. Unlike HLSL, WGSL doesn't have `and()` and `or()`.
-            if (as<IRBasicType>(inst->getDataType()))
+            auto vecType = as<IRVectorType>(inst->getDataType());
+            if (!vecType)
                 return false;
 
+            // The function signature for `select` in WGSL is different from others:
+            // @const @must_use fn select(f: T, t: T, cond: bool) -> T
             if (inst->getOp() == kIROp_And)
             {
-                m_writer->emit("select(");
-                emitType(inst->getDataType());
-                m_writer->emit("(false), ");
-                emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
-                m_writer->emit(", ");
+                m_writer->emit(" select(vec");
+                m_writer->emit(getIntVal(vecType->getElementCount()));
+                m_writer->emit("<bool>(false), ");
                 emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
-                m_writer->emit(")");
+                m_writer->emit(", ");
+                emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+                m_writer->emit(") ");
             }
             else
             {
-                m_writer->emit("select(");
-                emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
-                m_writer->emit(", ");
-                emitType(inst->getDataType());
-                m_writer->emit("(true), ");
+                m_writer->emit(" select(");
                 emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
-                m_writer->emit(")");
+                m_writer->emit(", vec");
+                m_writer->emit(getIntVal(vecType->getElementCount()));
+                m_writer->emit("<bool>(true), ");
+                emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+                m_writer->emit(") ");
             }
             return true;
         }

source/slang/slang-emit.cpp

@@ -52,6 +52,7 @@
 #include "slang-ir-insts.h"
 #include "slang-ir-layout.h"
 #include "slang-ir-legalize-array-return-type.h"
+#include "slang-ir-legalize-binary-operator.h"
 #include "slang-ir-legalize-global-values.h"
 #include "slang-ir-legalize-image-subscript.h"
 #include "slang-ir-legalize-mesh-outputs.h"
@@ -1460,6 +1461,9 @@ Result linkAndOptimizeIR(
         floatNonUniformResourceIndex(irModule, NonUniformResourceIndexFloatMode::Textual);
     }
 
+    if (isD3DTarget(targetRequest) || isKhronosTarget(targetRequest) || isWGPUTarget(targetRequest))
+        legalizeLogicalAndOr(irModule->getModuleInst());
+
     // Legalize non struct parameters that are expected to be structs for HLSL.
     if (isD3DTarget(targetRequest))
         legalizeNonStructParameterToStructForHLSL(irModule);

source/slang/slang-ir-insts.h

@@ -4518,6 +4518,9 @@ struct IRBuilder
     IRInst* emitShr(IRType* type, IRInst* op0, IRInst* op1);
     IRInst* emitShl(IRType* type, IRInst* op0, IRInst* op1);
 
+    IRInst* emitAnd(IRType* type, IRInst* left, IRInst* right);
+    IRInst* emitOr(IRType* type, IRInst* left, IRInst* right);
+
     IRSPIRVAsmOperand* emitSPIRVAsmOperandLiteral(IRInst* literal);
     IRSPIRVAsmOperand* emitSPIRVAsmOperandInst(IRInst* inst);
     IRSPIRVAsmOperand* createSPIRVAsmOperandInst(IRInst* inst);

source/slang/slang-ir-legalize-binary-operator.cpp

@@ -118,4 +118,101 @@ void legalizeBinaryOp(IRInst* inst)
     }
 }
 
+void legalizeLogicalAndOr(IRInst* inst)
+{
+    switch (inst->getOp())
+    {
+    case kIROp_And:
+    case kIROp_Or:
+        {
+            IRBuilder builder(inst);
+            builder.setInsertBefore(inst);
+
+            // Logical-AND and logical-OR takes boolean types as its operands.
+            // If they are not, legalize them by casting to boolean type.
+            //
+            SLANG_ASSERT(inst->getOperandCount() == 2);
+            for (UInt i = 0; i < 2; i++)
+            {
+                auto operand = inst->getOperand(i);
+                auto operandDataType = operand->getDataType();
+
+                if (auto vecType = as<IRVectorType>(operandDataType))
+                {
+                    if (!as<IRBoolType>(vecType->getElementType()))
+                    {
+                        // Cast operand to vector<bool,N>
+                        auto elemCount = vecType->getElementCount();
+                        auto vb = builder.getVectorType(builder.getBoolType(), elemCount);
+                        auto v = builder.emitCast(vb, operand);
+                        builder.replaceOperand(inst->getOperands() + i, v);
+                    }
+                }
+                else if (!as<IRBoolType>(operandDataType))
+                {
+                    // Cast operand to bool
+                    auto s = builder.emitCast(builder.getBoolType(), operand);
+                    builder.replaceOperand(inst->getOperands() + i, s);
+                }
+            }
+
+            // Legalize the return type; mostly for SPIRV.
+            // The return type of OpLogicalOr must be boolean type.
+            // If not, we need to recreate the instruction with boolean return type.
+            // Then, we have to cast it back to the original type so that other instrucitons that
+            // use have the matching types.
+            //
+            auto dataType = inst->getDataType();
+            auto lhs = inst->getOperand(0);
+            auto rhs = inst->getOperand(1);
+            IRInst* newInst = nullptr;
+
+            if (auto vecType = as<IRVectorType>(dataType))
+            {
+                if (!as<IRBoolType>(vecType->getElementType()))
+                {
+                    // Return type should be vector<bool,N>
+                    auto elemCount = vecType->getElementCount();
+                    auto vb = builder.getVectorType(builder.getBoolType(), elemCount);
+
+                    if (inst->getOp() == kIROp_And)
+                    {
+                        newInst = builder.emitAnd(vb, lhs, rhs);
+                    }
+                    else
+                    {
+                        newInst = builder.emitOr(vb, lhs, rhs);
+                    }
+                    newInst = builder.emitCast(dataType, newInst);
+                }
+            }
+            else if (!as<IRBoolType>(dataType))
+            {
+                // Return type should be bool
+                if (inst->getOp() == kIROp_And)
+                {
+                    newInst = builder.emitAnd(builder.getBoolType(), lhs, rhs);
+                }
+                else
+                {
+                    newInst = builder.emitOr(builder.getBoolType(), lhs, rhs);
+                }
+                newInst = builder.emitCast(dataType, newInst);
+            }
+
+            if (newInst && inst != newInst)
+            {
+                inst->replaceUsesWith(newInst);
+                inst->removeAndDeallocate();
+            }
+        }
+        break;
+    }
+
+    for (auto child : inst->getModifiableChildren())
+    {
+        legalizeLogicalAndOr(child);
+    }
+}
+
 } // namespace Slang

source/slang/slang-ir-legalize-binary-operator.h

@@ -13,4 +13,9 @@ struct IRInst;
 //   signed operand is converted to unsigned.
 void legalizeBinaryOp(IRInst* inst);
 
+// The logical binary operators such as AND and OR takes boolean types are its input.
+// If they are in integer type, as an example, we need to explicitly cast to bool type.
+// Also the return type from the logical operators should be a boolean type.
+void legalizeLogicalAndOr(IRInst* inst);
+
 } // namespace Slang

source/slang/slang-ir.cpp

@@ -6015,6 +6015,20 @@ IRInst* IRBuilder::emitShl(IRType* type, IRInst* left, IRInst* right)
     return inst;
 }
 
+IRInst* IRBuilder::emitAnd(IRType* type, IRInst* left, IRInst* right)
+{
+    auto inst = createInst<IRInst>(this, kIROp_And, type, left, right);
+    addInst(inst);
+    return inst;
+}
+
+IRInst* IRBuilder::emitOr(IRType* type, IRInst* left, IRInst* right)
+{
+    auto inst = createInst<IRInst>(this, kIROp_Or, type, left, right);
+    addInst(inst);
+    return inst;
+}
+
 IRInst* IRBuilder::emitGetNativePtr(IRInst* value)
 {
     auto valueType = value->getDataType();

tests/compute/logic-no-short-circuit-evaluation.slang

@@ -1,10 +1,11 @@
 //TEST(compute):SIMPLE(filecheck=SM5):-target hlsl -profile cs_5_1 -entry computeMain
 //TEST(compute):SIMPLE(filecheck=SM6):-target hlsl -profile cs_6_0 -entry computeMain
 //TEST(compute):SIMPLE(filecheck=WGSL):-target wgsl -stage compute -entry computeMain
-//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHK):-slang -compute -shaderobj -xslang -Wno-30056
-//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHK):-vk -compute  -shaderobj -xslang -Wno-30056
-//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHK):-cuda -compute -shaderobj -xslang -Wno-30056
-//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHK):-cpu -compute -shaderobj -xslang -Wno-30056
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHK):-slang -compute -shaderobj -output-using-type -xslang -Wno-30056
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHK):-vk -compute  -shaderobj -output-using-type -xslang -Wno-30056
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHK):-mtl -compute  -shaderobj -output-using-type -xslang -Wno-30056
+//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHK):-cuda -compute -shaderobj -output-using-type -xslang -Wno-30056
+//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHK):-cpu -compute -shaderobj -output-using-type -xslang -Wno-30056
 
 // Testnig logical-AND, logical-OR and ternary operator with non-scalar operands
 
@@ -15,7 +16,7 @@ static int result = 0;
 
 bool2 assignFunc(int index)
 {
-    result++;
+    result += 10;
     return bool2(true);
 }
 
@@ -24,34 +25,38 @@ void computeMain(int3 dispatchThreadID : SV_DispatchThreadID)
 {
     int index = dispatchThreadID.x;
 
+    // No short-circuiting for vector types
+
     //SM5:(all({{.*}}&&
     //SM6:(all(and(
-    //WGSL:(all(select(vec2<bool>(false),
+    //WGSL:(all( select(vec2<bool>(false),
     if (all(bool2(index >= 1) && assignFunc(index)))
     {
         result++;
     }
 
+    // Intentionally using non-boolean type for testing.
+
     //SM5:(all({{.*}}||
-    //SM6:(all(or(
-    //WGSL:(all(select({{.*}}vec2<bool>(true),
-    if (all(bool2(index >= 2) || !assignFunc(index)))
+    //SM6:(or(vector<bool,2>(
+    //WGSL:( select({{.*}}, vec2<bool>(true), vec2<bool>(
+    if (all(int2(index >= 2) || !assignFunc(index)))
     {
         result++;
     }
 
     //SM5:(all({{.*}}?{{.*}}:
-    //SM6:(all(select({{ *}}
-    //WGSL:(all(select(
+    //SM6:(all(select(
+    //WGSL:(all(select(vec2<bool>(false),
     if (all(bool2(index >= 3) ? assignFunc(index) : bool2(false)))
     {
         result++;
     }
 
     outputBuffer[index] = result;
 
-    //CHK:3
-    //CHK-NEXT:4
-    //CHK-NEXT:5
-    //CHK-NEXT:6
+    //CHK:30
+    //CHK-NEXT:31
+    //CHK-NEXT:32
+    //CHK-NEXT:33
 }

tests/compute/logic-short-circuit-evaluation.slang

@@ -1,5 +1,6 @@
 //TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHK):-dx12 -compute -shaderobj
 //TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHK):-vk -compute  -shaderobj
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHK):-mtl -compute  -shaderobj
 //TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHK):-cuda -compute -shaderobj
 //TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHK):-cpu -compute -compile-arg -O3 -shaderobj
 //TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHK):-slang -compute -shaderobj