Cleanup atomic intrinsics. (#5324)

* Cleanup atomic intrinsics.

* Fix.

* Fix glsl.

* Remove hacky intrinsic expansion logic for glsl image atomics.

* Fix all tests.

* Fix.

* Add `InterlockedAddF16Emulated`.

* Fix glsl intrinsic.

* Fix.

1 files changed, 198 insertions, 28 deletions

diff --git a/source/slang/slang-emit-glsl.cpp b/source/slang/slang-emit-glsl.cpp
index ca5569602..7f8bc14b4 100644
--- a/source/slang/slang-emit-glsl.cpp
+++ b/source/slang/slang-emit-glsl.cpp
@@ -2153,8 +2153,50 @@ bool GLSLSourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOu
     return false;
 }
 
+static IRImageSubscript* isTextureAccess(IRInst* inst)
+{
+    return as<IRImageSubscript>(getRootAddr(inst->getOperand(0)));
+}
+
+void GLSLSourceEmitter::emitAtomicImageCoord(IRImageSubscript* inst)
+{
+    emitOperand(inst->getImage(), getInfo(EmitOp::General));
+    m_writer->emit(", ");
+    if (auto vecType = as<IRVectorType>(inst->getCoord()->getDataType()))
+    {
+        m_writer->emit("ivec");
+        m_writer->emit(getIntVal(vecType->getElementCount()));
+    }
+    else
+    {
+        m_writer->emit("int");
+    }
+    m_writer->emit("(");
+    emitOperand(inst->getCoord(), getInfo(EmitOp::General));
+    m_writer->emit(")");
+    if (inst->hasSampleCoord())
+    {
+        m_writer->emit(", ");
+        emitOperand(inst->getSampleCoord(), getInfo(EmitOp::General));
+    }
+}
+
 bool GLSLSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
 {
+    auto requireAtomicExtIfNeeded = [&]()
+        {
+            if (isFloatingType(inst->getDataType()))
+            {
+                _requireGLSLExtension(toSlice("GL_EXT_shader_atomic_float"));
+            }
+            if (isIntegralType(inst->getDataType()))
+            {
+                if (getIntTypeInfo(inst->getDataType()).width == 64)
+                {
+                    _requireGLSLExtension(toSlice("GL_EXT_shader_atomic_int64"));
+                }
+            }
+        };
     switch (inst->getOp())
     {
     case kIROp_StructuredBufferGetDimensions:
@@ -2176,24 +2218,52 @@ bool GLSLSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
     case kIROp_AtomicLoad:
     {
         emitInstResultDecl(inst);
-        emitDereferenceOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        if (auto imageSubscript = isTextureAccess(inst))
+        {
+            m_writer->emit("imageLoad(");
+            emitAtomicImageCoord(imageSubscript);
+            m_writer->emit(")");
+        }
+        else
+        {
+            emitDereferenceOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        }
         m_writer->emit(";\n");
         return true;
     }
     case kIROp_AtomicStore:
     {
-        emitInstResultDecl(inst);
-        emitDereferenceOperand(inst->getOperand(0), getInfo(EmitOp::General));
-        m_writer->emit(" = ");
-        emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
-        m_writer->emit(";\n");
+        if (auto imageSubscript = isTextureAccess(inst))
+        {
+            m_writer->emit("imageStore(");
+            emitAtomicImageCoord(imageSubscript);
+            m_writer->emit(", ");
+            emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
+            m_writer->emit(")");
+        }
+        else
+        {
+            emitDereferenceOperand(inst->getOperand(0), getInfo(EmitOp::General));
+            m_writer->emit(" = ");
+            emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
+            m_writer->emit(";\n");
+        }
         return true;
     }
     case kIROp_AtomicExchange:
     {
+        requireAtomicExtIfNeeded();
         emitInstResultDecl(inst);
-        m_writer->emit("atomicExchange(");
-        emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        if (auto imageSubscript = isTextureAccess(inst))
+        {
+            m_writer->emit("imageAtomicExchange(");
+            emitAtomicImageCoord(imageSubscript);
+        }
+        else
+        {
+            m_writer->emit("atomicExchange(");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        }
         m_writer->emit(", ");
         emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
         m_writer->emit(");\n");
@@ -2201,9 +2271,19 @@ bool GLSLSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
     }
     case kIROp_AtomicCompareExchange:
     {
+        requireAtomicExtIfNeeded();
+
         emitInstResultDecl(inst);
-        m_writer->emit("atomicCompSwap(");
-        emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        if (auto imageSubscript = isTextureAccess(inst))
+        {
+            m_writer->emit("imageAtomicCompSwap(");
+            emitAtomicImageCoord(imageSubscript);
+        }
+        else
+        {
+            m_writer->emit("atomicCompSwap(");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        }
         m_writer->emit(", ");
         emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
         m_writer->emit(", ");
@@ -2213,9 +2293,19 @@ bool GLSLSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
     }
     case kIROp_AtomicAdd:
     {
+        requireAtomicExtIfNeeded();
+
         emitInstResultDecl(inst);
-        m_writer->emit("atomicAdd(");
-        emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        if (auto imageSubscript = isTextureAccess(inst))
+        {
+            m_writer->emit("imageAtomicAdd(");
+            emitAtomicImageCoord(imageSubscript);
+        }
+        else
+        {
+            m_writer->emit("atomicAdd(");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        }
         m_writer->emit(", ");
         emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
         m_writer->emit(");\n");
@@ -2223,9 +2313,19 @@ bool GLSLSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
     }
     case kIROp_AtomicSub:
     {
+        requireAtomicExtIfNeeded();
+
         emitInstResultDecl(inst);
-        m_writer->emit("atomicAdd(");
-        emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        if (auto imageSubscript = isTextureAccess(inst))
+        {
+            m_writer->emit("imageAtomicAdd(");
+            emitAtomicImageCoord(imageSubscript);
+        }
+        else
+        {
+            m_writer->emit("atomicAdd(");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        }
         m_writer->emit(", -(");
         emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
         m_writer->emit("));\n");
@@ -2233,9 +2333,19 @@ bool GLSLSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
     }
     case kIROp_AtomicAnd:
     {
+        requireAtomicExtIfNeeded();
+
         emitInstResultDecl(inst);
-        m_writer->emit("atomicAnd(");
-        emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        if (auto imageSubscript = isTextureAccess(inst))
+        {
+            m_writer->emit("imageAtomicAnd(");
+            emitAtomicImageCoord(imageSubscript);
+        }
+        else
+        {
+            m_writer->emit("atomicAnd(");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        }
         m_writer->emit(", ");
         emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
         m_writer->emit(");\n");
@@ -2243,9 +2353,19 @@ bool GLSLSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
     }
     case kIROp_AtomicOr:
     {
+        requireAtomicExtIfNeeded();
+
         emitInstResultDecl(inst);
-        m_writer->emit("atomicOr(");
-        emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        if (auto imageSubscript = isTextureAccess(inst))
+        {
+            m_writer->emit("imageAtomicOr(");
+            emitAtomicImageCoord(imageSubscript);
+        }
+        else
+        {
+            m_writer->emit("atomicOr(");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        }
         m_writer->emit(", ");
         emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
         m_writer->emit(");\n");
@@ -2253,9 +2373,19 @@ bool GLSLSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
     }
     case kIROp_AtomicXor:
     {
+        requireAtomicExtIfNeeded();
+
         emitInstResultDecl(inst);
-        m_writer->emit("atomicXor(");
-        emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        if (auto imageSubscript = isTextureAccess(inst))
+        {
+            m_writer->emit("imageAtomicXor(");
+            emitAtomicImageCoord(imageSubscript);
+        }
+        else
+        {
+            m_writer->emit("atomicXor(");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        }
         m_writer->emit(", ");
         emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
         m_writer->emit(");\n");
@@ -2263,9 +2393,19 @@ bool GLSLSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
     }
     case kIROp_AtomicMin:
     {
+        requireAtomicExtIfNeeded();
+
         emitInstResultDecl(inst);
-        m_writer->emit("atomicMin(");
-        emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        if (auto imageSubscript = isTextureAccess(inst))
+        {
+            m_writer->emit("imageAtomicMin(");
+            emitAtomicImageCoord(imageSubscript);
+        }
+        else
+        {
+            m_writer->emit("atomicMin(");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        }
         m_writer->emit(", ");
         emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
         m_writer->emit(");\n");
@@ -2273,9 +2413,19 @@ bool GLSLSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
     }
     case kIROp_AtomicMax:
     {
+        requireAtomicExtIfNeeded();
+
         emitInstResultDecl(inst);
-        m_writer->emit("atomicMax(");
-        emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        if (auto imageSubscript = isTextureAccess(inst))
+        {
+            m_writer->emit("imageAtomicMax(");
+            emitAtomicImageCoord(imageSubscript);
+        }
+        else
+        {
+            m_writer->emit("atomicMax(");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        }
         m_writer->emit(", ");
         emitOperand(inst->getOperand(1), getInfo(EmitOp::General));
         m_writer->emit(");\n");
@@ -2283,9 +2433,19 @@ bool GLSLSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
     }
     case kIROp_AtomicInc:
     {
+        requireAtomicExtIfNeeded();
+
         emitInstResultDecl(inst);
-        m_writer->emit("atomicAdd(");
-        emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        if (auto imageSubscript = isTextureAccess(inst))
+        {
+            m_writer->emit("imageAtomicAdd(");
+            emitAtomicImageCoord(imageSubscript);
+        }
+        else
+        {
+            m_writer->emit("atomicAdd(");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        }
         m_writer->emit(", ");
         emitType(inst->getDataType());
         m_writer->emit("(1)");
@@ -2294,9 +2454,19 @@ bool GLSLSourceEmitter::tryEmitInstStmtImpl(IRInst* inst)
     }
     case kIROp_AtomicDec:
     {
+        requireAtomicExtIfNeeded();
+
         emitInstResultDecl(inst);
-        m_writer->emit("atomicAdd(");
-        emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        if (auto imageSubscript = isTextureAccess(inst))
+        {
+            m_writer->emit("imageAtomicAdd(");
+            emitAtomicImageCoord(imageSubscript);
+        }
+        else
+        {
+            m_writer->emit("atomicAdd(");
+            emitOperand(inst->getOperand(0), getInfo(EmitOp::General));
+        }
         m_writer->emit(", ");
         emitType(inst->getDataType());
         m_writer->emit("(-1)");