IR: add support for `discard` statement (#261)

- Add definition of `discard` instruction
  - A `discard` is a terminator instruction, just like `returnVoid`

- Lower `DiscardStmt` in AST to a `discard` instruction in the IR

- Emit `discard` instruction as a `discard;` statement when emitting HLSL/GLSL

- Add a test case using the "graphics compute" mode that tests discard. The test writes to one entry in a UAV before doing a conditional (always true at runtime) discard, and then writes to another entry; we expect to see the results of the first write, but not the second.

7 files changed, 122 insertions, 0 deletions

diff --git a/source/slang/emit.cpp b/source/slang/emit.cpp
index 332125fc6..2b7030897 100644
--- a/source/slang/emit.cpp
+++ b/source/slang/emit.cpp
@@ -5296,6 +5296,10 @@ emitDeclImpl(decl, nullptr);
             emit(";\n");
             break;
 
+        case kIROp_discard:
+            emit("discard;\n");
+            break;
+
         case kIROp_swizzleSet:
             {
                 auto ii = (IRSwizzleSet*)inst;
@@ -5425,6 +5429,7 @@ emitDeclImpl(decl, nullptr);
 
             case kIROp_ReturnVal:
             case kIROp_ReturnVoid:
+            case kIROp_discard:
                 emitIRInst(ctx, terminator);
                 return;
 
diff --git a/source/slang/ir-inst-defs.h b/source/slang/ir-inst-defs.h
index d2e2c12e8..a8118f714 100644
--- a/source/slang/ir-inst-defs.h
+++ b/source/slang/ir-inst-defs.h
@@ -181,6 +181,8 @@ INST(if, if, 3, 0)
 INST(ifElse, ifElse, 4, 0)
 INST(loopTest, loopTest, 3, 0)
 
+INST(discard, discard, 0, 0)
+
 INST(Add, add, 2, 0)
 INST(Sub, sub, 2, 0)
 INST(Mul, mul, 2, 0)
diff --git a/source/slang/ir-insts.h b/source/slang/ir-insts.h
index f70ab46ff..06ccf2921 100644
--- a/source/slang/ir-insts.h
+++ b/source/slang/ir-insts.h
@@ -135,6 +135,9 @@ struct IRReturnVal : IRReturn
 struct IRReturnVoid : IRReturn
 {};
 
+struct IRDiscard : IRTerminatorInst
+{};
+
 struct IRBlock;
 
 struct IRUnconditionalBranch : IRTerminatorInst
@@ -464,6 +467,8 @@ struct IRBuilder
 
     IRInst* emitReturn();
 
+    IRInst* emitDiscard();
+
     IRInst* emitBranch(
         IRBlock*    block);
 
diff --git a/source/slang/ir.cpp b/source/slang/ir.cpp
index b1af6521c..8fbe20aa6 100644
--- a/source/slang/ir.cpp
+++ b/source/slang/ir.cpp
@@ -166,6 +166,7 @@ namespace Slang
         case kIROp_if:
         case kIROp_ifElse:
         case kIROp_loopTest:
+        case kIROp_discard:
             return true;
         }
     }
@@ -1150,6 +1151,17 @@ namespace Slang
         return inst;
     }
 
+    IRInst* IRBuilder::emitDiscard()
+    {
+        auto inst = createInst<IRDiscard>(
+            this,
+            kIROp_discard,
+            nullptr);
+        addInst(inst);
+        return inst;
+    }
+
+
     IRInst* IRBuilder::emitBranch(
         IRBlock*    pBlock)
     {
diff --git a/source/slang/lower-to-ir.cpp b/source/slang/lower-to-ir.cpp
index 4fabd6a81..86f037435 100644
--- a/source/slang/lower-to-ir.cpp
+++ b/source/slang/lower-to-ir.cpp
@@ -1912,6 +1912,11 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
             getBuilder()->emitReturn();
         }
     }
+
+    void visitDiscardStmt(DiscardStmt* /*stmt*/)
+    {
+        getBuilder()->emitDiscard();
+    }
 };
 
 void lowerStmt(
diff --git a/tests/compute/discard-stmt.slang b/tests/compute/discard-stmt.slang
new file mode 100644
index 000000000..18ffc39e2
--- /dev/null
+++ b/tests/compute/discard-stmt.slang
@@ -0,0 +1,91 @@
+//TEST(compute):COMPARE_RENDER_COMPUTE:
+//TEST_INPUT: Texture2D(size=4, content = one) : dxbinding(0),glbinding(0)
+//TEST_INPUT: Sampler : dxbinding(0),glbinding(0)
+//TEST_INPUT: ubuffer(data=[0 0], stride=4):dxbinding(1),glbinding(0),out
+
+
+Texture2D tex;
+SamplerState samp;
+RWStructuredBuffer<float> outputBuffer;
+
+cbuffer Uniforms
+{
+	float4x4 modelViewProjection;
+}
+
+struct AssembledVertex
+{
+	float3	position;
+	float3	color;
+    float2  uv;
+};
+
+struct CoarseVertex
+{
+	float3	color;
+    float2  uv;
+};
+
+struct Fragment
+{
+	float4 color;
+};
+
+
+// Vertex  Shader
+
+struct VertexStageInput
+{
+	AssembledVertex assembledVertex	: A;
+};
+
+struct VertexStageOutput
+{
+	CoarseVertex	coarseVertex	: CoarseVertex;
+	float4			sv_position		: SV_Position;
+};
+
+VertexStageOutput vertexMain(VertexStageInput input)
+{
+	VertexStageOutput output;
+
+	float3 position = input.assembledVertex.position;
+	float3 color	= input.assembledVertex.color;
+
+	output.coarseVertex.color = color;
+	output.sv_position = mul(modelViewProjection, float4(position, 1.0));
+    output.coarseVertex.uv = input.assembledVertex.uv;
+	return output;
+}
+
+// Fragment Shader
+
+struct FragmentStageInput
+{
+	CoarseVertex	coarseVertex	: CoarseVertex;
+};
+
+struct FragmentStageOutput
+{
+	Fragment fragment	: SV_Target;
+};
+
+FragmentStageOutput fragmentMain(FragmentStageInput input)
+{
+	FragmentStageOutput output;
+
+	float3 color = input.coarseVertex.color;
+    float2 uv = input.coarseVertex.uv;
+	output.fragment.color = float4(color, 1.0);
+
+	float4 val = float4(color, 1.0);
+	val = val - 16*tex.Sample(samp, uv);
+
+	outputBuffer[0] = 1;
+
+	if(val.x < 0)
+		discard;
+
+	outputBuffer[1] = 1;
+	return output;
+}
\ No newline at end of file
diff --git a/tests/compute/discard-stmt.slang.expected.txt b/tests/compute/discard-stmt.slang.expected.txt
new file mode 100644
index 000000000..e0e43c4e8
--- /dev/null
+++ b/tests/compute/discard-stmt.slang.expected.txt
@@ -0,0 +1,2 @@
+3F800000
+0