1 files changed, 193 insertions, 0 deletions
diff --git a/extensions/dx/flexExt.cpp b/extensions/dx/flexExt.cpp
new file mode 100644
index 0000000..4dcc0e4
--- /dev/null
+++ b/extensions/dx/flexExt.cpp
@@ -0,0 +1,193 @@
+// This code contains NVIDIA Confidential Information and is disclosed to you
+// under a form of NVIDIA software license agreement provided separately to you.
+//
+// Notice
+// NVIDIA Corporation and its licensors retain all intellectual property and
+// proprietary rights in and to this software and related documentation and
+// any modifications thereto. Any use, reproduction, disclosure, or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA Corporation is strictly prohibited.
+//
+// ALL NVIDIA DESIGN SPECIFICATIONS, CODE ARE PROVIDED "AS IS.". NVIDIA MAKES
+// NO WARRANTIES, EXPRESSED, IMPLIED, STATUTORY, OR OTHERWISE WITH RESPECT TO
+// THE MATERIALS, AND EXPRESSLY DISCLAIMS ALL IMPLIED WARRANTIES OF NONINFRINGEMENT,
+// MERCHANTABILITY, AND FITNESS FOR A PARTICULAR PURPOSE.
+//
+// Information and code furnished is believed to be accurate and reliable.
+// However, NVIDIA Corporation assumes no responsibility for the consequences of use of such
+// information or for any infringement of patents or other rights of third parties that may
+// result from its use. No license is granted by implication or otherwise under any patent
+// or patent rights of NVIDIA Corporation. Details are subject to change without notice.
+// This code supersedes and replaces all information previously supplied.
+// NVIDIA Corporation products are not authorized for use as critical
+// components in life support devices or systems without express written approval of
+// NVIDIA Corporation.
+//
+// Copyright (c) 2013-2017 NVIDIA Corporation. All rights reserved.
+
+#include "core/core.h"
+#include "core/maths.h"
+
+#include "include/NvFlex.h"
+#include "include/NvFlexExt.h"
+
+#include "src/dx/context/context.h"
+#include "src/dx/context/device.h"
+
+#include "flexExt_dx_common.h"
+
+#include "shaders\flexExt.UpdateForceFields.h"
+
+
+struct NvFlexExtForceFieldCallback
+{
+	NvFlexExtForceFieldCallback(NvFlexSolver* solver) : mSolver(solver)
+	{
+		// force fields
+		mMaxForceFields = 0;
+		mNumForceFields = 0;
+		
+		mForceFieldsGpu = NULL;
+
+		mDevice = NULL;
+		mContext = NULL;
+
+		NvFlexLibrary* lib = NvFlexGetSolverLibrary(solver);
+		NvFlexGetDeviceAndContext(lib, (void**)&mDevice, (void**)&mContext);
+
+		{
+			// force field shader
+			NvFlex::ComputeShaderDesc desc{};
+			desc.cs = (void*)g_flexExt_UpdateForceFields;
+			desc.cs_length = sizeof(g_flexExt_UpdateForceFields);
+			desc.label = L"NvFlexExtForceFieldCallback";
+			desc.NvAPI_Slot = 0;
+
+			mShaderUpdateForceFields = mContext->createComputeShader(&desc);
+		}
+
+		{
+			// constant buffer
+			NvFlex::ConstantBufferDesc desc;
+			desc.stride = sizeof(int);
+			desc.dim = 4;
+			desc.uploadAccess = true;
+			
+			mConstantBuffer = mContext->createConstantBuffer(&desc);
+		}
+	}
+
+	~NvFlexExtForceFieldCallback()
+	{
+		// force fields
+		delete mForceFieldsGpu;
+		delete mConstantBuffer;
+		delete mShaderUpdateForceFields;
+	}
+	
+	NvFlex::Buffer* mForceFieldsGpu;
+
+	// DX Specific
+	NvFlex::ComputeShader* mShaderUpdateForceFields;
+	NvFlex::ConstantBuffer* mConstantBuffer;
+
+	int mMaxForceFields;
+	int mNumForceFields;
+
+	// D3D device and context wrappers for the solver library
+	NvFlex::Device* mDevice;
+	NvFlex::Context* mContext;
+
+	NvFlexSolver* mSolver;
+
+};
+
+NvFlexExtForceFieldCallback* NvFlexExtCreateForceFieldCallback(NvFlexSolver* solver)
+{
+	return new NvFlexExtForceFieldCallback(solver);	
+}
+
+void NvFlexExtDestroyForceFieldCallback(NvFlexExtForceFieldCallback* callback)
+{
+	delete callback;
+}
+
+void ApplyForceFieldsCallback(NvFlexSolverCallbackParams params)
+{
+	// callbacks always have the correct CUDA device set so we can safely launch kernels without acquiring
+
+	NvFlexExtForceFieldCallback* c = (NvFlexExtForceFieldCallback*)params.userData;
+
+	if (params.numActive && c->mNumForceFields)
+	{
+		const unsigned int numThreadsPerBlock = 256;
+		const unsigned int kNumBlocks = (params.numActive + numThreadsPerBlock - 1) / numThreadsPerBlock;
+
+		NvFlex::Buffer* particles = (NvFlex::Buffer*)params.particles;
+		NvFlex::Buffer* velocities = (NvFlex::Buffer*)params.velocities;
+
+		// Init constant buffer
+		{
+			FlexExtConstParams constBuffer;
+
+			constBuffer.kNumParticles = params.numActive;
+			constBuffer.kNumForceFields = c->mNumForceFields;
+			constBuffer.kDt = params.dt;
+
+			memcpy(c->mContext->map(c->mConstantBuffer), &constBuffer, sizeof(FlexExtConstParams));
+			c->mContext->unmap(c->mConstantBuffer);
+		}
+
+		{
+			NvFlex::DispatchParams params = {};
+			params.shader = c->mShaderUpdateForceFields;
+			params.readWrite[0] = velocities->getResourceRW();
+			params.readOnly[0] = particles->getResource();
+			params.readOnly[1] = c->mForceFieldsGpu->getResource();
+			params.gridDim = { kNumBlocks , 1, 1 };
+			params.rootConstantBuffer = c->mConstantBuffer;
+
+			c->mContext->dispatch(&params);
+		}
+	}
+}
+
+void NvFlexExtSetForceFields(NvFlexExtForceFieldCallback* c, const NvFlexExtForceField* forceFields, int numForceFields)
+{
+	// re-alloc if necessary
+	if (numForceFields > c->mMaxForceFields)
+	{
+		delete c->mForceFieldsGpu;
+
+		NvFlex::BufferDesc desc {};
+		desc.dim = numForceFields;
+		desc.stride = sizeof(NvFlexExtForceField);
+		desc.bufferType = NvFlex::eBuffer | NvFlex::eUAV_SRV | NvFlex::eStructured | NvFlex::eStage; 
+		desc.format = NvFlexFormat::eNvFlexFormat_unknown;
+		desc.data = NULL;
+	
+		c->mForceFieldsGpu = c->mContext->createBuffer(&desc);
+
+		c->mMaxForceFields = numForceFields;
+	}
+	c->mNumForceFields = numForceFields;
+
+	if (numForceFields > 0)
+	{
+		// update staging buffer
+		void* dstPtr = c->mContext->mapUpload(c->mForceFieldsGpu);
+		memcpy(dstPtr, forceFields, numForceFields*sizeof(NvFlexExtForceField));
+		c->mContext->unmap(c->mForceFieldsGpu);
+
+		// upload to device buffer
+		c->mContext->upload(c->mForceFieldsGpu, 0, numForceFields*sizeof(NvFlexExtForceField));
+
+	}
+
+	NvFlexSolverCallback callback;
+	callback.function = ApplyForceFieldsCallback;
+	callback.userData = c;
+
+	// register a callback to calculate the forces at the end of the time-step
+	NvFlexRegisterSolverCallback(c->mSolver, callback, eNvFlexStageUpdateEnd);
+}