Native call marshalling for ComPtr parameters and return values. (#2305)

Co-authored-by: Yong He <yhe@nvidia.com>

7 files changed, 180 insertions, 7 deletions

diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index 10ac2ca84..9bd999087 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -1544,8 +1544,11 @@ void CLikeSourceEmitter::emitCallExpr(IRCall* inst, EmitOpInfo outerPrec)
     {
         const auto operand0TypeOp = funcValue->getOperand(0)->getDataType()->getOp();
 
-        if (operand0TypeOp == kIROp_ComPtrType || operand0TypeOp == kIROp_PtrType)
+        switch (operand0TypeOp)
         {
+        case kIROp_ComPtrType:
+        case kIROp_PtrType:
+        case kIROp_NativePtrType:
             emitComInterfaceCallExpr(inst, outerPrec);
             return;
         }
@@ -1825,7 +1828,40 @@ void CLikeSourceEmitter::defaultEmitInstExpr(IRInst* inst, const EmitOpInfo& inO
             m_writer->emit(")");
             break;
         }
-
+    case kIROp_GetNativePtr:
+    {
+        auto prec = getInfo(EmitOp::Postfix);
+        needClose = maybeEmitParens(outerPrec, prec);
+        emitOperand(inst->getOperand(0), leftSide(outerPrec, prec));
+        m_writer->emit(".get()");
+        break;
+    }
+    case kIROp_GetManagedPtrWriteRef:
+    {
+        auto prec = getInfo(EmitOp::Postfix);
+        needClose = maybeEmitParens(outerPrec, prec);
+        emitDereferenceOperand(inst->getOperand(0), leftSide(outerPrec, prec));
+        m_writer->emit(".writeRef()");
+        break;
+    }
+    case kIROp_ManagedPtrAttach:
+    {
+        auto prec = getInfo(EmitOp::Postfix);
+        needClose = maybeEmitParens(outerPrec, prec);
+        emitDereferenceOperand(inst->getOperand(0), leftSide(outerPrec, prec));
+        m_writer->emit(".attach(");
+        emitOperand(inst->getOperand(1), EmitOpInfo());
+        m_writer->emit(")");
+        break;
+    }
+    case kIROp_ManagedPtrDetach:
+    {
+        auto prec = getInfo(EmitOp::Postfix);
+        needClose = maybeEmitParens(outerPrec, prec);
+        emitDereferenceOperand(inst->getOperand(0), leftSide(outerPrec, prec));
+        m_writer->emit(".detach()");
+        break;
+    }
     case kIROp_getElement:
     case kIROp_getElementPtr:
     case kIROp_ImageSubscript:
diff --git a/source/slang/slang-emit-cpp.cpp b/source/slang/slang-emit-cpp.cpp
index 31f09e400..b52407412 100644
--- a/source/slang/slang-emit-cpp.cpp
+++ b/source/slang/slang-emit-cpp.cpp
@@ -522,6 +522,14 @@ SlangResult CPPSourceEmitter::calcTypeName(IRType* type, CodeGenTarget target, S
             out << "*";
             return SLANG_OK;
         }
+        case kIROp_NativePtrType:
+        case kIROp_PtrType:
+        {
+            auto elementType = (IRType*)type->getOperand(0);
+            SLANG_RETURN_ON_FAIL(calcTypeName(elementType, target, out));
+            out << "*";
+            return SLANG_OK;
+        }
         case kIROp_RTTIType:
         {
             out << "TypeInfo";
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index 6eae710c0..b4530d655 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -112,6 +112,8 @@ INST(Nop, nop, 0, 0)
 
     // A ComPtr<T> type is treated as a opaque type that represents a reference-counted handle to a COM object.
     INST(ComPtrType, ComPtr, 1, 0)
+    // A NativePtr<T> type represents a native pointer to a managed resource.
+    INST(NativePtrType, NativePtr, 1, 0)
 
     /* SamplerStateTypeBase */
         INST(SamplerStateType, SamplerState, 0, 0)
@@ -314,6 +316,18 @@ INST(getNativeStr, getNativeStr, 1, 0)
 // Make String from a NativeString.
 INST(makeString, makeString, 1, 0)
 
+// Get a native ptr from a ComPtr or RefPtr
+INST(GetNativePtr, getNativePtr, 1, 0)
+
+// Get a write reference to a managed ptr var (operand must be Ptr<ComPtr<T>> or Ptr<RefPtr<T>>).
+INST(GetManagedPtrWriteRef, getManagedPtrWriteRef, 1, 0)
+
+// Attach a managedPtr var to a NativePtr without changing its ref count.
+INST(ManagedPtrAttach, ManagedPtrAttach, 1, 0)
+
+// Attach a managedPtr var to a NativePtr without changing its ref count.
+INST(ManagedPtrDetach, ManagedPtrDetach, 1, 0)
+
 // "Subscript" an image at a pixel coordinate to get pointer
 INST(ImageSubscript, imageSubscript, 2, 0)
 
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index 521570b22..7d0609645 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -1425,6 +1425,18 @@ struct IRGetElementPtr : IRInst
     IRInst* getIndex() { return getOperand(1); }
 };
 
+struct IRGetNativePtr : IRInst
+{
+    IR_LEAF_ISA(GetNativePtr);
+    IRInst* getElementType() { return getOperand(0); }
+};
+
+struct IRGetManagedPtrWriteRef : IRInst
+{
+    IR_LEAF_ISA(GetManagedPtrWriteRef);
+    IRInst* getPtrToManagedPtr() { return getOperand(0); }
+};
+
 struct IRGetAddress : IRInst
 {
     IR_LEAF_ISA(getAddr);
@@ -2175,6 +2187,7 @@ public:
     IRBasicType* getCharType();
     IRStringType* getStringType();
     IRNativeStringType* getNativeStringType();
+    IRNativePtrType* getNativePtrType(IRType* valueType);
 
     IRType* getCapabilitySetType();
 
@@ -2506,6 +2519,14 @@ public:
         return emitWrapExistential(type, value, slotArgCount, slotArgVals.getBuffer());
     }
 
+    IRInst* emitManagedPtrAttach(IRInst* managedPtrVar, IRInst* value);
+
+    IRInst* emitManagedPtrDetach(IRInst* managedPtrVar);
+
+    IRInst* emitGetNativePtr(IRInst* value);
+
+    IRInst* emitGetManagedPtrWriteRef(IRInst* ptrToManagedPtr);
+
     IRInst* emitGpuForeach(List<IRInst*> args);
 
     IRUndefined* emitUndefined(IRType* type);
diff --git a/source/slang/slang-ir-marshal-native-call.cpp b/source/slang/slang-ir-marshal-native-call.cpp
index b8e2edb2f..8e342cb26 100644
--- a/source/slang/slang-ir-marshal-native-call.cpp
+++ b/source/slang/slang-ir-marshal-native-call.cpp
@@ -12,6 +12,10 @@ namespace Slang
         {
         case kIROp_StringType:
             return builder.getNativeStringType();
+        case kIROp_InterfaceType:
+            return builder.getNativePtrType(type);
+        case kIROp_ComPtrType:
+            return builder.getNativePtrType((IRType*)as<IRComPtrType>(type)->getOperand(0));
         default:
             return type;
         }
@@ -30,10 +34,14 @@ namespace Slang
         IRType* returnType = declaredFuncType->getResultType();
         if (auto resultType = as<IRResultType>(declaredFuncType->getResultType()))
         {
-            nativeParamTypes.add(builder.getPtrType(resultType->getValueType()));
+            auto nativeResultType = getNativeType(builder, resultType->getValueType());
+            nativeParamTypes.add(builder.getPtrType(nativeResultType));
             returnType = resultType->getErrorType();
         }
-        getNativeType(builder, declaredFuncType->getResultType());
+        else
+        {
+            returnType = getNativeType(builder, returnType);
+        }
         auto funcType = builder.getFuncType(
             nativeParamTypes.getCount(), (IRType**)nativeParamTypes.getBuffer(), returnType);
 
@@ -43,9 +51,18 @@ namespace Slang
     void NativeCallMarshallingContext::marshalRefManagedValueToNativeValue(
         IRBuilder& builder, IRInst* originalArg, List<IRInst*>& args)
     {
-        SLANG_UNUSED(builder);
-        SLANG_UNUSED(originalArg);
-        args.add(originalArg);
+        auto ptrTypeBase = as<IRPtrTypeBase>(originalArg->getDataType());
+        SLANG_RELEASE_ASSERT(ptrTypeBase);
+        switch (ptrTypeBase->getValueType()->getOp())
+        {
+        case kIROp_InterfaceType:
+        case kIROp_ComPtrType:
+            args.add(builder.emitGetManagedPtrWriteRef(originalArg));
+            break;
+        default:
+            args.add(originalArg);
+            break;
+        }
     }
 
     void NativeCallMarshallingContext::marshalManagedValueToNativeValue(
@@ -63,6 +80,12 @@ namespace Slang
                 args.add(nativeStr);
             }
             break;
+        case kIROp_InterfaceType:
+            {
+                auto nativePtr = builder.emitGetNativePtr(originalArg);
+                args.add(nativePtr);
+            }
+            break;
         default:
             args.add(originalArg);
             break;
@@ -79,6 +102,21 @@ namespace Slang
                 return builder.emitMakeString(nativeVal);
             }
             break;
+        case kIROp_NativePtrType:
+            {
+                SLANG_RELEASE_ASSERT(nativeVal->getDataType()->getOperand(0)->getOp() == kIROp_InterfaceType);
+                auto comPtrVar = builder.emitVar(builder.getComPtrType((IRType*)nativeVal->getDataType()->getOperand(0)));
+                builder.emitManagedPtrAttach(comPtrVar, nativeVal);
+                return builder.emitLoad(comPtrVar);
+            }
+            break;
+        case kIROp_InterfaceType:
+            {
+                auto comPtrVar = builder.emitVar(nativeVal->getDataType());
+                builder.emitManagedPtrAttach(comPtrVar, nativeVal);
+                return builder.emitLoad(comPtrVar);
+            }
+            break;
         default:
             return nativeVal;
             break;
diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index dd71ae782..081b9103f 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -2539,6 +2539,11 @@ namespace Slang
         return (IRNativeStringType*)getType(kIROp_NativeStringType);
     }
 
+    IRNativePtrType* IRBuilder::getNativePtrType(IRType* valueType)
+    {
+        return (IRNativePtrType*)getType(kIROp_NativePtrType, (IRInst*)valueType);
+    }
+
 
     IRType* IRBuilder::getCapabilitySetType()
     {
@@ -4441,6 +4446,56 @@ namespace Slang
         return inst;
     }
 
+    IRInst* IRBuilder::emitGetNativePtr(IRInst* value)
+    {
+        auto valueType = value->getDataType();
+        SLANG_RELEASE_ASSERT(valueType);
+        switch (valueType->getOp())
+        {
+        case kIROp_InterfaceType:
+            return emitIntrinsicInst(
+                getNativePtrType((IRType*)valueType), kIROp_GetNativePtr, 1, &value);
+            break;
+        case kIROp_ComPtrType:
+            return emitIntrinsicInst(
+                getNativePtrType((IRType*)valueType->getOperand(0)), kIROp_GetNativePtr, 1, &value);
+            break;
+        default:
+            SLANG_UNEXPECTED("invalid operand type for `getNativePtr`.");
+            UNREACHABLE_RETURN(nullptr);
+        }
+    }
+
+    IRInst* IRBuilder::emitManagedPtrAttach(IRInst* managedPtrVar, IRInst* value)
+    {
+        IRInst* args[] = { managedPtrVar, value };
+        return emitIntrinsicInst(getVoidType(), kIROp_ManagedPtrAttach, 2, args);
+    }
+
+    IRInst* IRBuilder::emitManagedPtrDetach(IRInst* managedPtrVar)
+    {
+        return emitIntrinsicInst(getVoidType(), kIROp_ManagedPtrDetach, 1, &managedPtrVar);
+    }
+
+    IRInst* IRBuilder::emitGetManagedPtrWriteRef(IRInst* ptrToManagedPtr)
+    {
+        auto type = ptrToManagedPtr->getDataType();
+        auto ptrType = as<IRPtrTypeBase>(type);
+        SLANG_RELEASE_ASSERT(ptrType);
+        auto managedPtrType = ptrType->getValueType();
+        switch (managedPtrType->getOp())
+        {
+        case kIROp_InterfaceType:
+        case kIROp_ComPtrType:
+            return emitIntrinsicInst(
+                getPtrType(getNativePtrType((IRType*)managedPtrType->getOperand(0))), kIROp_GetManagedPtrWriteRef, 1, &ptrToManagedPtr);
+            break;
+        default:
+            SLANG_UNEXPECTED("invalid operand type for `getNativePtr`.");
+            UNREACHABLE_RETURN(nullptr);
+        }
+    }
+
     IRInst* IRBuilder::emitGpuForeach(List<IRInst*> args)
     {
         auto inst = createInst<IRInst>(
diff --git a/source/slang/slang-ir.h b/source/slang/slang-ir.h
index 1ed7e52c7..0a27eb4cd 100644
--- a/source/slang/slang-ir.h
+++ b/source/slang/slang-ir.h
@@ -1300,6 +1300,7 @@ SIMPLE_IR_TYPE(OutType, OutTypeBase)
 SIMPLE_IR_TYPE(InOutType, OutTypeBase)
 
 SIMPLE_IR_TYPE(ComPtrType, Type)
+SIMPLE_IR_TYPE(NativePtrType, Type)
 
 struct IRPseudoPtrType : public IRPtrTypeBase
 {