diff --git a/include/madrona/block_list.hpp b/include/madrona/block_list.hpp deleted file mode 100644 index 307f5592..00000000 --- a/include/madrona/block_list.hpp +++ /dev/null @@ -1,92 +0,0 @@ -#pragma once - -#include -#include - -namespace madrona { - -template -class BlockList { - struct BlockBase {}; - - struct Metadata { - BlockBase *next; - CountT numElems; - }; - - template - struct Block : BlockBase { - T arr[num_elems]; - Metadata metadata; - }; - - static constexpr inline CountT computeElemsPerBlock() - { - static_assert(desired_elems_per_block >= 0); - - if constexpr (desired_elems_per_block != 0) { - return desired_elems_per_block; - } - - constexpr CountT default_block_size = 1024; - - sizeof(T) - - constexpr CountT num_elems = (default_block_size - sizeof(Metadata)) / sizeof(T); - - using TestT = Block - - if constexpr (num_elems == 0) { - return 1; - } - - return num_elems; - } - - static constexpr inline CountT per_block_ = computeElemsPerBlock(); - - struct Block { - T arr[per_block_]; - Metadata metadata; - }; - - -public: - BlockList() - : head_(nullptr) - {} - - class Iter { - public: - - private: - Block *cur_block_; - CountT cur_offset_; - - friend class BlockList; - }; - - Iter begin() - { - } - - Iter end() - { - return Iter { - }; - } - -private: - static inline constexpr CountT num_elems_ = - block_size / - - union Block { - Block *next; - CountT numElems; - }; - - Block *head_; -}; - -} diff --git a/include/madrona/broadphase.hpp b/include/madrona/broadphase.hpp deleted file mode 100644 index 2028fd96..00000000 --- a/include/madrona/broadphase.hpp +++ /dev/null @@ -1,119 +0,0 @@ -#pragma once - -#include -#include - -namespace madrona::phys { - -struct ObjectManager; - -} - -namespace madrona::phys::broadphase { - -struct LeafID { - int32_t id; -}; - -class BVH { -public: - BVH(const ObjectManager *obj_mgr, - CountT max_leaves, - float leaf_velocity_expansion, - float leaf_accel_expansion); - - inline LeafID reserveLeaf(Entity e, base::ObjectID obj_id); - inline math::AABB getLeafAABB(LeafID leaf_id) const; - - template - inline void findIntersecting(const math::AABB &aabb, Fn &&fn) const; - - template - inline void findLeafIntersecting(LeafID leaf_id, Fn &&fn) const; - - Entity traceRay(math::Vector3 o, - math::Vector3 d, - float *out_hit_t, - math::Vector3 *out_hit_normal, - float t_max = float(INFINITY)); - - void updateLeafPosition(LeafID leaf_id, - const math::Vector3 &pos, - const math::Quat &rot, - const math::Diag3x3 &scale, - const math::Vector3 &linear_vel, - const math::AABB &obj_aabb); - - math::AABB expandLeaf(LeafID leaf_id, - const math::Vector3 &linear_vel); - - void refitLeaf(LeafID leaf_id, const math::AABB &leaf_aabb); - - inline void rebuildOnUpdate(); - void updateTree(); - - inline void clearLeaves(); - -private: - static constexpr int32_t sentinel_ = 0xFFFF'FFFF_i32; - - struct Node { - float minX[4]; - float minY[4]; - float minZ[4]; - float maxX[4]; - float maxY[4]; - float maxZ[4]; - int32_t children[4]; - int32_t parentID; - - inline bool isLeaf(CountT child) const; - inline int32_t leafIDX(CountT child) const; - - inline void setLeaf(CountT child, int32_t idx); - inline void setInternal(CountT child, int32_t internal_idx); - inline bool hasChild(CountT child) const; - inline void clearChild(CountT child); - }; - - // FIXME: evaluate whether storing this in-line in the tree - // makes sense or if we should force a lookup through the entity ID - struct LeafTransform { - math::Vector3 pos; - math::Quat rot; - math::Diag3x3 scale; - }; - - inline CountT numInternalNodes(CountT num_leaves) const; - - void rebuild(); - void refit(LeafID *leaf_ids, CountT num_moved); - - bool traceRayIntoLeaf(int32_t leaf_idx, - math::Vector3 world_ray_o, - math::Vector3 world_ray_d, - float t_min, - float t_max, - float *hit_t, - math::Vector3 *hit_normal); - - Node *nodes_; - CountT num_nodes_; - const CountT num_allocated_nodes_; - Entity *leaf_entities_; - const ObjectManager *obj_mgr_; - base::ObjectID *leaf_obj_ids_; - math::AABB *leaf_aabbs_; // FIXME: remove this, it's duplicated data - LeafTransform *leaf_transforms_; - uint32_t *leaf_parents_; - int32_t *sorted_leaves_; - AtomicI32 num_leaves_; - int32_t num_allocated_leaves_; - float leaf_velocity_expansion_; - float leaf_accel_expansion_; - bool force_rebuild_; -}; - -} - -#include "broadphase.inl" diff --git a/include/madrona/broadphase.inl b/include/madrona/broadphase.inl deleted file mode 100644 index 494b35e1..00000000 --- a/include/madrona/broadphase.inl +++ /dev/null @@ -1,108 +0,0 @@ -namespace madrona::phys::broadphase { - -LeafID BVH::reserveLeaf(Entity e, base::ObjectID obj_id) -{ - int32_t leaf_idx = num_leaves_.fetch_add_relaxed(1); - assert(leaf_idx < num_allocated_leaves_); - - leaf_entities_[leaf_idx] = e; - leaf_obj_ids_[leaf_idx] = obj_id; - - return LeafID { - leaf_idx, - }; -} - -math::AABB BVH::getLeafAABB(LeafID leaf_id) const -{ - return leaf_aabbs_[leaf_id.id]; -} - -template -void BVH::findIntersecting(const math::AABB &aabb, Fn &&fn) const -{ - int32_t stack[32]; - stack[0] = 0; - CountT stack_size = 1; - - while (stack_size > 0) { - int32_t node_idx = stack[--stack_size]; - const Node &node = nodes_[node_idx]; - for (int i = 0; i < 4; i++) { - if (!node.hasChild(i)) { - continue; // Technically this could be break? - }; - - madrona::math::AABB child_aabb { - /* .pMin = */ { - node.minX[i], - node.minY[i], - node.minZ[i], - }, - /* .pMax = */ { - node.maxX[i], - node.maxY[i], - node.maxZ[i], - }, - }; - - if (aabb.overlaps(child_aabb)) { - if (node.isLeaf(i)) { - Entity e = leaf_entities_[node.leafIDX(i)]; - fn(e); - } else { - stack[stack_size++] = node.children[i]; - } - } - } - } -} - -template -void BVH::findLeafIntersecting(LeafID leaf_id, Fn &&fn) const -{ - math::AABB leaf_aabb = leaf_aabbs_[leaf_id.id]; - findIntersecting(leaf_aabb, std::forward(fn)); -} - -void BVH::rebuildOnUpdate() -{ - force_rebuild_ = true; -} - -void BVH::clearLeaves() -{ - num_leaves_.store_relaxed(0); -} - -bool BVH::Node::isLeaf(CountT child) const -{ - return children[child] & 0x80000000; -} - -int32_t BVH::Node::leafIDX(CountT child) const -{ - return children[child] & ~0x80000000; -} - -void BVH::Node::setLeaf(CountT child, int32_t idx) -{ - children[child] = 0x80000000 | idx; -} - -void BVH::Node::setInternal(CountT child, int32_t internal_idx) -{ - children[child] = internal_idx; -} - -bool BVH::Node::hasChild(CountT child) const -{ - return children[child] != sentinel_; -} - -void BVH::Node::clearChild(CountT child) -{ - children[child] = sentinel_; -} - -} diff --git a/include/madrona/custom_context.hpp b/include/madrona/custom_context.hpp index 1ef8f375..6e60defd 100644 --- a/include/madrona/custom_context.hpp +++ b/include/madrona/custom_context.hpp @@ -23,7 +23,6 @@ class CustomContext : public Context { private: using WorldDataT = DataT; -friend class JobManager; friend class TaskGraph; }; diff --git a/include/madrona/fwd.hpp b/include/madrona/fwd.hpp index f205451e..11358a2f 100644 --- a/include/madrona/fwd.hpp +++ b/include/madrona/fwd.hpp @@ -9,7 +9,6 @@ namespace madrona { -class JobManager; class StateCache; class StateManager; class ECSRegistry; diff --git a/include/madrona/job.hpp b/include/madrona/job.hpp deleted file mode 100644 index 3339bf70..00000000 --- a/include/madrona/job.hpp +++ /dev/null @@ -1,272 +0,0 @@ -/* - * Copyright 2021-2022 Brennan Shacklett and contributors - * - * Use of this source code is governed by an MIT-style - * license that can be found in the LICENSE file or at - * https://opensource.org/licenses/MIT. - */ -#pragma once - -#include -#include -#include -#include -#include - -#include -#include -#include -#include - -namespace madrona { - -enum class JobPriority { - High, - Normal, - IO, -}; - -struct JobID { - uint32_t gen; - int32_t id; - - static constexpr inline JobID none(); -}; - -struct JobContainerBase { - JobID id; - uint32_t jobSize; -#ifdef MADRONA_MW_MODE - uint32_t worldID; -#endif - uint32_t numDependencies; - - template struct DepsArray; -}; - -template -struct JobContainer : public JobContainerBase { - [[no_unique_address]] DepsArray dependencies; - [[no_unique_address]] Fn fn; - - template - inline JobContainer(uint32_t job_size, MADRONA_MW_COND(uint32_t world_id,) - Fn &&fn, DepTs ...deps); -}; - -struct Job { - void (*func)(); - JobContainerBase *data; - uint32_t invocationOffset; - uint32_t numInvocations; -}; - -class JobManager { -public: - template struct EntryConfig; - - template - static EntryConfig makeEntry( - StartFn &&start_fn); - - template - static EntryConfig makeEntry(StartFn &&start_fn, - UpdateFn &&update_fn); - - template - JobManager(const EntryConfig &entry_cfg, - int desired_num_workers, - int num_io, - StateManager *state_mgr, - bool pin_workers = true); - - ~JobManager(); - - inline JobID reserveProxyJobID(int thread_idx, JobID parent_id); - inline void relinquishProxyJobID(int thread_idx, JobID job_id); - - template - JobID queueJob(int thread_idx, - Fn &&fn, - uint32_t num_invocations, - JobID parent_id, - MADRONA_MW_COND(uint32_t world_id,) - JobPriority prio = JobPriority::Normal, - DepTs ...deps); - -#if 0 - JobID queueJobs(int thread_idx, JobID parent_id, - const Job *jobs, uint32_t num_jobs, - const JobID *deps, uint32_t num_dependencies, - JobPriority prio = JobPriority::Normal); -#endif - - void waitForAllFinished(); - - // Custom allocator that recycles small arenas out of a large chunk of - // preallocated memory. - class Alloc { - public: - struct Arena { - // Doubles as next pointer in freelist or - // count of num bytes freed before being put on freelist - AtomicU32 metadata; - }; - - struct SharedState { - void *memoryBase; - void *jobMemory; - Arena *arenas; - AtomicU32 freeHead; - }; - - static constexpr size_t maxJobSize = 1024; - static constexpr size_t maxJobAlignment = 128; - - Alloc(SharedState &shared); - void * alloc(SharedState &shared, uint32_t num_bytes, uint32_t alignment); - void dealloc(SharedState &shared, void *ptr, uint32_t num_bytes); - - // FIXME: fix InitAlloc ownership - static SharedState makeSharedState(InitAlloc alloc, - uint32_t num_arenas); - - private: - static constexpr size_t arena_size_ = 4096; - - uint32_t cur_arena_; - uint32_t next_arena_; - uint32_t arena_offset_; - uint32_t arena_used_bytes_; - }; - - struct RunQueue { - AtomicU32 head; - AtomicU32 correction; - AtomicU32 auth; - char pad[MADRONA_CACHE_LINE - sizeof(AtomicU32) * 3]; - AtomicU32 tail; - }; - -private: - struct SchedulerState { - uint32_t numWaiting; - uint32_t numSleepingWorkers; - SpinLock lock; - }; - - template - static void singleInvokeEntry(Context *ctx_base, JobContainerBase *data); - - template - static void multiInvokeEntry(Context *ctx_base, JobContainerBase *data, - uint64_t invocation_offset, - uint64_t num_invocations, - RunQueue *thread_queue); - - using SingleInvokeFn = decltype(&singleInvokeEntry); - using MultiInvokeFn = decltype(&multiInvokeEntry); - - JobManager(uint32_t num_ctx_userdata_bytes, - uint32_t ctx_userdata_alignment, - void (*ctx_init_fn)(void *, void *, WorkerInit &&), - uint32_t num_ctx_bytes, - uint32_t ctx_alignment, - void (*start_fn)(Context *, void *), - void *start_fn_data, - void (*update_fn)(Context *, void *), - void *update_fn_data, - int desired_num_workers, - int num_io, - StateManager *state_mgr, - bool pin_workers); - - struct Init; - JobManager(const Init &init); - - inline void * allocJob(int worker_idx, uint32_t num_bytes, - uint32_t alignment); - inline void deallocJob(int worker_idx, void *ptr, uint32_t num_bytes); - - JobID queueJob(int thread_idx, void (*job_func)(), - JobContainerBase *job_data, uint32_t num_invocations, - uint32_t parent_job_idx, - JobPriority prio = JobPriority::Normal); - - JobID reserveProxyJobID(int thread_idx, uint32_t parent_job_idx); - void relinquishProxyJobID(int thread_idx, uint32_t job_idx); - - void markInvocationsFinished(int thread_idx, - JobContainerBase *job_data, - int32_t job_idx, - uint32_t num_invocations); - - inline JobID getNewJobID(int thread_idx, uint32_t parent_job_idx, - uint32_t num_invocations); - - template - inline void addToRunQueue(int thread_idx, JobPriority prio, Fn &&add_cb); - - inline void addToWaitQueue(int thread_idx, void (*job_func)(), - JobContainerBase *job_data, uint32_t num_invocations, - JobPriority prio); - - enum class WorkerControl : uint64_t; - inline WorkerControl schedule(int thread_idx, Job *run_job); - - inline bool isQueueEmpty(uint32_t head, uint32_t correction, - uint32_t tail) const; - - inline uint32_t dequeueJobIndex(RunQueue *run_queue); - - inline WorkerControl getNextJob(void *queue_base, int thread_idx, - int init_search_idx, - bool run_scheduler, - Job *job); - - inline WorkerControl tryScheduling(JobManager::WorkerControl default_ctrl, - int thread_idx, Job *next_job); - - inline bool shouldSplitJob(RunQueue *queue) const; - - void splitJob(MultiInvokeFn fn_ptr, JobContainerBase *job_data, - uint32_t invocation_offset, uint32_t num_invocations, - RunQueue *run_queue); - - inline void runJob(const int thread_idx, Context *ctx, - void (*generic_fn)(), JobContainerBase *job_data, - uint32_t invocation_offset, uint32_t num_invocations); - - void workerThread(const int thread_idx, - void *context_base, - uint32_t num_context_bytes); - - void ioThread(const int thread_idx, - void *context_base, - uint32_t num_context_bytes); - - HeapArray threads_; - - Alloc::SharedState alloc_state_; - HeapArray job_allocs_; - - SchedulerState scheduler_; - void *const state_ptr_; - void *const high_base_; - void *const normal_base_; - void *const io_base_; - void *const tracker_base_; - void *const tracker_cache_base_; - void *const worker_base_; - void *const log_base_; - void *const waiting_jobs_; - uint32_t num_compute_workers_; - - std::counting_semaphore<> io_sema_; - alignas(MADRONA_CACHE_LINE) AtomicU32 num_high_; -}; - -} - -#include "job.inl" diff --git a/include/madrona/job.inl b/include/madrona/job.inl deleted file mode 100644 index 3df94267..00000000 --- a/include/madrona/job.inl +++ /dev/null @@ -1,304 +0,0 @@ -/* - * Copyright 2021-2022 Brennan Shacklett and contributors - * - * Use of this source code is governed by an MIT-style - * license that can be found in the LICENSE file or at - * https://opensource.org/licenses/MIT. - */ -#pragma once - -#include - -namespace madrona { - -constexpr JobID JobID::none() -{ - return JobID { - 0xFFFF'FFFF_u32, - 0xFFFF'FFFF_i32, - }; -} - -template -struct JobContainerBase::DepsArray { - JobID dependencies[N]; - - template - inline DepsArray(DepTs ...deps) - : dependencies { deps ... } - {} -}; - -template <> struct JobContainerBase::DepsArray<0> { - template - inline DepsArray(DepTs...) {} -}; - -template -template -JobContainer::JobContainer(uint32_t job_size, - MADRONA_MW_COND(uint32_t world_id,) - Fn &&func, - DepTs ...deps) - : JobContainerBase { - .id = JobID::none(), // Assigned in JobManager::queueJob - .jobSize = job_size, - MADRONA_MW_COND(.worldID = world_id,) - .numDependencies = N, - }, - dependencies(deps...), - fn(std::forward(func)) -{} - -bool JobManager::isQueueEmpty(uint32_t head, - uint32_t correction, - uint32_t tail) const -{ - auto checkGEWrapped = [](uint32_t a, uint32_t b) { - return a - b <= (1u << 31u); - }; - - return checkGEWrapped(head - correction, tail); -} - -template -struct JobManager::EntryConfig { - uint32_t numUserdataBytes; - uint32_t userdataAlignment; - void (*ctxInitCB)(void *, void *, WorkerInit &&); - uint32_t numCtxBytes; - uint32_t ctxAlignment; - StartFn startFnData; - void (*startWrapper)(Context *, void *); - UpdateFn updateFnData; - void (*updateLoop)(Context *, void *); -}; - -template -JobManager::EntryConfig - JobManager::makeEntry(StartFn &&start_fn) -{ - return makeEntry( - std::forward(start_fn), nullptr); -} - -template -JobManager::EntryConfig JobManager::makeEntry( - StartFn &&start_fn, UpdateFn &&update_fn) -{ - static_assert(std::is_trivially_destructible_v, - "Context types with custom destructors are not supported"); - - void (*start_wrapper)(Context *, void *); - if constexpr (!std::is_same_v) { - start_wrapper = [](Context *ctx_base, void *data) { - auto &ctx = *static_cast(ctx_base); - auto fn_ptr = (StartFn *)data; - - ctx.submit([fn = StartFn(*fn_ptr)](ContextT &ctx) { - fn(ctx); - }, false, ctx.currentJobID()); - }; - } else { - start_wrapper = start_fn; - start_fn = nullptr; - } - - void (*update_wrapper)(Context *, void *); - if constexpr (!std::is_same_v) { - update_wrapper = [](Context *ctx_base, void *data) { - auto &ctx = *static_cast(ctx_base); - auto fn_ptr = (UpdateFn *)data; - - ctx.submit([fn = UpdateFn(*fn_ptr)](ContextT &ctx) { - fn(ctx); - }, false, ctx.currentJobID()); - }; - } else { - update_wrapper = update_fn; - update_fn = nullptr; - } - - using DataT = typename ContextT::WorldDataT; - - return { - sizeof(DataT), - alignof(DataT), - [](void *ctx, void *data, WorkerInit &&init) { - new (ctx) ContextT((DataT *)data, std::forward(init)); - }, - sizeof(ContextT), - std::alignment_of_v, - std::forward(start_fn), - start_wrapper, - std::forward(update_fn), - update_wrapper, - }; -} - -template -JobManager::JobManager(const EntryConfig &entry_cfg, - int desired_num_workers, - int num_io, - StateManager *state_mgr, - bool pin_workers) - : JobManager(entry_cfg.numUserdataBytes, - entry_cfg.userdataAlignment, - entry_cfg.ctxInitCB, - entry_cfg.numCtxBytes, - entry_cfg.ctxAlignment, - entry_cfg.startWrapper, - [&entry_cfg]() { - if constexpr (std::is_same_v) { - (void)entry_cfg; - return nullptr; - } else { - return (void *)&entry_cfg.startFnData; - } - }(), - entry_cfg.updateLoop, - [&entry_cfg]() { - if constexpr (std::is_same_v) { - (void)entry_cfg; - return nullptr; - } else { - return (void *)&entry_cfg.updateFnData; - } - }(), - desired_num_workers, - num_io, - state_mgr, - pin_workers) -{} - -JobID JobManager::reserveProxyJobID(int thread_idx, JobID parent_id) -{ - return reserveProxyJobID(thread_idx, parent_id.id); -} - -void JobManager::relinquishProxyJobID(int thread_idx, JobID job_id) -{ - return markInvocationsFinished(thread_idx, nullptr, job_id.id, 1); -} - -bool JobManager::shouldSplitJob(RunQueue *queue) const -{ - uint32_t cur_tail = queue->tail.load_relaxed(); - uint32_t cur_correction = queue->correction.load_relaxed(); - uint32_t cur_head = queue->head.load_relaxed(); - - return isQueueEmpty(cur_head, cur_correction, cur_tail); -} - -template -void JobManager::singleInvokeEntry(Context *ctx_base, - JobContainerBase *data) -{ - ContextT &ctx = *static_cast(ctx_base); - auto container = static_cast(data); - JobManager *job_mgr = ctx.job_mgr_; - - container->fn(ctx); - - job_mgr->markInvocationsFinished(ctx.worker_idx_, data, data->id.id, 1); -} - -template -void JobManager::multiInvokeEntry(Context *ctx_base, - JobContainerBase *data, - uint64_t invocation_offset, - uint64_t num_invocations, - RunQueue *thread_queue) -{ - ContextT &ctx = *static_cast(ctx_base); - auto container = static_cast(data); - JobManager *job_mgr = ctx.job_mgr_; - - // This loop is never called with num_invocations == 0 - uint64_t invocation_idx = invocation_offset; - uint64_t remaining_invocations = num_invocations; - do { - uint64_t cur_invocation = invocation_idx++; - remaining_invocations -= 1; - - if (remaining_invocations > 0 && - job_mgr->shouldSplitJob(thread_queue)) { - job_mgr->splitJob(&multiInvokeEntry, data, - invocation_idx, remaining_invocations, thread_queue); - remaining_invocations = 0; - } - - container->fn(ctx, cur_invocation); - } while (remaining_invocations > 0); - - job_mgr->markInvocationsFinished(ctx.worker_idx_, data, data->id.id, - invocation_idx - invocation_offset); -} - -template -JobID JobManager::queueJob(int thread_idx, - Fn &&fn, - uint32_t num_invocations, - JobID parent_id, - MADRONA_MW_COND(uint32_t world_id,) - JobPriority prio, - DepTs ...deps) -{ - static constexpr uint32_t num_deps = sizeof...(DepTs); - using ContainerT = JobContainer; - static_assert(std::is_trivially_destructible_v); - -#ifdef MADRONA_GCC -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Winvalid-offsetof" -#endif - static_assert(num_deps == 0 || - offsetof(ContainerT, dependencies) == sizeof(JobContainerBase), - "Dependencies at incorrect offset in container type"); -#ifdef MADRONA_GCC -#pragma GCC diagnostic pop -#endif - - static constexpr uint64_t job_size = sizeof(ContainerT); - static constexpr uint64_t job_alignment = alignof(ContainerT); - static_assert(job_size <= JobManager::Alloc::maxJobSize, - "Job lambda capture is too large"); - static_assert(job_alignment <= JobManager::Alloc::maxJobAlignment, - "Job lambda capture has too large an alignment requirement"); - static_assert(utils::isPower2(job_alignment)); - - void *store = allocJob(thread_idx, job_size, job_alignment); - - auto container = new (store) ContainerT( - job_size, MADRONA_MW_COND(world_id,) std::forward(fn), deps...); - - void (*entry)(); - if constexpr (single_invoke) { - SingleInvokeFn fn_ptr = &singleInvokeEntry; - entry = (void (*)())fn_ptr; - } else { - MultiInvokeFn fn_ptr = &multiInvokeEntry; - entry = (void (*)())fn_ptr; - } - - return queueJob(thread_idx, entry, container, num_invocations, - parent_id.id, prio); -} - -void * JobManager::allocJob(int worker_idx, uint32_t num_bytes, - uint32_t alignment) -{ - return job_allocs_[worker_idx].alloc(alloc_state_, num_bytes, - alignment); -} - -void JobManager::deallocJob(int worker_idx, void *ptr, uint32_t num_bytes) -{ - job_allocs_[worker_idx].dealloc(alloc_state_, ptr, num_bytes); -} - -} diff --git a/include/madrona/mw_cpu.hpp b/include/madrona/mw_cpu.hpp deleted file mode 100644 index c10a989a..00000000 --- a/include/madrona/mw_cpu.hpp +++ /dev/null @@ -1,114 +0,0 @@ -/* - * Copyright 2021-2023 Brennan Shacklett and contributors - * - * Use of this source code is governed by an MIT-style - * license that can be found in the LICENSE file or at - * https://opensource.org/licenses/MIT. - */ -#pragma once - -#include -#include -#include - -namespace madrona { - -// Base class for TaskGraphExecutor below, don't use directly -class ThreadPoolExecutor { -public: - struct Config { - // Batch size for the backend - uint32_t numWorlds; - // Number of exported ECS components - uint32_t numExportedBuffers; - // Number of worker threads - uint32_t numWorkers = 0; - }; - - struct Job { - void (*fn)(void *); - void *data; - }; - - ThreadPoolExecutor(const Config &cfg); - ThreadPoolExecutor(ThreadPoolExecutor &&o); - - ~ThreadPoolExecutor(); - void run(Job *jobs, CountT num_jobs); - - // Get the base pointer of the component data exported with - // ECSRegister::exportColumn - void * getExported(CountT slot) const; - -protected: - void initializeContexts( - Context & (*init_fn)(void *, const WorkerInit &, CountT), - void *init_data, CountT num_worlds); - - ECSRegistry getECSRegistry(); - - void initExport(); - -private: - struct Impl; - std::unique_ptr impl_; -}; - -// The TaskGraphExecutor class is the entry point for the CPU backend. -// Use as follows: -// using MyCPUBackend = TaskGraphExecutor< -// MyContextSubclass, MyPerWorldState, MyConfig, MyPerWorldInit>; -// -// MyCPUBackend backend({ -// .numWorlds = 1024, -// .numExportedBuffers = 5, // Make sure this is set correctly! -// .numWorkers = 0, // (Autodetect number of CPU cores) -// }, MyConfig {}, my_world_inits); -// -// backend.run(); // Take one step -// -// The above code will initialize the simulation state with -// 1024 copies of the MyPerWorldState class, passing MyConfig and the -// appropriate my_world_inits reference to the MyPerWorldState constructor -template -class TaskGraphExecutor : private ThreadPoolExecutor { -public: - TaskGraphExecutor( - const Config &cfg, - const ConfigT &user_cfg, - const InitT *user_inits, - CountT num_taskgraphs); - - // Run one invocation of the task graph across all worlds (one step) - template - inline void runTaskGraph(EnumT taskgraph_id); - - inline void runTaskGraph(uint32_t taskgraph_idx); - - inline void run(); - - // Get the base pointer of the component data exported with - // ECSRegister::exportColumn - using ThreadPoolExecutor::getExported; - - // Get a reference to the per world data class - inline WorldT & getWorldData(CountT world_idx); - - inline ContextT & getWorldContext(CountT idx); - -private: - struct JobData { - Context *ctx; - TaskGraph taskgraph; - }; - - HeapArray contexts_; - HeapArray world_datas_; - HeapArray job_datas_; - HeapArray jobs_; - uint32_t num_taskgraphs_; -}; - -} - -#include "mw_cpu.inl" diff --git a/include/madrona/mw_cpu.inl b/include/madrona/mw_cpu.inl deleted file mode 100644 index af9faba9..00000000 --- a/include/madrona/mw_cpu.inl +++ /dev/null @@ -1,113 +0,0 @@ -#pragma once - -namespace madrona { - -template -TaskGraphExecutor::TaskGraphExecutor( - const Config &cfg, - const ConfigT &user_cfg, - const InitT *user_inits, - CountT num_taskgraphs) - : ThreadPoolExecutor(cfg), - contexts_(cfg.numWorlds), - world_datas_(cfg.numWorlds), - job_datas_((CountT)cfg.numWorlds * num_taskgraphs), - jobs_((CountT)cfg.numWorlds * num_taskgraphs), - num_taskgraphs_((uint32_t)num_taskgraphs) -{ - auto ecs_reg = getECSRegistry(); - WorldT::registerTypes(ecs_reg, user_cfg); - - HeapArray taskgraph_mgrs(cfg.numWorlds); - - auto ctx_init_cb = [&](const WorkerInit &worker_init, - CountT world_idx) -> Context & { - WorldT *world_data_ptr = &world_datas_[world_idx]; - new (&contexts_[world_idx]) Context(world_data_ptr, worker_init); - - taskgraph_mgrs.emplace(world_idx, num_taskgraphs, worker_init); - - return contexts_[world_idx]; - }; - - using CBPtrT = decltype(&ctx_init_cb); - initializeContexts([](void *ptr_raw, - const WorkerInit &worker_init, - CountT world_idx) -> Context & { - return (*(CBPtrT)ptr_raw)(worker_init, world_idx); - }, &ctx_init_cb, cfg.numWorlds); - - for (CountT world_idx = 0; world_idx < (CountT)cfg.numWorlds; - world_idx++) { - world_datas_.emplace(world_idx, contexts_[world_idx], - user_cfg, user_inits[world_idx]); - } - - HeapArray> built_graphs(cfg.numWorlds); - for (CountT world_idx = 0; world_idx < (CountT)cfg.numWorlds; - world_idx++) { - WorldT::setupTasks(taskgraph_mgrs[world_idx], user_cfg); - - built_graphs.emplace(world_idx, - taskgraph_mgrs[world_idx].constructGraphs()); - } - - for (CountT taskgraph_idx = 0; taskgraph_idx < num_taskgraphs; - taskgraph_idx++) { - for (CountT world_idx = 0; world_idx < (CountT)cfg.numWorlds; - world_idx++) { - CountT job_idx = taskgraph_idx * cfg.numWorlds + world_idx; - - job_datas_.emplace(job_idx, JobData { - .ctx = &contexts_[world_idx], - .taskgraph = std::move(built_graphs[world_idx][taskgraph_idx]), - }); - - jobs_[job_idx].fn = [](void *ptr) { - auto job_data = (JobData *)ptr; - job_data->taskgraph.run(job_data->ctx); - }; - jobs_[job_idx].data = &job_datas_[job_idx]; - } - } - - initExport(); -} - -template -template -void TaskGraphExecutor::runTaskGraph(EnumT taskgraph_id) -{ - runTaskGraph(static_cast(taskgraph_id)); -} - -template -void TaskGraphExecutor::runTaskGraph(uint32_t taskgraph_idx) -{ - CountT offset = taskgraph_idx * world_datas_.size(); - ThreadPoolExecutor::run(jobs_.data() + offset, world_datas_.size()); -} - -template -void TaskGraphExecutor::run() -{ - for (uint32_t i = 0; i < (uint32_t)num_taskgraphs_; i++) { - runTaskGraph(i); - } -} - -template -WorldT & TaskGraphExecutor::getWorldData( - CountT world_idx) -{ - return world_datas_[world_idx]; -} - -template -ContextT & TaskGraphExecutor::getWorldContext( - CountT world_idx) -{ - return contexts_[world_idx]; -} - -} diff --git a/include/madrona/navmesh.hpp b/include/madrona/navmesh.hpp deleted file mode 100644 index 93d4dac5..00000000 --- a/include/madrona/navmesh.hpp +++ /dev/null @@ -1,79 +0,0 @@ -#pragma once - -#include -#include -#include - -namespace madrona { - -struct Navmesh { - struct PathFindQueue { - float *costs; - uint32_t *heap; - uint32_t *heapIndex; - CountT heapSize; - - void add(uint32_t poly, float cost); - uint32_t removeMin(); - void decreaseCost(uint32_t poly, float cost); - }; - - struct AliasEntry { - float tau; - uint32_t alias; - }; - - math::Vector3 *vertices; - uint32_t *triIndices; - uint32_t *triAdjacency; - AliasEntry *triSampleAliasTable; - uint32_t numVerts; - uint32_t numTris; - - inline math::Vector3 samplePointAndPoly(RandKey rnd, uint32_t *out_poly); - inline math::Vector3 samplePoint(RandKey rnd); - - inline void getTriangleVertices(uint32_t tri_idx, - math::Vector3 *out_a, - math::Vector3 *out_b, - math::Vector3 *out_c); - - struct BFSState { - uint32_t *queue; - bool *visited; - }; - - template - inline void bfsFromPoly( - uint32_t poly, - BFSState bfs_state, - Fn &&fn); - - struct DijkstrasState { - float *distances; - math::Vector3 *entryPoints; - uint32_t *heap; - uint32_t *heapIndex; - }; - - template - inline void dijkstrasFromPoly( - uint32_t start_poly, - math::Vector3 start_pos, - DijkstrasState dijkstras_state, - Fn &&fn); - - static Navmesh initFromPolygons( - math::Vector3 *poly_vertices, - uint32_t *poly_idxs, - uint32_t *poly_idx_offsets, - uint32_t *poly_sizes, - uint32_t num_verts, - uint32_t num_polys); - - static constexpr inline uint32_t sentinel = 0xFFFF'FFFF; -}; - -} - -#include "navmesh.inl" diff --git a/include/madrona/navmesh.inl b/include/madrona/navmesh.inl deleted file mode 100644 index 4ee7eab8..00000000 --- a/include/madrona/navmesh.inl +++ /dev/null @@ -1,156 +0,0 @@ -#include - -namespace madrona { - -math::Vector3 Navmesh::samplePointAndPoly(RandKey rnd, uint32_t *out_poly) -{ - using namespace math; - - RandKey tbl_row_rnd = rand::split_i(rnd, 0); - RandKey alias_p_rnd = rand::split_i(rnd, 1); - RandKey bary_rnd = rand::split_i(rnd, 2); - - uint32_t tbl_row_idx = rand::sampleI32(tbl_row_rnd, 0, numTris); - float p = rand::sampleUniform(alias_p_rnd); - - AliasEntry tbl_row = triSampleAliasTable[tbl_row_idx]; - - uint32_t tri_idx = p < tbl_row.tau ? tbl_row_idx : tbl_row.alias; - *out_poly = tri_idx; - - Vector3 a, b, c; - getTriangleVertices(tri_idx, &a, &b, &c); - - Vector2 uv = rand::sample2xUniform(bary_rnd); - - if (uv.x + uv.y > 1.f) { - uv.x = 1.f - uv.x; - uv.y = 1.f - uv.y; - } - - float w = 1.f - uv.x - uv.y; - - return a * uv.x + b * uv.y + c * w; -} - -math::Vector3 Navmesh::samplePoint(RandKey rnd) -{ - uint32_t poly; - return samplePointAndPoly(rnd, &poly); -} - -void Navmesh::getTriangleVertices(uint32_t tri_idx, - math::Vector3 *out_a, - math::Vector3 *out_b, - math::Vector3 *out_c) -{ - *out_a = vertices[triIndices[3 * tri_idx]]; - *out_b = vertices[triIndices[3 * tri_idx + 1]]; - *out_c = vertices[triIndices[3 * tri_idx + 2]]; -} - -template -void Navmesh::bfsFromPoly(uint32_t start_poly, - BFSState bfs_state, - Fn &&fn) -{ - ArrayQueue bfs_queue(bfs_state.queue, numTris); - bool *visited = bfs_state.visited; - - utils::zeroN(visited, numTris); - - bfs_queue.add(start_poly); - visited[start_poly] = true; - - while (!bfs_queue.isEmpty()) { - uint32_t poly = bfs_queue.remove(); - - bool accept = fn(poly); - if (!accept) { - continue; - } - - MADRONA_UNROLL - for (CountT i = 0; i < 3; i++) { - uint32_t adjacent = triAdjacency[3 * poly + i]; - - if (adjacent != sentinel && !visited[adjacent]) { - bfs_queue.add(adjacent); - visited[adjacent] = true; - } - } - } -} - - -template -inline void Navmesh::dijkstrasFromPoly( - uint32_t start_poly, - math::Vector3 start_pos, - DijkstrasState dijkstras_state, - Fn &&fn) -{ - using namespace math; - - float *distances = dijkstras_state.distances; - - PathFindQueue prio_queue { - .costs = distances, - .heap = dijkstras_state.heap, - .heapIndex = dijkstras_state.heapIndex, - .heapSize = 0, - }; - utils::fillN(prio_queue.heapIndex, sentinel, numTris); - utils::fillN(distances, FLT_MAX, numTris); - - Vector3 *entry_points = dijkstras_state.entryPoints; - entry_points[start_poly] = start_pos; - - prio_queue.add(start_poly, 0.f); - while (prio_queue.heapSize > 0) { - uint32_t min_poly = prio_queue.removeMin(); - Vector3 cur_pos = entry_points[min_poly]; - float dist_so_far = distances[min_poly]; - - fn(min_poly, cur_pos, dist_so_far); - - Vector3 edge_midpoints[3]; - { - Vector3 a, b, c; - getTriangleVertices(min_poly, &a, &b, &c); - - edge_midpoints[0] = (a + b) / 2.f; - edge_midpoints[1] = (b + c) / 2.f; - edge_midpoints[2] = (c + a) / 2.f; - } - - MADRONA_UNROLL - for (CountT i = 0; i < 3; i++) { - uint32_t adjacent = triAdjacency[3 * min_poly + i]; - if (adjacent == Navmesh::sentinel) { - continue; - } - - Vector3 edge_midpoint = edge_midpoints[i]; - - float dist_to_edge = cur_pos.distance(edge_midpoint); - float new_dist = dist_so_far + dist_to_edge; - float prev_dist = distances[adjacent]; - - if (new_dist >= prev_dist) { - continue; - } - - entry_points[adjacent] = edge_midpoint; - - uint32_t prio_queue_idx = prio_queue.heapIndex[adjacent]; - if (prio_queue_idx == sentinel) { - prio_queue.add(adjacent, new_dist); - } else { - prio_queue.decreaseCost(adjacent, new_dist); - } - } - } -} - -} diff --git a/include/madrona/physics.hpp b/include/madrona/physics.hpp deleted file mode 100644 index 1fbee69c..00000000 --- a/include/madrona/physics.hpp +++ /dev/null @@ -1,232 +0,0 @@ -#pragma once -#include -#include -#include -#include -#include - -#include -#include - -namespace madrona::phys { - -struct ExternalForce : math::Vector3 { - ExternalForce(math::Vector3 v) - : Vector3(v) - {} -}; - -struct ExternalTorque : math::Vector3 { - ExternalTorque(math::Vector3 v) - : Vector3(v) - {} -}; - -enum class ResponseType : uint32_t { - Dynamic, - Kinematic, - Static, -}; - -struct Velocity { - math::Vector3 linear; - math::Vector3 angular; -}; - -struct SolverBundleAlias {}; - -struct RigidBody : Bundle< - base::ObjectInstance, - ResponseType, - broadphase::LeafID, - Velocity, - ExternalForce, - ExternalTorque, - SolverBundleAlias -> {}; - -struct CandidateCollision { - Loc a; - Loc b; - uint32_t aPrim; - uint32_t bPrim; -}; - -struct ContactConstraint { - Loc ref; - Loc alt; - math::Vector4 points[4]; - int32_t numPoints; - math::Vector3 normal; -}; - -struct JointConstraint { - enum class Type { - Fixed, - Hinge - }; - - struct Fixed { - math::Quat attachRot1; - math::Quat attachRot2; - float separation; - }; - - struct Hinge { - math::Vector3 a1Local; - math::Vector3 a2Local; - math::Vector3 b1Local; - math::Vector3 b2Local; - }; - - Entity e1; - Entity e2; - Type type; - - union { - Fixed fixed; - Hinge hinge; - }; - - math::Vector3 r1; - math::Vector3 r2; -}; - -struct CollisionEvent { - Entity a; - Entity b; -}; - -struct CollisionEventTemporary : Archetype {}; - -// Per object state -struct RigidBodyMassData { - float invMass; - math::Vector3 invInertiaTensor; - math::Vector3 toCenterOfMass; - math::Quat toInteriaFrame; -}; - -struct RigidBodyFrictionData { - float muS; - float muD; -}; - -struct RigidBodyMetadata { - RigidBodyMassData mass; - RigidBodyFrictionData friction; -}; - -struct CollisionPrimitive { - enum class Type : uint32_t { - Sphere = 1 << 0, - Hull = 1 << 1, - Plane = 1 << 2, - }; - - struct Sphere { - float radius; - }; - - struct Hull { - geo::HalfEdgeMesh halfEdgeMesh; - }; - - struct Plane {}; - - Type type; - union { - Sphere sphere; - Plane plane; - Hull hull; - }; -}; - -struct ObjectManager { - CollisionPrimitive *collisionPrimitives; - math::AABB *primitiveAABBs; - - math::AABB *rigidBodyAABBs; - uint32_t *rigidBodyPrimitiveOffsets; - uint32_t *rigidBodyPrimitiveCounts; - RigidBodyMetadata *metadata; -}; - -struct ObjectData { - ObjectManager *mgr; -}; - -namespace PhysicsSystem { - enum class Solver : uint32_t { - XPBD, - TGS, - }; - - void init(Context &ctx, - ObjectManager *obj_mgr, - float delta_t, - CountT num_substeps, - math::Vector3 gravity, - CountT max_dynamic_objects, - Solver solver = Solver::XPBD); - - void reset(Context &ctx); - broadphase::LeafID registerEntity(Context &ctx, - Entity e, - base::ObjectID obj_id); - - template - void findEntitiesWithinAABB(Context &ctx, - math::AABB aabb, - Fn &&fn); - - bool checkEntityAABBOverlap(Context &ctx, - math::AABB aabb, - Entity e); - - Entity makeFixedJoint(Context &ctx, - Entity e1, Entity e2, - math::Quat attach_rot1, math::Quat attach_rot2, - math::Vector3 r1, math::Vector3 r2, - float separation); - - Entity makeHingeJoint(Context &ctx, - Entity e1, Entity e2, - math::Vector3 a1_local, math::Vector3 a2_local, - math::Vector3 b1_local, math::Vector3 b2_local, - math::Vector3 r1, math::Vector3 r2); - - - void registerTypes(ECSRegistry ®istry, - Solver solver = Solver::XPBD); - - TaskGraphNodeID setupBroadphaseTasks( - TaskGraphBuilder &builder, - Span deps); - - TaskGraphNodeID setupPhysicsStepTasks( - TaskGraphBuilder &builder, - Span deps, - CountT num_substeps, - Solver solver = Solver::XPBD); - - TaskGraphNodeID setupCleanupTasks( - TaskGraphBuilder &builder, - Span deps); - - // Use the below two functions if you just want to use the broadphase without - // the rest of the physics system - - TaskGraphNodeID setupStandaloneBroadphaseOverlapTasks( - TaskGraphBuilder &builder, - Span deps); - - TaskGraphNodeID setupStandaloneBroadphaseCleanupTasks( - TaskGraphBuilder &builder, - Span deps); - -}; - -} - -#include "physics.inl" diff --git a/include/madrona/physics.inl b/include/madrona/physics.inl deleted file mode 100644 index b0654029..00000000 --- a/include/madrona/physics.inl +++ /dev/null @@ -1,28 +0,0 @@ -#pragma once - -namespace madrona::phys { - -namespace PhysicsSystem { - -template -void findEntitiesWithinAABB(Context &ctx, - math::AABB aabb, - Fn &&fn) -{ - using namespace madrona::base; - using namespace madrona::math; - - auto &bvh = ctx.singleton(); - - bvh.findIntersecting(aabb, [&](Entity e) { - bool overlap = checkEntityAABBOverlap( - ctx, aabb, e); - if (overlap) { - fn(e); - } - }); -} - -} - -} diff --git a/include/madrona/physics_assets.hpp b/include/madrona/physics_assets.hpp deleted file mode 100644 index d88954b7..00000000 --- a/include/madrona/physics_assets.hpp +++ /dev/null @@ -1,67 +0,0 @@ -#pragma once - -#include -#include -#include -#include - -#include - -namespace madrona::phys { - -struct SourceCollisionPrimitive { - struct HullInput { - uint32_t hullIDX; - }; - - CollisionPrimitive::Type type; - union { - CollisionPrimitive::Sphere sphere; - CollisionPrimitive::Plane plane; - HullInput hullInput; - }; -}; - -struct SourceCollisionObject { - Span prims; - float invMass; - RigidBodyFrictionData friction; -}; - -struct RigidBodyAssets { - struct HullData { - geo::HalfEdge *halfEdges; - uint32_t *faceBaseHalfEdges; - geo::Plane *facePlanes; - math::Vector3 *vertices; - - uint32_t numHalfEdges; - uint32_t numFaces; - uint32_t numVerts; - } hullData; - - // Per Primitive Data - CollisionPrimitive *primitives; - math::AABB *primitiveAABBs; - - // Per Object Data - RigidBodyMetadata *metadatas; - math::AABB *objAABBs; - uint32_t *primOffsets; - uint32_t *primCounts; - - uint32_t numConvexHulls; - uint32_t totalNumPrimitives; - uint32_t numObjs; - - static void * processRigidBodyAssets( - Span convex_hull_meshes, - Span collision_objs, - bool build_convex_hulls, - StackAlloc &tmp_alloc, - RigidBodyAssets *out_assets, - CountT *out_num_bytes); -}; - - -} diff --git a/include/madrona/physics_loader.hpp b/include/madrona/physics_loader.hpp deleted file mode 100644 index 0e7c84a8..00000000 --- a/include/madrona/physics_loader.hpp +++ /dev/null @@ -1,25 +0,0 @@ -#pragma once - -#include -#include - -#include - -namespace madrona::phys { - -class PhysicsLoader { -public: - PhysicsLoader(ExecMode exec_mode, CountT max_objects); - ~PhysicsLoader(); - PhysicsLoader(PhysicsLoader &&o); - - CountT loadRigidBodies(const RigidBodyAssets &assets); - - ObjectManager & getObjectManager(); - -private: - struct Impl; - std::unique_ptr impl_; -}; - -} diff --git a/include/madrona/py/bindings.hpp b/include/madrona/py/bindings.hpp index 9584a259..9c09adb6 100644 --- a/include/madrona/py/bindings.hpp +++ b/include/madrona/py/bindings.hpp @@ -16,44 +16,6 @@ namespace nb = nanobind; namespace madrona::py { -class JAXInterface { -public: - // Returns a function that registers custom_call_name with XLA - // to call step_fn (or async_step_fn in GPU mode) and returns - // the python function implementing the custom call. - template - static auto buildEntry(); - -private: - template - static void cpuEntryFn(void **out, void **in); - -#ifdef MADRONA_CUDA_SUPPORT - template - static void gpuEntryFn(cudaStream_t strm, void **buffers, - const char *opaque, size_t opaque_len); -#endif - - static nb::dict setup(const TrainInterface &iface, - nb::object sim_obj, - void *sim_ptr, - void *init_fn, - void *step_fn, - void *save_ckpts_fn, - void *restore_ckpts_fn, - bool xla_gpu); -}; - void setupMadronaSubmodule(nb::module_ parent_mod); } - -#include "bindings.inl" diff --git a/include/madrona/py/bindings.inl b/include/madrona/py/bindings.inl deleted file mode 100644 index c10dc784..00000000 --- a/include/madrona/py/bindings.inl +++ /dev/null @@ -1,123 +0,0 @@ -#include - -#include -#include -#include - -namespace madrona::py { - -template -auto JAXInterface::buildEntry() -{ - using SimT = - typename utils::ExtractClassFromMemberPtr::type; - - return [](nb::object sim, bool xla_gpu) { - void *init_fn; - void *step_fn; - void *save_ckpts_fn; - void *restore_ckpts_fn; - if (xla_gpu) { -#ifdef MADRONA_CUDA_SUPPORT - if constexpr (gpu_init_fn != nullptr && - gpu_step_fn != nullptr) { - auto init_wrapper = - &JAXInterface::gpuEntryFn; - init_fn = std::bit_cast(init_wrapper); - - auto step_wrapper = - &JAXInterface::gpuEntryFn; - step_fn = std::bit_cast(step_wrapper); - } else { - init_fn = nullptr; - step_fn = nullptr; - } - - if constexpr (gpu_save_ckpts_fn != nullptr && - gpu_restore_ckpts_fn != nullptr) { - auto gpu_save_ckpts_wrapper = - &JAXInterface::gpuEntryFn< - SimT, gpu_save_ckpts_fn>; - save_ckpts_fn = std::bit_cast( - gpu_save_ckpts_wrapper); - - auto gpu_restore_ckpts_wrapper = - &JAXInterface::gpuEntryFn< - SimT, gpu_restore_ckpts_fn>; - restore_ckpts_fn = std::bit_cast( - gpu_restore_ckpts_wrapper); - } else { - save_ckpts_fn = nullptr; - restore_ckpts_fn = nullptr; - } -#else - init_fn = nullptr; - step_fn = nullptr; - save_ckpts_fn = nullptr; - restore_ckpts_fn = nullptr; -#endif - } else { - auto init_wrapper = - &JAXInterface::cpuEntryFn; - init_fn = std::bit_cast(init_wrapper); - - auto step_wrapper = - &JAXInterface::cpuEntryFn; - step_fn = std::bit_cast(step_wrapper); - - if constexpr (cpu_save_ckpts_fn != nullptr && - cpu_restore_ckpts_fn != nullptr) { - auto cpu_save_ckpts_wrapper = - &JAXInterface::cpuEntryFn< - SimT, cpu_save_ckpts_fn>; - save_ckpts_fn = std::bit_cast( - cpu_save_ckpts_wrapper); - - auto cpu_restore_ckpts_wrapper = - &JAXInterface::cpuEntryFn< - SimT, cpu_restore_ckpts_fn>; - restore_ckpts_fn = std::bit_cast( - cpu_restore_ckpts_wrapper); - } else { - save_ckpts_fn = nullptr; - restore_ckpts_fn = nullptr; - } - } - assert(init_fn != nullptr && step_fn != nullptr); - - SimT *sim_ptr = nb::inst_ptr(sim); - TrainInterface iface = std::invoke(iface_fn, *sim_ptr); - return setup(iface, sim, (void *)sim_ptr, init_fn, step_fn, - save_ckpts_fn, restore_ckpts_fn, xla_gpu); - }; -} - -template -void JAXInterface::cpuEntryFn(void **out, void **in) -{ - SimT *sim = *(SimT **)in[0]; - std::invoke(fn, *sim, in + 2, out); -} - -#ifdef MADRONA_CUDA_SUPPORT -template -void JAXInterface::gpuEntryFn(cudaStream_t strm, void **buffers, - const char *opaque, size_t) -{ - SimT *sim = *(SimT **)opaque; - - // The first buffer entry is used by the CPU backend, skip - // The scond buffer entry is a token JAX uses for ordering, skip - std::invoke(fn, *sim, strm, buffers + 2); -} -#endif - -} diff --git a/include/madrona/py/utils.hpp b/include/madrona/py/utils.hpp index 9305528a..91a02071 100644 --- a/include/madrona/py/utils.hpp +++ b/include/madrona/py/utils.hpp @@ -16,45 +16,6 @@ namespace madrona::py { -#ifdef MADRONA_CUDA_SUPPORT -class CudaSync final { -public: - CudaSync(cudaExternalSemaphore_t sema); - void wait(uint64_t strm); - -private: -#ifdef MADRONA_LINUX - // These classes have to be virtual on linux so a unique typeinfo - // is emitted. Otherwise every user of this class gets a weak symbol - // reference and nanobind can't map the types correctly - virtual void key_(); -#endif - - cudaExternalSemaphore_t sema_; -}; -#endif - -// Need to wrap the actual enum class because macos -// RTTI for enum classes isn't consistent across libraries -class PyExecMode final { -public: - inline PyExecMode(ExecMode v) - : v_(v) - {} - - inline operator ExecMode() const - { - return v_; - } - -private: -#ifdef MADRONA_LINUX - virtual void key_(); -#endif - - ExecMode v_; -}; - enum class TensorElementType { UInt8, Int8, @@ -142,65 +103,4 @@ class Tensor final { std::array dimensions_; }; -struct NamedTensor { - const char *name; - Tensor tensor; -}; - -struct TrainStepInputInterface { - Span actions; - Tensor resets; - Tensor simCtrl; - Span pbt = {}; -}; - -struct TrainStepOutputInterface { - Span observations; - Tensor rewards; - Tensor dones; - Span stats = {}; - Span pbt = {}; -}; - -struct TrainCheckpointingInterface { - Tensor checkpointData; -}; - -class TrainInterface { -public: - TrainInterface(); - TrainInterface(TrainStepInputInterface step_inputs, - TrainStepOutputInterface step_outputs, - Optional checkpointing = - Optional::none()); - TrainInterface(TrainInterface &&o); - ~TrainInterface(); - - TrainInterface & operator=(TrainInterface &&o); - - TrainStepInputInterface stepInputs() const; - TrainStepOutputInterface stepOutputs() const; - Optional checkpointing() const; - - void cpuCopyStepInputs(void **buffers); - void cpuCopyObservations(void **buffers); - void cpuCopyStepOutputs(void **buffers); - -#ifdef MADRONA_CUDA_SUPPORT - void ** cudaCopyStepInputs(cudaStream_t strm, void **buffers); - void cudaCopyObservations(cudaStream_t strm, void **buffers); - void cudaCopyStepOutputs(cudaStream_t strm, void **buffers); -#endif - -private: - struct Impl; - -#ifdef MADRONA_LINUX - virtual void key_(); -#endif - - std::unique_ptr impl_; -}; - - } diff --git a/include/madrona/render/common.hpp b/include/madrona/render/common.hpp index 021ad0c5..cb0d5190 100644 --- a/include/madrona/render/common.hpp +++ b/include/madrona/render/common.hpp @@ -8,13 +8,6 @@ struct APILib {}; struct APIBackend {}; struct GPUDevice {}; -// If voxel generation is to happen -struct VoxelConfig { - uint32_t xLength; - uint32_t yLength; - uint32_t zLength; -}; - inline float srgbToLinear(float srgb); inline math::Vector4 srgb8ToFloat(uint8_t r, uint8_t g, uint8_t b); diff --git a/include/madrona/render/ecs.hpp b/include/madrona/render/ecs.hpp index 4e9eb5fb..c3b91390 100644 --- a/include/madrona/render/ecs.hpp +++ b/include/madrona/render/ecs.hpp @@ -196,7 +196,6 @@ namespace RenderingSystem { bool update_visual_properties = false); void init(Context &ctx, const RenderECSBridge *bridge); - uint32_t * getVoxelPtr(Context &ctx); void makeEntityRenderable(Context &ctx, Entity e); void disableEntityRenderable(Context &ctx, Entity e); void attachEntityToView(Context &ctx, diff --git a/include/madrona/render/render_mgr.hpp b/include/madrona/render/render_mgr.hpp index 3d428cca..f2c047c6 100644 --- a/include/madrona/render/render_mgr.hpp +++ b/include/madrona/render/render_mgr.hpp @@ -37,8 +37,6 @@ class RenderManager { uint32_t maxLightsPerWorld; uint32_t maxInstancesPerWorld; ExecMode execMode; - - VoxelConfig voxelCfg; }; RenderManager(APIBackend *render_backend, diff --git a/scripts/parse_device_tracing.py b/scripts/parse_device_tracing.py deleted file mode 100644 index 74148171..00000000 --- a/scripts/parse_device_tracing.py +++ /dev/null @@ -1,466 +0,0 @@ -import os -import sys -import argparse -import pandas as pd -from PIL import Image, ImageDraw, ImageFont - -HIDE_SEEK = False -SPLIT_LINES = False -MAX_BLOCKS_PER_SM = 6 - - -def parse_device_logs(events): - LOG_STEPS = {} - STEP = -1 - - def new_step(num_warps, num_blocks, num_sms): - nonlocal STEP - STEP += 1 - LOG_STEPS[STEP] = { - "events": {}, - "SMs": {}, - "mapping": {}, - # "final_cycles": {}, - "start_timestamp": 0xFFFFFFFFFFFFFFFF, - "final_timestamp": 0, - "num_warps": num_warps, - "num_blocks": num_blocks, - "num_sms": num_sms, - "configs": {}, - } - - for i in range(0, len(events), 40): - array = events[i : i + 40] - - event = int.from_bytes(array[:4], byteorder="little") - funcID = int.from_bytes(array[4:8], byteorder="little") - numInvocations = int.from_bytes(array[8:12], byteorder="little") - nodeID = int.from_bytes(array[12:16], byteorder="little") - warpID = int.from_bytes(array[16:20], byteorder="little") - blockID = int.from_bytes(array[20:24], byteorder="little") - smID = int.from_bytes(array[24:28], byteorder="little") - logIndex = int.from_bytes(array[28:32], byteorder="little") - cycleCount = int.from_bytes(array[32:40], byteorder="little") - - if STEP != -1: - # to make a unique warp id - warpID += blockID * LOG_STEPS[STEP]["num_warps"] - - # print("event: {}, funcID: {}, numInvocations: {}, nodeID: {}, warpID: {}, blockID: {}, smID: {}, cycleCount: {}, logIndex: {}".format(event, funcID, numInvocations, nodeID, warpID, blockID, smID, cycleCount, logIndex)) - if event == 0: - # for calibration event, we log kernel config instead of node - # when logIndex == 0, it is the first node of the whole step - if logIndex == 0: - new_step(num_warps=funcID, num_blocks=numInvocations, num_sms=nodeID) - LOG_STEPS[STEP]["start_timestamp"] = cycleCount - else: - assert LOG_STEPS[STEP]["num_warps"] == funcID and LOG_STEPS[STEP]["num_sms"] == nodeID - LOG_STEPS[STEP]["num_blocks"] = numInvocations - - elif event in [1, 2]: - if nodeID not in LOG_STEPS[STEP]["mapping"]: - LOG_STEPS[STEP]["mapping"][nodeID] = (funcID, numInvocations) - else: - assert LOG_STEPS[STEP]["mapping"][nodeID] == (funcID, numInvocations) - - if nodeID not in LOG_STEPS[STEP]["events"]: - LOG_STEPS[STEP]["events"][nodeID] = {event: (smID, warpID, cycleCount)} - else: - assert event not in LOG_STEPS[STEP]["events"][nodeID] - LOG_STEPS[STEP]["events"][nodeID][event] = (smID, warpID, cycleCount) - if nodeID not in LOG_STEPS[STEP]["configs"]: - LOG_STEPS[STEP]["configs"][nodeID] = {"num_blocks": LOG_STEPS[STEP]["num_blocks"]} - - elif event in [3, 4]: - if smID not in LOG_STEPS[STEP]["SMs"]: - assert event == 3 - LOG_STEPS[STEP]["SMs"][smID] = {(numInvocations, nodeID, warpID): [cycleCount]} - else: - if (numInvocations, nodeID, warpID) in LOG_STEPS[STEP]["SMs"][smID]: - assert event == 4 - LOG_STEPS[STEP]["SMs"][smID][(numInvocations, nodeID, warpID)].append(cycleCount) - else: - assert event == 3 - LOG_STEPS[STEP]["SMs"][smID][(numInvocations, nodeID, warpID)] = [cycleCount] - - elif event == 5: - # assert warpID not in LOG_STEPS[STEP]["final_cycles"] - # LOG_STEPS[STEP]["final_cycles"][warpID] = cycleCount - LOG_STEPS[STEP]["final_timestamp"] = max(LOG_STEPS[STEP]["final_timestamp"], cycleCount) - else: - assert False & "event {} not supported".format(event) - - # drop the last step which might be corrupted - # del LOG_STEPS[STEP] - print("At the end, complete traces for {} steps are generated".format(STEP)) - - for s in LOG_STEPS: - LOG_STEPS[s]["final_timestamp"] -= LOG_STEPS[s]["start_timestamp"] - # for b in LOG_STEPS[s]["final_cycles"]: - # LOG_STEPS[s]["final_cycles"][b] -= LOG_STEPS[s]["start_timestamp"] - - return LOG_STEPS - - -def serialized_analysis(step_log, nodes_map): - - def calibrate(timestamp, reference): - return timestamp[2] - reference - - for i in range(max(step_log["events"]) + 1): - # skipped nodes - if i not in step_log["events"]: - continue - nodes_map[i] = { - "nodeID": i, - "funcID": step_log["mapping"][i][0], - "invocations": step_log["mapping"][i][1], - "start": calibrate(step_log["events"][i][1], step_log["start_timestamp"]), - "end": calibrate(step_log["events"][i][2], step_log["start_timestamp"]), - "SM utilization": [], - } - nodes_map[i]["duration (ns)"] = nodes_map[i]["end"] - nodes_map[i]["start"] - - total_exec_time = sum(v["duration (ns)"] for v in nodes_map.values()) - for i in nodes_map: - nodes_map[i]["percentage (%)"] = nodes_map[i]["duration (ns)"] / total_exec_time * 100 - - -def block_analysis(step_log, nodes_map): - sm_execution = {k: [] for k in step_log["SMs"].keys()} - block_exec_time = { - "blocks": {k: {} for k in step_log["SMs"].keys()}, # horizontal - "nodes": {k: {} for k in step_log["SMs"].keys()}, # vertical - } - - for sm in step_log["SMs"]: - for (_, nodeID, warpID), time_stamps in step_log["SMs"][sm].items(): - assert len(time_stamps) == 2 - start = time_stamps[0] - end = max(time_stamps[1:]) - # confirm clock does proceed within an SM - assert end >= start - assert start > step_log["start_timestamp"] - - start, end = [i - step_log["start_timestamp"] for i in [start, end]] - sm_execution[sm].append((start, end)) - - if warpID not in block_exec_time["blocks"][sm]: - block_exec_time["blocks"][sm][warpID] = [(start, end, nodeID)] - else: - assert start > block_exec_time["blocks"][sm][warpID][-1][1] - block_exec_time["blocks"][sm][warpID].append((start, end, nodeID)) - - if nodeID not in block_exec_time["nodes"][sm]: - block_exec_time["nodes"][sm][nodeID] = [] - block_exec_time["nodes"][sm][nodeID].append((start, end, warpID)) - - block_exec_time["blocks"][sm] = {k: sorted(v) for k, v in block_exec_time["blocks"][sm].items()} - block_exec_time["nodes"][sm] = {k: sorted(v) for k, v in block_exec_time["nodes"][sm].items()} - - for s in sm_execution: - intervals = [] - for start, end in sm_execution[s]: - if len(intervals) == 0: - intervals.append((start, end)) - continue - p = 0 - while p < len(intervals): - if start > intervals[p][1]: - p += 1 - continue - elif end < intervals[p][0]: - intervals.insert(p, (start, end)) - break - else: - intervals[p] = ( - min(start, intervals[p][0]), - max(end, intervals[p][1]), - ) - break - else: - intervals.insert(p, (start, end)) - - i_pointer = 0 - for k, v in nodes_map.items(): - occupied_time = 0 - while i_pointer < len(intervals): - i_start, i_end = intervals[i_pointer] - if v["end"] < i_start: - break - elif v["start"] <= i_start and v["end"] >= i_end: - occupied_time += i_end - i_start - i_pointer += 1 - continue - else: - assert False and "no intersections" - nodes_map[k]["SM utilization"].append(occupied_time / nodes_map[k]["duration (ns)"]) - - for i in nodes_map: - assert len(nodes_map[i]["SM utilization"]) == len(sm_execution) - nodes_map[i]["SM utilization"] = sum(nodes_map[i]["SM utilization"]) / len(sm_execution) - - print( - "For each SM on average, {:.3f}% of the time there is at least one block is running on".format( - sum(v["SM utilization"] * v["percentage (%)"] for v in nodes_map.values()) - ) - ) - - return block_exec_time - - -COLORS = [ - "blue", - "orange", - "red", - "green", - "purple", - "cyan", - "pink", - "magenta", - "olive", - "navy", - "teal", - "maroon", - "yellow", - "black", -] - - -def plot_events(step_log, nodes_map, blocks, file_name, args): - # todo: here we have an assumption that each SM has the same number of blocks, which are expected to be true - num_sms = step_log["num_sms"] - # num_block_per_sm = step_log["num_blocks"] - num_warp_per_sm = step_log["num_warps"] * MAX_BLOCKS_PER_SM - # num_block_per_sm = 8 - num_pixel_per_warp = 1 - sm_interval_pixel = num_pixel_per_warp * 3 - num_pixel_per_sm = num_warp_per_sm * num_pixel_per_warp + sm_interval_pixel - y_blank = num_pixel_per_warp * num_warp_per_sm * (num_sms // 8) - y_limit = num_sms * num_pixel_per_sm + y_blank - x_limit = y_limit * args.aspect_ratio - print("the figure size will be {}x{}".format(x_limit, y_limit)) - - top_nodes = sorted( - [ - i[0] - for i in sorted(nodes_map.items(), key=lambda item: item[1]["duration (ns)"])[ - len(nodes_map) - args.num_highlight_nodes : - ] - ] - ) - - colors = {} - if HIDE_SEEK: - colors = { - # narrowphase - 28: (0, 76, 153), - # solvers - 20: (0, 102, 204), - 22: (0, 128, 255), - 24: (102, 178, 255), - 26: (178, 216, 255), - # broadphase - 30: (199, 31, 102), - 34: (230, 96, 152), - 36: (248, 181, 209), - # Collect Observations System - 52: (139, 235, 219), - # Compute Visibility System - 54: (90, 184, 168), - # LIDAR System - 56: (148, 112, 206), - } - else: - for n in top_nodes: - func = nodes_map[n]["funcID"] - if func not in colors and len(colors) < len(COLORS): - colors[func] = COLORS[len(colors)] - print("Color mapping for functions:", colors) - - img = Image.new("RGB", (x_limit, y_limit), "white") - draw = ImageDraw.Draw(img) - - def cast_coor(timestamp, limit=x_limit): - assert timestamp <= step_log["final_timestamp"] - if args.fixed_scale: - return int(timestamp / args.tpp) - else: - return int(timestamp / step_log["final_timestamp"] * limit) - - color_span = {} - for n in nodes_map: - if nodes_map[n]["funcID"] not in colors: - continue - left, right = cast_coor(nodes_map[n]["start"]), cast_coor(nodes_map[n]["end"]) - color_span[(left, right)] = colors[nodes_map[n]["funcID"]] - sorted_color_span = sorted(color_span.keys()) - - num_stamps = active_warps = 0 - - for sm, b in blocks.items(): - node_pixels_sm = {} - y = (sm + 1) * num_pixel_per_sm - vertical_pixels = {} - for warp, events in b.items(): - # to avoid duplication - last_end_pixel = -1 - for e in events: - start = max(last_end_pixel + 1, cast_coor(e[0])) - end = cast_coor(e[1]) - if e[2] not in node_pixels_sm: - node_pixels_sm[e[2]] = [start, end] - else: - node_pixels_sm[e[2]] = [ - min(node_pixels_sm[e[2]][0], start), - max(node_pixels_sm[e[2]][1], end), - ] - for i in range(start, end + 1): - if i not in vertical_pixels: - vertical_pixels[i] = 1 * MAX_BLOCKS_PER_SM / step_log["configs"][e[2]]["num_blocks"] - else: - vertical_pixels[i] += 1 * MAX_BLOCKS_PER_SM / step_log["configs"][e[2]]["num_blocks"] - last_end_pixel = end - - n_pointer = 0 - for node_start, node_end in sorted(node_pixels_sm.values()): - for p in range(node_start, node_end + 1): - if p not in vertical_pixels: - vertical_pixels[p] = 0 - num_stamps += num_warp_per_sm - active_warps += vertical_pixels[p] - - while n_pointer < len(color_span): - left, right = sorted_color_span[n_pointer] - if p > right: - n_pointer += 1 - continue - elif p < left: - bar_color = (0, 0, 0) - break - else: - bar_color = color_span[(left, right)] - break - else: - bar_color = (0, 0, 0) - assert vertical_pixels[p] <= num_warp_per_sm - # if vertical_pixels[p] != 0: - draw.line( - (p, y, p, y - int(vertical_pixels[p] * num_pixel_per_warp)), - fill=bar_color, - width=1 if vertical_pixels[p] != 0 else 0, - ) - y_low = y - int(vertical_pixels[p] * num_pixel_per_warp) - y_low -= 1 if vertical_pixels[p] != 0 else 0 - y_high = y - num_warp_per_sm * num_pixel_per_warp - if y_low <= y_high: - pass - else: - draw.line((p, y_low, p, y_high), fill=(211, 211, 211), width=1) - - if SPLIT_LINES: - # indicate the start of each splitted kernel - nodes = sorted(nodes_map.keys()) - last_node = nodes[0] - for n in nodes[1:]: - if step_log["configs"][n]["num_blocks"] != step_log["configs"][last_node]["num_blocks"]: - node_start = cast_coor(nodes_map[n]["start"]) - draw.line( - (node_start, 0, node_start, y_limit - y_blank / 2), - fill="plum", - width=2, - ) - last_node = n - - print("Percentage of active warps is {:.2f}%".format(active_warps / num_stamps * 100)) - - fontsize = 65 - font = ImageFont.load_default(size=fontsize) - - if not HIDE_SEEK: - # mark the start and the end of major nodes_map - y_shift = 0.9 - for n in top_nodes: - # for n, v in nodes_map.items(): - left, right = cast_coor(nodes_map[n]["start"]), cast_coor(nodes_map[n]["end"]) - draw.line((left, 0, left, y_limit), fill="red", width=1) - draw.line((right, 0, right, y_limit), fill="green", width=1) - draw.text( - (left, y_limit - y_blank * y_shift), - " f: {}\n t: {:.3f}ms\n {:.1f}%".format( - nodes_map[n]["funcID"], - (nodes_map[n]["duration (ns)"]) / 1000000, - nodes_map[n]["percentage (%)"], - ), - fill=(0, 0, 0), - font=font, - ) - y_shift = 1.3 - y_shift - - img.save(file_name) - - -def step_analysis(step_log, file_name, tabular_data, args=None): - nodes_map = {} - serialized_analysis(step_log, nodes_map) - - block_exec_time = block_analysis(step_log, nodes_map) - if args is not None: - plot_events(step_log, nodes_map, block_exec_time["blocks"], file_name, args) - - for n in nodes_map: - tabular_data = pd.concat([tabular_data, pd.DataFrame({k: [v] for k, v in nodes_map[n].items()})]) - return tabular_data - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--trace_file", type=str, required=True) - parser.add_argument("--start_step", type=int, default=10, help="analysis start from which step") - parser.add_argument("--num_steps", type=int, default=5, help="number of steps to be analyzed") - parser.add_argument("--num_highlight_nodes", type=int, default=16) - parser.add_argument("--aspect_ratio", type=float, default=2) - parser.add_argument("--fixed_scale", action="store_true") - parser.add_argument("--tpp", type=int, default=8000, help="time(ns) per pixel") - - args = parser.parse_args() - - with open(args.trace_file, "rb") as f: - events = bytearray(f.read()) - assert len(events) % 40 == 0 - print("{} events were logged in total".format(len(events) // 40)) - LOG_STEPS = parse_device_logs(events) - - steps = args.num_steps - start_from = args.start_step - - dir_path = args.trace_file + "_megakernel_events" - isExist = os.path.exists(dir_path) - if not isExist: - os.mkdir(dir_path) - # todo: limit - assert start_from < len(LOG_STEPS) - end_at = min(start_from + steps, len(LOG_STEPS)) - - tabular_data = [ - pd.DataFrame( - { - "nodeID": [], - "funcID": [], - "duration (ns)": [], - "invocations": [], - "percentage (%)": [], - "SM utilization": [], - } - ) - for _ in range(start_from, end_at) - ] - - with pd.ExcelWriter(dir_path + "/metrics.xlsx") as writer: - for s in range(start_from, end_at): - step_analysis( - LOG_STEPS[s], - dir_path + "/step{}.png".format(s), - tabular_data[s - start_from], - args, - ).to_excel(writer, sheet_name="step{}".format(s), index=False) diff --git a/scripts/parse_host_tracing.py b/scripts/parse_host_tracing.py deleted file mode 100644 index 9f3b9b6a..00000000 --- a/scripts/parse_host_tracing.py +++ /dev/null @@ -1,97 +0,0 @@ -import sys -import numpy as np -import matplotlib.pyplot as plt -from PIL import Image, ImageDraw - - -def read_binary_file(file_name): - with open(file_name, "rb") as f: - # Read the contents of the file into a NumPy array - events, time_stamps = np.fromfile(f, dtype=np.int64).reshape(2, -1) - # set the time stamp of first event to be 0 - time_stamps = [i - time_stamps[0] for i in time_stamps] - return events, time_stamps - - -# apt to change -event_dict = {0: 1, 2: 3, 4: 5} -color_dict = {1: "r", 3: "g", 5: "b"} - - -def get_steps(events, time_stamps): - steps = [] - for i in range(len(events)): - if events[i : i + 4].tolist() == [2, 3, 4, 5]: - steps.append(time_stamps[i : i + 4]) - return steps - - -def plot_events(events, time_stamps, file_name, exclude_init=True, drop_warmup=2, display_steps=20): - - steps = get_steps(events, time_stamps) - steps = steps[drop_warmup : drop_warmup + display_steps] - steps = [[i - s[0] for i in s] for s in steps] - - max_time = max(max(s) for s in steps) - - num_steps = len(steps) - x_limit = 200 - pixel_per_step = 10 - bar_interval = 2 - y_limit = num_steps * (pixel_per_step + bar_interval) - - img = Image.new("RGB", (x_limit, y_limit), "white") - draw = ImageDraw.Draw(img) - - def cast_coor(timestamp, limit=x_limit): - assert timestamp <= max_time - return int(timestamp / max_time * limit) - - for i, s in enumerate(steps): - s = [i - s[0] for i in s] - y = pixel_per_step / 2 + i * (pixel_per_step + bar_interval) - draw.line((0, y, cast_coor(s[1]), y), fill="green", width=pixel_per_step) - draw.line((cast_coor(s[2]), y, cast_coor(s[3]), y), fill="blue", width=pixel_per_step) - # print(s[1]/1000000, (s[3] - s[2])/1000000, s[1] / (s[3] - s[2] + s[1]) ) - - img.save(file_name + "_events.png") - - # num_events = len(events) // 2 - # event_stack = [0] - # event_starts = [] - # event_ends = [] - # event_colors = [] - # for i in range(1, len(events)): - # if events[i] in event_dict: - # event_stack.append(i) - # else: - # assert (events[i] == event_dict[events[event_stack[-1]]]) - # if exclude_init and events[i] == 1: - # num_events -= 1 - # continue - # event_starts.append(time_stamps[event_stack[-1]]) - # event_ends.append(time_stamps[i]) - # event_colors.append(color_dict[events[i]]) - # event_stack.pop() - # assert num_events == len(event_starts) == len(event_ends) - - # num_events -= drop_warmup - # event_starts = event_starts[drop_warmup:] - # event_ends = event_ends[drop_warmup:] - # event_colors = event_colors[drop_warmup:] - - # plt.barh(range(num_events), - # [event_ends[i] - event_starts[i] for i in range(num_events)], - # left=event_starts, - # color=event_colors) - # plt.xlim(min(event_starts), max(event_ends)) - # plt.savefig(file_name + '_events.png') - - -if __name__ == "__main__": - if len(sys.argv) != 2: - print("python parse_tracing.py [file_name]") - exit() - - events, time_stamps = read_binary_file(sys.argv[1]) - plot_events(events, time_stamps, sys.argv[1]) diff --git a/scripts/perf_benchmark/README.md b/scripts/perf_benchmark/README.md deleted file mode 100644 index 6bb7cc6f..00000000 --- a/scripts/perf_benchmark/README.md +++ /dev/null @@ -1,442 +0,0 @@ -# Performance Benchmark Suite - -This directory contains a comprehensive performance benchmarking suite for evaluating rendering and simulation performance across different renderers, configurations, and hardware setups. - -## Overview - -The benchmark suite is designed to systematically test performance across multiple dimensions: -- **Renderers**: Madrona, Omniverse, PyRender, ManiSkill -- **Rendering Modes**: Rasterizer vs Raytracer -- **Batch Sizes**: Multiple environment counts -- **Resolutions**: Various image resolutions -- **Assets**: Different MJCF/URDF models - -## Directory Structure - -``` -perf_benchmark/ -├── README.md # This file -├── batch_benchmark.py # Main batch execution script -├── benchmark_configs.py # Configuration parser -├── benchmark_report_generator.py # Report and visualization generator -├── benchmark_profiler.py # Performance profiling utilities -├── benchmark_madrona.py # Madrona renderer benchmark -├── benchmark_omni.py # Omniverse renderer benchmark -├── benchmark_pyrender.py # PyRender benchmark -├── benchmark_maniskill.py # ManiSkill renderer benchmark -├── process_xml.py # XML asset preprocessing utility -├── configs/ # Configuration files -│ ├── benchmark_config_smoke_test.yml # Quick test configuration -│ ├── benchmark_config_madrona.yml # Madrona-specific config -│ ├── benchmark_config_omni.yml # Omniverse-specific config -│ ├── benchmark_config_maniskill.yml # ManiSkill-specific config -│ └── benchmark_config_full.yml # Comprehensive test config -``` - -## Quick Start - - -### 1. Optional steps -To enable benchmarking with IsaacLab and ManiSkill, follow these optional setup steps: - -- Install IsaacLab - - Download and install IsaacLab from [NVIDIA IsaacLab Documentation](https://isaac-sim.github.io/IsaacLab/main/source/setup/installation/index.html). - - Add the IsaacLab installation directory to your system `PATH`. -- Install ManiSkill - - Install ManiSkill2 by following [ManiSkill Documentation](https://maniskill.readthedocs.io/en/latest/user_guide/). -- Set Environment Variables - - Both IsaacLab and ManiSkill use the `ASSET_DIR` environment variable to locate the Genesis assets directory. - - Use `genesis.utils.misc.asset_dir()` in Genesis to retrieve the exact directory path and then set the environment variable: - ``` - export ASSET_DIR=/path/to/genesis/asset_dir - ``` -- Preprocess Required Assets - - IsaacLab benchmarks require MJCF assets to be preprocessed for compatibility with Omniverse. - ```bash - python process_xml.py --file ./configs/benchmark_config_omni.yml - ``` - -### 2. Run a Quick Smoke Test - -```bash -python batch_benchmark.py -f benchmark_config_smoke_test.yml -``` - -### 3. Run a Full Benchmark Suite - -```bash -python batch_benchmark.py -f benchmark_config_full.yml -``` - -### 4. Continue from a Previous Run - -```bash -python batch_benchmark.py -f benchmark_config_full.yml -c /name/of/previous/run/folder -``` - -## Configuration Files - -Configuration files are YAML-based and define the test parameters. Here's an example structure: - -```yaml -# List of MJCF/URDF files to test -mjcf_list: - - xml/franka_emika_panda/panda.xml - -# Renderer configurations -renderer_list: - - renderer: madrona - benchmark_script: benchmark_madrona.py - timeout: 120 - - renderer: omniverse - benchmark_script: benchmark_omni.py - timeout: 300 - - renderer: maniskill - benchmark_script: benchmark_maniskill.py - timeout: 180 - -# Test rasterizer and raytracer modes -rasterizer_list: - - true # Rasterizer - - false # Raytracer - -# Batch sizes to test -batch_size_list: - - 256 - - 512 - - 1024 - -# Resolutions to test (width x height) -resolution_list: - - [128, 128] - - [256, 256] - - [512, 512] - -# Raytracer settings -raytracer: - max_bounce: 2 - spp: 1 - -# Simulation settings -simulation: - n_steps: 1000 - -# Camera settings -camera: - position: [1.5, 0.5, 1.5] - lookat: [0.0, 0.0, 0.5] - fov: 45.0 - -# Display settings -display: - gui: false - -# Performance comparisons to generate -comparison_list: - - - renderer: madrona - rasterizer: true - - renderer: madrona - rasterizer: false -``` - -## Core Components - -### batch_benchmark.py - -The main orchestration script that: -- Parses configuration files -- Creates test combinations -- Executes benchmarks in parallel -- Handles failures and timeouts -- Generates reports - -**Key Features:** -- Hierarchical test execution (renderer → rasterizer → mjcf → batch_size → resolution) -- Automatic failure handling (skips larger resolutions if smaller ones fail) -- Resume capability from previous runs -- Timeout management per renderer - -### benchmark_configs.py - -Configuration parser that loads and validates YAML configuration files. - -**Supported Configuration Sections:** -- `mjcf_list`: List of asset files to test -- `renderer_list`: Renderer configurations with timeouts -- `rasterizer_list`: Boolean flags for rasterizer/raytracer modes -- `batch_size_list`: Environment counts to test -- `resolution_list`: Image resolutions to test -- `raytracer`: Raytracing parameters (max_bounce, spp) -- `simulation`: Simulation parameters (n_steps) -- `camera`: Camera positioning and settings -- `display`: Display options (gui) -- `comparison_list`: Performance comparison definitions - -### benchmark_profiler.py - -High-precision performance profiling using CUDA events and CPU timing. - -**Profiling Capabilities:** -- GPU timing using CUDA events -- CPU timing using high-resolution timers -- Per-step detailed timing -- Per-environment timing calculations -- FPS calculations (total and per-environment) - -**Key Methods:** -- `on_simulation_start()`: Start simulation timing -- `on_rendering_start()`: Start rendering timing -- `on_rendering_end()`: End rendering timing -- `get_total_rendering_gpu_time()`: Total GPU rendering time -- `get_rendering_fps()`: Overall FPS -- `get_rendering_fps_per_env()`: FPS per environment - -### benchmark_report_generator.py - -Generates comprehensive performance reports and visualizations. - -**Output Types:** -- Individual performance plots per MJCF/rasterizer combination -- Comparison plots between renderers -- HTML report with embedded plots and tables -- Performance summary tables - -**Visualization Features:** -- FPS vs batch size plots -- Resolution-based color coding -- Performance comparison charts -- Interactive HTML reports - -### process_xml.py - -Utility script for preprocessing MJCF/XML assets to ensure compatibility with different renderers. - -**Features:** -- Wraps visual geometry elements in body tags -- Generates unique identifiers for mesh elements -- Handles collision geometry appropriately -- Creates processed XML files with `_new.xml` suffix - -**Usage:** -```bash -python process_xml.py --file ./genesis/assets/xml/franka_emika_panda/panda.xml -``` - -## Renderer-Specific Scripts - -### benchmark_madrona.py - -Benchmarks the Madrona renderer using the Genesis framework. - -**Features:** -- GPU/CPU fallback handling -- Batch rendering support -- Configurable camera and lighting -- Performance profiling integration - -### benchmark_omni.py - -Benchmarks the Omniverse renderer using Isaac Sim. - -**Features:** -- Omniverse-specific settings optimization -- Path tracing and rasterizer modes -- GPU memory management -- Performance-optimized rendering settings - -### benchmark_pyrender.py - -Benchmarks the PyRender renderer. - -**Features:** -- OpenGL-based rendering -- CPU-based simulation -- Cross-platform compatibility - -### benchmark_maniskill.py - -Benchmarks the ManiSkill renderer using the ManiSkill framework. - -**Features:** -- Custom ManiSkill environment implementation -- Support for multiple robot types (Panda, Unitree Go2, Unitree G1) -- Configurable camera modes (minimal for rasterizer, rt-fast for raytracer) -- GPU memory optimization for large batch sizes -- Image saving capabilities for debugging - -**Supported Robots:** -- `panda.xml` → Panda robot -- `go2.xml` → Unitree Go2 robot -- `g1.xml` → Unitree G1 robot - -**Camera Modes:** -- `minimal`: Fast rasterizer mode -- `rt-fast`: Fast raytracer mode - -## Output Structure - -Benchmark runs create the following structure: - -``` -benchmark_reports/ -└── perf_benchmark_YYYYMMDD_HHMMSS/ - ├── perf_data.csv # Raw benchmark data - ├── plots/ # Generated plots - │ ├── panda_plot.png - │ ├── panda_comparison_plot.png - │ └── ... - ├── images/ # Rendered images (ManiSkill) - │ ├── step00_env00_minimal_panda.png - │ └── ... - └── report.html # Interactive HTML report -``` - -### CSV Data Format - -The `perf_data.csv` file contains the following columns: -- `result`: "succeeded" or "failed" -- `mjcf`: Asset file path -- `renderer`: Renderer name -- `rasterizer`: Boolean rasterizer flag -- `n_envs`: Number of environments -- `n_steps`: Number of simulation steps -- `resX`, `resY`: Image resolution -- `camera_posX/Y/Z`: Camera position -- `camera_lookatX/Y/Z`: Camera lookat point -- `camera_fov`: Camera field of view -- `time_taken_gpu`: Total GPU time (seconds) -- `time_taken_per_env_gpu`: GPU time per environment (seconds) -- `time_taken_cpu`: Total CPU time (seconds) -- `time_taken_per_env_cpu`: CPU time per environment (seconds) -- `fps`: Overall FPS -- `fps_per_env`: FPS per environment - -## Advanced Usage - -### Custom Configuration - -Create a custom configuration file: - -```yaml -# custom_config.yml -mjcf_list: - - path/to/your/model.xml - -renderer_list: - - renderer: madrona - benchmark_script: benchmark_madrona.py - timeout: 180 - - renderer: maniskill - benchmark_script: benchmark_maniskill.py - timeout: 240 - -rasterizer_list: [true, false] -batch_size_list: [64, 128, 256] -resolution_list: [[256, 256], [512, 512]] - -raytracer: - max_bounce: 4 - spp: 2 - -simulation: - n_steps: 500 - -camera: - position: [2.0, 1.0, 2.0] - lookat: [0.0, 0.0, 0.0] - fov: 60.0 -``` - -Run with custom config: -```bash -python batch_benchmark.py -f custom_config.yml -``` - -### Asset Preprocessing - -Some renderers may require specific XML formatting. Use the preprocessing utility: - -```bash -# Process a single file -python process_xml.py --file ./genesis/assets/xml/franka_emika_panda/panda.xml - -# Process multiple files -for file in ./genesis/assets/xml/*/*.xml; do - python process_xml.py --file "$file" -done -``` - -### Performance Analysis - -The benchmark suite provides detailed performance analysis: - -1. **Raw Data**: CSV files with detailed timing information -2. **Visualizations**: Plots showing FPS vs batch size relationships -3. **Comparisons**: Side-by-side renderer performance comparisons -4. **HTML Reports**: Interactive reports with embedded plots and tables -5. **Rendered Images**: Sample images for visual verification (ManiSkill) - -### Failure Handling - -The benchmark suite includes robust failure handling: - -- **Timeout Management**: Each renderer has configurable timeouts -- **Progressive Skipping**: If a resolution fails, larger resolutions are skipped -- **Resume Capability**: Can continue from previous runs -- **Error Logging**: Failed runs are recorded with "failed" status - -## Dependencies - -### Required Python Packages -- `pandas`: Data analysis and CSV handling -- `matplotlib`: Plot generation -- `numpy`: Numerical computations -- `pyyaml`: Configuration file parsing -- `torch`: CUDA event handling (for profiling) - -### Renderer-Specific Dependencies -- **Madrona**: Genesis framework -- **Omniverse**: Isaac Sim, Omniverse Kit -- **PyRender**: PyRender, OpenGL -- **ManiSkill**: ManiSkill, SAPIEN, Gymnasium - -## Troubleshooting - -### Common Issues - -1. **CUDA Out of Memory**: Reduce batch sizes or resolutions -2. **Timeout Errors**: Increase timeout values in configuration -3. **Missing Assets**: Ensure MJCF/URDF files are in the correct paths -4. **Renderer Failures**: Check renderer-specific dependencies -5. **XML Compatibility**: Use `process_xml.py` to preprocess assets if needed - -### Debug Mode - -Enable GUI mode for debugging: -```yaml -display: - gui: true -``` - -### Verbose Logging - -Check the console output for detailed timing information and error messages. - -### ManiSkill-Specific Issues - -- **Robot Compatibility**: Ensure robot XML files match supported robot types -- **Memory Configuration**: Adjust GPU memory settings for large batch sizes -- **Image Saving**: Check image output directory permissions - -## Contributing - -To add a new renderer: - -1. Create a new benchmark script (e.g., `benchmark_newrenderer.py`) -2. Implement the required interface (see `benchmark_madrona.py` for reference) -3. Add the renderer to your configuration file -4. Test with the smoke test configuration first - -## License - -This benchmark suite is part of the Genesis project. See the main LICENSE file for details. \ No newline at end of file diff --git a/scripts/perf_benchmark/batch_benchmark.py b/scripts/perf_benchmark/batch_benchmark.py deleted file mode 100644 index 576ebdb4..00000000 --- a/scripts/perf_benchmark/batch_benchmark.py +++ /dev/null @@ -1,385 +0,0 @@ -import argparse -import subprocess -import os -from datetime import datetime - -import pandas as pd - -from benchmark_report_generator import generate_report -from benchmark_configs import BenchmarkConfigs - -# Example command: -# python batch_benchmark.py -f benchmark_config_smoke_test.yml - - -# Create a struct to store the arguments -class BenchmarkArgs: - def __init__( - self, - renderer, - rasterizer, - n_envs, - n_steps, - resX, - resY, - camera_posX, - camera_posY, - camera_posZ, - camera_lookatX, - camera_lookatY, - camera_lookatZ, - camera_fov, - mjcf, - benchmark_result_file, - benchmark_config_file, - max_bounce, - spp, - gui=False, - benchmark_script=None, - renderer_timeout=None, - ): - self.renderer = renderer - self.rasterizer = rasterizer - self.n_envs = n_envs - self.n_steps = n_steps - self.resX = resX - self.resY = resY - self.camera_posX = camera_posX - self.camera_posY = camera_posY - self.camera_posZ = camera_posZ - self.camera_lookatX = camera_lookatX - self.camera_lookatY = camera_lookatY - self.camera_lookatZ = camera_lookatZ - self.camera_fov = camera_fov - self.mjcf = mjcf - self.benchmark_result_file = benchmark_result_file - self.benchmark_config_file = benchmark_config_file - self.max_bounce = max_bounce - self.spp = spp - self.gui = gui - self.benchmark_script = benchmark_script - self.renderer_timeout = renderer_timeout - - @staticmethod - def parse_benchmark_args(): - parser = argparse.ArgumentParser() - parser.add_argument("-d", "--renderer", required=True, type=str) - parser.add_argument("-r", "--rasterizer", action="store_true", default=False) - parser.add_argument("-n", "--n_envs", required=True, type=int) - parser.add_argument("-x", "--resX", required=True, type=int) - parser.add_argument("-y", "--resY", required=True, type=int) - parser.add_argument("-f", "--mjcf", required=True, type=str) - parser.add_argument("-g", "--benchmark_result_file", required=True, type=str) - parser.add_argument("-c", "--benchmark_config_file", required=True, type=str) - args = parser.parse_args() - benchmark_config = BenchmarkConfigs(args.benchmark_config_file) - benchmark_args = BenchmarkArgs( - renderer=args.renderer, - rasterizer=args.rasterizer, - n_envs=args.n_envs, - n_steps=benchmark_config.n_steps, - resX=args.resX, - resY=args.resY, - camera_posX=benchmark_config.camera_pos[0], - camera_posY=benchmark_config.camera_pos[1], - camera_posZ=benchmark_config.camera_pos[2], - camera_lookatX=benchmark_config.camera_lookat[0], - camera_lookatY=benchmark_config.camera_lookat[1], - camera_lookatZ=benchmark_config.camera_lookat[2], - camera_fov=benchmark_config.camera_fov, - mjcf=args.mjcf, - benchmark_result_file=args.benchmark_result_file, - benchmark_config_file=args.benchmark_config_file, - max_bounce=benchmark_config.max_bounce, - spp=benchmark_config.spp, - gui=benchmark_config.gui, - ) - print(f"Benchmark with args:") - print(f" renderer: {benchmark_args.renderer}") - print(f" rasterizer: {benchmark_args.rasterizer}") - print(f" n_envs: {benchmark_args.n_envs}") - print(f" n_steps: {benchmark_args.n_steps}") - print(f" resolution: {benchmark_args.resX}x{benchmark_args.resY}") - print( - f" camera_pos: ({benchmark_args.camera_posX}, {benchmark_args.camera_posY}, {benchmark_args.camera_posZ})" - ) - print( - f" camera_lookat: ({benchmark_args.camera_lookatX}, {benchmark_args.camera_lookatY}, {benchmark_args.camera_lookatZ})" - ) - print(f" camera_fov: {benchmark_args.camera_fov}") - print(f" mjcf: {benchmark_args.mjcf}") - print(f" benchmark_result_file: {benchmark_args.benchmark_result_file}") - print(f" benchmark_config_file: {benchmark_args.benchmark_config_file}") - print(f" max_bounce: {benchmark_args.max_bounce}") - print(f" spp: {benchmark_args.spp}") - print(f" gui: {benchmark_args.gui}") - return benchmark_args - - -class BatchBenchmarkArgs: - def __init__(self, config_file, continue_from): - self.config_file = config_file - self.continue_from = continue_from - - def parse_batch_benchmark_args(): - parser = argparse.ArgumentParser() - parser.add_argument("-f", "--config_file", type=str, default="benchmark_config_smoke_test.yml") - parser.add_argument("-c", "--continue_from", type=str, default=None) - args = parser.parse_args() - return BatchBenchmarkArgs(config_file=args.config_file, continue_from=args.continue_from) - - -def create_batch_args(benchmark_result_file, config_file): - # Ensure the directory exists - os.makedirs(os.path.dirname(benchmark_result_file), exist_ok=True) - - # Load configuration - config = BenchmarkConfigs(config_file) - mjcf_list = config.mjcf_list - renderer_list = config.renderer_list - rasterizer_list = config.rasterizer_list - batch_size_list = config.batch_size_list - resolution_list = config.resolution_list - n_steps = config.n_steps - camera_pos = config.camera_pos - camera_lookat = config.camera_lookat - camera_fov = config.camera_fov - max_bounce = config.max_bounce - spp = config.spp - gui = config.gui - - # Batch data for resolution and batch size needs to be sorted in ascending order of resX x resY - # so that if one resolution fails, all the resolutions, which are larger, will be skipped. - resolution_list.sort(key=lambda x: x[0] * x[1]) - - # Create a hierarchical dictionary to store all combinations - batch_args_dict = {} - - # Build hierarchical structure - for renderer_info in renderer_list: - renderer = renderer_info["renderer"] - benchmark_script = renderer_info["benchmark_script"] - renderer_timeout = renderer_info["timeout"] - batch_args_dict[renderer] = {} - for rasterizer in rasterizer_list: - batch_args_dict[renderer][rasterizer] = {} - for mjcf in mjcf_list: - batch_args_dict[renderer][rasterizer][mjcf] = {} - for batch_size in batch_size_list: - batch_args_dict[renderer][rasterizer][mjcf][batch_size] = {} - for resolution in resolution_list: - resX, resY = resolution - # Create benchmark args for this combination - args = BenchmarkArgs( - renderer=renderer, - rasterizer=rasterizer, - n_envs=batch_size, - n_steps=n_steps, - resX=resX, - resY=resY, - camera_posX=camera_pos[0], - camera_posY=camera_pos[1], - camera_posZ=camera_pos[2], - camera_lookatX=camera_lookat[0], - camera_lookatY=camera_lookat[1], - camera_lookatZ=camera_lookat[2], - camera_fov=camera_fov, - mjcf=mjcf, - benchmark_result_file=benchmark_result_file, - benchmark_config_file=config_file, - max_bounce=max_bounce, - spp=spp, - gui=gui, - benchmark_script=benchmark_script, - renderer_timeout=renderer_timeout, - ) - batch_args_dict[renderer][rasterizer][mjcf][batch_size][(resX, resY)] = args - - return batch_args_dict - - -def create_benchmark_result_file(continue_from): - benchmark_report_root = os.path.join(os.path.dirname(os.path.abspath(__file__)), "benchmark_reports") - if continue_from is not None: - continue_from_file = os.path.join(benchmark_report_root, continue_from, "perf_data.csv") - if not os.path.exists(continue_from_file): - raise FileNotFoundError(f"Continue from file not found: {continue_from_file}") - print(f"Continuing from file: {continue_from_file}") - return continue_from_file - else: - # Create benchmark result data file with header - benchmark_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - - # Get the benchmark project root directory - benchmark_data_directory = os.path.join(benchmark_report_root, f"perf_benchmark_{benchmark_timestamp}") - - if not os.path.exists(benchmark_data_directory): - os.makedirs(benchmark_data_directory) - - benchmark_result_file = f"{benchmark_data_directory}/perf_data.csv" - with open(benchmark_result_file, "w") as f: - f.write( - "result,mjcf,renderer,rasterizer,n_envs,n_steps,resX,resY,camera_posX,camera_posY,camera_posZ,camera_lookatX,camera_lookatY,camera_lookatZ,camera_fov,time_taken_gpu,time_taken_per_env_gpu,time_taken_cpu,time_taken_per_env_cpu,fps,fps_per_env\n" - ) - print(f"Created new benchmark result file: {benchmark_result_file}") - return benchmark_result_file - - -def write_benchmark_result_file(args: BenchmarkArgs, performance_results: dict): - os.makedirs(os.path.dirname(args.benchmark_result_file), exist_ok=True) - with open(args.benchmark_result_file, "a") as f: - f.write( - f"succeeded,{args.mjcf},{args.renderer}," - f"{args.rasterizer},{args.n_envs},{args.n_steps}," - f"{args.resX},{args.resY}," - f"{args.camera_posX},{args.camera_posY},{args.camera_posZ}," - f"{args.camera_lookatX},{args.camera_lookatY},{args.camera_lookatZ}," - f"{args.camera_fov}," - f"{performance_results['time_taken_gpu']},{performance_results['time_taken_per_env_gpu']},{performance_results['time_taken_cpu']}," - f"{performance_results['time_taken_per_env_cpu']},{performance_results['fps']},{performance_results['fps_per_env']}\n" - ) - - -def get_previous_runs(continue_from_file): - if continue_from_file is None: - return [] - - # Read the existing benchmark data file - df = pd.read_csv(continue_from_file) - - # Create a list of tuples containing run info and status - previous_runs = [] - - for _, row in df.iterrows(): - run_info = ( - row["mjcf"], - row["renderer"], - row["rasterizer"], - row["n_envs"], - (row["resX"], row["resY"]), - row["result"], # 'succeeded' or 'failed' - ) - previous_runs.append(run_info) - - return previous_runs - - -def run_batch_benchmark(batch_args_dict, previous_runs=None): - if previous_runs is None: - previous_runs = [] - - for renderer in batch_args_dict: - print(f"Running benchmark for {renderer}") - for rasterizer in batch_args_dict[renderer]: - for mjcf in batch_args_dict[renderer][rasterizer]: - for batch_size in batch_args_dict[renderer][rasterizer][mjcf]: - last_resolution_failed = False - for resolution in batch_args_dict[renderer][rasterizer][mjcf][batch_size]: - if last_resolution_failed: - break - - # Check if this run was in a previous execution - run_info = (mjcf, renderer, rasterizer, batch_size, resolution) - skip_this_run = False - - for prev_run in previous_runs: - if run_info == prev_run[:5]: # Compare only the run parameters, not the status - skip_this_run = True - if prev_run[4] == "failed": - # Skip this and subsequent resolutions if it failed before - last_resolution_failed = True - break - - if skip_this_run: - continue - - # Run the benchmark - batch_args = batch_args_dict[renderer][rasterizer][mjcf][batch_size][resolution] - - # launch a process to run the benchmark - current_dir = os.path.dirname(os.path.abspath(__file__)) - benchmark_script_path = os.path.join(current_dir, batch_args.benchmark_script) - if not os.path.exists(benchmark_script_path): - raise FileNotFoundError(f"Benchmark script not found: {benchmark_script_path}") - cmd = ["python3", benchmark_script_path] - if batch_args.rasterizer: - cmd.append("--rasterizer") - cmd.extend( - [ - "--renderer", - batch_args.renderer, - "--n_envs", - str(batch_args.n_envs), - "--resX", - str(batch_args.resX), - "--resY", - str(batch_args.resY), - "--mjcf", - batch_args.mjcf, - "--benchmark_result_file", - batch_args.benchmark_result_file, - "--benchmark_config_file", - batch_args.benchmark_config_file, - ] - ) - try: - # Read timeout from config - process = subprocess.Popen(cmd) - try: - # Hack to avoid omniverse runs to take forever. - timeout = batch_args.renderer_timeout - return_code = process.wait(timeout=timeout) - if return_code != 0: - raise subprocess.CalledProcessError(return_code, cmd) - except subprocess.TimeoutExpired: - process.kill() - process.wait() # Wait for the process to be killed - raise TimeoutError(f"Process did not complete within {timeout} seconds") - except Exception as e: - print(f"Error running benchmark: {str(e)}") - if isinstance(e, subprocess.CalledProcessError): - last_resolution_failed = True - # Write failed result without timing data - with open(batch_args.benchmark_result_file, "a") as f: - f.write( - f"failed,{batch_args.mjcf},{batch_args.renderer},{batch_args.rasterizer},{batch_args.n_envs},{batch_args.n_steps},{batch_args.resX},{batch_args.resY},{batch_args.camera_posX},{batch_args.camera_posY},{batch_args.camera_posZ},{batch_args.camera_lookatX},{batch_args.camera_lookatY},{batch_args.camera_lookatZ},{batch_args.camera_fov},,,,,,\n" - ) - - -def sort_and_dedupe_benchmark_result_file(benchmark_result_file): - # Sort by mjcf asc, renderer asc, rasterizer desc, n_envs asc, resX asc, resY asc, n_envs asc - df = pd.read_csv(benchmark_result_file) - df = df.sort_values( - by=["mjcf", "renderer", "rasterizer", "resX", "resY", "n_envs", "result"], - ascending=[True, True, False, True, True, True, False], - ) - - # Deduplicate by keeping the first occurrence of each unique combination of mjcf, renderer, rasterizer, resX, resY, n_envs - # Keep succeeded runs if there are multiple runs for the same combination. - df = df.drop_duplicates( - subset=["mjcf", "renderer", "rasterizer", "resX", "resY", "n_envs"], - keep="first", - ) - df.to_csv(benchmark_result_file, index=False) - - -def main(): - batch_benchmark_args = BatchBenchmarkArgs.parse_batch_benchmark_args() - benchmark_result_file = create_benchmark_result_file(batch_benchmark_args.continue_from) - - # Get list of previous runs if continuing from a previous run - previous_runs = get_previous_runs(benchmark_result_file) - - # Run benchmark in batch - batch_args_dict = create_batch_args(benchmark_result_file, config_file=batch_benchmark_args.config_file) - run_batch_benchmark(batch_args_dict, previous_runs) - - # Sort benchmark result file - sort_and_dedupe_benchmark_result_file(benchmark_result_file) - - # Generate plots - generate_report(benchmark_result_file, config_file=batch_benchmark_args.config_file) - - -if __name__ == "__main__": - main() diff --git a/scripts/perf_benchmark/benchmark_assets/plane_urdf/checker.png b/scripts/perf_benchmark/benchmark_assets/plane_urdf/checker.png deleted file mode 100644 index 7c57d041..00000000 Binary files a/scripts/perf_benchmark/benchmark_assets/plane_urdf/checker.png and /dev/null differ diff --git a/scripts/perf_benchmark/benchmark_assets/plane_urdf/plane.mtl b/scripts/perf_benchmark/benchmark_assets/plane_urdf/plane.mtl deleted file mode 100644 index 12fb0655..00000000 --- a/scripts/perf_benchmark/benchmark_assets/plane_urdf/plane.mtl +++ /dev/null @@ -1,14 +0,0 @@ -newmtl Material - Ns 10.0000 - Ni 1.5000 - d 1.0000 - Tr 0.0000 - Tf 1.0000 1.0000 1.0000 - illum 2 - Ka 0.0000 0.0000 0.0000 - Kd 1.0000 1.0000 1.0000 - Ks 0.0000 0.0000 0.0000 - Ke 0.0000 0.0000 0.0000 - map_Ka cube.tga - map_Kd checker.png - diff --git a/scripts/perf_benchmark/benchmark_assets/plane_urdf/plane.urdf b/scripts/perf_benchmark/benchmark_assets/plane_urdf/plane.urdf deleted file mode 100644 index 448eabee..00000000 --- a/scripts/perf_benchmark/benchmark_assets/plane_urdf/plane.urdf +++ /dev/null @@ -1,26 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/scripts/perf_benchmark/benchmark_assets/plane_urdf/plane100.obj b/scripts/perf_benchmark/benchmark_assets/plane_urdf/plane100.obj deleted file mode 100644 index 3a74f590..00000000 --- a/scripts/perf_benchmark/benchmark_assets/plane_urdf/plane100.obj +++ /dev/null @@ -1,22 +0,0 @@ -# Blender v2.66 (sub 1) OBJ File: '' -# www.blender.org -mtllib plane.mtl -o Plane -v 100.000000 -100.000000 0.000000 -v 100.000000 100.000000 0.000000 -v -100.000000 100.000000 0.000000 -v -100.000000 -100.000000 0.000000 - -vt 100.000000 0.000000 -vt 100.000000 100.000000 -vt 0.000000 100.000000 -vt 0.000000 0.000000 - - - -usemtl Material -s off -f 1/1 2/2 3/3 -f 1/1 3/3 4/4 - - diff --git a/scripts/perf_benchmark/benchmark_assets/plane_usd/.asset_hash b/scripts/perf_benchmark/benchmark_assets/plane_usd/.asset_hash deleted file mode 100644 index c2288d4a..00000000 --- a/scripts/perf_benchmark/benchmark_assets/plane_usd/.asset_hash +++ /dev/null @@ -1 +0,0 @@ -b3f970c53c90c6563970b4d47b8e0bab \ No newline at end of file diff --git a/scripts/perf_benchmark/benchmark_assets/plane_usd/config.yaml b/scripts/perf_benchmark/benchmark_assets/plane_usd/config.yaml deleted file mode 100644 index 22459db1..00000000 --- a/scripts/perf_benchmark/benchmark_assets/plane_usd/config.yaml +++ /dev/null @@ -1,18 +0,0 @@ -asset_path: genesis/assets/urdf/plane/plane.urdf -usd_dir: null -usd_file_name: null -force_usd_conversion: true -make_instanceable: true -fix_base: true -root_link_name: null -link_density: 0.0 -merge_fixed_joints: true -convert_mimic_joints_to_normal_joints: false -joint_drive: null -collider_type: convex_hull -self_collision: false -replace_cylinders_with_capsules: false -collision_from_visuals: false -## -# Generated by UrdfConverter on 2025-06-10 at 18:41:48. -## diff --git a/scripts/perf_benchmark/benchmark_assets/plane_usd/configuration/materials/textures/checker.png b/scripts/perf_benchmark/benchmark_assets/plane_usd/configuration/materials/textures/checker.png deleted file mode 100644 index 7c57d041..00000000 Binary files a/scripts/perf_benchmark/benchmark_assets/plane_usd/configuration/materials/textures/checker.png and /dev/null differ diff --git a/scripts/perf_benchmark/benchmark_assets/plane_usd/configuration/plane_base.usd b/scripts/perf_benchmark/benchmark_assets/plane_usd/configuration/plane_base.usd deleted file mode 100644 index c73a418d..00000000 Binary files a/scripts/perf_benchmark/benchmark_assets/plane_usd/configuration/plane_base.usd and /dev/null differ diff --git a/scripts/perf_benchmark/benchmark_assets/plane_usd/configuration/plane_physics.usd b/scripts/perf_benchmark/benchmark_assets/plane_usd/configuration/plane_physics.usd deleted file mode 100644 index e3d61822..00000000 Binary files a/scripts/perf_benchmark/benchmark_assets/plane_usd/configuration/plane_physics.usd and /dev/null differ diff --git a/scripts/perf_benchmark/benchmark_assets/plane_usd/configuration/plane_sensor.usd b/scripts/perf_benchmark/benchmark_assets/plane_usd/configuration/plane_sensor.usd deleted file mode 100644 index 5f117bb3..00000000 Binary files a/scripts/perf_benchmark/benchmark_assets/plane_usd/configuration/plane_sensor.usd and /dev/null differ diff --git a/scripts/perf_benchmark/benchmark_assets/plane_usd/plane.usd b/scripts/perf_benchmark/benchmark_assets/plane_usd/plane.usd deleted file mode 100644 index 44550cfc..00000000 Binary files a/scripts/perf_benchmark/benchmark_assets/plane_usd/plane.usd and /dev/null differ diff --git a/scripts/perf_benchmark/benchmark_configs.py b/scripts/perf_benchmark/benchmark_configs.py deleted file mode 100644 index 78a9f434..00000000 --- a/scripts/perf_benchmark/benchmark_configs.py +++ /dev/null @@ -1,45 +0,0 @@ -import os -import yaml - - -class BenchmarkConfigs: - def __init__(self, config_file): - self.load_from_config_file(config_file) - - def load_from_config_file(self, config_file): - self.config_path = os.path.join(os.path.dirname(__file__), "configs", config_file) - if not os.path.exists(self.config_path): - raise FileNotFoundError(f"Config file not found: {self.config_path}") - with open(self.config_path, "r") as f: - config = yaml.safe_load(f) - - self.mjcf_list = config["mjcf_list"] - self.rasterizer_list = config["rasterizer_list"] - self.batch_size_list = config["batch_size_list"] - self.resolution_list = config["resolution_list"] - self.gui = config.get("gui", False) - - # Get renderer list with defaults - self.renderer_list = config["renderer_list"] - - # Get raytracer config with defaults - raytracer_config = config.get("raytracer", {}) - self.max_bounce = raytracer_config.get("max_bounce", 2) - self.spp = raytracer_config.get("spp", 1) - - # Get simulation config with defaults - simulation_config = config.get("simulation", {}) - self.n_steps = simulation_config.get("n_steps", 1) - - # Get camera config with defaults - camera_config = config.get("camera", {}) - self.camera_pos = camera_config.get("position", [1.5, 0.5, 1.5]) - self.camera_lookat = camera_config.get("lookat", [0.0, 0.0, 0.5]) - self.camera_fov = camera_config.get("fov", 45.0) - - # Get display config with defaults - display_config = config.get("display", {}) - self.gui = display_config.get("gui", False) - - # Get comparison list with defaults - self.comparison_list = config.get("comparison_list", []) diff --git a/scripts/perf_benchmark/benchmark_madrona.py b/scripts/perf_benchmark/benchmark_madrona.py deleted file mode 100644 index dbadb69e..00000000 --- a/scripts/perf_benchmark/benchmark_madrona.py +++ /dev/null @@ -1,141 +0,0 @@ -import os - -import genesis as gs -from genesis.utils.image_exporter import FrameImageExporter - -from batch_benchmark import BenchmarkArgs, write_benchmark_result_file -from benchmark_profiler import BenchmarkProfiler - - -def init_gs(benchmark_args): - ########################## init ########################## - try: - gs.init(backend=gs.gpu) - except Exception as e: - print(f"Failed to initialize GPU backend: {e}") - print("Falling back to CPU backend") - gs.init(backend=gs.cpu) - - ########################## create a scene ########################## - scene = gs.Scene( - viewer_options=gs.options.ViewerOptions( - camera_pos=( - benchmark_args.camera_posX, - benchmark_args.camera_posY, - benchmark_args.camera_posZ, - ), - camera_lookat=( - benchmark_args.camera_lookatX, - benchmark_args.camera_lookatY, - benchmark_args.camera_lookatZ, - ), - camera_fov=benchmark_args.camera_fov, - ), - show_viewer=False, - renderer=gs.options.renderers.BatchRenderer( - use_rasterizer=benchmark_args.rasterizer, - ), - ) - - ########################## entities ########################## - plane = scene.add_entity( - gs.morphs.Plane(), - ) - franka = scene.add_entity( - gs.morphs.MJCF(file=benchmark_args.mjcf), - visualize_contact=False, - ) - - ########################## cameras ########################## - cam_0 = scene.add_camera( - res=(benchmark_args.resX, benchmark_args.resY), - pos=( - benchmark_args.camera_posX, - benchmark_args.camera_posY, - benchmark_args.camera_posZ, - ), - lookat=( - benchmark_args.camera_lookatX, - benchmark_args.camera_lookatY, - benchmark_args.camera_lookatZ, - ), - fov=benchmark_args.camera_fov, - ) - scene.add_light( - pos=(0.0, 0.0, 1.5), - dir=(1.0, 1.0, -2.0), - directional=True, - castshadow=False, - cutoff=45.0, - intensity=0.5, - ) - scene.add_light( - pos=(4, -4, 4), - dir=(-1, 1, -1), - directional=False, - castshadow=False, - cutoff=45.0, - intensity=0.5, - ) - ########################## build ########################## - scene.build(n_envs=benchmark_args.n_envs) - return scene - - -def run_benchmark(scene, benchmark_args): - try: - n_envs = benchmark_args.n_envs - n_steps = benchmark_args.n_steps - - # warmup - scene.step() - rgb, depth, _, _ = scene.render_all_cameras(rgb=True, depth=True) - - # Profiler - profiler = BenchmarkProfiler(n_steps, n_envs) - output_dir = os.path.dirname(benchmark_args.benchmark_result_file) - os.makedirs(output_dir, exist_ok=True) - image_dirname = f"{benchmark_args.renderer}-{benchmark_args.rasterizer}-{benchmark_args.n_envs}-{benchmark_args.resX}" - image_dir = os.path.join(output_dir, image_dirname) - if n_steps < 10: - exporter = FrameImageExporter(image_dir) - - for i in range(n_steps): - profiler.on_simulation_start() - scene.step() - profiler.on_rendering_start() - rgb, depth, _, _ = scene.render_all_cameras(rgb=True, depth=True) - profiler.on_rendering_end() - if n_steps < 10: - exporter.export_frame_all_cameras(i, rgb=rgb) - profiler.end() - profiler.print_summary() - - performance_results = { - "time_taken_gpu": profiler.get_total_rendering_gpu_time(), - "time_taken_cpu": profiler.get_total_rendering_cpu_time(), - "time_taken_per_env_gpu": profiler.get_total_rendering_gpu_time_per_env(), - "time_taken_per_env_cpu": profiler.get_total_rendering_cpu_time_per_env(), - "fps": profiler.get_rendering_fps(), - "fps_per_env": profiler.get_rendering_fps_per_env(), - } - write_benchmark_result_file(benchmark_args, performance_results) - - except Exception as e: - print(f"Error during benchmark: {e}") - raise - - -def main(): - ######################## Parse arguments ####################### - benchmark_args = BenchmarkArgs.parse_benchmark_args() - - ######################## Initialize scene ####################### - scene = init_gs(benchmark_args) - - ######################## Run benchmark ####################### - run_benchmark(scene, benchmark_args) - - -if __name__ == "__main__": - main() diff --git a/scripts/perf_benchmark/benchmark_maniskill.py b/scripts/perf_benchmark/benchmark_maniskill.py deleted file mode 100644 index aa283e46..00000000 --- a/scripts/perf_benchmark/benchmark_maniskill.py +++ /dev/null @@ -1,209 +0,0 @@ -from typing import Dict -import numpy as np -import torch -import os -from PIL import Image - -import gymnasium as gym -import sapien -from mani_skill.envs.sapien_env import BaseEnv -from mani_skill.utils.structs.types import GPUMemoryConfig, SimConfig -from mani_skill.sensors.camera import CameraConfig -from mani_skill.utils import sapien_utils -from mani_skill.utils.registration import register_env -from mani_skill.utils.wrappers.flatten import FlattenActionSpaceWrapper - -from batch_benchmark import BenchmarkArgs, write_benchmark_result_file -from benchmark_profiler import BenchmarkProfiler, get_utilization_percentages, print_system_utilization - - -# Get asset directory from environment variable -asset_dir = os.path.abspath(os.getenv("ASSET_DIR")) -benchmark_dir = os.path.abspath(os.path.dirname(__file__)) - - -@register_env("SingleRobotBenchmark-v1") -class SingleRobotBenchmarkEnv(BaseEnv): - SUPPORTED_REWARD_MODES = ["none"] - SUPPORTED_ROBOTS = ["panda", "unitree_go2", "unitree_g1"] - - def __init__( - self, - *args, - robot_uid="panda", - camera_mode="minimal", - camera_width=128, - camera_height=128, - camera_fov=0.7854, # math.radian(45) - camera_pos=(1.5, 0.5, 1.5), - camera_lookat=(0.0, 0.0, 0.5), - **kwargs, - ): - self.camera_mode = camera_mode - self.camera_width = camera_width - self.camera_height = camera_height - self.camera_pos = camera_pos - self.camera_lookat = camera_lookat - self.camera_fov = camera_fov - super().__init__(*args, robot_uids=robot_uid, **kwargs) - - @property - def _default_sensor_configs(self): - return [ - CameraConfig( - uid="render_camera", - pose=sapien_utils.look_at(self.camera_pos, self.camera_lookat), - width=self.camera_width, - height=self.camera_height, - fov=self.camera_fov, - shader_pack=self.camera_mode, - ), - ] - - @property - def _default_human_render_camera_configs(self): - return dict() - - def _load_agent(self, options: dict): - super()._load_agent(options, sapien.Pose(p=[0, 0, 0], q=[1, 0, 0, 0])) - - def _load_scene(self, options: dict): - ground_path = os.path.join(benchmark_dir, "benchmark_assets/plane_urdf/plane.urdf") - urdf_loader = self.scene.create_urdf_loader() - urdf_loader.fix_root_link = True - ground_actor = urdf_loader.parse(ground_path)["actor_builders"][0] - ground_actor._auto_inertial = True # Force ground urdf to be static! - self.ground = ground_actor.build_static(name="ground") - - def _initialize_episode(self, env_idx: torch.Tensor, options: dict): - with torch.device(self.device): - qpos = np.zeros(len(self.agent.robot.get_active_joints())) - self.agent.robot.set_qpos(qpos) - - def _load_lighting(self, options: Dict): - self.scene.add_directional_light( # norm([1.0, 1.0, -2.0] - [0, 0, 1.5]) - [0.26490647, 0.26490647, -0.92717265], - [3.0, 3.0, 3.0], - shadow=False, - ) - - def _get_obs_extra(self, info: Dict): - return dict() - - def evaluate(self): - return {} - - -def main(): - args = BenchmarkArgs.parse_benchmark_args() - - env_id = "SingleRobotBenchmark-v1" - n_envs = args.n_envs - n_steps = args.n_steps - sim_config = SimConfig( - gpu_memory_config=GPUMemoryConfig( - max_rigid_contact_count=n_envs * max(1024, n_envs) * 80, - max_rigid_patch_count=n_envs * max(1024, n_envs) * 4, - found_lost_pairs_capacity=2**26, - ), - sim_freq=100, # dt = 0.01 - control_freq=100, # substep = 1 - spacing=10.0, - ) - - if args.mjcf.endswith("panda.xml"): - robot_uid = "panda" - elif args.mjcf.endswith("go2.xml"): - robot_uid = "unitree_go2" - elif args.mjcf.endswith("g1.xml"): - robot_uid = "unitree_g1" - else: - raise Exception(f"Invalid robot: {args.mjcf}") - - if args.rasterizer: - camera_mode = "minimal" - else: - camera_mode = "rt-fast" - obs_mode = "rgbd" # Or "rgb" - # render_mode = "sensors" - render_mode = "rgb_array" # "human for GUI" - env = gym.make( - env_id, - num_envs=n_envs, - obs_mode=obs_mode, - render_mode=render_mode, - control_mode="pd_joint_delta_pos", - sim_config=sim_config, - robot_uid=robot_uid, - camera_mode=camera_mode, - camera_width=args.resX, - camera_height=args.resY, - # parallel_in_single_scene=True, # This actually combines all environments into a single env, is it batch rendering? - ) - if isinstance(env.action_space, gym.spaces.Dict): - env = FlattenActionSpaceWrapper(env) - base_env: BaseEnv = env.unwrapped - base_env.print_sim_details() - - # Build image output directory similar to _omni - output_dir = os.path.dirname(args.benchmark_result_file) - os.makedirs(output_dir, exist_ok=True) - image_dirname = f"{args.renderer}-{args.rasterizer}-{args.n_envs}-{args.resX}" - image_dir = os.path.join(output_dir, image_dirname) - image_tiles = [] - - profiler = BenchmarkProfiler(n_steps, n_envs) - with torch.inference_mode(): - env.reset(seed=2022) - for i in range(3): - env.step(None) # warmup step - env.render() - env.reset(seed=2022) - for i in range(n_steps): - # obs, rew, terminated, truncated, info = env.step(actions) - print(f"Step: {i}") - profiler.on_simulation_start() - profiler.on_rendering_start() - obs = env.step(None)[0] - # When render_mode="sensors", speed of env.render() suffers a great decline in large batch size. - profiler.on_rendering_end() - if render_mode == "human": - viewer = env.render() - image_tile = None - else: - image_tile = env.render() - image_tile = obs["sensor_data"]["render_camera"]["rgb"] - - if n_steps < 10: - image_tiles.append(image_tile) - - if i % 10 == 0: - system_analysis = get_utilization_percentages() - print_system_utilization(system_analysis) - - profiler.end() - profiler.print_summary() - if n_steps < 10: - os.makedirs(image_dir, exist_ok=True) - image_tiles = [image_tile.cpu().numpy() for image_tile in image_tiles] - for i in range(n_steps): - for j in range(n_envs): - image_pil = Image.fromarray(image_tiles[i][j]) - image_path = os.path.join(image_dir, f"step{i:02d}_env{j:02d}_{camera_mode}_{robot_uid}.png") - print(f"Image saved: {image_path}") - image_pil.save(image_path) - - env.close() - performance_results = { - "time_taken_gpu": profiler.get_total_rendering_gpu_time(), - "time_taken_cpu": profiler.get_total_rendering_cpu_time(), - "time_taken_per_env_gpu": profiler.get_total_rendering_gpu_time_per_env(), - "time_taken_per_env_cpu": profiler.get_total_rendering_cpu_time_per_env(), - "fps": profiler.get_rendering_fps(), - "fps_per_env": profiler.get_rendering_fps_per_env(), - } - write_benchmark_result_file(args, performance_results) - - -if __name__ == "__main__": - main() diff --git a/scripts/perf_benchmark/benchmark_omni.py b/scripts/perf_benchmark/benchmark_omni.py deleted file mode 100644 index 36b6a5a4..00000000 --- a/scripts/perf_benchmark/benchmark_omni.py +++ /dev/null @@ -1,330 +0,0 @@ -######################## Parse arguments ####################### -from batch_benchmark import BenchmarkArgs, write_benchmark_result_file -benchmark_args = BenchmarkArgs.parse_benchmark_args() -######################## Launch app ####################### -from isaaclab.app import AppLauncher -app = AppLauncher( - headless=not benchmark_args.gui, - enable_cameras=True, - device="cuda:0", -).app - -import os -import math -from scipy.spatial.transform import Rotation as R -import torch -from pxr import PhysxSchema -from PIL import Image - -import carb -import omni.replicator.core as rep -import isaaclab.sim as sim_utils -import isaacsim.core.utils.stage as stage_utils -import isaaclab.assets as asset_utils -import isaaclab_assets.robots as asset_robots -from isaaclab.scene.interactive_scene import InteractiveScene -from isaaclab.sensors import TiledCameraCfg -from isaaclab.utils.math import create_rotation_matrix_from_view, quat_from_matrix -from isaaclab.utils import configclass -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim.converters import MjcfConverter, MjcfConverterCfg -from isaacsim.core.utils.extensions import enable_extension -enable_extension("isaacsim.asset.importer.mjcf") -import isaacsim.asset.importer.mjcf - -from benchmark_profiler import BenchmarkProfiler, get_utilization_percentages, print_system_utilization - - -# Get asset directory from environment variable -asset_dir = os.path.abspath(os.getenv("ASSET_DIR")) -benchmark_dir = os.path.abspath(os.path.dirname(__file__)) - - -def load_mjcf(mjcf_path: str) -> str: - return MjcfConverter( - MjcfConverterCfg( - asset_path=mjcf_path, - fix_base=True, - force_usd_conversion=True - ) - ).usd_path - - -def get_robot_config() -> asset_utils.AssetBaseCfg: - robot_name = f"{os.path.splitext(benchmark_args.mjcf)[0]}_new.xml" - robot_path = load_mjcf(os.path.abspath(os.path.join(asset_dir, robot_name))) - print("Robot asset:", robot_path) - - if benchmark_args.mjcf.endswith("g1.xml"): - robot_cfg = asset_utils.AssetBaseCfg( - spawn=asset_robots.unitree.G1_CFG.spawn.copy() - ) - elif benchmark_args.mjcf.endswith("go2.xml"): - robot_cfg = asset_utils.AssetBaseCfg( - spawn=asset_robots.unitree.UNITREE_GO2_CFG.spawn.copy() - ) - elif benchmark_args.mjcf.endswith("panda.xml"): - robot_cfg = asset_utils.AssetBaseCfg( - spawn=asset_robots.franka.FRANKA_PANDA_CFG.spawn.copy() - ) - else: - raise Exception(f"Invalid robot: {benchmark_args.mjcf}") - robot_cfg.spawn.usd_path = robot_path - return robot_cfg.replace(prim_path="{ENV_REGEX_NS}/Robot") - - -def get_dir_light_config() -> asset_utils.AssetBaseCfg: - dir_light_pos = torch.Tensor([[0.0, 0.0, 1.5]]) - dir_light_quat = quat_from_matrix( - create_rotation_matrix_from_view( - dir_light_pos, - torch.Tensor([[1.0, 1.0, -2.0]]), - stage_utils.get_stage_up_axis())) - dir_light_pos = tuple(dir_light_pos.detach().cpu().squeeze().numpy()) - dir_light_quat = tuple(dir_light_quat.detach().cpu().squeeze().numpy()) - dir_light_cfg = asset_utils.AssetBaseCfg( - prim_path="/World/direct_light", - spawn=sim_utils.DistantLightCfg(intensity=500.0, angle=45.0), - init_state=asset_utils.AssetBaseCfg.InitialStateCfg( - pos=dir_light_pos, rot=dir_light_quat - ) - ) - return dir_light_cfg - - -@configclass -class RobotSceneCfg(InteractiveSceneCfg): - """Configuration for a cart-pole scene.""" - ground = asset_utils.AssetBaseCfg( - # prim_path="{ENV_REGEX_NS}/ground", # Each environment should have a ground - prim_path="/World/ground", # All environment shares a ground - spawn=sim_utils.UsdFileCfg( - usd_path=os.path.abspath(os.path.join(benchmark_dir, "benchmark_assets/plane_usd/plane.usd")) - ), - ) - robot: asset_utils.ArticulationCfg = get_robot_config() - dir_light = get_dir_light_config() - - -def apply_benchmark_physics_settings(): - stage = stage_utils.get_current_stage() - physxSceneAPI = PhysxSchema.PhysxSceneAPI.Apply(stage.GetPrimAtPath("/physicsScene")) - physxSceneAPI.CreateGpuTempBufferCapacityAttr(16 * 1024 * 1024 * 2) - physxSceneAPI.CreateGpuHeapCapacityAttr(64 * 1024 * 1024 * 2) - physxSceneAPI.CreateGpuMaxRigidPatchCountAttr(8388608) - physxSceneAPI.CreateGpuMaxRigidContactCountAttr(16777216) - - -def print_render_settings(settings): - print("Render mode:", settings.get("/rtx/rendermode")) - print("Sample per pixel:", settings.get("/rtx/pathtracing/spp")) - print("Total spp:", settings.get("/rtx/pathtracing/totalSpp")) - print("Clamp spp:", settings.get("/rtx/pathtracing/clampSpp")) - print("Max bounce:", settings.get("/rtx/pathtracing/maxBounces")) - print("Optix Denoiser", settings.get("/rtx/pathtracing/optixDenoiser/enabled")) - print("Shadows", settings.get("/rtx/shadows/enabled")) - print("dlss/enabled:", settings.get("/rtx/post/dlss/enabled")) - print("dlss/auto:", settings.get("/rtx/post/dlss/auto")) - print("upscaling/enabled:", settings.get("/rtx/post/upscaling/enabled")) - print("aa/denoiser/enabled:", settings.get("/rtx/post/aa/denoiser/enabled")) - print("aa/taa/enabled:", settings.get("/rtx/post/aa/taa/enabled")) - print("motionBlur/enabled:", settings.get("/rtx/post/motionBlur/enabled")) - print("dof/enabled:", settings.get("/rtx/post/dof/enabled")) - print("bloom/enabled:", settings.get("/rtx/post/bloom/enabled")) - print("tonemap/enabled:", settings.get("/rtx/post/tonemap/enabled")) - print("exposure/enabled:", settings.get("/rtx/post/exposure/enabled")) - print("vsync:", settings.get("/app/window/vsync")) - - -def apply_benchmark_carb_settings(print_changes: bool = False) -> None: - # rep.settings.set_render_rtx_realtime() # Keep default pipeline; explicitly set below - settings = carb.settings.get_settings() - - # Print settings before applying the settings - if print_changes: - print("Before settings:") - print_render_settings(settings) - - # Options: https://docs.omniverse.nvidia.com/materials-and-rendering/latest/rtx-renderer_pt.html - if benchmark_args.rasterizer: - settings.set("/rtx/rendermode", "RayTracedLighting") - else: - settings.set("/rtx/rendermode", "PathTracing") - settings.set("/rtx/shadows/enabled", False) - - # Path tracing settings - settings.set("/rtx/pathtracing/spp", benchmark_args.spp) - settings.set("/rtx/pathtracing/totalSpp", benchmark_args.spp) - settings.set("/rtx/pathtracing/clampSpp", benchmark_args.spp) - settings.set("/rtx/pathtracing/maxBounces", benchmark_args.max_bounce) - settings.set("/rtx/pathtracing/optixDenoiser/enabled", False) - settings.set("/rtx/pathtracing/adaptiveSampling/enabled", False) - - # Disable DLSS & upscaling - settings.set("/rtx-transient/dlssg/enabled", False) - settings.set("/rtx/post/dlss/enabled", False) - settings.set("/rtx/post/dlss/auto", False) - settings.set("/rtx/post/upscaling/enabled", False) - - # Disable post-processing - settings.set("/rtx/post/aa/denoiser/enabled", False) - settings.set("/rtx/post/aa/taa/enabled", False) - settings.set("/rtx/post/motionBlur/enabled", False) - settings.set("/rtx/post/dof/enabled", False) - settings.set("/rtx/post/bloom/enabled", False) - settings.set("/rtx/post/tonemap/enabled", False) - settings.set("/rtx/post/exposure/enabled", False) - - # Disable VSync - settings.set("/app/window/vsync", False) - - # Print settings after applying the settings - if print_changes: - print("After settings:") - print_render_settings(settings) - - -def create_scene(): - """Create simulation and scene with camera and physics/render settings applied.""" - sim_cfg = sim_utils.SimulationCfg( - device="cuda:0", dt=0.01, use_fabric=False, - ) - sim = sim_utils.SimulationContext(sim_cfg) - scene_cfg = RobotSceneCfg(num_envs=benchmark_args.n_envs, env_spacing=10.0) - - apply_benchmark_physics_settings() - apply_benchmark_carb_settings(True) - - camera_fov = math.radians(benchmark_args.camera_fov) - camera_aperture = 20.955 - camera_fol = camera_aperture / (2 * math.tan(camera_fov / 2)) - - camera_pos = torch.tensor(( - benchmark_args.camera_posX, - benchmark_args.camera_posY, - benchmark_args.camera_posZ - )).reshape(-1, 3) - camera_lookat = torch.tensor(( - benchmark_args.camera_lookatX, - benchmark_args.camera_lookatY, - benchmark_args.camera_lookatZ - )).reshape(-1, 3) - camera_quat = quat_from_matrix( - create_rotation_matrix_from_view( - camera_lookat, camera_pos, stage_utils.get_stage_up_axis() - ) @ R.from_euler('z', 180, degrees=True).as_matrix() - ) - camera_pos = tuple(camera_pos.detach().cpu().squeeze().numpy()) - camera_quat = tuple(camera_quat.detach().cpu().squeeze().numpy()) - camera_cfg = TiledCameraCfg( - prim_path="{ENV_REGEX_NS}/tiled_camera", - update_period=0, - height=benchmark_args.resY, - width=benchmark_args.resX, - offset=TiledCameraCfg.OffsetCfg( - pos=camera_pos, - rot=camera_quat, - convention="ros" - ), - data_types=["rgb", "depth"], - spawn=sim_utils.PinholeCameraCfg( - focal_length=camera_fol, - horizontal_aperture=camera_aperture, - ), - ) - setattr(scene_cfg, "tiled_camera", camera_cfg) - scene = InteractiveScene(scene_cfg) - return sim, scene - - -def run_simulator( - sim: sim_utils.SimulationContext, - scene: InteractiveScene, -) -> None: - """Run the simulator with all cameras, and return timing analytics. Visualize if desired.""" - n_envs = benchmark_args.n_envs - n_steps = benchmark_args.n_steps - camera = scene["tiled_camera"] - camera_data_types = ["rgb", "depth"] - - # Initialize timing variables - system_utilization_analytics = get_utilization_percentages() - print_system_utilization(system_utilization_analytics) - - sim.reset() - dt = sim.get_physics_dt() - n_warm_steps = 3 - for i in range(n_warm_steps): - print(f"Warm up step {i}.") - sim.step() - camera.update(dt) - _ = camera.data - - print("Warm up finished.") - output_dir = os.path.dirname(benchmark_args.benchmark_result_file) - os.makedirs(output_dir, exist_ok=True) - image_dirname = f'{benchmark_args.renderer}-{benchmark_args.rasterizer}-{benchmark_args.n_envs}-{benchmark_args.resX}' - image_dir = os.path.join(output_dir, image_dirname) - - profiler = BenchmarkProfiler(n_steps, n_envs) - for i in range(n_steps): - print(f"Step {i}:") - get_utilization_percentages() - - # Measure the total simulation step time - profiler.on_simulation_start() - sim.step(render=False) - profiler.on_rendering_start() - sim.render() - - # Update cameras and process vision data within the simulation step - # Loop through all camera lists and their data_types - camera.update(dt=dt) - rgb_tiles = camera.data.output.get("rgb") - depth_tiles = camera.data.output.get("depth") - profiler.on_rendering_end() - - if n_steps < 10: - os.makedirs(image_dir, exist_ok=True) - rgb_tiles = rgb_tiles.detach().cpu().numpy() - for j in range(n_envs): - rgb_image = rgb_tiles[j] - rgb_image = Image.fromarray(rgb_image) - - image_name = f"rgb_step{i}_env{j}.png" - image_path = os.path.join(image_dir, image_name) - rgb_image.save(image_path) - print("Image saved:", image_path) - # End timing for the step - - profiler.end() - profiler.print_summary() - - system_utilization_analytics = get_utilization_percentages() - print_system_utilization(system_utilization_analytics) - - performance_results = { - "time_taken_gpu": profiler.get_total_rendering_gpu_time(), - "time_taken_cpu": profiler.get_total_rendering_cpu_time(), - "time_taken_per_env_gpu": profiler.get_total_rendering_gpu_time_per_env(), - "time_taken_per_env_cpu": profiler.get_total_rendering_cpu_time_per_env(), - "fps": profiler.get_rendering_fps(), - "fps_per_env": profiler.get_rendering_fps_per_env(), - } - write_benchmark_result_file(benchmark_args, performance_results) - - print("App closing..") - # app.close() - print("App closed!") - - -def main() -> None: - """Entry point for running the benchmark scene and simulator.""" - sim, scene = create_scene() - run_simulator(sim=sim, scene=scene) - - -if __name__ == "__main__": - # run the main function - main() - # simulation_app.close() diff --git a/scripts/perf_benchmark/benchmark_profiler.py b/scripts/perf_benchmark/benchmark_profiler.py deleted file mode 100644 index d288401d..00000000 --- a/scripts/perf_benchmark/benchmark_profiler.py +++ /dev/null @@ -1,270 +0,0 @@ -import time - -import numpy as np -import psutil -import torch -import pynvml - - -def get_utilization_percentages(reset: bool = False, max_values: list[float] = [0.0, 0.0, 0.0, 0.0]) -> list[float]: - """Get the maximum CPU, RAM, GPU utilization (processing), and - GPU memory usage percentages since the last time reset was true.""" - if reset: - max_values[:] = [0, 0, 0, 0] # Reset the max values - - # CPU utilization - cpu_usage = psutil.cpu_percent(interval=0.1) - max_values[0] = max(max_values[0], cpu_usage) - - # RAM utilization - memory_info = psutil.virtual_memory() - ram_usage = memory_info.percent - max_values[1] = max(max_values[1], ram_usage) - - # GPU utilization using pynvml - if torch.cuda.is_available(): - pynvml.nvmlInit() # Initialize NVML - for i in range(torch.cuda.device_count()): - handle = pynvml.nvmlDeviceGetHandleByIndex(i) - - # GPU Utilization - gpu_utilization = pynvml.nvmlDeviceGetUtilizationRates(handle) - gpu_processing_utilization_percent = gpu_utilization.gpu # GPU core utilization - max_values[2] = max(max_values[2], gpu_processing_utilization_percent) - - # GPU Memory Usage - memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle) - gpu_memory_total = memory_info.total - gpu_memory_used = memory_info.used - gpu_memory_utilization_percent = (gpu_memory_used / gpu_memory_total) * 100 - max_values[3] = max(max_values[3], gpu_memory_utilization_percent) - - pynvml.nvmlShutdown() # Shutdown NVML after usage - return max_values - - -def print_system_utilization(analytics: list[float]) -> None: - print( - f"| CPU: {analytics[0]}% |" - f" RAM: {analytics[1]}% |" - f" GPU Compute: {analytics[2]}% |" - f" GPU Memory: {analytics[3]:.2f}% |" - ) - - -class BenchmarkProfiler: - def __init__(self, n_steps, n_envs): - self.reset(n_steps) - self.n_envs = n_envs - - def reset(self, n_steps): - self.n_steps = n_steps - # Create arrays of CUDA events for each step - # Each step has 3 events: simulation_start, render_start, render_end - self.events = [] - # CPU timing arrays - self.cpu_times = [] - for _ in range(n_steps): - step_events = { - "simulation_start": torch.cuda.Event(enable_timing=True), - "render_start": torch.cuda.Event(enable_timing=True), - "render_end": torch.cuda.Event(enable_timing=True), - } - self.events.append(step_events) - # Initialize CPU timing structure for each step - self.cpu_times.append({"simulation_start": 0.0, "render_start": 0.0, "render_end": 0.0}) - self.current_step = 0 - - # Synchronize all previous GPU events - torch.cuda.synchronize() - self.is_synchronized = False - - ######################## Profiling Events ####################### - def on_simulation_start(self): - """Record the start of simulation for current step""" - if self.current_step >= self.n_steps: - raise Exception("All steps have been profiled") - self.events[self.current_step]["simulation_start"].record() - self.cpu_times[self.current_step]["simulation_start"] = time.time() - - def on_rendering_start(self): - """Record the start of rendering for current step""" - if self.current_step >= self.n_steps: - raise Exception("All steps have been profiled") - self.events[self.current_step]["render_start"].record() - self.cpu_times[self.current_step]["render_start"] = time.time() - - def on_rendering_end(self): - """Record the end of rendering for current step""" - if self.current_step >= self.n_steps: - raise Exception("All steps have been profiled") - self.events[self.current_step]["render_end"].record() - self.cpu_times[self.current_step]["render_end"] = time.time() - self.current_step += 1 - - def end(self): - """End the profiler""" - self._synchronize() - - def _synchronize(self): - """Synchronize GPU to ensure all events are recorded""" - torch.cuda.synchronize() - self.is_synchronized = True - - ######################## Simulation Performance ####################### - def get_total_simulation_gpu_time(self): - """Calculate total simulation GPU time across all steps in seconds""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - total_time = 0.0 - for step in range(self.current_step): - events = self.events[step] - total_time += events["simulation_start"].elapsed_time(events["render_start"]) - return total_time / 1000.0 - - def get_total_simulation_cpu_time(self): - """Calculate total simulation CPU time across all steps in seconds""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - total_time = 0.0 - for step in range(self.current_step): - cpu_times = self.cpu_times[step] - total_time += (cpu_times["render_start"] - cpu_times["simulation_start"]) * 1000 # Convert to ms - return total_time / 1000.0 - - def get_simulation_fps(self): - """Get the FPS for the current step""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - return self.n_envs * self.n_steps / self.get_total_simulation_gpu_time() - - def get_simulation_fps_per_env(self): - """Get the FPS per env for the current step""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - return self.n_steps / self.get_total_simulation_gpu_time() - - ######################## Rendering Performance ####################### - def get_total_rendering_gpu_time(self): - """Calculate total rendering GPU time across all steps in seconds""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - total_time = 0.0 - for step in range(self.current_step): - events = self.events[step] - total_time += events["render_start"].elapsed_time(events["render_end"]) - return total_time / 1000.0 - - def get_total_rendering_cpu_time(self): - """Calculate total rendering CPU time across all steps in seconds""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - total_time = 0.0 - for step in range(self.current_step): - cpu_times = self.cpu_times[step] - total_time += (cpu_times["render_end"] - cpu_times["render_start"]) * 1000 # Convert to ms - return total_time / 1000.0 - - def get_rendering_fps(self): - """Get the FPS for the current step""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - return self.n_envs * self.n_steps / self.get_total_rendering_gpu_time() - - def get_rendering_fps_per_env(self): - """Get the FPS per env for the current step""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - return self.n_steps / self.get_total_rendering_gpu_time() - - def get_total_rendering_gpu_time_per_env(self): - """Get the total rendering GPU time per env for the current step in seconds""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - return self.get_total_rendering_gpu_time() / self.n_envs - - def get_total_rendering_cpu_time_per_env(self): - """Get the total rendering CPU time per env for the current step in seconds""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - return self.get_total_rendering_cpu_time() / self.n_envs - - ######################## Total Performance ####################### - def get_total_gpu_time(self): - """Calculate total GPU time across all steps in seconds""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - total_time = 0.0 - for step in range(self.current_step): - events = self.events[step] - total_time += events["simulation_start"].elapsed_time(events["render_end"]) - return total_time / 1000.0 - - def get_total_cpu_time(self): - """Calculate total CPU time across all steps in seconds""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - total_time = 0.0 - for step in range(self.current_step): - cpu_times = self.cpu_times[step] - total_time += (cpu_times["render_end"] - cpu_times["simulation_start"]) * 1000 # Convert to ms - return total_time / 1000.0 - - def get_total_gpu_time_per_env(self): - """Get the total GPU time per env for the current step in seconds""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - return self.get_total_gpu_time() / self.n_envs - - def get_total_cpu_time_per_env(self): - """Get the total CPU time per env for the current step in seconds""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - return self.get_total_cpu_time() / self.n_envs - - def get_step_times(self, step_idx): - """Get detailed timing for a specific step in seconds""" - if not self.is_synchronized: - raise Exception("GPU profiler is not synchronized") - if step_idx >= self.current_step: - raise Exception(f"Step {step_idx} has not been profiled yet") - - events = self.events[step_idx] - cpu_times = self.cpu_times[step_idx] - - return { - "simulation": { - "gpu_ms": events["simulation_start"].elapsed_time(events["render_start"]), - "cpu_ms": (cpu_times["render_start"] - cpu_times["simulation_start"]) * 1000, - }, - "rendering": { - "gpu_ms": events["render_start"].elapsed_time(events["render_end"]), - "cpu_ms": (cpu_times["render_end"] - cpu_times["render_start"]) * 1000, - }, - "total": { - "gpu_ms": events["simulation_start"].elapsed_time(events["render_end"]), - "cpu_ms": (cpu_times["render_end"] - cpu_times["simulation_start"]) * 1000, - }, - } - - ######################## Print Summary ####################### - def print_rendering_summary(self): - """Print a summary of the profiler""" - print(f"Total rendering GPU time: {self.get_total_rendering_gpu_time()} seconds") - print(f"Total rendering CPU time: {self.get_total_rendering_cpu_time()} seconds") - print(f"Total rendering GPU time per env: {self.get_total_rendering_gpu_time_per_env()} seconds") - print(f"Total rendering CPU time per env: {self.get_total_rendering_cpu_time_per_env()} seconds") - print(f"Rendering FPS: {self.get_rendering_fps()}") - print(f"Rendering FPS per env: {self.get_rendering_fps_per_env()}") - - def print_simulation_summary(self): - """Print a summary of the profiler""" - print(f"Total simulation GPU time: {self.get_total_simulation_gpu_time()} seconds") - print(f"Total simulation CPU time: {self.get_total_simulation_cpu_time()} seconds") - print(f"Simulation FPS: {self.get_simulation_fps()}") - print(f"Simulation FPS per env: {self.get_simulation_fps_per_env()}") - - def print_summary(self): - """Print a summary of the profiler""" - self.print_rendering_summary() - self.print_simulation_summary() diff --git a/scripts/perf_benchmark/benchmark_pyrender.py b/scripts/perf_benchmark/benchmark_pyrender.py deleted file mode 100644 index 3bed018b..00000000 --- a/scripts/perf_benchmark/benchmark_pyrender.py +++ /dev/null @@ -1,136 +0,0 @@ -import os - -import genesis as gs - -from batch_benchmark import BenchmarkArgs -from benchmark_profiler import BenchmarkProfiler - - -def init_gs(benchmark_args): - ########################## init ########################## - try: - gs.init(backend=gs.gpu) - except Exception as e: - print(f"Failed to initialize GPU backend: {e}") - print("Falling back to CPU backend") - gs.init(backend=gs.cpu) - - ########################## create a scene ########################## - scene = gs.Scene( - viewer_options=gs.options.ViewerOptions( - camera_pos=( - benchmark_args.camera_posX, - benchmark_args.camera_posY, - benchmark_args.camera_posZ, - ), - camera_lookat=( - benchmark_args.camera_lookatX, - benchmark_args.camera_lookatY, - benchmark_args.camera_lookatZ, - ), - camera_fov=benchmark_args.camera_fov, - ), - vis_options=gs.options.VisOptions( - lights=[ - { - "type": "directional", - "dir": (1.0, 1.0, -2.0), - "color": (1.0, 1.0, 1.0), - "intensity": 0.5, - }, - { - "type": "point", - "pos": (4, -4, 4), - "color": (1.0, 1.0, 1.0), - "intensity": 1, - }, - ], - ), - show_viewer=False, - rigid_options=gs.options.RigidOptions( - # constraint_solver=gs.constraint_solver.Newton, - ), - renderer=benchmark_args.rasterizer and gs.options.renderers.Rasterizer() or gs.options.renderers.RayTracer(), - ) - - ########################## entities ########################## - plane = scene.add_entity( - gs.morphs.Plane(), - ) - franka = scene.add_entity( - gs.morphs.MJCF(file=benchmark_args.mjcf), - visualize_contact=False, - ) - - ########################## cameras ########################## - cam_0 = scene.add_camera( - res=(benchmark_args.resX, benchmark_args.resY), - pos=( - benchmark_args.camera_posX, - benchmark_args.camera_posY, - benchmark_args.camera_posZ, - ), - lookat=( - benchmark_args.camera_lookatX, - benchmark_args.camera_lookatY, - benchmark_args.camera_lookatZ, - ), - fov=benchmark_args.camera_fov, - ) - ########################## build ########################## - scene.build() - return scene - - -def run_benchmark(scene, benchmark_args): - try: - n_envs = benchmark_args.n_envs - n_steps = benchmark_args.n_steps - - # warmup - scene.step() - rgb, depth, _, _ = scene.visualizer.cameras[0].render(rgb=True, depth=True) - - # Profiler - profiler = BenchmarkProfiler(n_steps, n_envs) - for i in range(n_steps): - profiler.on_simulation_start() - scene.step() - profiler.on_rendering_start() - rgb, depth, _, _ = scene.visualizer.cameras[0].render(rgb=True, depth=True) - profiler.on_rendering_end() - - profiler.end() - profiler.print_summary() - - time_taken_gpu = profiler.get_total_rendering_gpu_time() - time_taken_cpu = profiler.get_total_rendering_cpu_time() - time_taken_per_env_gpu = profiler.get_total_rendering_gpu_time_per_env() - time_taken_per_env_cpu = profiler.get_total_rendering_cpu_time_per_env() - fps = profiler.get_rendering_fps() - fps_per_env = profiler.get_rendering_fps_per_env() - - # Append a line with all args and results in csv format - os.makedirs(os.path.dirname(benchmark_args.benchmark_result_file), exist_ok=True) - with open(benchmark_args.benchmark_result_file, "a") as f: - f.write( - f"succeeded,{benchmark_args.mjcf},{benchmark_args.renderer},{benchmark_args.rasterizer},{benchmark_args.n_envs},{benchmark_args.n_steps},{benchmark_args.resX},{benchmark_args.resY},{benchmark_args.camera_posX},{benchmark_args.camera_posY},{benchmark_args.camera_posZ},{benchmark_args.camera_lookatX},{benchmark_args.camera_lookatY},{benchmark_args.camera_lookatZ},{benchmark_args.camera_fov},{time_taken_gpu},{time_taken_per_env_gpu},{time_taken_cpu},{time_taken_per_env_cpu},{fps},{fps_per_env}\n" - ) - except Exception as e: - print(f"Error during benchmark: {e}") - raise - - -def main(): - ######################## Parse arguments ####################### - benchmark_args = BenchmarkArgs.parse_benchmark_args() - - ######################## Initialize scene ####################### - scene = init_gs(benchmark_args) - - ######################## Run benchmark ####################### - run_benchmark(scene, benchmark_args) - - -if __name__ == "__main__": - main() diff --git a/scripts/perf_benchmark/benchmark_report_generator.py b/scripts/perf_benchmark/benchmark_report_generator.py deleted file mode 100644 index f285d9e6..00000000 --- a/scripts/perf_benchmark/benchmark_report_generator.py +++ /dev/null @@ -1,489 +0,0 @@ -import glob -import os -import html -import argparse - -import pandas as pd -import matplotlib.pyplot as plt -import numpy as np -from benchmark_configs import BenchmarkConfigs - - -def generate_table_html(plot_table_data): - # Add CSS styling for the table - html_table = """ - - \n""" - - # Get all batch sizes and renderers across all plots - all_batch_sizes = [] - all_renderers = [] - for renderer, renderer_data in plot_table_data.items(): - all_renderers.append(renderer) - for batch_size in renderer_data.keys(): - if batch_size not in all_batch_sizes: - all_batch_sizes.append(batch_size) - - sorted_batch_sizes = sorted(all_batch_sizes) - - # Header row with batch sizes - html_table += "" - for batch_size in sorted_batch_sizes: - html_table += f"" - html_table += "\n" - - # Data rows - renderer_data = [] - for renderer in all_renderers: - html_table += f"" - row_data = [] - for batch_size in sorted_batch_sizes: - if renderer not in plot_table_data or batch_size not in plot_table_data[renderer]: - row_data.append(None) - html_table += "" - else: - fps = plot_table_data[renderer][batch_size] - row_data.append(fps) - html_table += f"" - html_table += "\n" - renderer_data.append(row_data) - - # Add speedup row for every two renderers - if len(renderer_data) % 2 == 0: - html_table += f"" - last_renderer_data = [None, None] - for i in range(len(sorted_batch_sizes)): - if renderer_data[-2][i] is not None and renderer_data[-1][i] is not None: - ratio = renderer_data[-1][i] / renderer_data[-2][i] - last_renderer_data[-2] = renderer_data[-2][i] - last_renderer_data[-1] = renderer_data[-1][i] - html_table += f"" - elif renderer_data[-2][i] is not None and renderer_data[-1][i] is None: - ratio = last_renderer_data[-1] / renderer_data[-2][i] - last_renderer_data[-2] = renderer_data[-2][i] - html_table += f"" - elif renderer_data[-2][i] is None and renderer_data[-1][i] is not None: - ratio = renderer_data[-1][i] / last_renderer_data[-2] - last_renderer_data[-1] = renderer_data[-1][i] - html_table += f"" - else: - html_table += "" - html_table += "\n" - - html_table += "
Renderer{batch_size}
{html.escape(renderer)}N/A{fps:.1f}
Speedup{ratio:.1f}x{ratio:.1f}x{ratio:.1f}xN/A
" - return html_table - - -def generatePlotHtml(plots_dir, all_plot_table_data): - # Generate an html page to display all the plots - - # Get all plot files - plot_files = glob.glob(f"{plots_dir}/*.png") - if len(plot_files) == 0: - print(f"No plot files found in {plots_dir}") - return - - # Separate regular plots from comparison charts - regular_plot_files = [p for p in plot_files if p.endswith("_plot.png") and not p.endswith("_comparison_plot.png")] - - # Group regular plots by MJCF file - plot_groups = {} - for plot_file in regular_plot_files: - basename = os.path.basename(plot_file) - mjcf_name = basename.split("_")[0] - if mjcf_name not in plot_groups: - plot_groups[mjcf_name] = [] - plot_groups[mjcf_name].append(plot_file) - - # Sort plot groups by mjcf name and plot file name - plot_groups = sorted(plot_groups.items(), key=lambda x: (x[0], x[1][0])) - - # Group comparison plots by resolution - comparison_plot_files = {} - for plot_file in plot_files: - if plot_file.endswith("_comparison_plot.png"): - # Extract resolution from filename (e.g., "128x128" from "..._128x128_comparison_plot.png") - resolution = plot_file.split("_")[-3] # Get the resolution part - if resolution not in comparison_plot_files: - comparison_plot_files[resolution] = [] - comparison_plot_files[resolution].append(plot_file) - - # Sort resolutions by their dimensions - def get_resolution_dims(res): - width, height = map(int, res.split("x")) - return width * height # Sort by total pixels - - sorted_resolutions = sorted(comparison_plot_files.keys(), key=get_resolution_dims) - - # Create HTML file - html_content = """ - - - - Benchmark Results - - - -

Benchmark Results

- """ - - # Add comparison plots sections by resolution - if comparison_plot_files: - html_content += "
\n" - html_content += "

Performance Comparison Plots

\n" - for resolution in sorted_resolutions: - html_content += f"

Resolution: {resolution}

\n" - html_content += "
\n" - for plot in comparison_plot_files[resolution]: - html_content += generate_table_html(all_plot_table_data[plot]) - html_content += f"{html.escape(os.path.basename(plot))}
\n" - html_content += "
\n" - html_content += "
\n" - - # Add regular plots section - html_content += "
\n" - html_content += "

Performance Plots

\n" - for mjcf_name, plots in plot_groups: - html_content += f"
\n" - for plot in plots: - html_content += f"

{html.escape(mjcf_name)} - {os.path.basename(plot)}

\n" - html_content += ( - f"{html.escape(os.path.basename(plot))}
\n" - ) - html_content += "
\n" - html_content += "
\n" - - html_content += """ - - - """ - - # Write HTML file - with open(f"{plots_dir}/index.html", "w") as f: - f.write(html_content) - - -def get_comparison_data_list(config_file): - config = BenchmarkConfigs(config_file) - return config.comparison_list - - -def generate_report(data_file_path, config_file, width=20, height=15): - # Load the log file as csv - # For each mjcf, rasterizer (rasterizer or not(=raytracer)), generate a plot image and save it to a directory. - # The plot image has batch size on the x-axis and fps on the y-axis. - # Each resolution has a different color. - # The plot image has a legend for the resolution. - # The plot image has a title for the mjcf. - # The plot image has a x-axis label for the batch size. - # The plot image has a y-axis label for the fps. - - # Read CSV file - df = pd.read_csv(data_file_path) - - # Create plots directory if it doesn't exist - plots_dir = os.path.dirname(data_file_path) - if not os.path.exists(plots_dir): - os.makedirs(plots_dir) - # Generate individual plots for each mjcf/rasterizer combination - generate_individual_plots(df, plots_dir, width, height) - - # Generate difference plots for specific aspect ratios - all_plot_table_data = dict() - for aspect_ratio in ["1:1", "4:3", "16:9"]: - for comparison_list in get_comparison_data_list(config_file): - plot_table_data = generate_comparison_plots(df, plots_dir, width, height, comparison_list, aspect_ratio) - all_plot_table_data.update(plot_table_data) - - # Generate an html page to display all the plots - generatePlotHtml(plots_dir, all_plot_table_data) - - -def generate_individual_plots(df, plots_dir, width, height): - # Get unique combinations of mjcf and rasterizer - for mjcf in df["mjcf"].unique(): - for renderer in df[df["mjcf"] == mjcf]["renderer"].unique(): - for rasterizer in df[(df["mjcf"] == mjcf) & (df["renderer"] == renderer)]["rasterizer"].unique(): - # Filter data for this mjcf and rasterizer - data = df[(df["mjcf"] == mjcf) & (df["renderer"] == renderer) & (df["rasterizer"] == rasterizer)] - - # continue if there is no data - if len(data) == 0: - print(f"No data found for {mjcf} and {renderer} and rasterizer:{rasterizer}") - continue - - # Create new figure - plt.figure(figsize=(width, height)) - - # Group data by resolution - resolutions = sorted(data.groupby(["resX", "resY"]), key=lambda x: (x[0][0], x[0][1])) - - # Get all batch sizes - all_batch_sizes = sorted(data["n_envs"].unique()) - - # Create bar chart - x = np.arange(len(all_batch_sizes)) - bar_width = 0.8 / len(resolutions) - - # Plot bars for each resolution - for i, (resolution, res_data) in enumerate(resolutions): - # Create mapping from batch size to index - batch_to_idx = {batch: idx for idx, batch in enumerate(all_batch_sizes)} - - # Create array of FPS for all batch sizes - fps_array = np.zeros(len(all_batch_sizes)) - for batch, fps in zip(res_data["n_envs"], res_data["fps"]): - fps_array[batch_to_idx[batch]] = fps - - # Plot bars - bars = plt.bar( - x + i * bar_width, - fps_array, - bar_width, - label=f"{resolution[0]}x{resolution[1]}", - ) - - # Add value labels on top of bars - for bar in bars: - bar_height = bar.get_height() - if bar_height > 0: # Only add label if there's a value - plt.annotate( - f"{bar_height:.1f}", - xy=(bar.get_x() + bar.get_width() / 2, bar_height), - xytext=(0, 3), # 3 points vertical offset - textcoords="offset points", - ha="center", - va="bottom", - fontsize=8, - ) - - # Customize plot - plt.title( - f'Performance for {os.path.basename(mjcf)}\n{renderer} {"Rasterizer" if rasterizer else "Raytracer"}' - ) - plt.xlabel("Batch Size") - plt.ylabel("FPS") - plt.xticks(x + bar_width * (len(resolutions) - 1) / 2, all_batch_sizes) - plt.legend(title="Resolution") - plt.grid(True, axis="y") - - # Save plot - plot_filename = f"{plots_dir}/{os.path.splitext(os.path.basename(mjcf))[0]}_{renderer}_{'rasterizer' if rasterizer else 'raytracer'}_plot.png" - plt.savefig(plot_filename) - plt.close() - - -def generate_comparison_plots(df, plots_dir, width, height, comparison_list, aspect_ratio=None): - renderer_array = [comparison_info["renderer"] for comparison_info in comparison_list] - renderer_is_rasterizer_array = [comparison_info["rasterizer"] for comparison_info in comparison_list] - rasterizer_str_array = [ - "rasterizer" if renderer_is_rasterizer else "raytracer" - for renderer_is_rasterizer in renderer_is_rasterizer_array - ] - - # Filter by aspect ratio if specified - if aspect_ratio: - if aspect_ratio == "1:1": - df = df[df["resX"] == df["resY"]] - elif aspect_ratio == "4:3": - df = df[df["resX"] * 3 == df["resY"] * 4] - elif aspect_ratio == "16:9": - df = df[df["resX"] * 9 == df["resY"] * 16] - else: - raise ValueError(f"Unsupported aspect ratio: {aspect_ratio}") - - plot_table_data = dict() - - plt.clf() - plt.cla() - - # Generate plots showing fps comparison between renderer_1 and renderer_2 - for mjcf in df["mjcf"].unique(): - mjcf_data = df[df["mjcf"] == mjcf] - - # Get resolutions available for both renderer_1 and renderer_2 - for comparison in comparison_list: - renderer = comparison["renderer"] - renderer_is_rasterizer = comparison["rasterizer"] - renderer_resolutions = [ - set( - zip( - mjcf_data[ - (mjcf_data["renderer"] == renderer) & (mjcf_data["rasterizer"] == renderer_is_rasterizer) - ]["resX"], - mjcf_data[ - (mjcf_data["renderer"] == renderer) & (mjcf_data["rasterizer"] == renderer_is_rasterizer) - ]["resY"], - ) - ) - ] - print(f"renderer: {renderer}, renderer_is_rasterizer: {renderer_is_rasterizer}") - print(f"renderer_resolutions: {renderer_resolutions}") - common_res = set.intersection(*renderer_resolutions) - - # continue if there is no data - if len(common_res) == 0: - print(f"No data found for {mjcf}") - continue - - # Plot comparison for each resolution - for resX, resY in sorted(common_res, key=lambda x: x[0] * x[1]): - plt.figure(figsize=(width, height)) - renderer_data_array = [] - for comparison in comparison_list: - renderer = comparison["renderer"] - renderer_is_rasterizer = comparison["rasterizer"] - renderer_data = mjcf_data[ - (mjcf_data["result"] == "succeeded") - & (mjcf_data["renderer"] == renderer) - & (mjcf_data["rasterizer"] == renderer_is_rasterizer) - & (mjcf_data["resX"] == resX) - & (mjcf_data["resY"] == resY) - ] - renderer_data_array.append(renderer_data) - - # Match batch sizes and calculate difference - batch_sizes = set.union(*[set(renderer_data["n_envs"]) for renderer_data in renderer_data_array]) - sorted_batch_sizes = sorted(list(batch_sizes)) - - # Create bar chart - def add_labels(bars): - for bar in bars: - bar_height = bar.get_height() - plt.annotate( - f"{bar_height:.1f}", - xy=(bar.get_x() + bar.get_width() / 2, bar_height), - xytext=(0, 3), # 3 points vertical offset - textcoords="offset points", - ha="center", - va="bottom", - fontsize=8, - ) - - # Plot bars - bar_width = 0.8 / len(comparison_list) - fps_array = [ - renderer_data[renderer_data["n_envs"].isin(sorted_batch_sizes)]["fps"].values - for renderer_data in renderer_data_array - ] - for i, (fps, renderer, rasterizer_str) in enumerate(zip(fps_array, renderer_array, rasterizer_str_array)): - x = np.arange(len(fps)) - bars = plt.bar( - x + i * bar_width, - fps, - bar_width, - label=f"{renderer} {rasterizer_str}", - ) - add_labels(bars) - - # Customize plot - renderer_str_array = [ - f"{renderer} {rasterizer_str}" for renderer, rasterizer_str in zip(renderer_array, rasterizer_str_array) - ] - renderer_str_array_str = ", ".join(renderer_str_array) - plt.title(f"FPS Comparison: {renderer_str_array_str}\n{os.path.basename(mjcf)} - Resolution: {resX}x{resY}") - plt.xlabel("Batch Size") - plt.ylabel("FPS") - plt.xticks(np.arange(len(sorted_batch_sizes)), sorted_batch_sizes) - plt.legend() - plt.grid(True, axis="y") - - # Save plot - renderer_str_array_str_for_filename = renderer_str_array_str.replace(",", "_") - plot_filename = f"{plots_dir}/{os.path.splitext(os.path.basename(mjcf))[0]}_{renderer_str_array_str_for_filename}_{resX}x{resY}_comparison_plot.png" - plt.savefig(plot_filename, dpi=100) # Added dpi parameter for better quality - plt.close() - - # Create a table of the data in plot_table_data, the key is the plot_filename, the value is a nested dict - # The key of the outer dict is "{renderer} - {rasterizer_str}" - # The key of the inner dict is "batch_size" - # The value of the inner dict is the fps - plot_table_data[plot_filename] = { - f"{renderer} - {rasterizer_str}": { - batch_size: fps for batch_size, fps in zip(sorted_batch_sizes, fps_array[i]) - } - for i, (renderer, rasterizer_str) in enumerate(zip(renderer_array, rasterizer_str_array)) - } - - return plot_table_data - - -def main(): - import sys - import os - - print("Script arguments:", sys.argv) # Debug print - - parser = argparse.ArgumentParser() - parser.add_argument( - "-d", - "--data_file_path", - type=str, - default="scripts/perf_benchmark/benchmark_reports/Perf_Run_Name/perf_data.csv", - help="Path to the benchmark data CSV file", - ) - parser.add_argument( - "-c", - "--config_file", - type=str, - default="benchmark_config_smoke_test.yml", - help="Path to the benchmark config file", - ) - parser.add_argument("-w", "--width", type=int, default=20, help="Width of the plot in inches") - parser.add_argument("-y", "--height", type=int, default=8, help="Height of the plot in inches") - - # If no arguments provided, try to get from environment variables - if len(sys.argv) == 1: - data_file = os.environ.get("BENCHMARK_DATA_FILE") - if data_file: - sys.argv.extend(["-d", data_file]) - - args = parser.parse_args() - print("Parsed arguments:", args) # Debug print - generate_report(args.data_file_path, args.config_file, args.width, args.height) - - -if __name__ == "__main__": - main() diff --git a/scripts/perf_benchmark/configs/benchmark_config_full.yml b/scripts/perf_benchmark/configs/benchmark_config_full.yml deleted file mode 100644 index 0df95f9f..00000000 --- a/scripts/perf_benchmark/configs/benchmark_config_full.yml +++ /dev/null @@ -1,100 +0,0 @@ -mjcf_list: - - xml/franka_emika_panda/panda.xml - - xml/unitree_g1/g1.xml - - xml/unitree_go2/go2.xml - -renderer_list: - - renderer: madrona - benchmark_script: benchmark_madrona.py - timeout: 120 - -rasterizer_list: - - true - - false - -batch_size_list: - - 1 - - 2 - - 4 - - 8 - - 16 - - 32 - - 64 - - 128 - - 256 - - 512 - - 768 - - 1024 - - 1536 - - 2048 - - 3072 - - 4096 - - 6144 - - 8192 - - 12288 - - 16384 - -resolution_list: - #square: - - [64, 64] - - [128, 128] - - [256, 256] - - [512, 512] - - [1024, 1024] - - [2048, 2048] - - [4096, 4096] - - [8192, 8192] - - #four_three: - - [320, 240] - - [640, 480] - - [800, 600] - - [1024, 768] - - [1280, 960] - - [1600, 1200] - - [1920, 1440] - - [2048, 1536] - - [2560, 1920] - - [3200, 2400] - - [4096, 3072] - - [8192, 6144] - - #sixteen_nine: - - [320, 180] - - [640, 360] - - [800, 450] - - [1024, 576] - - [1280, 720] - - [1600, 900] - - [1920, 1080] - - [2048, 1152] - - [2560, 1440] - - [3200, 1800] - - [4096, 2304] - - [8192, 4608] - -comparison_list: - - - renderer: madrona - rasterizer: true - - renderer: madrona - rasterizer: false - -# Configurations shared betwen batch_benchmark.py and benchmark_*.py -# Raytracer configuration -raytracer: - max_bounce: 2 - spp: 1 - -# Simulation configuration -simulation: - n_steps: 1000 - -# Camera configuration -camera: - position: [1.5, 0.5, 1.5] # [x, y, z] - lookat: [0.0, 0.0, 0.5] # [x, y, z] - fov: 45.0 # degrees - -# Display configuration -display: - gui: false # Enable/disable GUI mode diff --git a/scripts/perf_benchmark/configs/benchmark_config_madrona.yml b/scripts/perf_benchmark/configs/benchmark_config_madrona.yml deleted file mode 100644 index b9782c05..00000000 --- a/scripts/perf_benchmark/configs/benchmark_config_madrona.yml +++ /dev/null @@ -1,68 +0,0 @@ -mjcf_list: - - xml/franka_emika_panda/panda.xml - -renderer_list: - - renderer: madrona - benchmark_script: benchmark_madrona.py - timeout: 120 - -rasterizer_list: - - true - - false - -batch_size_list: - - 1 - - 2 - - 4 - - 8 - - 16 - - 32 - - 64 - - 128 - - 256 - - 512 - - 768 - - 1024 - - 1536 - - 2048 - - 3072 - - 4096 - - 6144 - - 8192 - - 12288 - - 16384 - -resolution_list: - #square: - - [128, 128] - - [256, 256] - - #four_three: - - #sixteen_nine: - -comparison_list: - - - renderer: madrona - rasterizer: true - - renderer: madrona - rasterizer: false - -# Configurations shared betwen batch_benchmark.py and benchmark_*.py -# Raytracer configuration -raytracer: - max_bounce: 2 - spp: 1 - -# Simulation configuration -simulation: - n_steps: 100 - -# Camera configuration -camera: - position: [1.5, 0.5, 1.5] # [x, y, z] - lookat: [0.0, 0.0, 0.5] # [x, y, z] - fov: 45.0 # degrees - -# Display configuration -display: - gui: false # Enable/disable GUI mode diff --git a/scripts/perf_benchmark/configs/benchmark_config_maniskill.yml b/scripts/perf_benchmark/configs/benchmark_config_maniskill.yml deleted file mode 100644 index 1c37780f..00000000 --- a/scripts/perf_benchmark/configs/benchmark_config_maniskill.yml +++ /dev/null @@ -1,58 +0,0 @@ -mjcf_list: - - xml/franka_emika_panda/panda.xml - -renderer_list: - - renderer: maniskill - benchmark_script: benchmark_maniskill.py - timeout: 120 - -rasterizer_list: - - true - -batch_size_list: - - 1 - - 2 - - 4 - - 8 - - 16 - - 32 - - 64 - - 128 - - 256 - - 512 - - 768 - - 1024 - - 2048 - -resolution_list: - #square: - - [128, 128] - - [256, 256] - - #four_three: - - #sixteen_nine: - -comparison_list: - - - renderer: maniskill - rasterizer: true - -# Configurations shared betwen batch_benchmark.py and benchmark_*.py -# Raytracer configuration -raytracer: - max_bounce: 2 - spp: 1 - -# Simulation configuration -simulation: - n_steps: 100 - -# Camera configuration -camera: - position: [1.5, 0.5, 1.5] # [x, y, z] - lookat: [0.0, 0.0, 0.5] # [x, y, z] - fov: 45.0 # degrees - -# Display configuration -display: - gui: false # Enable/disable GUI mode diff --git a/scripts/perf_benchmark/configs/benchmark_config_omni.yml b/scripts/perf_benchmark/configs/benchmark_config_omni.yml deleted file mode 100644 index dd0386e0..00000000 --- a/scripts/perf_benchmark/configs/benchmark_config_omni.yml +++ /dev/null @@ -1,61 +0,0 @@ -mjcf_list: - - xml/franka_emika_panda/panda.xml - -renderer_list: - - renderer: omniverse - benchmark_script: benchmark_omni.py - timeout: 120 - -rasterizer_list: - - true - - false - -batch_size_list: - - 1 - - 2 - - 4 - - 8 - - 16 - - 32 - - 64 - - 128 - - 256 - - 512 - - 768 - - 1024 - - 2048 - -resolution_list: - #square: - - [128, 128] - - [256, 256] - - #four_three: - - #sixteen_nine: - -comparison_list: - - - renderer: omniverse - rasterizer: true - - renderer: omniverse - rasterizer: false - -# Configurations shared betwen batch_benchmark.py and benchmark_*.py -# Raytracer configuration -raytracer: - max_bounce: 2 - spp: 1 - -# Simulation configuration -simulation: - n_steps: 100 - -# Camera configuration -camera: - position: [1.5, 0.5, 1.5] # [x, y, z] - lookat: [0.0, 0.0, 0.5] # [x, y, z] - fov: 45.0 # degrees - -# Display configuration -display: - gui: false # Enable/disable GUI mode diff --git a/scripts/perf_benchmark/configs/benchmark_config_smoke_test.yml b/scripts/perf_benchmark/configs/benchmark_config_smoke_test.yml deleted file mode 100644 index da28d074..00000000 --- a/scripts/perf_benchmark/configs/benchmark_config_smoke_test.yml +++ /dev/null @@ -1,49 +0,0 @@ -mjcf_list: - - xml/franka_emika_panda/panda.xml - -renderer_list: - - renderer: madrona - benchmark_script: benchmark_madrona.py - timeout: 120 - -rasterizer_list: - - true - - false - -batch_size_list: - - 256 - -resolution_list: - #square: - - [128, 128] - - [256, 256] - - #four_three: - - #sixteen_nine: - -comparison_list: - - - renderer: madrona - rasterizer: true - - renderer: madrona - rasterizer: false - -# Configurations shared betwen batch_benchmark.py and benchmark_*.py -# Raytracer configuration -raytracer: - max_bounce: 2 - spp: 1 - -# Simulation configuration -simulation: - n_steps: 1000 - -# Camera configuration -camera: - position: [1.5, 0.5, 1.5] # [x, y, z] - lookat: [0.0, 0.0, 0.5] # [x, y, z] - fov: 45.0 # degrees - -# Display configuration -display: - gui: false # Enable/disable GUI mode diff --git a/scripts/perf_benchmark/example_report/index.html b/scripts/perf_benchmark/example_report/index.html deleted file mode 100644 index 2cc6b936..00000000 --- a/scripts/perf_benchmark/example_report/index.html +++ /dev/null @@ -1,123 +0,0 @@ - - - - - Benchmark Results - - - -

Benchmark Results

-
-

Performance Comparison Plots

-

Resolution: 128x128

-
- - - - - - - -
Renderer12481632641282565127681024153620483072
madrona - rasterizer550.01077.72099.94002.87210.412448.219639.127706.233683.333811.828812.228596.128518.228085.528065.9
madrona - raytracer526.11025.61980.43860.07022.613338.621824.832294.740519.249755.435023.335532.336919.237179.037510.0
Speedup1.0x1.0x0.9x1.0x1.0x1.1x1.1x1.2x1.2x1.5x1.2x1.2x1.3x1.3x1.3x
panda_madrona rasterizer_ madrona raytracer_128x128_comparison_plot.png
-
-

Resolution: 256x256

-
- - - - - - - -
Renderer12481632641282565127681024153620483072
madrona - rasterizer547.61066.12031.83771.66762.310946.516284.721487.813058.213682.113859.313990.213842.013674.613512.6
madrona - raytracer480.0939.41748.33265.65704.68972.412211.415166.110396.810825.110909.511107.511148.211062.111146.5
Speedup0.9x0.9x0.9x0.9x0.8x0.8x0.7x0.7x0.8x0.8x0.8x0.8x0.8x0.8x0.8x
panda_madrona rasterizer_ madrona raytracer_256x256_comparison_plot.png
-
-
-
-

Performance Plots

-
-

panda - panda_madrona_rasterizer_plot.png

-panda_madrona_rasterizer_plot.png
-

panda - panda_madrona_raytracer_plot.png

-panda_madrona_raytracer_plot.png
-
-
- - - - \ No newline at end of file diff --git a/scripts/perf_benchmark/example_report/panda_madrona rasterizer_ madrona raytracer_128x128_comparison_plot.png b/scripts/perf_benchmark/example_report/panda_madrona rasterizer_ madrona raytracer_128x128_comparison_plot.png deleted file mode 100644 index 9d3e9e77..00000000 Binary files a/scripts/perf_benchmark/example_report/panda_madrona rasterizer_ madrona raytracer_128x128_comparison_plot.png and /dev/null differ diff --git a/scripts/perf_benchmark/example_report/panda_madrona rasterizer_ madrona raytracer_128x128_comparison_table.png b/scripts/perf_benchmark/example_report/panda_madrona rasterizer_ madrona raytracer_128x128_comparison_table.png deleted file mode 100644 index eb00e7b8..00000000 Binary files a/scripts/perf_benchmark/example_report/panda_madrona rasterizer_ madrona raytracer_128x128_comparison_table.png and /dev/null differ diff --git a/scripts/perf_benchmark/example_report/panda_madrona_rasterizer_plot.png b/scripts/perf_benchmark/example_report/panda_madrona_rasterizer_plot.png deleted file mode 100644 index 792e6be6..00000000 Binary files a/scripts/perf_benchmark/example_report/panda_madrona_rasterizer_plot.png and /dev/null differ diff --git a/scripts/perf_benchmark/example_report/panda_madrona_raytracer_plot.png b/scripts/perf_benchmark/example_report/panda_madrona_raytracer_plot.png deleted file mode 100644 index f044f5d8..00000000 Binary files a/scripts/perf_benchmark/example_report/panda_madrona_raytracer_plot.png and /dev/null differ diff --git a/scripts/perf_benchmark/example_report/perf_data.csv b/scripts/perf_benchmark/example_report/perf_data.csv deleted file mode 100644 index 752bb4ea..00000000 --- a/scripts/perf_benchmark/example_report/perf_data.csv +++ /dev/null @@ -1,71 +0,0 @@ -result,mjcf,renderer,rasterizer,n_envs,n_steps,resX,resY,camera_posX,camera_posY,camera_posZ,camera_lookatX,camera_lookatY,camera_lookatZ,camera_fov,time_taken_gpu,time_taken_per_env_gpu,time_taken_cpu,time_taken_per_env_cpu,fps,fps_per_env -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,1,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.1818104321956634,0.1818104321956634,0.1930997371673584,0.1930997371673584,550.0234436073531,550.0234436073531 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,2,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.1855801277160644,0.0927900638580322,0.2083349227905273,0.1041674613952636,1077.7015969403674,538.8507984701837 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,4,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.1904850237369537,0.0476212559342384,0.2134573459625244,0.0533643364906311,2099.902617816147,524.9756544540368 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,8,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.199859935760498,0.0249824919700622,0.2249636650085449,0.0281204581260681,4002.80324796401,500.35040599550126 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,16,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.221900832414627,0.0138688020259141,0.2477030754089355,0.0154814422130584,7210.428111465401,450.6517569665876 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,32,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.257064959526062,0.0080332799851894,0.2882359027862549,0.0090073719620704,12448.215446787,389.00673271209376 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,64,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.3258807048797607,0.0050918860137462,0.3639366626739502,0.0056865103542804,19639.08848902665,306.8607576410414 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,128,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.4619897890090942,0.0036092952266335,0.5255374908447266,0.0041057616472244,27706.240060963843,216.45500047628 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,256,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.7600212163925171,0.0029688328765332,0.8242590427398682,0.0032197618857026,33683.27021383932,131.57527427280985 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,512,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,1.5142652807235717,0.0029575493764132,1.5475430488586426,0.003022545017302,33811.777006162854,66.03862696516182 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,768,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,2.6655333137512205,0.0034707465022802,2.7337162494659424,0.0035595263664921,28812.24541587849,37.51594455192512 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,1024,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,3.5809054374694824,0.0034969779662787,3.6493477821350098,0.0035638161934912,28596.119553596192,27.92589800155878 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,1536,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,5.386033821105957,0.0035065324356158,5.450999736785889,0.0035488279536366,28518.20190918521,18.566537701292454 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,2048,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,7.29202025604248,0.0035605567656457,7.355396747589111,0.0035915023181587,28085.49521379812,13.71362070986236 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,3072,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,10.945659706115723,0.0035630402689178,11.008342027664185,0.0035834446704636,28065.919117543606,9.136041379408724 -failed,xml/franka_emika_panda/panda.xml,madrona,True,4096,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,,,,,, -failed,xml/franka_emika_panda/panda.xml,madrona,True,6144,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,,,,,, -failed,xml/franka_emika_panda/panda.xml,madrona,True,8192,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,,,,,, -failed,xml/franka_emika_panda/panda.xml,madrona,True,12288,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,,,,,, -failed,xml/franka_emika_panda/panda.xml,madrona,True,16384,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,,,,,, -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,1,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.1826125756502151,0.1826125756502151,0.1937849521636963,0.1937849521636963,547.607412271238,547.607412271238 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,2,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.1875923197269439,0.0937961598634719,0.2091264724731445,0.1045632362365722,1066.1417284626389,533.0708642313194 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,4,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.1968726726770401,0.04921816816926,0.2214667797088623,0.0553666949272155,2031.770049955995,507.9425124889988 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,8,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.2121126413345337,0.0265140801668167,0.2369122505187988,0.0296140313148498,3771.580962674823,471.4476203343529 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,16,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.2366069767475128,0.0147879360467195,0.2632186412811279,0.0164511650800704,6762.268898382426,422.6418061489017 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,32,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.2923321604728698,0.0091353800147771,0.3365225791931152,0.0105163305997848,10946.452127688424,342.07662899026326 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,64,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.3930070729255676,0.0061407355144619,0.4455823898315429,0.0069622248411178,16284.694197379262,254.44834683405097 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,128,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.5956879329681396,0.0046538119763135,0.6594650745391846,0.0051520708948373,21487.76111045481,167.8731336754282 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,256,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,1.9604556407928464,0.007658029846847,2.0290844440460205,0.0079261111095547,13058.188855345312,51.008550216192624 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,512,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,3.74211616897583,0.0073088206425309,3.809530258178711,0.0074404887855052,13682.09795956516,26.722847577275704 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,768,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,5.541409248352051,0.0072153766254583,5.605247259140015,0.0072984990353385,13859.290400332624,18.04595104209977 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,1024,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,7.319408500671387,0.0071478598639369,7.3833396434783936,0.0072102926205843,13990.20152934587,13.662306181001826 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,1536,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,11.096647178649905,0.0072243796736001,11.15586543083191,0.0072629332231978,13842.01890238778,9.011731056242043 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,2048,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,14.976678344726562,0.007312831223011,15.03558588027954,0.0073415946681052,13674.594278250766,6.677047987427132 -succeeded,xml/franka_emika_panda/panda.xml,madrona,True,3072,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,22.73427993774414,0.0074004817505677,22.7915723323822,0.0074191316186139,13512.633821754664,4.398643822185763 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,1,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.1900674557685852,0.1900674557685852,0.2022051811218261,0.2022051811218261,526.1289977057095,526.1289977057095 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,2,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.1949984003305435,0.0974992001652717,0.2171480655670166,0.1085740327835083,1025.649439487597,512.8247197437985 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,4,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.2019746559858322,0.050493663996458,0.2255463600158691,0.0563865900039672,1980.446497346967,495.11162433674167 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,8,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.2072516167163849,0.0259064520895481,0.2336342334747314,0.0292042791843414,3860.042264928463,482.50528311605785 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,16,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.2278363840579986,0.0142397740036249,0.2552666664123535,0.015954166650772,7022.583362246039,438.9114601403774 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,32,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.2399055359363555,0.0074970479980111,0.2729389667510986,0.0085293427109718,13338.583403298066,416.83073135306455 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,64,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.2932444808483124,0.0045819450132548,0.3325328826904297,0.0051958262920379,21824.792683176023,341.01238567462536 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,128,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.3963496654033661,0.0030964817609637,0.4526624679565429,0.0035364255309104,32294.716300500484,252.30247109766003 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,256,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.6317999358177185,0.0024679684992879,0.6983375549316406,0.0027278810739517,40519.15574645752,158.27795213459967 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,512,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,1.029033890724182,0.0020098318178206,1.0975844860076904,0.0021437196992337,49755.40695162919,97.17852920240075 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,768,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,2.1928288612365723,0.0028552459130684,2.2612593173980717,0.0029443480695287,35023.25300328783,45.603194014697685 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,1024,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,2.8818820476531983,0.0028143379371613,2.948338508605957,0.0028792368248105,35532.33557334081,34.699546458340635 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,1536,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,4.1604327316284175,0.0027086150596539,4.226066112518311,0.0027513451253374,36919.23650929457,24.035961269071983 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,2048,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,5.508481456756591,0.0026896882113069,5.573500156402588,0.0027214356232434,37179.03048376363,18.15382347840021 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,3072,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,8.189823631286622,0.0026659582133094,8.253234386444092,0.0026865997351706,37509.96527281002,12.21027515390951 -failed,xml/franka_emika_panda/panda.xml,madrona,False,4096,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,,,,,, -failed,xml/franka_emika_panda/panda.xml,madrona,False,6144,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,,,,,, -failed,xml/franka_emika_panda/panda.xml,madrona,False,8192,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,,,,,, -failed,xml/franka_emika_panda/panda.xml,madrona,False,12288,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,,,,,, -failed,xml/franka_emika_panda/panda.xml,madrona,False,16384,100,128,128,1.5,0.5,1.5,0.0,0.0,0.5,45.0,,,,,, -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,1,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.2083193914890289,0.2083193914890289,0.2212691307067871,0.2212691307067871,480.0321241590535,480.0321241590535 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,2,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.2129035832881927,0.1064517916440963,0.2365458011627197,0.1182729005813598,939.3923620781616,469.6961810390808 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,4,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.2287999367713928,0.0571999841928482,0.2518947124481201,0.06297367811203,1748.2522313791678,437.063057844792 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,8,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.2449816331863403,0.0306227041482925,0.2697319984436035,0.0337164998054504,3265.55092965478,408.19386620684753 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,16,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.2804776637554169,0.0175298539847135,0.3079962730407715,0.0192497670650482,5704.554075989586,356.5346297493491 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,32,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.3566481597423553,0.0111452549919486,0.3974072933197021,0.0124189779162406,8972.428183315731,280.3883807286166 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,64,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.5240999989509583,0.0081890624836087,0.5747663974761963,0.0089807249605655,12211.410060695056,190.80328219836025 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,128,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,0.8439869766235352,0.0065936482548713,0.907806634902954,0.0070922393351793,15166.110798543172,118.48524061361852 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,256,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,2.462288450241089,0.0096183142587542,2.5315020084381104,0.0098886797204613,10396.832262886765,40.61262602690142 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,512,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,4.729760364532471,0.0092378132119774,4.79765510559082,0.009370420128107,10825.072742361026,21.14272019992388 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,768,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,7.039768127441406,0.0091663647492726,7.104008436203003,0.0092500109846393,10909.450227576295,14.205013317156634 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,1024,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,9.219011352539065,0.0090029407739639,9.28283953666687,0.0090652729850262,11107.48171188632,10.847150109263984 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,1536,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,13.77803616333008,0.008970075627168,13.83764934539795,0.0090088862925767,11148.178026183645,7.25792840246331 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,2048,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,18.51371530151367,0.0090399000495672,18.572023153305054,0.0090683706803247,11062.069210022672,5.401400981456384 -succeeded,xml/franka_emika_panda/panda.xml,madrona,False,3072,100,256,256,1.5,0.5,1.5,0.0,0.0,0.5,45.0,27.56010620117188,0.0089713887373606,27.61528587341309,0.0089893508702516,11146.54630710159,3.6284330426762983 diff --git a/scripts/perf_benchmark/process_xml.py b/scripts/perf_benchmark/process_xml.py deleted file mode 100644 index a71f40b5..00000000 --- a/scripts/perf_benchmark/process_xml.py +++ /dev/null @@ -1,104 +0,0 @@ -import xml.etree.ElementTree as ET -import argparse -import os -import copy -import yaml - -# Get asset directory from environment variable -asset_dir = os.getenv("ASSET_DIR") - - -def process_mjcf_geoms(input_file, output_file): - tree = ET.parse(input_file) - root = tree.getroot() - link_geoms = {} - - def wrap_geom(cur, cur_link_name=None): - replace_elems = [] - remove_elems = [] - for i, elem in enumerate(cur): - if elem.tag == "geom": - if cur_link_name is None: - continue - - if elem.get("class", None) == "collision" or \ - elem.get("conaffinity", 0) > 0 or \ - elem.get("contype", 0) > 0: - remove_elems.append(elem) - continue - - body = ET.Element("body") - body_idx = link_geoms.get(cur_link_name, 0) - link_geoms[cur_link_name] = body_idx + 1 - body.set("name", f"{cur_link_name}_geom{body_idx}") - body.append(copy.deepcopy(elem)) - replace_elems.append((elem, body)) - else: - nex_link_name = None - if elem.tag == "body": - nex_link_name = elem.get("name", None) - # Recurse into children - wrap_geom(elem, nex_link_name if nex_link_name else cur_link_name) - - for elem, body in replace_elems: - idx = list(cur).index(elem) - cur.remove(elem) - cur.insert(idx, body) - for elem in remove_elems: - cur.remove(elem) - - wrap_geom(root) - tree.write(output_file, encoding="utf-8", xml_declaration=True) - - -def process_config_file(config_file): - """Process all robot files mentioned in a specific benchmark configuration file.""" - if not os.path.exists(config_file): - print(f"Configuration file not found: {config_file}") - return - - print(f"Processing configuration file: {config_file}") - - try: - with open(config_file, 'r') as f: - config = yaml.safe_load(f) - - if 'mjcf_list' not in config: - print(f"No mjcf_list found in {config_file}") - return - - mjcf_list = config['mjcf_list'] - print(f"Found {len(mjcf_list)} robot files in {config_file}") - - processed_count = 0 - for mjcf_path in mjcf_list: - # Construct full path using asset directory - full_mjcf_path = os.path.join(asset_dir, mjcf_path) - - if not os.path.exists(full_mjcf_path): - print(f"Warning: Robot file not found: {full_mjcf_path}") - continue - - print(f"Processing robot file: {full_mjcf_path}") - output_file = f"{os.path.splitext(full_mjcf_path)[0]}_new.xml" - process_mjcf_geoms(full_mjcf_path, output_file) - processed_count += 1 - print(f"Created processed file: {output_file}") - - print(f"Total files processed: {processed_count}") - - except Exception as e: - print(f"Error processing {config_file}: {e}") - - -def main(): - parser = argparse.ArgumentParser(description="Process MJCF robot files from a benchmark configuration file") - parser.add_argument("--file", type=str, help="Path to the benchmark configuration YAML file") - args = parser.parse_args() - - # Process robot files from the specified configuration file - process_config_file(args.file) - - -if __name__ == "__main__": - main() diff --git a/scripts/profile.py b/scripts/profile.py deleted file mode 100644 index b4df90e1..00000000 --- a/scripts/profile.py +++ /dev/null @@ -1,106 +0,0 @@ -import os -import sys -import json -import subprocess -import pandas as pd - -NUM_THREADS_PER_BLOCK = 256 -NUM_SMS = 82 -BASE_STEP = 10 -DIR_PATH = "/tmp/profile_blocks__megakernel_events" -# only take action when the change of config for certain node can contribute to at least 0.1% overall acceleration -THRESHOLD = 1000 - - -def profile_madrona(bench_cmd, block_config=range(1, 7), cache="/tmp/madcache"): - try: - # os.remove(cache) - pass - except: - pass - - for config in block_config: - profile_command = "MADRONA_MWGPU_ENABLE_PGO=1 MADRONA_MWGPU_TRACE_NAME=profile_{block}_block MADRONA_MWGPU_EXEC_CONFIG_OVERRIDE={thread},{block},{sm} MADRONA_MWGPU_KERNEL_CACHE={cache} {bench_cmd}".format( - thread=NUM_THREADS_PER_BLOCK, - block=config, - sm=NUM_SMS, - cache=cache, - bench_cmd=bench_cmd, - ) - subprocess.run(profile_command, shell=True, text=True) - - -def parse_traces(block_config=range(1, 7)): - from parse_device_tracing import parse_device_logs, step_analysis - - tabular_data = [ - pd.DataFrame( - { - "nodeID": [], - "funcID": [], - "duration (ns)": [], - "invocations": [], - "percentage (%)": [], - "SM utilization": [], - } - ) - for _ in block_config - ] - - isExist = os.path.exists(DIR_PATH) - if not isExist: - os.mkdir(DIR_PATH) - for i in block_config: - path_to_trace = "/tmp/profile_{}_block_madrona_device_tracing.bin".format(i) - - with open(path_to_trace, "rb") as f: - events = bytearray(f.read()) - assert len(events) % 40 == 0 - print("{} events were logged in total".format(len(events) // 40)) - log_steps = parse_device_logs(events) - assert len(log_steps) > BASE_STEP - - with pd.ExcelWriter(DIR_PATH + "/block_{}_metrics.xlsx".format(i)) as writer: - step_analysis( - log_steps[BASE_STEP], - DIR_PATH + "/block_{}.png".format(i), - tabular_data[i - 1], - ).to_excel(writer, index=False) - - -def generate_json(block_config=range(1, 7)): - tabular_data = [pd.read_excel(DIR_PATH + "/block_{}_metrics.xlsx".format(i)) for i in block_config] - tabular_data = ["paddding"] + tabular_data - - overall_durations = {i: sum(tabular_data[i]["duration (ns)"]) for i in block_config} - base_config, base_duration = min(overall_durations.items(), key=lambda x: x[1]) - - duration_deduction = 0 - split_nodes = {} - for i, node in enumerate(tabular_data[base_config]["nodeID"]): - min_duration = tabular_data[base_config]["duration (ns)"][i] - for b in block_config: - if min_duration - tabular_data[b]["duration (ns)"][i] > base_duration / THRESHOLD: - duration_deduction += min_duration - tabular_data[b]["duration (ns)"][i] - min_duration = tabular_data[b]["duration (ns)"][i] - split_nodes[node] = b - - print( - "with default {} block per SM config and following exceptions:".format(base_config), - split_nodes, - ) - print("the estimated acceleration will be {:.3f}%".format(duration_deduction / base_duration * 100)) - - with open(DIR_PATH + "/node_blocks.json", "w") as f: - json.dump(split_nodes, f) - - -if len(sys.argv) < 2: - print(f"{sys.argv[0]} benchmark command and args", file=sys.stderr) - sys.exit(1) - -bench_cmd = " ".join(sys.argv[1:]) - -profile_madrona(bench_cmd) -parse_traces() -generate_json() diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 6f70ecee..d98501cf 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -2,7 +2,6 @@ add_subdirectory(common) add_subdirectory(core) add_subdirectory(importer) -add_subdirectory(physics) add_subdirectory(render) diff --git a/src/bridge/CMakeLists.txt b/src/bridge/CMakeLists.txt index 07ade7ec..2f0c3faf 100644 --- a/src/bridge/CMakeLists.txt +++ b/src/bridge/CMakeLists.txt @@ -11,7 +11,6 @@ target_link_libraries(madgs_cpu_impl madrona_mw_core PRIVATE madrona_common - madrona_mw_physics madrona_rendering_system ) @@ -34,10 +33,8 @@ target_link_libraries(madgs_mgr PRIVATE madrona_python_utils madgs_cpu_impl - madrona_mw_cpu madrona_common madrona_importer - madrona_physics_loader madrona_render madrona_render_asset_processor ) diff --git a/src/bridge/bindings.cpp b/src/bridge/bindings.cpp index afc702cb..89f14825 100644 --- a/src/bridge/bindings.cpp +++ b/src/bridge/bindings.cpp @@ -205,10 +205,6 @@ NB_MODULE(_gs_madrona_batch_renderer, m) { reinterpret_cast(render_options.data()) ); }) - .def("instance_positions_tensor", &Manager::instancePositionsTensor) - .def("instance_rotations_tensor", &Manager::instanceRotationsTensor) - .def("camera_positions_tensor", &Manager::cameraPositionsTensor) - .def("camera_rotations_tensor", &Manager::cameraRotationsTensor) .def("rgb_tensor", &Manager::rgbTensor) .def("depth_tensor", &Manager::depthTensor) .def("normal_tensor", &Manager::normalTensor) diff --git a/src/bridge/mgr.cpp b/src/bridge/mgr.cpp index f5259aa9..83b7d5a8 100644 --- a/src/bridge/mgr.cpp +++ b/src/bridge/mgr.cpp @@ -4,9 +4,7 @@ #include #include -#include #include -#include #include #include @@ -23,7 +21,6 @@ using namespace madrona; using namespace madrona::math; -using namespace madrona::phys; using namespace madrona::py; using namespace madrona::imp; @@ -104,7 +101,6 @@ static inline Optional initRenderManager( .maxLightsPerWorld = gs_model.numLights, .maxInstancesPerWorld = max_instances_per_world, .execMode = ExecMode::CUDA, - .voxelCfg = {}, }); } @@ -826,50 +822,6 @@ void Manager::render(const math::Vector3 *geom_pos, const math::Quat *geom_rot, impl_->render(geom_pos, geom_rot, cam_pos, cam_rot, render_options); } -Tensor Manager::instancePositionsTensor() const -{ - return impl_->exportTensor(ExportID::InstancePositions, - TensorElementType::Float32, - { - impl_->cfg.numWorlds, - impl_->numGeoms, - sizeof(Vector3) / sizeof(float), - }); -} - -Tensor Manager::instanceRotationsTensor() const -{ - return impl_->exportTensor(ExportID::InstanceRotations, - TensorElementType::Float32, - { - impl_->cfg.numWorlds, - impl_->numGeoms, - sizeof(Quat) / sizeof(float), - }); -} - -Tensor Manager::cameraPositionsTensor() const -{ - return impl_->exportTensor(ExportID::CameraPositions, - TensorElementType::Float32, - { - impl_->cfg.numWorlds, - impl_->numCams, - sizeof(Vector3) / sizeof(float), - }); -} - -Tensor Manager::cameraRotationsTensor() const -{ - return impl_->exportTensor(ExportID::CameraRotations, - TensorElementType::Float32, - { - impl_->cfg.numWorlds, - impl_->numCams, - sizeof(Quat) / sizeof(float), - }); -} - Tensor Manager::rgbTensor() const { const uint8_t *rgb_ptr = impl_->getRGBOut(); diff --git a/src/bridge/mgr.hpp b/src/bridge/mgr.hpp index fed8c015..db4d36be 100644 --- a/src/bridge/mgr.hpp +++ b/src/bridge/mgr.hpp @@ -116,11 +116,6 @@ class Manager { // These functions export Tensor objects that link the ECS // simulation state to the python bindings / PyTorch tensors (src/bindings.cpp) // - MGR_EXPORT madrona::py::Tensor instancePositionsTensor() const; - MGR_EXPORT madrona::py::Tensor instanceRotationsTensor() const; - MGR_EXPORT madrona::py::Tensor cameraPositionsTensor() const; - MGR_EXPORT madrona::py::Tensor cameraRotationsTensor() const; - MGR_EXPORT madrona::py::Tensor rgbTensor() const; MGR_EXPORT madrona::py::Tensor depthTensor() const; MGR_EXPORT madrona::py::Tensor normalTensor() const; diff --git a/src/bridge/sim.cpp b/src/bridge/sim.cpp index ad116609..cb9e7608 100644 --- a/src/bridge/sim.cpp +++ b/src/bridge/sim.cpp @@ -4,7 +4,6 @@ using namespace madrona; using namespace madrona::math; -using namespace madrona::phys; namespace RenderingSystem = madrona::render::RenderingSystem; diff --git a/src/bridge/types.hpp b/src/bridge/types.hpp index a5408bb3..cf19e641 100644 --- a/src/bridge/types.hpp +++ b/src/bridge/types.hpp @@ -3,7 +3,6 @@ #include #include #include -#include #include namespace madGS { diff --git a/src/common/CMakeLists.txt b/src/common/CMakeLists.txt index 453905ea..0120357c 100644 --- a/src/common/CMakeLists.txt +++ b/src/common/CMakeLists.txt @@ -149,12 +149,6 @@ target_link_libraries(madrona_bvh_builder madrona_embree ) -add_library(madrona_navmesh STATIC - ${MADRONA_INC_DIR}/navmesh.hpp ${MADRONA_INC_DIR}/navmesh.inl navmesh.cpp -) - -target_link_libraries(madrona_navmesh PUBLIC madrona_common) - if (CUDAToolkit_FOUND) add_library(madrona_cuda STATIC ${MADRONA_INC_DIR}/cuda_utils.hpp ${MADRONA_INC_DIR}/cuda_utils.inl diff --git a/src/common/mesh_bvh_builder.cpp b/src/common/mesh_bvh_builder.cpp index 3907f4bf..30ed7780 100644 --- a/src/common/mesh_bvh_builder.cpp +++ b/src/common/mesh_bvh_builder.cpp @@ -1,6 +1,5 @@ #include -#include #include #include diff --git a/src/common/navmesh.cpp b/src/common/navmesh.cpp deleted file mode 100644 index bd47f87e..00000000 --- a/src/common/navmesh.cpp +++ /dev/null @@ -1,318 +0,0 @@ -#include -#include -#include - -namespace madrona { - -using namespace math; - -static inline CountT heapParent(CountT idx) -{ - return (idx - 1) / 2; -} - -static inline CountT heapChildOffset(CountT idx) -{ - return 2 * idx + 1; -} - -static inline void heapMoveUp(CountT moved_idx, - uint32_t moved_poly, - float moved_cost, - uint32_t *heap, - uint32_t *heap_index, - float *costs) -{ - while (moved_idx != 0) { - CountT parent_idx = heapParent(moved_idx); - uint32_t parent_poly = heap[parent_idx]; - if (costs[parent_poly] <= moved_cost) { - break; - } - - heap[moved_idx] = parent_poly; - heap_index[parent_poly] = moved_idx; - - moved_idx = parent_idx; - } - - heap[moved_idx] = moved_poly; - heap_index[moved_poly] = moved_idx; -} - -void Navmesh::PathFindQueue::add(uint32_t poly, float cost) -{ - costs[poly] = cost; - - CountT new_idx = heapSize++; - heapMoveUp(new_idx, poly, cost, heap, heapIndex, costs); -} - -uint32_t Navmesh::PathFindQueue::removeMin() -{ - uint32_t root_poly = heap[0]; - - uint32_t moved_poly = heap[--heapSize]; - float moved_cost = costs[moved_poly]; - - CountT moved_idx = 0; - CountT child_offset; - while ((child_offset = heapChildOffset(moved_idx)) < heapSize) { - CountT child_idx = child_offset; - uint32_t child_poly = heap[child_idx]; - float child_cost = costs[child_poly]; - { - // Pick the lowest cost child - CountT right_idx = child_idx + 1; - if (right_idx < heapSize) { - uint32_t right_poly = heap[right_idx]; - float right_cost = costs[right_poly]; - if (right_cost < child_cost) { - child_idx = right_idx; - child_poly = right_poly; - child_cost = right_cost; - } - } - } - - // moved_idx is now a valid position for moved_poly in the heap - if (moved_cost < child_cost) { - break; - } - - heap[moved_idx] = child_poly; - heapIndex[child_poly] = moved_idx; - - moved_idx = child_idx; - } - - heap[moved_idx] = moved_poly; - heapIndex[moved_poly] = moved_idx; - - heapIndex[root_poly] = Navmesh::sentinel; - return root_poly; -} - -void Navmesh::PathFindQueue::decreaseCost(uint32_t poly, float cost) -{ - costs[poly] = cost; - - CountT cur_idx = (CountT)heapIndex[poly]; - - heapMoveUp(cur_idx, poly, cost, heap, heapIndex, costs); -} - -static inline uint32_t hashNavmeshEdge(uint32_t a, uint32_t b) -{ - // MurmurHash2 Finalizer - - const uint32_t m = 0x5bd1e995; - - a ^= b >> 18; - a *= m; - b ^= a >> 22; - b *= m; - a ^= b >> 17; - a *= m; - b ^= a >> 19; - b *= m; - - return b; -} - -Navmesh Navmesh::initFromPolygons( - Vector3 *poly_vertices, - uint32_t *poly_idxs, - uint32_t *poly_idx_offsets, - uint32_t *poly_sizes, - uint32_t num_verts, - uint32_t num_polys) -{ - Vector3 *out_vertices = (Vector3 *)rawAlloc(sizeof(Vector3) * num_verts); - utils::copyN(out_vertices, poly_vertices, num_verts); - - uint32_t num_tris = 0; - for (CountT i = 0; i < (CountT)num_polys; i++) { - uint32_t poly_size = poly_sizes[i]; - - num_tris += poly_size - 2; - } - - uint32_t *tri_indices = - (uint32_t *)rawAlloc(sizeof(uint32_t) * 3 * num_tris); - uint32_t *tri_adjacency = - (uint32_t *)rawAlloc(sizeof(uint32_t) * 3 * num_tris); - - AliasEntry *alias_tbl = - (AliasEntry *)rawAlloc(sizeof(AliasEntry) * num_tris); - - // Temporary data - float *tri_weights = (float *)rawAlloc(sizeof(float) * num_tris); - uint32_t *alias_stack = - (uint32_t *)rawAlloc(sizeof(uint32_t) * num_tris * 2); - uint32_t *under_stack = alias_stack; - uint32_t *over_stack = alias_stack + num_tris; - uint32_t under_stack_size = 0; - uint32_t over_stack_size = 0; - - // Triangulate the input polygons - float tri_weight_sum = 0.f; - uint32_t cur_tri = 0; - for (CountT i = 0; i < (CountT)num_polys; i++) { - uint32_t poly_size = poly_sizes[i]; - uint32_t poly_idx_base = poly_idx_offsets[i]; - for (uint32_t tri_offset = 1; tri_offset < poly_size - 1; tri_offset++) { - uint32_t idx_a = poly_idxs[poly_idx_base]; - uint32_t idx_b = poly_idxs[poly_idx_base + tri_offset]; - uint32_t idx_c = poly_idxs[poly_idx_base + tri_offset + 1]; - - tri_indices[3 * cur_tri] = idx_a; - tri_indices[3 * cur_tri + 1] = idx_b; - tri_indices[3 * cur_tri + 2] = idx_c; - - Vector3 a = poly_vertices[idx_a]; - Vector3 b = poly_vertices[idx_b]; - Vector3 c = poly_vertices[idx_c]; - - Vector3 ab = b - a; - Vector3 ac = c - a; - float tri_area_x2 = cross(ab, ac).length(); - tri_weights[cur_tri] = tri_area_x2; - tri_weight_sum += tri_area_x2; - - cur_tri++; - } - } - - for (uint32_t tri_idx = 0; tri_idx < num_tris; tri_idx++) { - float normalized_weight = - tri_weights[tri_idx] * float(num_tris) / tri_weight_sum; - tri_weights[tri_idx] = normalized_weight; - - if (normalized_weight < 1.f) { - under_stack[under_stack_size++] = tri_idx; - } else { - over_stack[over_stack_size++] = tri_idx; - } - } - - while (under_stack_size != 0 && over_stack_size != 0) { - uint32_t under_idx = under_stack[--under_stack_size]; - uint32_t over_idx = over_stack[--over_stack_size]; - - alias_tbl[under_idx] = { - .tau = tri_weights[under_idx], - .alias = over_idx, - }; - - float new_over_weight = - (tri_weights[over_idx] + tri_weights[under_idx]) - 1.f; - tri_weights[over_idx] = new_over_weight; - - if (new_over_weight < 1.f) { - under_stack[under_stack_size++] = over_idx; - } else { - over_stack[over_stack_size++] = over_idx; - } - } - - for (uint32_t i = 0; i < under_stack_size; i++) { - uint32_t idx = under_stack[i]; - - alias_tbl[idx] = { - .tau = 1.f, - .alias = idx, // never accessed - }; - } - - for (uint32_t i = 0; i < over_stack_size; i++) { - uint32_t idx = over_stack[i]; - - alias_tbl[idx] = { - .tau = 1.f, - .alias = idx, // never accessed - }; - } - - rawDealloc(alias_stack); - rawDealloc(tri_weights); - - struct EdgeEntry { - uint32_t vertA; - uint32_t vertB; - uint32_t firstTriIdx; - uint32_t firstTriEdgeOffset; - }; - - uint32_t max_edges = num_tris * 3; - auto *edge_tbl = (EdgeEntry *)rawAlloc(sizeof(EdgeEntry) * max_edges); - - for (uint32_t i = 0; i < max_edges; i++) { - tri_adjacency[i] = sentinel; - edge_tbl[i] = { sentinel, sentinel, sentinel, 0 }; - } - - // This uses a super simple open addressing hash table. The max_edges bound - // is quite pessimistic, so shouldn't get too close to full on the hash - // table. - auto recordEdge = [edge_tbl, tri_adjacency, max_edges, out_vertices]( - uint32_t tri_idx, - uint32_t tri_edge_offset, - uint32_t a, uint32_t b) - { - if (b < a) { - std::swap(a, b); - } - - // Use Lemire fast modulo replacement trick - uint32_t edge_hash = utils::u32mulhi(hashNavmeshEdge(a, b), max_edges); - - while (edge_tbl[edge_hash].vertA != sentinel && ( - edge_tbl[edge_hash].vertA != a || - edge_tbl[edge_hash].vertB != b)) { - if (edge_hash == max_edges - 1) { - edge_hash = 0; - } else { - edge_hash += 1; - } - } - - EdgeEntry &entry = edge_tbl[edge_hash]; - entry.vertA = a; - entry.vertB = b; - - if (entry.firstTriIdx == sentinel) { - entry.firstTriIdx = tri_idx; - entry.firstTriEdgeOffset = tri_edge_offset; - } else { - uint32_t other_tri_idx = entry.firstTriIdx; - uint32_t other_tri_edge_offset = entry.firstTriEdgeOffset; - - tri_adjacency[3 * tri_idx + tri_edge_offset] = other_tri_idx; - tri_adjacency[3 * other_tri_idx + other_tri_edge_offset] = tri_idx; - } - }; - - for (uint32_t tri_idx = 0; tri_idx < num_tris; tri_idx++) { - uint32_t a_idx = tri_indices[3 * tri_idx]; - uint32_t b_idx = tri_indices[3 * tri_idx + 1]; - uint32_t c_idx = tri_indices[3 * tri_idx + 2]; - - recordEdge(tri_idx, 0, a_idx, b_idx); - recordEdge(tri_idx, 1, b_idx, c_idx); - recordEdge(tri_idx, 2, c_idx, a_idx); - } - - rawDealloc(edge_tbl); - - return Navmesh { - .vertices = out_vertices, - .triIndices = tri_indices, - .triAdjacency = tri_adjacency, - .triSampleAliasTable = alias_tbl, - .numVerts = num_verts, - .numTris = num_tris, - }; -} - -} diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt index fc3daf17..b0d147b3 100644 --- a/src/core/CMakeLists.txt +++ b/src/core/CMakeLists.txt @@ -1,7 +1,6 @@ set(MADRONA_CORE_SRCS ${MADRONA_INC_DIR}/fwd.hpp ${MADRONA_INC_DIR}/exec_mode.hpp - #${MADRONA_INC_DIR}/job.hpp ${MADRONA_INC_DIR}/job.inl job.cpp ${MADRONA_INC_DIR}/state.hpp ${MADRONA_INC_DIR}/state.inl state.cpp ${MADRONA_INC_DIR}/context.hpp ${MADRONA_INC_DIR}/context.inl context.cpp ${MADRONA_INC_DIR}/components.hpp base.cpp @@ -11,20 +10,6 @@ set(MADRONA_CORE_SRCS # platform_utils.cpp ) -add_library(madrona_core STATIC - ${MADRONA_CORE_SRCS} -) - -target_link_libraries(madrona_core - PUBLIC - madrona_common -) - -target_compile_definitions(madrona_core - PUBLIC - -DMADRONA_USE_TASK_GRAPH=1 -) - add_library(madrona_mw_core STATIC ${MADRONA_CORE_SRCS} ${MADRONA_INC_DIR}/taskgraph.hpp ${MADRONA_INC_DIR}/taskgraph.inl taskgraph.cpp diff --git a/src/core/job.cpp b/src/core/job.cpp deleted file mode 100644 index 40df7f15..00000000 --- a/src/core/job.cpp +++ /dev/null @@ -1,1668 +0,0 @@ -/* - * Copyright 2021-2022 Brennan Shacklett and contributors - * - * Use of this source code is governed by an MIT-style - * license that can be found in the LICENSE file or at - * https://opensource.org/licenses/MIT. - */ -#include -#include -#include -#include - -#include "worker_init.hpp" - -#if defined(__linux__) or defined(__APPLE__) -#include -#include -#endif - -#include - -#if defined(MADRONA_X64) -#include -#elif defined(MADRONA_ARM) -#endif - -using std::atomic_thread_fence; - -namespace madrona { -namespace { - -struct JobTracker { - uint32_t parent; - uint32_t remainingInvocations; - uint32_t numOutstandingJobs; -}; - -template -struct JobTrackerMapStore -{ - inline T & operator[](uint32_t idx); - inline const T & operator[](uint32_t idx) const; - JobTrackerMapStore(uint32_t) {} - uint32_t expand(uint32_t) - { - FATAL("Out of job IDs\n"); - } - - static constexpr uint64_t pastMapOffset(); -}; - -using JobTrackerMap = IDMap; - -template -constexpr uint64_t JobTrackerMapStore::pastMapOffset() -{ - return sizeof(JobTrackerMap) - offsetof(JobTrackerMap, store_); -} - -template -T & JobTrackerMapStore::operator[](uint32_t idx) -{ - return ((T *)((char *)this + pastMapOffset()))[idx]; -} - -template -const T & JobTrackerMapStore::operator[](uint32_t idx) const -{ - return ((const T *)((const char *)this + pastMapOffset()))[idx]; -} - -struct LogEntry { - enum class Type : uint32_t { - JobFinished, - WaitingJobQueued, - JobCreated, - }; - - struct JobFinished { - JobContainerBase *jobData; - int32_t jobIdx; - uint32_t numCompleted; - }; - - struct WaitingJobQueued { - void (*fnPtr)(); - JobContainerBase *jobData; - uint32_t numInvocations; - }; - - struct JobCreated { - uint32_t parentID; - }; - - Type type; - union { - JobFinished finished; - WaitingJobQueued waiting; - JobCreated created; - }; -}; - -#if 0 -#include -std::vector globalLog; -void printGlobalLog() -{ - for (LogEntry &entry : globalLog) { - switch (entry.type) { - case LogEntry::Type::JobFinished: { - printf("F: (%u %u) %u %u\n", entry.finished.jobIdx.id, entry.finished.jobIdx.gen, entry.finished.numCompleted, entry.finished.threadID); - } break; - case LogEntry::Type::WaitingJobQueued: { - printf("W: (%u %u) %u\n", entry.waiting.id.id, entry.waiting.id.gen, entry.waiting.numInvocations); - } break; - case LogEntry::Type::JobCreated: { - printf("C: (%u %u) %u %u\n", entry.created.curID.id, entry.created.curID.gen, entry.created.parentID, entry.created.numInvocations); - } break; - } - } -} -#endif - -struct alignas(MADRONA_CACHE_LINE) WorkerState { - uint32_t numIdleLoops; - uint32_t numConsecutiveSchedulerCalls; - AtomicU32 logHead; // Only modified by scheduler - uint32_t logTailCache; // Scheduler's last read tail value - // logTail below is only modified by worker thread. - // Still has to be atomic for memory ordering guarantees (worker - // releases updates to log tail to guarantee the log entries themselves are - // visible after an acquire fence by the scheduler - alignas(MADRONA_CACHE_LINE) AtomicU32 logTail; - alignas(MADRONA_CACHE_LINE) AtomicU32 wakeUp; -}; - -// Plan: -// - High level goal: remove all contended atomics in job dependency system and outstanding job tracking system. -// - Centralize job system control into the "Scheduler" -// - When a worker thread can't find work, it attempts to -// wait acquire the Scheduler lock. If this fails, another -// thread is running the scheduler, so it will go to sleep -// on WorkerWakeup primitive. -// - Q: Is this going to create massive contention on the Scheduler lock, or is -// there a way to have a relaxed load of the scheduler lock be safe / accurate (double checked locking maybe) -// - WorkerWakeup primitive: futex / WaitOnAddress / __ulock_wait + spin. -// - Futex may not be strictly necessary here, since the assumption is likely that if we're about to wait, we're actually going to be put to sleep by the kernel. On OSX the only alternative seems to be posix cond vars though, which seem really slow. Regardless, the idea would be a couple of rounds of spinning (possibly using the PAUSE instruction?) where the worker makes sure it isn't just about to be woken up. -// - Q: Does futex already provide this spin? A: Possibly, but it's irrelevant. This spin should actually be continually looking for new jobs after each PAUSE. -// - Worker Logs: each worker thread gets a log. This records every job the worker has completed, including job ID and number of completed invocations. -// - Scheduler: the scheduler has 2 jobs: -// - Run through all worker logs and use this to update job dependency info, moving any jobs from wait queue to run queue that have fulfilled dependencies. (Possible optimization: workers could check job dependency info and if they've fully satisified a requirement, just immediately unblock a job in order to skip needing to go through the scheduler). -// - Wake up sleeping workers based on # of jobs that are ready for execution -// - Related change: when worker decides to split or queue immediately runnable job, it should similarly trigger worker wakeups. Solution could be to wake up 1 worker, which immediately runs and wakes up others. Want to avoid a really heavy search over all threads each time a split occurs. Futex2 / WaitForMultipleObjects are linux / windows options as well (one thread local futex + one "Wake them all" futex). -// - Pitfall: scheduler running is potentially a big synchronization point -// - Option: Scheduler early outs after finding runnable N jobs -// - Option: Some way to parallelize scheduler? -// - MW tick synchronization: -// - Scheduler runs: finds no runnable work, *and all other workers are sleeping*. At this point it does 3 things: -// - Check should_exit flag: if this is true, it wakes all other workers and exits (can imagine more efficient impls here, this will cause N threads worth of scheduler entries - but whatever).sizeof(WaitQueue) + -// - Signal external "frame done" futex. -// - Point A: Wait on external "Launch next frame" futex. Upon wakeup, rerun scheduler -// - JobManager updateLoop inserts dependency on JobID 0 rather than ctx.currentJobID() to update loop resubmission. The idea is that before signaling the launch next frame futex, external code updates JobID 0 (ez option, JobID 0 always has 1 outstanding job, just increase generation). This means after the futex is signaled and the worker at Point A runs, the "next frame" job will be ready and moved to the run queue by the scheduler. -// - Alternative: no resubmission of updateJob by itself. Instead, before signalling the "launch next frame" futex, the system manually queues updateJob requests into each worker run queue and signals all workers. -// - Issue: Entire above strategy is bugged if you have background work you want to keep going. Can not depend on workers being stopped as termination condition. Similarly, cannot guarantee the system is idle in order to slide updateJobs in manually. Need to use fake dependency strategy + an additional job dependent on ctx.currentJobID that when run increments an atomic. When that atomic reaches # worlds, we signal the "frame done" futex in order to wake up the external thread. -// - How do we handle waking up the worker threads in this model? Update the job dependency and then check if the scheduler is locked? If not, run the scheduler ourselves from the external thread? Scheduler could have it's own run queue. Then worker thread model would be: check my run queue -> check scheduler run queue -> try stealing from all other run queues. -// - This system almost certainly has a missed wakeup type race where the scheduler incorrectly concludes that no threads need to be woken up, just as other threads fail to acquire the scheduler lock and then go to sleep. Does this mean we need to ensure each worker thread acquires the scheduler before it sleeps? - -namespace consts { - constexpr uint64_t jobQueueStartAlignment = MADRONA_CACHE_LINE; - - constexpr int waitQueueSizePerWorld = 1024; - - constexpr int runQueueSizePerThread = 65536; - constexpr uint32_t runQueueIndexMask = (uint32_t)runQueueSizePerThread - 1; - - template - constexpr uint64_t computeRunQueueBytes() - { - static_assert(offsetof(JobManager::RunQueue, tail) == - jobQueueStartAlignment); - - constexpr uint64_t bytes_per_thread = - sizeof(JobManager::RunQueue) + num_jobs * sizeof(Job); - - return utils::roundUp(bytes_per_thread, jobQueueStartAlignment); - } - - constexpr uint64_t runQueueBytesPerThread = - computeRunQueueBytes(); - - constexpr int logSizePerThread = 4096; // FIXME: should decrease this and add functionality to force scheduler run - constexpr int logSizeSafetyMargin = logSizePerThread >> 3; - constexpr int logSizeMaxSafeCapacity = logSizePerThread - logSizeSafetyMargin; - constexpr uint32_t logIndexMask = (uint32_t)logSizePerThread - 1; - constexpr uint64_t logBytesPerThread = logSizePerThread * sizeof(LogEntry); - - constexpr uint32_t jobQueueSentinel = 0xFFFFFFFF; - constexpr uint32_t jobAllocSentinel = 0xFFFFFFFF; - constexpr uint32_t numJobAllocArenas = 1024; -} - -inline void workerPause() -{ -#if defined(MADRONA_X64) - _mm_pause(); -#elif defined(MADRONA_ARM) -#if defined(MADRONA_GCC) or defined(MADRONA_CLANG) - asm volatile("yield"); -#elif defined(MADRONA_MSVC) - YieldProcessor(); -#endif -#endif -} - -inline void workerYield() -{ -#if defined(__linux__) or defined(__APPLE__) - sched_yield(); -#elif defined(_WIN32) - STATIC_UNIMPLEMENTED(); -#else - STATIC_UNIMPLEMENTED(); -#endif -} - -inline uint32_t acquireArena(JobManager::Alloc::SharedState &shared) -{ - uint32_t cur_head = shared.freeHead.load_acquire(); - uint32_t new_head, arena_idx; - do { - if (cur_head == consts::jobAllocSentinel) { - FATAL("Out of job memory"); - } - - arena_idx = cur_head & 0xFFFF; - new_head = shared.arenas[arena_idx].metadata.load_relaxed(); - - // Update the tag - new_head += ((uint32_t)1u << (uint32_t)16); - } while (!shared.freeHead.compare_exchange_weak< - sync::release, sync::acquire>(cur_head, new_head)); - - // Arena metadata field is reused for counting used bytes, need to 0 out - shared.arenas[arena_idx].metadata.store_release(0); - - return arena_idx; -} - -inline void releaseArena(JobManager::Alloc::SharedState &shared, - uint32_t arena_idx) -{ - uint32_t cur_head = shared.freeHead.load_relaxed(); - uint32_t new_head; - - do { - new_head = (cur_head & 0xFFFF0000) + ((uint32_t)1u << (uint32_t)16) + arena_idx; - shared.arenas[arena_idx].metadata.store(cur_head, sync::relaxed); - } while (!shared.freeHead.compare_exchange_weak< - sync::release, sync::relaxed>(cur_head, new_head)); -} - -void disableThreadSignals() -{ -#if defined(__linux__) or defined(__APPLE__) - sigset_t mask; - sigfillset(&mask); - sigdelset(&mask, SIGSEGV); - sigdelset(&mask, SIGILL); - sigdelset(&mask, SIGBUS); - sigdelset(&mask, SIGTRAP); - sigdelset(&mask, SIGFPE); - int res = pthread_sigmask(SIG_BLOCK, &mask, nullptr); - bool failed = res != 0; -#elif defined(_WIN32) - STATIC_UNIMPLEMENTED(); -#else - STATIC_UNIMPLEMENTED(); -#endif - - if (failed) { - FATAL("failed to block signals for fiber executor"); - } -} - -int getNumWorkers(int num_workers) -{ - if (num_workers != 0) { - return num_workers; - } - -#if defined(__linux__) or defined(__APPLE__) - int os_num_threads = sysconf(_SC_NPROCESSORS_ONLN); - - if (os_num_threads == -1) { - FATAL("Failed to get number of concurrent threads"); - } - - return os_num_threads; -#elif defined(_WIN32) -#else - STATIC_UNIMPLEMENTED(); -#endif -} - -void setThreadAffinity(int thread_idx) -{ -#if defined(__linux__) - cpu_set_t cpuset; - pthread_getaffinity_np(pthread_self(), sizeof(cpuset), &cpuset); - - const int max_threads = CPU_COUNT(&cpuset); - - CPU_ZERO(&cpuset); - - if (thread_idx > max_threads) [[unlikely]] { - FATAL("Tried setting thread affinity to %d when %d is max", - thread_idx, max_threads); - } - - CPU_SET(thread_idx, &cpuset); - - int res = pthread_setaffinity_np(pthread_self(), sizeof(cpuset), &cpuset); - - if (res != 0) { - FATAL("Failed to set thread affinity to %d", thread_idx); - } -#elif defined(__APPLE__) - (void)thread_idx; - // No thread affinity on macOS / iOS :( -#elif defined(_WIN32) - STATIC_UNIMPLEMENTED(); -#else - STATIC_UNIMPLEMENTED(); -#endif -} - -// 2 phases: workers decrement numRemaining (initialized to # threads) -// and then spin until it reaches 0. Next, all workers add 1 to numAcked, -// and the main thread spins until numAcked == # threads, at which point it -// knows all threads have finished initialization. Simply waiting for -// numRemaining to be 0 is insufficient, because worker threads may still -// be spinning, waiting to see that numRemaining is 0, when the ThreadPoolInit -// struct is freed -struct ThreadPoolInit { - AtomicI32 numRemaining; - AtomicI32 numAcked; - - inline void workerWait() - { - numRemaining.fetch_sub_release(1); - - while (numRemaining.load_acquire() != 0) { - workerYield(); - } - - numAcked.fetch_add_release(1); - } - - inline void mainWait(int num_threads) - { - while (numAcked.load_acquire() != num_threads) { - workerYield(); - } - } -}; - -inline WorkerState & getWorkerState(void *base, int thread_idx) -{ - return ((WorkerState *)base)[thread_idx]; -} - -inline LogEntry * getWorkerLog(void *base, int thread_idx) -{ - return (LogEntry *)((char *)base + consts::logBytesPerThread * thread_idx); -} - -inline JobManager::RunQueue * getRunQueue( - void *queue_base, const int thread_idx) -{ - return (JobManager::RunQueue *)((char *)queue_base + - thread_idx * consts::runQueueBytesPerThread); -} - -inline Job * getRunnableJobs(JobManager::RunQueue *queue) -{ - return (Job *)((char *)queue + sizeof(JobManager::RunQueue)); -} - -inline JobTrackerMap & getTrackerMap(void *base) -{ - return *(JobTrackerMap *)base; -} - -inline JobTrackerMap::Cache & getTrackerCache(void *tracker_cache_base, - int thread_idx) -{ - return ((JobTrackerMap::Cache *)tracker_cache_base)[thread_idx]; -} - -inline void decrementJobTracker(JobTrackerMap &tracker_map, - JobTrackerMap::Cache &tracker_cache, - int32_t job_id) -{ - - while (job_id != JobID::none().id) { - JobTracker &tracker = tracker_map.getRef(job_id); - - uint32_t num_outstanding = --tracker.numOutstandingJobs; - - if (num_outstanding == 0 && tracker.remainingInvocations == 0) { - uint32_t parent = tracker.parent; - - tracker_map.releaseID(tracker_cache, job_id); -#ifdef TSAN_ENABLED - tracker_map.releaseGen(job_id); -#endif - job_id = parent; - } else { - break; - } - } -} - -inline const JobID * getJobDependencies(JobContainerBase *job_base) -{ - return (const JobID *)((char *)job_base + sizeof(JobContainerBase)); -} - -inline bool isRunnable(JobTrackerMap &tracker_map, - JobContainerBase *job_data) -{ - int num_deps = job_data->numDependencies; - - if (num_deps == 0) { - return true; - } - - const JobID *dependencies = getJobDependencies(job_data); - for (int i = 0; i < num_deps; i++) { - JobID dependency = dependencies[i]; - - if (tracker_map.present(dependency)) { - return false; - } - } - - return true; -} - -template -inline uint32_t addToRunQueueImpl(JobManager::RunQueue *run_queue, - Fn &&add_cb) -{ - // No one modifies queue_tail besides this thread - uint32_t cur_tail = run_queue->tail.load_relaxed(); - Job *job_array = getRunnableJobs(run_queue); - - uint32_t num_added = add_cb(job_array, cur_tail); - - cur_tail += num_added; - run_queue->tail.store_release(cur_tail); - - return num_added; -} - -inline void addToLog(WorkerState &worker_state, LogEntry *worker_log, - const LogEntry &entry) -{ - uint32_t cur_tail = worker_state.logTail.load_relaxed(); - uint32_t new_idx = cur_tail & consts::logIndexMask; - - worker_log[new_idx] = entry; - - uint32_t new_tail = cur_tail + 1; - worker_state.logTail.store_release(new_tail); - - uint32_t log_head = worker_state.logHead.load_relaxed(); - if (new_tail - log_head >= consts::logSizePerThread) [[unlikely]] { - for (uint32_t i = log_head; i != new_tail; i++) { - LogEntry &debug_entry = worker_log[i & consts::logIndexMask]; - switch (debug_entry.type) { - case LogEntry::Type::JobFinished: { - printf("Job finished\n"); - } break; - case LogEntry::Type::WaitingJobQueued: { - printf("Waiting job queued\n"); - } break; - case LogEntry::Type::JobCreated: { - printf("JobCreated\n"); - } break; - } - } - FATAL("Worker filled up job system log"); - } -} - -} - -JobManager::Alloc::Alloc(SharedState &shared) - : cur_arena_(acquireArena(shared)), - next_arena_(acquireArena(shared)), - arena_offset_(0), - arena_used_bytes_(0) -{} - -void * JobManager::Alloc::alloc(SharedState &shared, - uint32_t num_bytes, - uint32_t alignment) -{ - // Get offset necessary to meet alignment requirements. - // Alignment must be less than maxJobAlignment (otherwise base address not - // guaranteed to meet alignment). - uint32_t new_offset = utils::roundUpPow2(arena_offset_, alignment); - - if (new_offset + num_bytes <= arena_size_) { - arena_offset_ = new_offset; - } else { - // Out of space in this arena, mark this arena as freeable - // and get a new one - - // Marking the arena as freeable just involves adding the total memory - // used in the arena to the arena's metadata value. Once all jobs in - // the arena have been freed these values will cancel out and the - // metadata value will be zero. - uint32_t post_metadata = - shared.arenas[cur_arena_].metadata.fetch_add_acq_rel( - arena_used_bytes_); - post_metadata += arena_used_bytes_; - - // Edge case, if post_metadata == 0, we can skip getting a new arena - // because there are no active jobs left in the current arena, so - // the cur_arena_ can immediately be reused by resetting offsets to 0 - if (post_metadata != 0) { - // Get next free arena. First check the cached arena in next_arena_ - if (next_arena_ != consts::jobAllocSentinel) { - cur_arena_ = next_arena_; - next_arena_ = consts::jobAllocSentinel; - } else { - cur_arena_ = acquireArena(shared); - } - } - - arena_offset_ = 0; - arena_used_bytes_ = 0; - } - - void *mem = - (char *)shared.jobMemory + arena_size_ * cur_arena_ + arena_offset_; - - arena_offset_ += num_bytes; - - // Need to track arena_used_bytes_ separately from arena_offset_, - // because deallocation code doesn't know how many extra bytes get added - // to each job for alignment padding reasons. - arena_used_bytes_ += num_bytes; - - return mem; -} - -void JobManager::Alloc::dealloc(SharedState &shared, - void *ptr, uint32_t num_bytes) -{ - size_t ptr_offset = (char *)ptr - (char *)shared.jobMemory; - uint32_t arena_idx = ptr_offset / arena_size_; - - Arena &arena = shared.arenas[arena_idx]; - - uint32_t post_metadata = arena.metadata.fetch_sub_acq_rel(num_bytes); - post_metadata -= num_bytes; - - if (post_metadata == 0) { - // If this thread doesn't have a cached free arena, store there, - // otherwise release to global free list - if (next_arena_ == consts::jobAllocSentinel) { - next_arena_ = arena_idx; - } else { - releaseArena(shared, arena_idx); - } - } -} - -JobManager::Alloc::SharedState JobManager::Alloc::makeSharedState( - InitAlloc alloc, uint32_t num_arenas) -{ - if (num_arenas > 65536) { - FATAL("Job allocator can only support up to 2^16 arenas."); - } - - uint64_t total_bytes = (maxJobAlignment - 1) + num_arenas * arena_size_ + - num_arenas * sizeof(Arena); - - void *mem = alloc.alloc(total_bytes); - - void *job_mem = - (void *)utils::roundUp((uintptr_t)mem, (uintptr_t)maxJobAlignment); - - Arena *arenas = (Arena *)((char *)job_mem + arena_size_ * num_arenas); - - // Build initial linear freelist - for (int i = 0; i < (int)num_arenas; i++) { - new (&arenas[i]) Arena { - (i < int(num_arenas - 1)) ? i + 1 : consts::jobQueueSentinel, - }; - } - - return SharedState { - mem, - job_mem, - arenas, - 0, - }; -} - -struct JobManager::Init { - uint32_t numCtxUserdataBytes; - void (*ctxInitFn)(void *, void *, WorkerInit &&); - uint32_t numCtxBytes; - void (*startFn)(Context *, void *); - void *startFnData; - void (*updateFn)(Context *, void *); - void *updateFnData; - int numWorkers; - int numIO; - int numThreads; - StateManager *stateMgr; - bool pinWorkers; - void *statePtr; - void *ctxBase; - void *ctxUserdataBase; - void *stateCacheBase; - void *highBase; - void *normalBase; - void *ioBase; - void *workerStateBase; - void *logBase; - void *waitingJobs; - void *trackerBase; - void *trackerCacheBase; - int numTrackerSlots; -}; - -JobManager::JobManager(uint32_t num_ctx_userdata_bytes, - uint32_t ctx_userdata_alignment, - void (*ctx_init_fn)(void *, void *, WorkerInit &&), - uint32_t num_ctx_bytes, - uint32_t ctx_alignment, - void (*start_fn)(Context *, void *), - void *start_fn_data, - void (*update_fn)(Context *, void *), - void *update_fn_data, - int desired_num_workers, - int num_io, - StateManager *state_mgr, - bool pin_workers) - : JobManager([num_ctx_userdata_bytes, - ctx_userdata_alignment, ctx_init_fn, - num_ctx_bytes, ctx_alignment, - start_fn, start_fn_data, update_fn, update_fn_data, - desired_num_workers, num_io, state_mgr, pin_workers]() { - int num_workers = getNumWorkers(desired_num_workers); - int num_threads = num_workers + num_io; - - uint64_t num_state_bytes = 0; - - uint64_t total_ctx_bytes = - (uint64_t)num_threads * (uint64_t)num_ctx_bytes; - uint64_t total_userdata_bytes = num_ctx_userdata_bytes; -#ifdef MADRONA_MW_MODE - uint64_t num_worlds = state_mgr->numWorlds(); - - total_ctx_bytes *= num_worlds; - total_userdata_bytes *= num_worlds; -#else - uint64_t num_worlds = 1; -#endif - - uint64_t ctx_offset = 0; - num_state_bytes = ctx_offset + total_ctx_bytes; - - uint64_t ctx_userdata_offset = utils::roundUp(num_state_bytes, - (uint64_t)ctx_userdata_alignment); - - num_state_bytes = ctx_userdata_offset + total_userdata_bytes; - - uint64_t state_cache_offset = utils::roundUp(num_state_bytes, - (uint64_t)alignof(StateCache)); - - num_state_bytes = - state_cache_offset + sizeof(StateCache) * num_threads; - - uint64_t high_offset = - utils::roundUp(num_state_bytes, consts::jobQueueStartAlignment); - num_state_bytes = - high_offset + num_threads * consts::runQueueBytesPerThread; - - uint64_t normal_offset = - utils::roundUp(num_state_bytes, consts::jobQueueStartAlignment); - num_state_bytes = - normal_offset + num_threads * consts::runQueueBytesPerThread; - - uint64_t io_offset = - utils::roundUp(num_state_bytes, consts::jobQueueStartAlignment); - num_state_bytes = - io_offset + num_threads * consts::runQueueBytesPerThread; - - uint64_t worker_state_offset = - utils::roundUp(num_state_bytes, (uint64_t)alignof(WorkerState)); - num_state_bytes = - worker_state_offset + num_threads * sizeof(WorkerState); - - uint64_t log_offset = - utils::roundUp(num_state_bytes, consts::jobQueueStartAlignment); - num_state_bytes = log_offset + - num_threads * consts::logSizePerThread * sizeof(LogEntry); - - uint64_t wait_offset = - utils::roundUp(num_state_bytes, consts::jobQueueStartAlignment); - num_state_bytes = wait_offset + - num_worlds * consts::waitQueueSizePerWorld * sizeof(Job); - - uint64_t tracker_cache_offset = - utils::roundUp(num_state_bytes, (uint64_t)alignof(JobTrackerMap::Cache)); - - num_state_bytes = - tracker_cache_offset + num_threads * sizeof(JobTrackerMap::Cache); - - int num_tracker_slots = num_threads * ( - consts::logSizePerThread + consts::runQueueSizePerThread); - - uint64_t tracker_offset = - utils::roundUp(num_state_bytes, consts::jobQueueStartAlignment); - - static_assert( - sizeof(JobTrackerMap) % alignof(JobTrackerMap::Node) == 0); - - num_state_bytes = tracker_offset + sizeof(JobTrackerMap) + - num_tracker_slots * sizeof(JobTrackerMap::Node); - - // Add padding so the base pointer can be aligned - num_state_bytes += ctx_alignment - 1; - - void *state_ptr = InitAlloc().alloc(num_state_bytes); - - char *base_ptr = (char *)utils::alignPtr(state_ptr, ctx_alignment); - - return Init { - .numCtxUserdataBytes = num_ctx_userdata_bytes, - .ctxInitFn = ctx_init_fn, - .numCtxBytes = num_ctx_bytes, - .startFn = start_fn, - .startFnData = start_fn_data, - .updateFn = update_fn, - .updateFnData = update_fn_data, - .numWorkers = num_workers, - .numIO = num_io, - .numThreads = num_threads, - .stateMgr = state_mgr, - .pinWorkers = pin_workers, - .statePtr = state_ptr, - .ctxBase = base_ptr + ctx_offset, - .ctxUserdataBase = base_ptr + ctx_userdata_offset, - .stateCacheBase = base_ptr + state_cache_offset, - .highBase = base_ptr + high_offset, - .normalBase = base_ptr + normal_offset, - .ioBase = base_ptr + io_offset, - .workerStateBase = base_ptr + worker_state_offset, - .logBase = base_ptr + log_offset, - .waitingJobs = base_ptr + wait_offset, - .trackerBase = base_ptr + tracker_offset, - .trackerCacheBase = base_ptr + tracker_cache_offset, - .numTrackerSlots = num_tracker_slots, - }; - }()) -{} - -JobManager::JobManager(const Init &init) - : threads_(init.numThreads, InitAlloc()), - alloc_state_(Alloc::makeSharedState(InitAlloc(), - consts::numJobAllocArenas)), - job_allocs_(threads_.size(), InitAlloc()), - scheduler_ { - .numWaiting = 0, - .numSleepingWorkers = 0, - .lock {}, - }, - state_ptr_(init.statePtr), - high_base_(init.highBase), - normal_base_(init.normalBase), - io_base_(init.ioBase), - tracker_base_(init.trackerBase), - tracker_cache_base_(init.trackerCacheBase), - worker_base_(init.workerStateBase), - log_base_(init.logBase), - waiting_jobs_(init.waitingJobs), - num_compute_workers_(init.numWorkers), - io_sema_(0), - num_high_(0) -{ - for (int i = 0, n = init.numThreads; i < n; i++) { - job_allocs_.emplace(i, alloc_state_); - } - - auto initQueue = [](void *queue_start, int thread_idx) { - RunQueue *queue = getRunQueue(queue_start, thread_idx); - - new (queue) RunQueue { - .head = 0, - .correction = 0, - .auth = 0, - .pad = {}, - .tail = 0, - }; - }; - - JobTrackerMap &tracker_map = getTrackerMap(tracker_base_); - new (&tracker_map) JobTrackerMap(init.numTrackerSlots); - - // Setup per-thread state and queues - for (int i = 0, n = threads_.size(); i < n; i++) { - initQueue(normal_base_, i); - initQueue(high_base_, i); - initQueue(io_base_, i); - - WorkerState &worker_state = getWorkerState(worker_base_, i); - new (&worker_state) WorkerState { - .numIdleLoops = 0, - .numConsecutiveSchedulerCalls = 0, - .logHead = 0, - .logTailCache = 0, - .logTail = 0, - .wakeUp = i + 1, - }; - - JobTrackerMap::Cache &cache = getTrackerCache(tracker_cache_base_, i); - new (&cache) JobTrackerMap::Cache(); - } - - struct StartWrapper { - void (*func)(Context *, void *); - void *data; - AtomicU32 remainingLaunches; - } start_wrapper { - init.startFn, - init.startFnData, -#ifdef MADRONA_MW_MODE - init.stateMgr->numWorlds(), -#else - 1, -#endif - }; - - struct StartJob : JobContainerBase { - StartWrapper *wrapper; - }; - - SingleInvokeFn entry = [](Context *ctx, JobContainerBase *ptr) { - auto &job = *(StartJob *)ptr; - auto &start = *(job.wrapper); - - start.func(ctx, start.data); - - uint32_t job_id = ptr->id.id; - start.remainingLaunches.fetch_sub_release(1); - - ctx->job_mgr_->markInvocationsFinished(ctx->worker_idx_, nullptr, - job_id, 1); - }; - - // Initial job - -#ifdef MADRONA_MW_MODE - int num_worlds = init.stateMgr->numWorlds(); - - HeapArray start_jobs(num_worlds); - - for (int i = 0; i < num_worlds; i++) { - start_jobs[i] = StartJob { - JobContainerBase { JobID::none(), sizeof(StartJob), (uint32_t)i, - 0 }, - &start_wrapper, - }; - - queueJob(i % init.numWorkers, (void (*)())entry, &start_jobs[i], 0, - JobID::none().id, JobPriority::Normal); - } -#else - StartJob start_job { - JobContainerBase { JobID::none(), sizeof(StartJob), 0 }, - &start_wrapper, - }; - - queueJob(0, (void (*)())entry, &start_job, 0, JobID::none().id, - JobPriority::Normal); -#endif - - ThreadPoolInit pool_init { init.numThreads, 0 }; - - for (int thread_idx = 0; thread_idx < init.numThreads; thread_idx++) { - // Find the proper state cache for this thread and initialize it before - // passing to context - StateCache *thread_state_cache = (StateCache *)( - (char *)init.stateCacheBase + thread_idx * sizeof(StateCache)); - new (thread_state_cache) StateCache(); - -#ifdef MADRONA_MW_MODE - void *ctx_store = (char *)init.ctxBase + (uint64_t)thread_idx * - (uint64_t)init.numCtxBytes * (uint64_t)num_worlds; - - for (int world_idx = 0; world_idx < num_worlds; world_idx++) { - void *cur_ctx = - (char *)ctx_store + world_idx * (uint64_t)init.numCtxBytes; - - void *cur_userdata = (char *)init.ctxUserdataBase + - world_idx * (uint64_t)init.numCtxUserdataBytes; - - init.ctxInitFn(cur_ctx, cur_userdata, WorkerInit { - .jobMgr = this, - .stateMgr = init.stateMgr, - .stateCache = thread_state_cache, - .workerIdx = thread_idx, - .worldID = (uint32_t)world_idx, - }); - } -#else - void *ctx_store = (char *)init.ctxBase + thread_idx * init.numCtxBytes; - init.ctxInitFn(ctx_store, init.ctxUserdataBase, WorkerInit { - .jobMgr = this, - .stateMgr = init.stateMgr, - .stateCache = thread_state_cache, - .workerIdx = thread_idx, - }); -#endif - threads_.emplace(thread_idx, [this]( - int thread_idx, - void *context_base, - uint32_t num_context_bytes, - int num_workers, - bool pin_workers, - ThreadPoolInit *pool_init) { - bool is_worker = thread_idx < num_workers; - - if (is_worker) { - disableThreadSignals(); - if (pin_workers) { - setThreadAffinity(thread_idx); - } - } - - pool_init->workerWait(); - - if (is_worker) { - workerThread(thread_idx, context_base, - num_context_bytes); - } else { - ioThread(thread_idx, context_base, - num_context_bytes); - } - }, thread_idx, ctx_store, init.numCtxBytes, init.numWorkers, - init.pinWorkers, &pool_init); - } - - pool_init.mainWait(init.numThreads); - - // Need to ensure start job has run at this point. - // Otherwise, the start function data can be freed / go out of scope - // before the job actually runs. - while (start_wrapper.remainingLaunches.load_acquire() != 0) { - workerYield(); - } -} - -JobManager::~JobManager() -{ - InitAlloc().dealloc(alloc_state_.memoryBase); - - InitAlloc().dealloc(state_ptr_); -} - -JobID JobManager::getNewJobID(int thread_idx, - uint32_t parent_job_idx, - uint32_t num_invocations) -{ - JobTrackerMap &tracker_map = getTrackerMap(tracker_base_); - JobTrackerMap::Cache &tracker_cache = - getTrackerCache(tracker_cache_base_, thread_idx); - WorkerState &worker_state = getWorkerState(worker_base_, thread_idx); - LogEntry *log = getWorkerLog(log_base_, thread_idx); - JobID new_id = tracker_map.acquireID(tracker_cache); - - JobTracker &tracker = tracker_map.getRef(new_id.id); - tracker.parent = parent_job_idx; - tracker.remainingInvocations = num_invocations; - tracker.numOutstandingJobs = 1; - - addToLog(worker_state, log, LogEntry { - .type = LogEntry::Type::JobCreated, - .created = { - .parentID = parent_job_idx, - }, - }); - - return new_id; -} - -JobID JobManager::queueJob(int thread_idx, - void (*job_func)(), - JobContainerBase *job_data, - uint32_t num_invocations, - uint32_t parent_job_idx, - JobPriority prio) -{ - JobTrackerMap &tracker_map = getTrackerMap(tracker_base_); - // num_invocations can be passed in as 0 here to signify a single - // invocation job, but for the purposes of dependency tracking it - // counts as a single invocation - JobID id = getNewJobID(thread_idx, parent_job_idx, std::max(num_invocations, 1u)); - - job_data->id = id; - - if (isRunnable(tracker_map, job_data)) { - atomic_thread_fence(sync::acquire); -#ifdef TSAN_ENABLED - { - const JobID *dependencies = getJobDependencies(job_data); - uint32_t num_dependencies = job_data->numDependencies; - for (int i = 0; i < (int)num_dependencies; i++) { - tracker_map.acquireGen(dependencies[i].id); - } - } -#endif - addToRunQueue(thread_idx, prio, - [=](Job *job_array, uint32_t cur_tail) { - job_array[cur_tail & consts::runQueueIndexMask] = Job { - .func = job_func, - .data = job_data, - .invocationOffset = 0, - .numInvocations = num_invocations, - }; - - return 1u; - }); - } else { - addToWaitQueue(thread_idx, job_func, job_data, num_invocations, - prio); - } - - return id; -} - -JobID JobManager::reserveProxyJobID(int thread_idx, uint32_t parent_job_idx) -{ - return getNewJobID(thread_idx, parent_job_idx, 1); -} - -void JobManager::markInvocationsFinished(int thread_idx, - JobContainerBase *job_data, - int32_t job_idx, - uint32_t num_invocations) -{ - WorkerState &worker_state = getWorkerState(worker_base_, thread_idx); - LogEntry *log = getWorkerLog(log_base_, thread_idx); - - addToLog(worker_state, log, LogEntry { - .type = LogEntry::Type::JobFinished, - .finished = { - .jobData = job_data, - .jobIdx = job_idx, - .numCompleted = num_invocations, - }, - }); -} - -template -void JobManager::addToRunQueue(int thread_idx, - JobPriority prio, - Fn &&add_cb) -{ - RunQueue *queue; - if (prio == JobPriority::High) { - queue = getRunQueue(high_base_, thread_idx); - } else if (prio == JobPriority::Normal) { - queue = getRunQueue(normal_base_, thread_idx); - } else { - queue = getRunQueue(io_base_, thread_idx); - } - - uint32_t num_added = addToRunQueueImpl(queue, std::forward(add_cb)); - - if (prio == JobPriority::High) { - num_high_.fetch_add_relaxed(num_added); - } - if (prio == JobPriority::IO) { - io_sema_.release(num_added); - } -} - -void JobManager::addToWaitQueue(int thread_idx, - void (*job_func)(), - JobContainerBase *job_data, - uint32_t num_invocations, - JobPriority prio) -{ - // FIXME Priority is dropped on jobs that need to wait - (void)prio; - - WorkerState &worker_state = getWorkerState(worker_base_, thread_idx); - LogEntry *log = getWorkerLog(log_base_, thread_idx); - - addToLog(worker_state, log, LogEntry { - .type = LogEntry::Type::WaitingJobQueued, - .waiting = { - job_func, - job_data, - num_invocations, - }, - }); -} - -#if 0 -JobID JobManager::queueJobs(int thread_idx, const Job *jobs, uint32_t num_jobs, - const JobID *deps, uint32_t num_dependencies, - JobPriority prio) -{ - (void)deps; - (void)num_dependencies; - - JobQueueTail *queue_tail; - if (prio == JobPriority::High) { - queue_tail = getQueueTail(getQueueHead(high_base_, thread_idx)); - } else if (prio == JobPriority::Normal) { - queue_tail = getQueueTail(getQueueHead(normal_base_, thread_idx)); - } else { - queue_tail = getQueueTail(getQueueHead(io_base_, thread_idx)); - } - - AtomicU32 &tail = queue_tail->tail; - - // No one modifies queue_tail besides this thread - uint32_t cur_tail = tail.load(sync::relaxed); - uint32_t wrapped_idx = (cur_tail & consts::jobQueueIndexMask); - - Job *job_array = getRunnableJobs(queue_tail); - - uint32_t num_remaining = consts::jobQueueSizePerThread - wrapped_idx; - uint32_t num_fit = std::min(num_remaining, num_jobs); - memcpy(job_array + wrapped_idx, jobs, num_fit * sizeof(Job)); - - if (num_remaining < num_jobs) { - uint32_t num_wrapped = num_jobs - num_remaining; - memcpy(job_array, jobs + num_remaining, num_wrapped * sizeof(Job)); - } - - cur_tail += num_jobs; - tail.store(cur_tail, sync::relaxed); - - if (prio == JobPriority::High) { - num_high_.fetch_add(num_jobs, sync::relaxed); - } - if (prio == JobPriority::IO) { - io_sema_.release(num_jobs); - } - - num_outstanding_.fetch_add(num_jobs, sync::relaxed); - - atomic_thread_fence(sync::release); - - return JobID(0); -} -#endif - -enum class JobManager::WorkerControl : uint64_t { - Run, - LoopIdle, - LoopBusy, - Sleep, - Exit, -}; - -JobManager::WorkerControl JobManager::schedule(int thread_idx, Job *run_job) -{ - JobTrackerMap &tracker_map = getTrackerMap(tracker_base_); - JobTrackerMap::Cache &tracker_cache = - getTrackerCache(tracker_cache_base_, thread_idx); - WorkerState &scheduling_worker = - getWorkerState(worker_base_, thread_idx); - scheduling_worker.numConsecutiveSchedulerCalls++; - - Job *waiting_jobs = (Job *)waiting_jobs_; - CountT cur_num_waiting = CountT(scheduler_.numWaiting); - - auto handleJobFinished = [&](const LogEntry::JobFinished &finished) { - JobTracker &tracker = - tracker_map.getRef(finished.jobIdx); - uint32_t remaining = tracker.remainingInvocations; - remaining -= finished.numCompleted; - tracker.remainingInvocations = remaining; - - if (remaining == 0) { - if (finished.jobData != nullptr) { - deallocJob(thread_idx, finished.jobData, - finished.jobData->jobSize); - } - - decrementJobTracker(tracker_map, tracker_cache, - finished.jobIdx); - } - }; - - auto handleWaitingQueued = [&](const LogEntry::WaitingJobQueued &waiting) { - waiting_jobs[cur_num_waiting++] = Job { - waiting.fnPtr, - waiting.jobData, - 0, - waiting.numInvocations, - }; - }; - - auto handleJobCreated = [&](const LogEntry::JobCreated &created) { - uint32_t parent_id = created.parentID; - if (parent_id != ~0u) { - JobTracker &parent_tracker = tracker_map.getRef(parent_id); - parent_tracker.numOutstandingJobs++; - } - }; - - // First, read all the log tails and cache them. This allows us to do - // a single acquire release barrier (dmb on arm) to ensure that log entries - // are consistent with job tails, as well as to release JobTracker - // generation updates in bulk. - - for (int64_t i = 0, n = threads_.size(); i < n; i++) { - WorkerState &worker_state = getWorkerState(worker_base_, i); - worker_state.logTailCache = worker_state.logTail.load_relaxed(); - TSAN_ACQUIRE(&worker_state.logTail); - } - - // Release half synchronizes all the releaseID calls under handleJobFinished - // to ensure that when isRunnable is called outside the scheduler, the - // job skipping the waitlist is synchronized-with the thread that finished - // the dependencies. - atomic_thread_fence(sync::acq_rel); - - // First, we read all the logs. - for (int64_t i = 0, n = threads_.size(); i != n; i++) { - int64_t offset = i + thread_idx; - int64_t worker_idx = offset < n ? offset : offset - n; - WorkerState &worker_state = getWorkerState(worker_base_, worker_idx); - LogEntry *log = getWorkerLog(log_base_, worker_idx); - - uint32_t log_tail = worker_state.logTailCache; - uint32_t log_head = worker_state.logHead.load_relaxed(); - - for (; log_head != log_tail; log_head++) { - LogEntry &entry = log[log_head & consts::logIndexMask]; - - switch (entry.type) { - case LogEntry::Type::JobFinished: { - handleJobFinished(entry.finished); - } break; - case LogEntry::Type::WaitingJobQueued: { - handleWaitingQueued(entry.waiting); - } break; - case LogEntry::Type::JobCreated: { - handleJobCreated(entry.created); - } break; - } - } - - worker_state.logHead.store(log_head, sync::relaxed); - } - - // Move all now runnable jobs to the scheduler's global run queue - - RunQueue *sched_run = getRunQueue(normal_base_, thread_idx); - - Job *sched_run_jobs = getRunnableJobs(sched_run); - uint32_t cur_run_tail = sched_run->tail.load_relaxed(); - int64_t num_new_invocations = 0; - int64_t compaction_offset = 0; - - bool first_found_job = true; - for (int64_t i = 0; i < cur_num_waiting; i++) { - Job &job = waiting_jobs[i]; - if (isRunnable(tracker_map, job.data)) { - uint32_t num_invocations = job.numInvocations; - - // num_invocations == 0 is a special case that indicates a one-off - // submission as opposed to a parallel for / multi invocation - // submission. For the scheduler's purpose, this counts as one - // invocation regardless - num_new_invocations += num_invocations > 0 ? num_invocations : 1; - - if (first_found_job) { - *run_job = job; - first_found_job = false; - } else { - sched_run_jobs[cur_run_tail & consts::runQueueIndexMask] = job; - cur_run_tail++; - } - } else { - int64_t cur_compaction_offset = compaction_offset++; - if (i != cur_compaction_offset) { - waiting_jobs[cur_compaction_offset] = job; - } - } - } - scheduler_.numWaiting = compaction_offset; - - if (num_new_invocations == 0) { - uint32_t sched_run_auth = sched_run->auth.load_relaxed(); - - if (sched_run_auth == cur_run_tail) { - if (scheduling_worker.numConsecutiveSchedulerCalls > 1) { - if (scheduling_worker.wakeUp.load_relaxed() != 0) { - scheduler_.numSleepingWorkers++; - } - - if (scheduler_.numWaiting == 0 && - scheduler_.numSleepingWorkers == num_compute_workers_) { - for (int64_t i = 0; i < num_compute_workers_; i++) { - WorkerState &worker_state = - getWorkerState(worker_base_, i); - worker_state.wakeUp.store(~0_u32, - sync::relaxed); - worker_state.wakeUp.notify_one(); - } - return WorkerControl::Exit; - } else { - getWorkerState(worker_base_, thread_idx) - .wakeUp.store(0, sync::relaxed); - return WorkerControl::Sleep; - } - } else { - return WorkerControl::LoopIdle; - } - } else { - return WorkerControl::LoopBusy; - } - } - - sched_run->tail.store_release(cur_run_tail); - - // Wake up compute workers based on # of jobs - int64_t num_compute_workers = num_compute_workers_; - int64_t num_wakeup = std::min(num_compute_workers, num_new_invocations); - - for (int64_t i = 0; num_wakeup > 0 && i < num_compute_workers; i++) { - WorkerState &worker_state = getWorkerState(worker_base_, i); - - if (worker_state.wakeUp.load_relaxed() == 0) { - worker_state.wakeUp.store_relaxed((uint32_t)thread_idx + 1); - worker_state.wakeUp.notify_one(); - - num_wakeup--; - scheduler_.numSleepingWorkers--; - } - } - - return WorkerControl::Run; -} - -uint32_t JobManager::dequeueJobIndex(RunQueue *job_queue) -{ - AtomicU32 &head = job_queue->head; - AtomicU32 &correction = job_queue->correction; - AtomicU32 &auth = job_queue->auth; - AtomicU32 &tail = job_queue->tail; - - uint32_t cur_tail = tail.load_relaxed(); - uint32_t cur_correction = correction.load_relaxed(); - uint32_t cur_head = head.load_relaxed(); - - if (isQueueEmpty(cur_head, cur_correction, cur_tail)) { - return consts::jobQueueSentinel; - } - - atomic_thread_fence(sync::acquire); - TSAN_ACQUIRE(&tail); - TSAN_ACQUIRE(&correction); - TSAN_ACQUIRE(&head); - - cur_head = head.fetch_add_relaxed(1); - cur_tail = tail.load_acquire(); - - if (isQueueEmpty(cur_head, cur_correction, cur_tail)) [[unlikely]] { - correction.fetch_add_release(1); - return consts::jobQueueSentinel; - } - - // Note, there is some non intuitive behavior here, where the value of idx - // can seem to be past cur_tail above. This isn't a case where too many - // items have been dequeued, instead, the producer has added another item - // to the queue and another consumer thread has come in and dequeued - // the item this thread was planning on dequeuing, so this thread picks - // up the later item. If tail is re-read after the fetch add below, - // everything would appear consistent. - return auth.fetch_add_acq_rel(1); -} - -JobManager::WorkerControl JobManager::tryScheduling( - JobManager::WorkerControl default_ctrl, int thread_idx, Job *next_job) { - if (scheduler_.lock.tryLock()) { - default_ctrl = schedule(thread_idx, next_job); - scheduler_.lock.unlock(); - } - return default_ctrl; -} - -JobManager::WorkerControl JobManager::getNextJob(void *const queue_base, - int thread_idx, - int init_search_idx, - bool run_scheduler, - Job *next_job) -{ - WorkerControl sched_ctrl = WorkerControl::LoopIdle; - - WorkerState &worker_state = getWorkerState(worker_base_, thread_idx); - uint32_t cur_tail = worker_state.logTail.load_relaxed(); - uint32_t log_head = worker_state.logHead.load_relaxed(); - // Determine if log capacity is too high (and we should try scheduling). - if (cur_tail - log_head > consts::logSizeMaxSafeCapacity) { - return tryScheduling(WorkerControl::LoopBusy, thread_idx, next_job); - } - - // First, check the current thread's queue - RunQueue *queue = getRunQueue(queue_base, init_search_idx); - uint32_t job_idx = dequeueJobIndex(queue); - - if (run_scheduler && job_idx == consts::jobQueueSentinel) { - sched_ctrl = - tryScheduling(WorkerControl::LoopIdle, thread_idx, next_job); - if (sched_ctrl != WorkerControl::LoopIdle) { - return sched_ctrl; - } - } - - // Try work stealing - if (job_idx == consts::jobQueueSentinel) { - int64_t num_queues = threads_.size(); - for (int64_t i = 1; i < num_queues; i++) { - int64_t unwrapped_idx = i + thread_idx; - int64_t queue_idx = unwrapped_idx < num_queues ? - unwrapped_idx : unwrapped_idx - num_queues; - - queue = getRunQueue(queue_base, queue_idx); - - job_idx = dequeueJobIndex(queue); - if (job_idx != consts::jobQueueSentinel) { - break; - } - } - } - - if (job_idx == consts::jobQueueSentinel) { - return WorkerControl::LoopIdle; - } - - *next_job = getRunnableJobs(queue)[job_idx & consts::runQueueIndexMask]; - - // There's no protection to prevent queueJob overwriting next_job - // in between job_idx being assigned and the job actually being - // read. If this happens it is a bug where way too many jobs are - // being created, or jobs are being processed too slowly, so we - // detect and crash with a fatal error (rather than silently - // dropping or reading corrupted jobs). - - uint32_t post_read_tail = queue->tail.load_acquire(); - - if (post_read_tail - job_idx > consts::runQueueSizePerThread) [[unlikely]] { - // Note, this is not ideal because it doesn't detect the source - // of the issue. The tradeoff is that we skip needing to read - // the head information when queueing jobs, whereas this - // code already has to read the tail once before. - FATAL("Job queue has overwritten readers. Detected by thread %d.\n" - "Job: %u, Tail: %u, Difference: %u, Queue: %p\n", - thread_idx, job_idx, post_read_tail, post_read_tail - job_idx, - queue); - } - - return WorkerControl::Run; -} - -void JobManager::splitJob(MultiInvokeFn fn_ptr, JobContainerBase *job_data, - uint32_t invocation_offset, uint32_t num_invocations, - RunQueue *run_queue) -{ - void (*generic_fn)() = (void (*)())fn_ptr; - if (num_invocations == 1) { - addToRunQueueImpl(run_queue, - [=](Job *job_array, uint32_t cur_tail) { - job_array[cur_tail & consts::runQueueIndexMask] = Job { - .func = generic_fn, - .data = job_data, - .invocationOffset = invocation_offset, - .numInvocations = 1, - }; - - return 1u; - }); - } else { - uint32_t b_num_invocations = num_invocations / 2; - uint32_t a_num_invocations = - num_invocations - b_num_invocations; - - uint32_t a_offset = invocation_offset; - uint32_t b_offset = a_offset + a_num_invocations; - - // FIXME, again priority issues here - addToRunQueueImpl(run_queue, - [=](Job *job_array, uint32_t cur_tail) { - uint32_t first_idx = - cur_tail & consts::runQueueIndexMask; - - uint32_t second_idx = - (cur_tail + 1) & consts::runQueueIndexMask; - - job_array[first_idx] = Job { - .func = generic_fn, - .data = job_data, - .invocationOffset = a_offset, - .numInvocations = a_num_invocations, - }; - - job_array[second_idx] = Job { - .func = generic_fn, - .data = job_data, - .invocationOffset = b_offset, - .numInvocations = b_num_invocations, - }; - - return 2u; - }); - } -} - -void JobManager::runJob(const int thread_idx, - Context *ctx, - void (*generic_fn)(), - JobContainerBase *job_data, - uint32_t invocation_offset, - uint32_t num_invocations) -{ - ctx->cur_job_id_ = job_data->id; - - if (num_invocations == 0) { - auto fn = (SingleInvokeFn)generic_fn; - fn(ctx, job_data); - return; - } else { - // FIXME, figure out relationship between different queue priorities - // Should the normal priority queue always be the work indicator here? - RunQueue *check_queue = getRunQueue(normal_base_, thread_idx); - - auto fn = (MultiInvokeFn)generic_fn; - fn(ctx, job_data, invocation_offset, num_invocations, check_queue); - } -} - -void JobManager::workerThread( - const int thread_idx, - void *context_base, - uint32_t num_context_bytes) -{ -#ifndef MADRONA_MW_MODE - (void)num_context_bytes; - Context *ctx = (Context *)context_base; -#endif - - Job cur_job; - - WorkerState &worker_state = - getWorkerState(worker_base_, thread_idx); - - auto runCurJob = [&]() MADRONA_ALWAYS_INLINE { - worker_state.numConsecutiveSchedulerCalls = 0; - -#ifdef MADRONA_MW_MODE - Context *ctx = (Context *)((char *)context_base + - (uint64_t)cur_job.data->worldID * (uint64_t)num_context_bytes); -#endif - - runJob(thread_idx, ctx, cur_job.func, cur_job.data, - cur_job.invocationOffset, cur_job.numInvocations); - }; - - while (true) { - WorkerControl worker_ctrl = WorkerControl::LoopIdle; - if (num_high_.load_relaxed() > 0) { - worker_ctrl = getNextJob(high_base_, thread_idx, thread_idx, - false, &cur_job); - - if (worker_ctrl == WorkerControl::Run) { - num_high_.fetch_sub_relaxed(1); - } - } - - if (worker_ctrl != WorkerControl::Run) [[likely]] { - worker_ctrl = getNextJob(normal_base_, thread_idx, thread_idx, - true, &cur_job); - } - - if (worker_ctrl == WorkerControl::Run) { - runCurJob(); - } else if (worker_ctrl == WorkerControl::LoopIdle) { - // No available work and couldn't run scheduler - workerPause(); - worker_state.numIdleLoops++; - } else if (worker_ctrl == WorkerControl::LoopBusy) { - continue; - } else if (worker_ctrl == WorkerControl::Sleep) [[unlikely]] { - worker_state.wakeUp.wait(0); - uint32_t wakeup_idx = worker_state.wakeUp.load_relaxed(); - if (wakeup_idx == ~0_u32) [[unlikely]] { - break; - } - - int wakeup_search_idx = (int)wakeup_idx - 1; - - worker_ctrl = getNextJob(normal_base_, thread_idx, - wakeup_search_idx, false, &cur_job); - - if (worker_ctrl == WorkerControl::Run) { - runCurJob(); - } - } else if (worker_ctrl == WorkerControl::Exit) [[unlikely]] { - break; - } - } -} - -void JobManager::ioThread( - const int thread_idx, - void *context_base, - uint32_t num_context_bytes) -{ -#ifndef MADRONA_MW_MODE - (void)num_context_bytes; - Context *ctx = (Context *)context_base; -#endif - - Job cur_job; - - while (true) { - WorkerControl worker_ctrl = getNextJob(io_base_, thread_idx, - thread_idx, false, &cur_job); - - if (worker_ctrl != WorkerControl::Run) { - io_sema_.acquire(); - } - -#ifdef MADRONA_MW_MODE - Context *ctx = (Context *)((char *)context_base + - (uint64_t)cur_job.data->worldID * (uint64_t)num_context_bytes); -#endif - - runJob(thread_idx, ctx, cur_job.func, cur_job.data, - cur_job.invocationOffset, cur_job.numInvocations); - } -} - -void JobManager::waitForAllFinished() -{ - for (int i = 0, n = threads_.size(); i < n; i++) { - threads_[i].join(); - } -} - -} diff --git a/src/importer/importer.cpp b/src/importer/importer.cpp index dc3d7186..f039bf47 100644 --- a/src/importer/importer.cpp +++ b/src/importer/importer.cpp @@ -12,14 +12,6 @@ #include "obj.hpp" -#ifdef MADRONA_GLTF_SUPPORT -#include "gltf.hpp" -#endif - -#ifdef MADRONA_USD_SUPPORT -#include "usd.hpp" -#endif - #ifdef MADRONA_CUDA_SUPPORT #include #endif @@ -33,14 +25,6 @@ struct AssetImporter::Impl { Optional objLoader; -#ifdef MADRONA_GLTF_SUPPORT - Optional gltfLoader; -#endif - -#ifdef MADRONA_USD_SUPPORT - Optional usdLoader; -#endif - static inline Impl * make(ImageImporter &&img_importer); inline Optional importFromDisk( @@ -53,12 +37,6 @@ AssetImporter::Impl * AssetImporter::Impl::make(ImageImporter &&img_importer) return new Impl { .imgImporter = std::move(img_importer), .objLoader = Optional::none(), -#ifdef MADRONA_GLTF_SUPPORT - .gltfLoader = Optional::none(), -#endif -#ifdef MADRONA_USD_SUPPORT - .usdLoader = Optional::none(), -#endif }; } @@ -103,33 +81,6 @@ Optional AssetImporter::Impl::importFromDisk( } load_success = objLoader->load(path, imported); - } else if (extension == "gltf" || extension == "glb") { -#ifdef MADRONA_GLTF_SUPPORT - if (!gltfLoader.has_value()) { - gltfLoader.emplace(imgImporter, err_buf); - } - - load_success = gltfLoader->load( - path, imported, one_object_per_asset, imgImporter); -#else - load_success = false; - snprintf(err_buf.data(), err_buf.size(), "Madrona not compiled with glTF support"); -#endif - } else if (extension == "usd" || - extension == "usda" || - extension == "usdc" || - extension == "usdz") { -#ifdef MADRONA_USD_SUPPORT - if (!usdLoader.has_value()) { - usdLoader.emplace(imgImporter, err_buf); - } - - load_success = usdLoader->load( - path, imported, one_object_per_asset, imgImporter); -#else - load_success = false; - snprintf(err_buf.data(), err_buf.size(), "Madrona not compiled with USD support"); -#endif } if (!load_success) { diff --git a/src/mw/CMakeLists.txt b/src/mw/CMakeLists.txt index 2eaaed41..7a430814 100644 --- a/src/mw/CMakeLists.txt +++ b/src/mw/CMakeLists.txt @@ -1,12 +1,3 @@ -add_library(madrona_mw_cpu STATIC - ${MADRONA_INC_DIR}/mw_cpu.hpp ${MADRONA_INC_DIR}/mw_cpu.inl cpu_exec.cpp -) - -target_link_libraries(madrona_mw_cpu - PUBLIC - madrona_mw_core -) - if (NOT CUDAToolkit_FOUND OR CUDAToolkit_VERSION_MAJOR LESS 12 OR (CUDAToolkit_VERSION_MAJOR EQUAL 12 AND CUDAToolkit_VERSION_MINOR LESS 8)) return() @@ -44,6 +35,11 @@ set(MADRONA_MW_GPU_COMPILE_FLAGS -DTHRUST_DEVICE_SYSTEM=THRUST_DEVICE_SYSTEM_CPP -DCCCL_DISABLE_BF16_SUPPORT=1 -DCUB_DISABLE_BF16_SUPPORT=1 + # CUB 2.8.2 (CUDA 12.8) marks its own Traits::CATEGORY deprecated yet still + # uses it internally in WarpScan/BlockScan; its self-suppression is a clang + # pragma that nvrtc/cicc ignores and cannot catch at our instantiation site. + # CCCL_IGNORE_DEPRECATED_API is CCCL's documented opt-out (drops the attribute). + -DCCCL_IGNORE_DEPRECATED_API=1 ) if (SKBUILD) @@ -54,13 +50,6 @@ endif() set(DEVICE_SRC_DIR "${MADRONA_SRC_DIR}/mw/device") -set(JOB_SYS_DEVICE_SOURCES - ${DEVICE_SRC_DIR}/job.cpp - ${DEVICE_SRC_DIR}/state.cpp - ${DEVICE_SRC_DIR}/crash.cpp - ${DEVICE_SRC_DIR}/const.cpp -) - set(BVH_DEVICE_SOURCES ${DEVICE_SRC_DIR}/bvh.cpp ${DEVICE_SRC_DIR}/memory.cpp @@ -80,23 +69,11 @@ set(TASK_GRAPH_DEVICE_SOURCES ${DEVICE_SRC_DIR}/ecs_render_bridge.cpp # FIXME ${MADRONA_SRC_DIR}/common/hashmap.cpp - ${MADRONA_SRC_DIR}/common/navmesh.cpp ${MADRONA_SRC_DIR}/core/base.cpp - ${MADRONA_SRC_DIR}/physics/physics.cpp ${MADRONA_SRC_DIR}/physics/geo.cpp - ${MADRONA_SRC_DIR}/physics/xpbd.cpp - ${MADRONA_SRC_DIR}/physics/tgs.cpp - ${MADRONA_SRC_DIR}/physics/narrowphase.cpp - ${MADRONA_SRC_DIR}/physics/broadphase.cpp ${MADRONA_SRC_DIR}/render/ecs_system.cpp ) -set(JOB_SYS_INTERNAL_SRC_LIST) -foreach(f ${JOB_SYS_DEVICE_SOURCES}) - set(JOB_SYS_INTERNAL_SRC_LIST "${JOB_SYS_INTERNAL_SRC_LIST}\"${f}\", ") -endforeach() -set(JOB_SYS_INTERNAL_SRC_LIST "${JOB_SYS_INTERNAL_SRC_LIST}") - set(TASK_GRAPH_INTERNAL_SRC_LIST) foreach(f ${TASK_GRAPH_DEVICE_SOURCES}) set(TASK_GRAPH_INTERNAL_SRC_LIST "${TASK_GRAPH_INTERNAL_SRC_LIST}\"${f}\", ") @@ -142,7 +119,6 @@ target_compile_definitions(madrona_mw_gpu PRIVATE MADRONA_NVRTC_INCLUDE_DIRS=${NVRTC_INCLUDE_DIRS_LIST} MADRONA_NVRTC_OPTIONS=${NVRTC_OPTIONS} - MADRONA_MW_GPU_JOB_SYS_INTERNAL_CPP=${JOB_SYS_INTERNAL_SRC_LIST} MADRONA_MW_GPU_TASK_GRAPH_INTERNAL_CPP=${TASK_GRAPH_INTERNAL_SRC_LIST} MADRONA_MW_GPU_BVH_INTERNAL_CPP=${BVH_INTERNAL_SRC_LIST} MADRONA_MW_GPU_DEVICE_SRC_DIR=\"${DEVICE_SRC_DIR}\" diff --git a/src/mw/cpu_exec.cpp b/src/mw/cpu_exec.cpp deleted file mode 100644 index 68b1b457..00000000 --- a/src/mw/cpu_exec.cpp +++ /dev/null @@ -1,239 +0,0 @@ -#include -#include "../core/worker_init.hpp" - -#if defined(MADRONA_LINUX) or defined(MADRONA_MACOS) -#include -#elif defined(MADRONA_WINDOWS) -#include -#endif - -namespace madrona { - -struct ThreadPoolExecutor::Impl { - HeapArray workers; - alignas(MADRONA_CACHE_LINE) AtomicI32 workerWakeup; - alignas(MADRONA_CACHE_LINE) AtomicI32 mainWakeup; - ThreadPoolExecutor::Job *currentJobs; - uint32_t numJobs; - alignas(MADRONA_CACHE_LINE) AtomicU32 nextJob; - alignas(MADRONA_CACHE_LINE) AtomicU32 numFinished; - StateManager stateMgr; - HeapArray stateCaches; - HeapArray exportPtrs; - - static Impl * make(const ThreadPoolExecutor::Config &cfg); - ~Impl(); - void run(Job *jobs, CountT num_jobs); - void workerThread(CountT worker_id); -}; - -static CountT getNumCores() -{ -#if defined(MADRONA_MACOS) - int os_num_threads = sysconf(_SC_NPROCESSORS_ONLN); - - if (os_num_threads <= 0) { - FATAL("Failed to get number of concurrent threads"); - } - - return os_num_threads; -#elif defined(MADRONA_LINUX) - cpu_set_t cpuset; - pthread_getaffinity_np(pthread_self(), sizeof(cpuset), &cpuset); - CountT num_active_threads = CPU_COUNT(&cpuset); - if (num_active_threads <= 0) { - FATAL("Failed to get number of concurrent threads"); - } - - return num_active_threads; -#elif defined(MADRONA_WINDOWS) - SYSTEM_INFO sys_info; - GetSystemInfo(&sys_info); - return sys_info.dwNumberOfProcessors; -#else - STATIC_UNIMPLEMENTED(); -#endif -} - -static inline void pinThread([[maybe_unused]] CountT worker_id) -{ -#ifdef MADRONA_LINUX - cpu_set_t cpu_set; - pthread_getaffinity_np(pthread_self(), sizeof(cpu_set), &cpu_set); - - const int max_threads = CPU_COUNT(&cpu_set); - - if (worker_id > max_threads) [[unlikely]] { - FATAL("Tried setting thread affinity to %d when %d is max", - worker_id, max_threads); - } - - cpu_set_t worker_set; - CPU_ZERO(&worker_set); - - // This is needed in case there was already cpu masking via - // a different call to setaffinity or via cgroup (SLURM) - for (CountT thread_idx = 0, available_threads = 0; - thread_idx < (CountT)CPU_SETSIZE; thread_idx++) { - if (CPU_ISSET(thread_idx, &cpu_set)) { - if ((available_threads++) == worker_id) { - CPU_SET(thread_idx, &worker_set); - - break; - } - } - } - - int res = pthread_setaffinity_np(pthread_self(), - sizeof(worker_set), - &worker_set); - - if (res != 0) { - FATAL("Failed to set thread affinity to %d", worker_id); - } -#endif -} - -ThreadPoolExecutor::Impl * ThreadPoolExecutor::Impl::make( - const ThreadPoolExecutor::Config &cfg) -{ - Impl *impl = new Impl { - .workers = HeapArray( - cfg.numWorkers == 0 ? getNumCores() : cfg.numWorkers), - .workerWakeup = 0, - .mainWakeup = 0, - .currentJobs = nullptr, - .numJobs = 0, - .nextJob = 0, - .numFinished = 0, - .stateMgr = StateManager(cfg.numWorlds), - .stateCaches = HeapArray(cfg.numWorlds), - .exportPtrs = HeapArray(cfg.numExportedBuffers), - }; - - for (CountT i = 0; i < (CountT)cfg.numWorlds; i++) { - impl->stateCaches.emplace(i); - } - - for (CountT i = 0; i < impl->workers.size(); i++) { - impl->workers.emplace(i, [](Impl *impl, CountT i) { - impl->workerThread(i); - }, impl, i); - } - - return impl; -} - -ThreadPoolExecutor::ThreadPoolExecutor(const Config &cfg) - : impl_(Impl::make(cfg)) -{} - -ThreadPoolExecutor::ThreadPoolExecutor(ThreadPoolExecutor &&o) = default; - -ThreadPoolExecutor::Impl::~Impl() -{ - workerWakeup.store_release(-1); - workerWakeup.notify_all(); - - for (CountT i = 0; i < workers.size(); i++) { - workers[i].join(); - } -} - -ThreadPoolExecutor::~ThreadPoolExecutor() = default; - -void ThreadPoolExecutor::Impl::run(Job *jobs, CountT num_jobs) -{ - stateMgr.copyInExportedColumns(); - - currentJobs = jobs; - numJobs = uint32_t(num_jobs); - nextJob.store_relaxed(0); - numFinished.store_relaxed(0); - workerWakeup.store_release(1); - workerWakeup.notify_all(); - - mainWakeup.wait(0); - mainWakeup.store_relaxed(0); - - stateMgr.copyOutExportedColumns(); -} - -void ThreadPoolExecutor::run(Job *jobs, CountT num_jobs) -{ - impl_->run(jobs, num_jobs); -} - -void * ThreadPoolExecutor::getExported(CountT slot) const -{ - return impl_->exportPtrs[slot]; -} - -void ThreadPoolExecutor::initializeContexts( - Context & (*init_fn)(void *, const WorkerInit &, CountT), - void *init_data, CountT num_worlds) -{ - for (CountT world_idx = 0; world_idx < num_worlds; world_idx++) { - WorkerInit worker_init { - &impl_->stateMgr, - &impl_->stateCaches[world_idx], - uint32_t(world_idx), - }; - - init_fn(init_data, worker_init, world_idx); - } -} - -ECSRegistry ThreadPoolExecutor::getECSRegistry() -{ - return ECSRegistry(&impl_->stateMgr, impl_->exportPtrs.data()); -} - -void ThreadPoolExecutor::initExport() -{ - impl_->stateMgr.copyOutExportedColumns(); -} - -void ThreadPoolExecutor::Impl::workerThread(CountT worker_id) -{ - pinThread(worker_id); - - while (true) { - workerWakeup.wait(0); - int32_t ctrl = workerWakeup.load_acquire(); - - if (ctrl == 0) { - continue; - } else if (ctrl == -1) { - break; - } - - while (true) { - uint32_t job_idx = nextJob.fetch_add_relaxed(1); - - if (job_idx == numJobs) { - workerWakeup.store_relaxed(0); - } - - assert(job_idx < 0xFFFF'FFFF); - - if (job_idx >= numJobs) { - break; - } - - currentJobs[job_idx].fn(currentJobs[job_idx].data); - - // This has to be acq_rel so the finishing thread has seen - // all the other threads' effects - uint32_t prev_finished = - numFinished.fetch_add_acq_rel(1); - - if (prev_finished == numJobs - 1) { - mainWakeup.store_release(1); - mainWakeup.notify_one(); - } - } - } -} - -} diff --git a/src/mw/cuda_exec.cpp b/src/mw/cuda_exec.cpp index abfb2424..256f6fdc 100644 --- a/src/mw/cuda_exec.cpp +++ b/src/mw/cuda_exec.cpp @@ -323,11 +323,6 @@ static constexpr uint32_t numEntryQueueThreads = 512; using GPUImplConsts = mwGPU::madrona::mwGPU::GPUImplConsts; -enum class ExecutorMode { - JobSystem, - TaskGraph, -}; - struct GPUCompileResults { CUmodule mod; std::string initECSName; @@ -405,8 +400,6 @@ struct GPUKernels { CUfunction initECS; CUfunction initWorlds; CUfunction initTasks; - CUfunction queueUserInit; - CUfunction queueUserRun; CUfunction initBVHParams; CUfunction destroyECS; }; @@ -480,94 +473,6 @@ struct MWCudaExecutor::Impl { std::vector timingGroups; }; -static void getUserEntries(const char *entry_class, CUmodule mod, - const char **compile_flags, - uint32_t num_compile_flags, - CUfunction *init_out, CUfunction *run_out) -{ - static const char mangle_code_postfix[] = R"__( -#include - -namespace madrona { namespace mwGPU { - -template __global__ void submitInit(uint32_t, void *) {} -template __global__ void submitRun(uint32_t) {} - -} } -)__"; - - static const char init_template[] = - "::madrona::mwGPU::submitInit<::"; - static const char run_template[] = - "::madrona::mwGPU::submitRun<::"; - - std::string_view entry_view(entry_class); - - // If user prefixed with ::, trim off as it will be added later - if (entry_view[0] == ':' && entry_view[1] == ':') { - entry_view = entry_view.substr(2); - } - - std::string fwd_declare; - - // Find all namespace separators - int num_namespaces = 0; - size_t prev_off = 0, off = 0; - while ((off = entry_view.find("::", prev_off)) != std::string_view::npos) { - auto ns_view = entry_view.substr(prev_off, off - prev_off); - - fwd_declare += "namespace "; - fwd_declare += ns_view; - fwd_declare += " { "; - - prev_off = off + 2; - num_namespaces++; - } - - auto class_view = entry_view.substr(prev_off); - if (class_view.size() == 0) { - FATAL("Invalid entry class name\n"); - } - - fwd_declare += "class "; - fwd_declare += class_view; - fwd_declare += "; "; - - for (int i = 0; i < num_namespaces; i++) { - fwd_declare += "} "; - } - - std::string mangle_code = std::move(fwd_declare); - mangle_code += mangle_code_postfix; - - std::string init_name = init_template; - init_name += entry_view; - init_name += ">"; - - std::string run_name = run_template; - run_name += entry_view; - run_name += ">"; - - nvrtcProgram prog; - REQ_NVRTC(CudaDynamicLoader::nvrtcCreateProgram(&prog, mangle_code.c_str(), "mangle.cpp", - 0, nullptr, nullptr)); - - REQ_NVRTC(CudaDynamicLoader::nvrtcAddNameExpression(prog, init_name.c_str())); - REQ_NVRTC(CudaDynamicLoader::nvrtcAddNameExpression(prog, run_name.c_str())); - - REQ_NVRTC(CudaDynamicLoader::nvrtcCompileProgram(prog, num_compile_flags, compile_flags)); - - const char *init_lowered; - REQ_NVRTC(CudaDynamicLoader::nvrtcGetLoweredName(prog, init_name.c_str(), &init_lowered)); - const char *run_lowered; - REQ_NVRTC(CudaDynamicLoader::nvrtcGetLoweredName(prog, run_name.c_str(), &run_lowered)); - - REQ_CU(CudaDynamicLoader::cuModuleGetFunction(init_out, mod, init_lowered)); - REQ_CU(CudaDynamicLoader::cuModuleGetFunction(run_out, mod, run_lowered)); - - REQ_NVRTC(CudaDynamicLoader::nvrtcDestroyProgram(&prog)); -} - static MegakernelCache loadMegakernelCache(const std::string &cache_path) { std::ifstream cache_file(cache_path, @@ -676,7 +581,7 @@ static GPUCompileResults compileCode( const MegakernelConfig *megakernel_cfgs, int64_t num_megakernel_cfgs, CompileConfig::OptMode opt_mode, - ExecutorMode exec_mode, bool verbose_compile) + bool verbose_compile) { const std::string cache_path = getenv("MADRONA_MWGPU_KERNEL_CACHE"); @@ -761,30 +666,10 @@ static GPUCompileResults compileCode( (char *)cubin.data(), cubin.size(), name)); }; - std::string megakernel_job_prefix = R"__(#include "megakernel_job_impl.inl" - -extern "C" { - -)__"; - std::string megakernel_taskgraph_prefix = R"__(#include "megakernel_impl.inl" extern "C" { -)__"; - - std::string megakernel_job_body = R"__(namespace madrona { -namespace mwGPU { - -static __attribute__((always_inline)) inline void dispatch( - uint32_t func_id, - madrona::JobContainerBase *data, - uint32_t *data_indices, - uint32_t *invocation_offsets, - uint32_t num_launches, - uint32_t grid) -{ - switch (func_id) { )__"; std::string megakernel_taskgraph_body = R"__(namespace madrona { @@ -798,31 +683,14 @@ static __attribute__((always_inline)) inline void dispatch( switch (func_id) { )__"; - std::string megakernel_prefix; - std::string megakernel_body; + std::string megakernel_prefix = megakernel_taskgraph_prefix; + std::string megakernel_body = megakernel_taskgraph_body; std::string megakernel_func_ids; - std::string_view entry_prefix; - std::string_view entry_postfix; - std::string_view entry_params; - std::string_view entry_args; - std::string_view id_prefix; - if (exec_mode == ExecutorMode::JobSystem) { - megakernel_prefix = megakernel_job_prefix; - megakernel_body = megakernel_job_body; - entry_prefix = ".weak .func _ZN7madrona5mwGPU8jobEntry"; - entry_postfix = "EvPNS_16JobContainerBaseEPj"; - entry_params = "(madrona::JobContainerBase *, uint32_t *, uint32_t *, uint32_t, uint32_t);\n"; - entry_args = "(data, data_indices, invocation_offsets, num_launches, grid);\n"; - id_prefix = "_ZN7madrona5mwGPU13JobFuncIDBase"; - } else if (exec_mode == ExecutorMode::TaskGraph) { - megakernel_prefix = megakernel_taskgraph_prefix; - megakernel_body = megakernel_taskgraph_body; - entry_prefix = ".weak .func _ZN7madrona5mwGPU9userEntry"; - entry_postfix = "EvPNS_8NodeBaseEi"; - entry_params = "(madrona::NodeBase *, int32_t);\n"; - entry_args = "(node_data, invocation_offset);\n"; - id_prefix = "_ZN7madrona5mwGPU14UserFuncIDBase"; - } + std::string_view entry_prefix = ".weak .func _ZN7madrona5mwGPU9userEntry"; + std::string_view entry_postfix = "EvPNS_8NodeBaseEi"; + std::string_view entry_params = "(madrona::NodeBase *, int32_t);\n"; + std::string_view entry_args = "(node_data, invocation_offset);\n"; + std::string_view id_prefix = "_ZN7madrona5mwGPU14UserFuncIDBase"; std::string init_ecs_name; std::string init_worlds_name; @@ -1382,7 +1250,6 @@ static BVHKernels buildBVHKernels(const CompileConfig &cfg, static GPUKernels buildKernels(const CompileConfig &cfg, Span megakernel_cfgs, - ExecutorMode exec_mode, int32_t num_sms, std::pair cuda_arch) { @@ -1396,23 +1263,12 @@ static GPUKernels buildKernels(const CompileConfig &cfg, using namespace std; - array job_sys_cpp_files { - MADRONA_MW_GPU_JOB_SYS_INTERNAL_CPP - }; - array task_graph_cpp_files { MADRONA_MW_GPU_TASK_GRAPH_INTERNAL_CPP }; - uint32_t num_exec_srcs = 0; - const char **exec_srcs = nullptr; - if (exec_mode == ExecutorMode::JobSystem) { - num_exec_srcs = job_sys_cpp_files.size(); - exec_srcs = job_sys_cpp_files.data(); - } else if (exec_mode == ExecutorMode::TaskGraph) { - num_exec_srcs = task_graph_cpp_files.size(); - exec_srcs = task_graph_cpp_files.data(); - } + uint32_t num_exec_srcs = task_graph_cpp_files.size(); + const char **exec_srcs = task_graph_cpp_files.data(); size_t num_srcs = num_exec_srcs + cfg.userSources.size(); HeapArray all_cpp_files(num_srcs); @@ -1489,11 +1345,7 @@ static GPUKernels buildKernels(const CompileConfig &cfg, compile_flags.push_back("-DMADRONA_MWGPU_LTO_MODE=1"); } - if (exec_mode == ExecutorMode::JobSystem) { - compile_flags.push_back("-DMARONA_MWGPU_JOB_SYSTEM=1"); - } else if (exec_mode == ExecutorMode::TaskGraph) { - compile_flags.push_back("-DMADRONA_MWGPU_TASKGRAPH=1"); - } + compile_flags.push_back("-DMADRONA_MWGPU_TASKGRAPH=1"); DynArray linker_flags { gpu_arch_flag.c_str(), @@ -1538,7 +1390,7 @@ static GPUKernels buildKernels(const CompileConfig &cfg, fast_compile_flags.data(), fast_compile_flags.size(), linker_flags.data(), linker_flags.size(), megakernel_cfgs.data(), megakernel_cfgs.size(), - opt_mode, exec_mode, verbose_compile); + opt_mode, verbose_compile); HeapArray megakernel_fns(megakernel_cfgs.size()); for (int64_t i = 0; i < megakernel_cfgs.size(); i++) { @@ -1557,8 +1409,6 @@ static GPUKernels buildKernels(const CompileConfig &cfg, .initECS = nullptr, .initWorlds = nullptr, .initTasks = nullptr, - .queueUserInit = nullptr, - .queueUserRun = nullptr, .initBVHParams = nullptr, .destroyECS = nullptr, }; @@ -1566,21 +1416,12 @@ static GPUKernels buildKernels(const CompileConfig &cfg, REQ_CU(CudaDynamicLoader::cuModuleGetFunction(&gpu_kernels.computeGPUImplConsts, gpu_kernels.mod, "madronaMWGPUComputeConstants")); - if (exec_mode == ExecutorMode::JobSystem) { - REQ_CU(CudaDynamicLoader::cuModuleGetFunction(&gpu_kernels.initECS, gpu_kernels.mod, - "madronaMWGPUInitialize")); - // FIXME: getUserEntries is broken - getUserEntries("", gpu_kernels.mod, compile_flags.data(), - compile_flags.size(), &gpu_kernels.queueUserInit, - &gpu_kernels.queueUserRun); - } else if (exec_mode == ExecutorMode::TaskGraph) { - REQ_CU(CudaDynamicLoader::cuModuleGetFunction(&gpu_kernels.initECS, gpu_kernels.mod, - compile_results.initECSName.c_str())); - REQ_CU(CudaDynamicLoader::cuModuleGetFunction(&gpu_kernels.initWorlds, gpu_kernels.mod, - compile_results.initWorldsName.c_str())); - REQ_CU(CudaDynamicLoader::cuModuleGetFunction(&gpu_kernels.initTasks, gpu_kernels.mod, - compile_results.initTasksName.c_str())); - } + REQ_CU(CudaDynamicLoader::cuModuleGetFunction(&gpu_kernels.initECS, gpu_kernels.mod, + compile_results.initECSName.c_str())); + REQ_CU(CudaDynamicLoader::cuModuleGetFunction(&gpu_kernels.initWorlds, gpu_kernels.mod, + compile_results.initWorldsName.c_str())); + REQ_CU(CudaDynamicLoader::cuModuleGetFunction(&gpu_kernels.initTasks, gpu_kernels.mod, + compile_results.initTasksName.c_str())); REQ_CU(CudaDynamicLoader::cuModuleGetFunction(&gpu_kernels.initBVHParams, gpu_kernels.mod, "initBVHParams")); @@ -1796,7 +1637,6 @@ static GPUEngineState initEngineAndUserState( const GPUKernels &gpu_kernels, const BVHKernels &bvh_kernels, const Optional &render_cfg, - ExecutorMode exec_mode, CUdevice cu_gpu, CUcontext cu_ctx, cudaStream_t strm) @@ -1975,23 +1815,12 @@ static GPUEngineState initEngineAndUserState( REQ_CU(CudaDynamicLoader::cuMemcpyHtoD(job_sys_consts_addr, gpu_consts_readback, job_sys_consts_size)); - if (exec_mode == ExecutorMode::JobSystem) { - launchKernel(gpu_kernels.initWorlds, 1, consts::numMegakernelThreads, no_args); - uint32_t num_queue_blocks = utils::divideRoundUp(num_worlds, consts::numEntryQueueThreads); - - launchKernel(gpu_kernels.queueUserInit, num_queue_blocks, - consts::numEntryQueueThreads, init_worlds_args); + launchKernel(gpu_kernels.initECS, 1, 1, init_ecs_args); + uint32_t num_init_blocks = utils::divideRoundUp(num_worlds, consts::numMegakernelThreads); - launchKernel(gpu_kernels.megakernels[0], 1, - consts::numMegakernelThreads, no_args); - } else if (exec_mode == ExecutorMode::TaskGraph) { - launchKernel(gpu_kernels.initECS, 1, 1, init_ecs_args); - uint32_t num_init_blocks = utils::divideRoundUp(num_worlds, consts::numMegakernelThreads); - - launchKernel(gpu_kernels.initWorlds, num_init_blocks, - consts::numMegakernelThreads, init_worlds_args); - launchKernel(gpu_kernels.initTasks, 1, 1, init_tasks_args); - } + launchKernel(gpu_kernels.initWorlds, num_init_blocks, + consts::numMegakernelThreads, init_worlds_args); + launchKernel(gpu_kernels.initTasks, 1, 1, init_tasks_args); REQ_CUDA(cudaStreamSynchronize(strm)); @@ -2058,56 +1887,6 @@ static GPUEngineState initEngineAndUserState( }; } -[[maybe_unused]] static CUgraphExec makeJobSysRunGraph( - CUfunction queue_run_kernel, - CUfunction job_sys_kernel, - uint32_t num_worlds) -{ - auto queue_args = makeKernelArgBuffer(num_worlds); - auto no_args = makeKernelArgBuffer(); - - uint32_t num_queue_blocks = utils::divideRoundUp(num_worlds, - consts::numEntryQueueThreads); - - CUgraph run_graph; - REQ_CU(CudaDynamicLoader::cuGraphCreate(&run_graph, 0)); - - CUDA_KERNEL_NODE_PARAMS kernel_node_params { - .func = queue_run_kernel, - .gridDimX = num_queue_blocks, - .gridDimY = 1, - .gridDimZ = 1, - .blockDimX = consts::numEntryQueueThreads, - .blockDimY = 1, - .blockDimZ = 1, - .sharedMemBytes = 0, - .kernelParams = nullptr, - .extra = queue_args.data(), - .kern = nullptr, - .ctx = nullptr, - }; - - CUgraphNode queue_node; - REQ_CU(CudaDynamicLoader::cuGraphAddKernelNode(&queue_node, run_graph, - nullptr, 0, &kernel_node_params)); - - kernel_node_params.func = job_sys_kernel; - kernel_node_params.gridDimX = 1; - kernel_node_params.blockDimX = consts::numMegakernelThreads; - kernel_node_params.extra = no_args.data(); - - CUgraphNode job_sys_node; - REQ_CU(CudaDynamicLoader::cuGraphAddKernelNode(&job_sys_node, run_graph, - &queue_node, 1, &kernel_node_params)); - - CUgraphExec run_graph_exec; - REQ_CU(CudaDynamicLoader::cuGraphInstantiate(&run_graph_exec, run_graph, 0)); - - REQ_CU(CudaDynamicLoader::cuGraphDestroy(run_graph)); - - return run_graph_exec; -} - static MegakernelConfig processExecConfigOverride(const char *override_str) { auto err = []() { @@ -2337,8 +2116,6 @@ MWCudaExecutor::MWCudaExecutor( // MADRONA_DEBUG_LOG("Kernel cache directory set to: %s\n", kernel_cache_path.c_str()); // MADRONA_DEBUG_LOG("BVH cache directory set to: %s\n", bvh_cache_path.c_str()); - const ExecutorMode exec_mode = ExecutorMode::TaskGraph; - auto strm = cu::makeStream(); CUdevice cu_gpu; @@ -2384,7 +2161,7 @@ MWCudaExecutor::MWCudaExecutor( } GPUKernels gpu_kernels = buildKernels(compile_cfg, megakernel_cfgs, - exec_mode, num_sms, cu_capability); + num_sms, cu_capability); GPUEngineState eng_state = initEngineAndUserState( state_cfg.numWorlds, state_cfg.numWorldDataBytes, @@ -2394,7 +2171,7 @@ MWCudaExecutor::MWCudaExecutor( state_cfg.numExportedBuffers, gpu_kernels, bvh_kernels, render_cfg, - exec_mode, cu_gpu, cu_ctx, strm); + cu_gpu, cu_ctx, strm); TaskGraphsState taskgraphs_state { .megakernels = std::move(gpu_kernels.megakernels), diff --git a/src/mw/device/include/madrona/job.hpp b/src/mw/device/include/madrona/job.hpp deleted file mode 100644 index a206a0cf..00000000 --- a/src/mw/device/include/madrona/job.hpp +++ /dev/null @@ -1,67 +0,0 @@ -/* - * Copyright 2021-2022 Brennan Shacklett and contributors - * - * Use of this source code is governed by an MIT-style - * license that can be found in the LICENSE file or at - * https://opensource.org/licenses/MIT. - */ -#pragma once - -#include -#include -#include - -namespace madrona { - -struct JobID { - uint32_t gen; - uint32_t id; - - static constexpr inline JobID none(); -}; - -struct JobContainerBase { - JobID jobID; - uint32_t worldID; - uint32_t numInvocations; - uint32_t numDependencies; - - template struct DepsArray; -}; - -template -struct JobContainer : public JobContainerBase { - [[no_unique_address]] DepsArray dependencies; - [[no_unique_address]] Fn fn; - - template - inline JobContainer(JobID job_id, uint32_t world_id, - uint32_t num_invocations, Fn &&fn, DepTs ...deps); -}; - -struct Job { - JobContainerBase *data; - uint32_t funcID; - uint32_t numCombinedJobs; - uint32_t numBytesPerJob; -}; - -class JobManager { -public: - uint32_t numOutstandingInvocations; - - AtomicU32 freeTrackerHead; - - bool startBlockIter(RunnableJob *out_job); - void finishBlockIter(); - - static inline JobManager * get(); - - template - static inline ContextT makeContext(JobID job_id, uint32_t grid_id, - uint32_t world_id, uint32_t lane_id); -}; - -} - -#include "job.inl" diff --git a/src/mw/device/include/madrona/job.inl b/src/mw/device/include/madrona/job.inl deleted file mode 100644 index b5e530f3..00000000 --- a/src/mw/device/include/madrona/job.inl +++ /dev/null @@ -1,109 +0,0 @@ -#include - -#include - -#include "mw_gpu/worker_init.hpp" -#include "mw_gpu/const.hpp" - -namespace madrona { - -namespace mwGPU { -namespace consts { - -inline constexpr uint32_t numWarpThreads = 32; -inline constexpr uint32_t numMegakernelThreads = 256; -inline constexpr uint32_t numMegakernelWarps = - numMegakernelThreads / numWarpThreads; - -} -} - -constexpr JobID JobID::none() -{ - return JobID { - ~0u, - ~0u, - }; -} - -template -struct JobContainerBase::DepsArray { - JobID dependencies[N]; - - template - inline DepsArray(DepTs ...deps) - : dependencies { deps ... } - {} -}; - -template <> struct JobContainerBase::DepsArray<0> { - template - inline DepsArray(DepTs...) {} -}; - -template -template -JobContainer::JobContainer(JobID job_id, - uint32_t world_id, - uint32_t num_invocations, - Fn &&func, - DepTs ...deps) - : JobContainerBase { - .jobID = job_id, - .worldID = world_id, - .numInvocations = num_invocations, - .numDependencies = N, - }, - dependencies(deps...), - fn(std::forward(func)) -{} - -JobManager * JobManager::get() -{ - return (JobManager *)GPUImplConsts::get().jobSystemAddr; -} - -mwGPU::SharedJobTracker * JobManager::getSharedJobTrackers() -{ - return (mwGPU::SharedJobTracker *)((char *)this + - mwGPU::GPUImplConsts::get().sharedJobTrackerOffset); -} - -mwGPU::UserJobTracker * JobManager::getUserJobTrackers() -{ - return (mwGPU::UserJobTracker *)((char *)this + - mwGPU::GPUImplConsts::get().userJobTrackerOffset); -} - -template -ContextT JobManager::makeContext(JobID job_id, uint32_t grid_id, - uint32_t world_id, uint32_t lane_id) -{ - using DataT = std::conditional_t, - WorldBase, typename ContextT::WorldDataT>; - - DataT *world_data; - - // If this is being called with the generic Context base class, - // we need to look up the size of the world data in constant memory - if constexpr (std::is_same_v) { - char *world_data_base = - (char *)mwGPU::GPUImplConsts::get().worldDataAddr; - world_data = (DataT *)(world_data_base + world_id * - mwGPU::GPUImplConsts::get().numWorldDataBytes); - } else { - DataT *world_data_base = - (DataT *)mwGPU::GPUImplConsts::get().worldDataAddr; - - world_data = world_data_base + world_id; - } - - return ContextT(world_data, WorkerInit { - .jobID = job_id, - .gridID = grid_id, - .worldID = world_id, - .laneID = lane_id, - }); -} - -} diff --git a/src/mw/device/include/madrona/mw_gpu.hpp b/src/mw/device/include/madrona/mw_gpu.hpp deleted file mode 100644 index 4870ffe1..00000000 --- a/src/mw/device/include/madrona/mw_gpu.hpp +++ /dev/null @@ -1,26 +0,0 @@ -/* - * Copyright 2021-2022 Brennan Shacklett and contributors - * - * Use of this source code is governed by an MIT-style - * license that can be found in the LICENSE file or at - * https://opensource.org/licenses/MIT. - */ -#pragma once - -#include - -#include "mw_gpu/entry.hpp" - -namespace madrona { - -template -class GPUJobEntry : mwGPU::EntryBase { -public: - static void submitInit(uint32_t invocation_idx, void *world_init_ptr); - static void submitRun(uint32_t invocation_idx); - -private: - static ContextT makeFakeContext(uint32_t invocation_idx); -}; - -#include "mw_gpu.inl" diff --git a/src/mw/device/include/madrona/mw_gpu.inl b/src/mw/device/include/madrona/mw_gpu.inl deleted file mode 100644 index a004cdc9..00000000 --- a/src/mw/device/include/madrona/mw_gpu.inl +++ /dev/null @@ -1,50 +0,0 @@ -/* - * Copyright 2021-2022 Brennan Shacklett and contributors - * - * Use of this source code is governed by an MIT-style - * license that can be found in the LICENSE file or at - * https://opensource.org/licenses/MIT. - */ -#pragma once - -#include "mw_gpu/worker_init.hpp" - -namespace madrona { - -template -void GPUEntry::submitInit(uint32_t invocation_idx, - void *world_init_ptr) -{ - ContextT ctx = makeFakeContext(invocation_idx); - - InitT *base_init = (InitT *)world_init_ptr; - InitT *init = base_init + ctx.worldID(); - - ctx.submit([init](ContextT &ctx) { - BaseT::init(ctx, *init); - }, false); -} - -template -void GPUEntry::submitRun(uint32_t invocation_idx) -{ - ContextT ctx = makeFakeContext(invocation_idx); - - ctx.submit([](ContextT &ctx) { - BaseT::run(ctx); - }, false); -} - -template -ContextT GPUEntry::makeFakeContext( - uint32_t invocation_idx) -{ - uint32_t lane_id = - invocation_idx % mwGPU::consts::numWarpThreads; - uint32_t world_id = lane_id; - - return JobManager::makeContext(JobID { 0, 0 }, 0, - world_id, lane_id); -} - -} diff --git a/src/mw/device/job.cpp b/src/mw/device/job.cpp deleted file mode 100644 index e2d01416..00000000 --- a/src/mw/device/job.cpp +++ /dev/null @@ -1,999 +0,0 @@ -/* - * Copyright 2021-2022 Brennan Shacklett and contributors - * - * Use of this source code is governed by an MIT-style - * license that can be found in the LICENSE file or at - * https://opensource.org/licenses/MIT. - */ -#include -#include -#include - -#include -#include -#include - -#include "mw_gpu/const.hpp" -#include "mw_gpu/cu_utils.hpp" - -// Considerations: -// - Enforcing end-user data safety constraints around iteration & creation / deletion: fundamentally, a job needs to not run until it can be run by all worlds. -// - This means a given launch should have all the union of all dependencies of all instances of the same launch across worlds. Practically speaking, this means all the children launches of a given parent need to be deferred until the parent has completed across all worlds. Then a merging step has to run, where children jobs that match across worlds need to be grouped together into single launches with shared dependencies. -// - Challenges: -// - Currently job "uniqueness" isn't enforced. This means job A can launch the same function pointer twice, and the second launch can depend on the first. This could be changed, but may be jarring given the current lambda API -// - The above issue implies that searching through the list of children, merging based on function ptr alone is not enough. One solution is (func_ptr, # of launches of ptr in this job). That's annoying because now you need to track per func ptr state while building -// - Advantages: -// - This scheme enforces maximum merging of jobs. If this wasn't the case, there would need to be some other mechanism to recover common jobs in conditional execution cases. -// - Disadvantages: -// - Lots of synchronization, less ability to overlap. Partly lose a major advantage of the megakernel which is that there is now no driver mandated synchronization -// - Sketch: -// - Every running job has a num_worlds wide log array. -// - -using std::atomic_uint32_t; -using std::memory_order; - -extern "C" { -__global__ void madronaMWGPUMegakernel(uint32_t func_id, madrona::JobContainerBase *data, uint32_t *data_indices, uint32_t *invocation_offsets, uint32_t num_launches, uint32_t grid); -} - -namespace madrona { -namespace mwGPU { - -namespace consts { - -static constexpr uint32_t allActive = 0xFFFFFFFF; -static constexpr uint32_t jobTrackerTerm = ~0u; - -static constexpr uint32_t jobsPerWaitQueue = 16384; -static constexpr uint32_t numWaitQueues = 64; - -} - -struct JobTracker { - uint32_t parent; - std::atomic_uint32_t numOutstanding; - std::atomic_uint32_t remainingInvocations; -}; - -struct ThreadblockData { - ChunkAllocator::Cache chunkCache; - std::atomic_uint32_t curChunkID; - std::atomic_uint32_t remainingChunkBytes; -}; - -static __shared__ ThreadblockData tbData; - -struct WaitQueue { - SpinLock lock; - uint32_t numWaiting; - Job waitingJobs[consts::jobsPerWaitQueue]; -}; - -static inline WaitQueue * getWaitQueues(JobManager *mgr, uint32_t idx) -{ - return (WaitQueue *)((char *)mgr + GPUImplConsts::get().waitQueueOffset); -} - -static inline Job *getBaseJobList(JobManager *mgr) -{ - return (Job *)((char *)mgr + GPUImplConsts::get().jobListOffset); -} - -static void initializeJobSystem(JobManager *mgr) -{ - if (threadIdx.x == 0) { - mgr->numOutstandingInvocations = 0; - - for (int i = 0; i < 8; i++) { - mgr->activeGrids[i] = 0; - } - - mgr->freeTrackerHead.store(JobID { 0, 0 }, std::memory_order_relaxed); - } - - __syncthreads(); - - auto grids = getGridInfo(mgr); - - for (int i = threadIdx.x; i < (int)consts::numJobGrids; i += blockDim.x) { - grids[i].waitJobHead.store(0, std::memory_order_relaxed); - grids[i].waitJobTail.store(0, std::memory_order_relaxed); - grids[i].waitQueueLock.store(0, std::memory_order_relaxed); - grids[i].numRunning.store(0, std::memory_order_relaxed); - } - - auto shared_trackers = getSharedJobTrackers(mgr); - - int num_shared_trackers = - GPUImplConsts::get().maxJobsPerGrid * consts::numJobGrids; - - for (int i = threadIdx.x; i < num_trackers; i += blockDim.x) { - JobTracker &tracker = trackers[i]; - tracker.gen.store(0, std::memory_order_relaxed); - if (i < num_trackers - 1) { - tracker.parent.store(i + 1, std::memory_order_relaxed); - } else { - tracker.parent.store(consts::jobTrackerTerm, std::memory_order_relaxed); - } - tracker.numOutstanding.store(0, std::memory_order_relaxed); - } - - std::atomic_thread_fence(std::memory_order_release); -} - -static inline const JobID * getJobDependencies(JobContainerBase *job_base) -{ - return (const JobID *)((char *)job_base + sizeof(JobContainerBase)); -} - -static inline bool isJobReady(JobManager *job_mgr, Job &job) -{ - const JobTracker *trackers = getJobTrackers(job_mgr); - for (int job_idx = 0; job_idx < job.numCombinedJobs; job_idx++) { - JobContainerBase *job_data = (JobContainerBase *)( - (char *)job.data + job.numBytesPerJob * job_idx); - - int num_deps = job_data->numDependencies; - - const JobID *dependencies = getJobDependencies(job_data); - for (int i = 0; i < num_deps; i++) { - JobID dependency = dependencies[i]; - - if (trackers[dependency.id].gen == dependency.gen) { - return false; - } - } - } - - return true; -} - -static inline void *allocateJobData(uint32_t total_bytes) -{ - // FIXME - return malloc(total_bytes); -} - -static inline void freeJobData(void *data) -{ - free(data); -} - -static inline bool checkGEWrapped(uint32_t a, uint32_t b) -{ - return uint32_t(a - b) <= (1u << 31u); -} - -static inline bool checkLTWrapped(uint32_t a, uint32_t b) -{ - return uint32_t(a - b) > (1u << 31u); -} - -static inline uint32_t wrapQueueIdx(uint32_t idx, uint32_t num_elems) -{ - // Assumes num_elems is a power of 2 - return idx & (num_elems - 1); -} - -// Only a single thread block can run this function -static void jobLoop() -{ - uint32_t lane_id = threadIdx.x % consts::numWarpThreads; - uint32_t warp_id = threadIdx.x / consts::numWarpThreads; - - constexpr uint32_t total_num_warps = consts::numJobSystemKernelThreads / - consts::numWarpThreads; - - constexpr uint32_t grids_per_warp = - consts::numJobGrids / total_num_warps; - - auto job_mgr = getJobManager(); - auto base_job_grids = getGridInfo(job_mgr); - auto base_job_list = getBaseJobList(job_mgr); - - const uint32_t max_jobs_per_grid = GPUImplConsts::get().maxJobsPerGrid; - - static __shared__ uint32_t active_grids_tmp[8]; - - cuda::barrier cpy_barrier; - init(&cpy_barrier, 1); - - auto findFirstReadyJob = [job_mgr, lane_id, max_jobs_per_grid]( - Job *job_list, uint32_t job_head, uint32_t job_tail) { - int first_job_idx = -1; - if (lane_id == 0) { - for (uint32_t job_idx = job_head; - checkLTWrapped(job_idx, job_tail); job_idx++) { - int wrapped_idx = - (int)wrapQueueIdx(job_idx, max_jobs_per_grid); - Job &cur_job = job_list[wrapped_idx]; - - if (isJobReady(job_mgr, cur_job)) { - first_job_idx = wrapped_idx; - break; - } - } - } - - return __shfl_sync(consts::allActive, first_job_idx, 0); - }; - - auto getFreeGrid = [lane_id]() { - int run_grid_idx = -1; - if (lane_id == 0) { -#pragma unroll - for (int bitfield_idx = 0; bitfield_idx < 8; bitfield_idx++) { - uint32_t *grid_bitfield_ptr = &active_grids_tmp[bitfield_idx]; - - uint32_t old_bitfield, set_bitfield; - do { - old_bitfield = *grid_bitfield_ptr; - - uint32_t inverse = ~old_bitfield; - - // All grids running - if (inverse == 0) { - run_grid_idx = -1; - break; - } - - uint32_t idx = 31 - __clz(inverse); - - uint32_t mask = 1 << idx; - - set_bitfield = old_bitfield | mask; - - run_grid_idx = idx; - } while (atomicCAS(grid_bitfield_ptr, old_bitfield, - set_bitfield)); - - if (run_grid_idx != -1) { - run_grid_idx = bitfield_idx * 32 + run_grid_idx; - - break; - } - } - } - - return __shfl_sync(consts::allActive, run_grid_idx, 0); - }; - - auto nextLoopSetup = [job_mgr, warp_id, lane_id]() { - std::atomic_thread_fence(std::memory_order_acquire); - - if (warp_id == 0 && lane_id == 0) { - for (int i = 0; i < 8; i++) { - active_grids_tmp[i] = job_mgr->activeGrids[i]; - } - } - - __syncthreads(); - }; - - nextLoopSetup(); - while (true) { - for (int wait_grid_offset = 0; wait_grid_offset < (int)grids_per_warp; - wait_grid_offset++) { - uint32_t wait_grid_idx = - warp_id * grids_per_warp + wait_grid_offset; - - JobGridInfo &wait_grid = base_job_grids[wait_grid_idx]; - Job *waiting_jobs = base_job_list + - wait_grid_idx * max_jobs_per_grid; - - // Relaxed is safe for head & tail, because - // nextLoopSetup() does an acquire barrier - uint32_t job_head = - wait_grid.waitJobHead.load(std::memory_order_relaxed); - // Cache the value of job tail, and use it across the warp, - // it can be incremented by other threads - uint32_t job_tail; - if (lane_id == 0) { - job_tail = - wait_grid.waitJobTail.load(std::memory_order_relaxed); - } - job_tail = __shfl_sync(consts::allActive, job_tail, 0); - - int first_job_idx = - findFirstReadyJob(waiting_jobs, job_head, job_tail); - if (first_job_idx == -1) { - continue; - } - - int run_grid_idx = getFreeGrid(); - if (run_grid_idx == -1) { - break; - } - - JobGridInfo &run_grid = base_job_grids[run_grid_idx]; - - uint32_t first_func_id = waiting_jobs[first_job_idx].funcID; - auto isJobMergable = [first_func_id, job_mgr](Job &job) { - return job.funcID == first_func_id && isJobReady(job_mgr, job); - }; - - uint32_t num_bytes_per_job = - waiting_jobs[first_job_idx].numBytesPerJob; - - uint32_t total_num_jobs = 0; - uint32_t total_num_invocations = 0; - - // Could start from the unwrapped version of first_job_idx, - // but would need to change findFirstReadyJob to return - // a separate failure boolean - for (uint32_t job_offset = job_head; - checkLTWrapped(job_offset, job_tail); - job_offset += consts::numWarpThreads) { - uint32_t job_idx = job_offset + lane_id; - - bool inbounds = checkLTWrapped(job_idx, job_tail); - - // Force out of bounds indices in bounds - if (!inbounds) { - job_idx = job_offset; - } - - Job cur_job = - waiting_jobs[wrapQueueIdx(job_idx, max_jobs_per_grid)]; - - bool merge_job = inbounds && isJobMergable(cur_job); - uint32_t merge_mask = - __ballot_sync(consts::allActive, merge_job); - - uint32_t cur_num_jobs = - merge_job ? cur_job.numCombinedJobs : 0_u32; - - uint32_t num_prior_jobs = total_num_jobs + - warpExclusiveScan(lane_id, cur_num_jobs); - - // Copy job data into grid's run buffer - if (merge_job) { - cuda::memcpy_async(run_grid.runData.buf + - num_prior_jobs * num_bytes_per_job, cur_job.data, - num_bytes_per_job * cur_num_jobs, cpy_barrier); - - assert((num_prior_jobs + cur_num_jobs) * num_bytes_per_job <= - 1024 * 1024); - } - - // FIXME: this is a potentially massive loop that one thread - // has to deal with - uint32_t combined_num_invocations = 0; - for (int job_idx = 0; job_idx != (int)cur_num_jobs; job_idx++) { - JobContainerBase *job_data = (JobContainerBase *)( - (char *)cur_job.data + job_idx * num_bytes_per_job); - - uint32_t num_invocations = job_data->numInvocations; - combined_num_invocations += num_invocations; - } - - uint32_t num_prior_invocations = total_num_invocations + - warpExclusiveScan(lane_id, combined_num_invocations); - - uint32_t num_setup_invocations = 0; - for (int job_idx = 0; job_idx != (int)cur_num_jobs; job_idx++) { - JobContainerBase *job_data = (JobContainerBase *)( - (char *)cur_job.data + job_idx * num_bytes_per_job); - - uint32_t num_invocations = job_data->numInvocations; - - for (int i = 0; i != (int)num_invocations; i++) { - run_grid.jobDataIndices[(int)num_prior_invocations + - num_setup_invocations] = num_prior_jobs + job_idx; - run_grid.jobInvocationOffsets[(int)num_prior_invocations + - num_setup_invocations] = i; - - num_setup_invocations++; - - assert(num_prior_invocations + num_setup_invocations < - 65536 * 16); - } - } - - total_num_invocations = - num_prior_invocations + combined_num_invocations; - total_num_jobs = num_prior_jobs + cur_num_jobs; - - // Get current running total of jobs merged together for launch - uint32_t top_merge_thread = getHighestSetBit(merge_mask); - total_num_invocations = __shfl_sync(consts::allActive, - total_num_invocations, top_merge_thread); - total_num_jobs = __shfl_sync(consts::allActive, - total_num_jobs, top_merge_thread); - } - - // Wait for all the async copies - cpy_barrier.arrive_and_wait(); - - __syncwarp(); - - uint32_t base_wait_coalesce_idx = job_tail - 1u; - - // Free job data for jobs that have been copied into the run - // data block, coalesce waiting job list - // Unfortunately, this loop is in reverse, because the list needs - // to be coalesced into the tail - for (uint32_t job_offset = job_tail - 1u; - checkGEWrapped(job_offset, job_head); - job_offset -= consts::numWarpThreads) { - - uint32_t job_idx = job_offset - lane_id; - - bool inbounds = checkGEWrapped(job_idx, job_head); - - if (!inbounds) { - job_idx = job_offset; - } - - Job cur_job = - waiting_jobs[wrapQueueIdx(job_idx, max_jobs_per_grid)]; - - bool mergable = isJobMergable(cur_job); - bool coalesceable = inbounds && !mergable; - - if (inbounds && mergable) { - freeJobData(cur_job.data); - } - - // The sync here also ensures all threads are done - // using cur_job before any pointers are overwritten - // when coalescing the list - uint32_t coalesce_mask = - __ballot_sync(consts::allActive, coalesceable); - - // Coalesce jobs that won't be launched to make the waiting - // list contiguous - if (coalesceable) { - // Lower threads in the warp are farther ahead in the - // array due to reading backwards - uint32_t coalesce_idx = base_wait_coalesce_idx - - getNumLowerSetBits(coalesce_mask, lane_id); - - if (coalesce_idx != job_idx) { - int wrapped_coalesce_idx = - wrapQueueIdx(coalesce_idx, max_jobs_per_grid); - waiting_jobs[wrapped_coalesce_idx] = cur_job; - } - - base_wait_coalesce_idx = coalesce_idx - 1u; - } - - if (coalesce_mask != 0) { - uint32_t top_coalesce_thread = - getHighestSetBit(coalesce_mask); - base_wait_coalesce_idx = __shfl_sync(consts::allActive, - base_wait_coalesce_idx, top_coalesce_thread); - } - } - - - if (lane_id == 0) { - wait_grid.waitJobHead.store(base_wait_coalesce_idx + 1, - std::memory_order_relaxed); - } - - std::atomic_thread_fence(std::memory_order_release); - - if (lane_id == 0) { - uint32_t num_blocks = utils::divideRoundUp(total_num_invocations, - consts::numJobLaunchKernelThreads); - - run_grid.numRunning.store(total_num_invocations, - std::memory_order_relaxed); - - madronaMWGPUMegakernel<<>>( - first_func_id, - (JobContainerBase *)run_grid.runData.buf, - run_grid.jobDataIndices, run_grid.jobInvocationOffsets, - total_num_invocations, run_grid_idx); - } - } - - __nanosleep(500); - - // Call this function at the end of the loop, in order to use - // the same __syncthreads / acquire call for reading the activeGrid - // bitfields and checking numOutstandingInvocations - nextLoopSetup(); - - if (job_mgr->numOutstandingInvocations == 0) { - break; - } - } -} - -static inline uint32_t computeMaxNumJobs(uint32_t num_worlds) -{ - // FIXME: scaling linearly like this probably doesn't make sense - return consts::maxNumJobsPerWorld * num_worlds; -} - -static inline JobID allocateJobTrackerSlot(JobManager *job_mgr, - JobTracker *trackers) -{ - JobID cur_head = - job_mgr->freeTrackerHead.load(std::memory_order_acquire); - - JobID new_head; - - do { - if (cur_head.id == consts::jobTrackerTerm) { - break; - } - - new_head.gen = cur_head.gen + 1; - new_head.id = trackers[cur_head.id].parents[0]; - } while (!job_mgr->freeTrackerHead.compare_exchange_weak( - cur_head, new_head, memory_order::release, - memory_order::acquire)); - - uint32_t job_id = cur_head.id; - - // FIXME - if (job_id == consts::jobTrackerTerm) { - assert(false); - } - - uint32_t gen = tracker.gen.load(std::memory_order_relaxed); - - return JobID { - gen, - job_id, - }; -} - -static inline void freeJobTrackerSlot(uint32_t job_id) -{ - auto job_mgr = getJobManager(); - JobTracker *trackers = getJobTrackers(job_mgr); - - JobTracker &tracker = trackers[job_id]; - tracker.gen = tracker.gen + 1; - - JobID new_head; - new_head.id = job_id; - - JobID cur_head = job_mgr->freeTrackerHead.load( - std::memory_order_relaxed); - - do { - new_head.gen = cur_head.gen + 1; - - tracker.parents[0] = cur_head.id; - } while (!job_mgr->freeTrackerHead.compare_exchange_weak( - cur_head, new_head, - memory_order::release, memory_order::relaxed)); -} - -static inline void decrementJobTracker(JobTracker *job_trackers, JobID job_id) -{ - uint32_t cur_id = job_id.id; - while (cur_id != consts::jobTrackerTerm) { - JobTracker &tracker = job_trackers[cur_id]; - - uint32_t prev_outstanding = - tracker.numOutstanding.fetch_sub(1, std::memory_order_acq_rel); - - if (prev_outstanding == 1 && - tracker.remainingInvocations.load(std::memory_order_relaxed) == 0) { - uint32_t parent = tracker.parent; - - freeJobTrackerSlot(cur_id); - - cur_id = parent; - } else { - break; - } - } -} - -// This function should only be called by the wave leader -static inline void queueMultiJobInWaitList( - uint32_t func_id, JobContainerBase *data, uint32_t grid_id, - uint32_t num_jobs, uint32_t total_num_invocations, - uint32_t num_bytes_per_job) -{ - Job job { - .data = data, - .funcID = func_id, - .numCombinedJobs = num_jobs, - .numBytesPerJob = num_bytes_per_job, - }; - - auto job_mgr = getJobManager(); - - const auto base_job_grids = getGridInfo(job_mgr); - const auto base_job_list = getBaseJobList(job_mgr); - const uint32_t max_jobs_per_grid = GPUImplConsts::get().maxJobsPerGrid; - - JobGridInfo &cur_grid = base_job_grids[grid_id]; - Job *job_list = base_job_list + grid_id * max_jobs_per_grid; - - // Get lock - while (cur_grid.waitQueueLock.exchange(1, std::memory_order_acq_rel)) {} - - uint32_t cur_job_pos = cur_grid.waitJobTail.load(std::memory_order_acquire); - job_list[wrapQueueIdx(cur_job_pos, max_jobs_per_grid)] = job; - - cur_grid.waitJobTail.fetch_add(1, std::memory_order_relaxed); - - cur_grid.waitQueueLock.store(0, std::memory_order_relaxed); - - atomicAdd(&job_mgr->numOutstandingInvocations, total_num_invocations); - - std::atomic_thread_fence(std::memory_order_release); -} - -} - -bool JobManager::startBlockIter(uint32_t block_idx, RunnableJob *out_job) -{ -} - -void JobManager::finishBlockIter(uint32_t block_idx) -{ - __syncthreads(); // Ensure entire threadblock has finished - uint32_t num_log_entries = - tbRunData.numLogEntries.load(memory_order::relaxed); - - for (int i = threadIdx.x; i < (int)num_log_entries; - i += mwGPU::consts::numMegakernelThreads) { - uint8_t offset_lp = tbRunData.dataTmpOffsets[i]; - int offset = (int)offset_lp * 8; - - LogEntry *log = (LogEntry *)(tbRunData.dataTmp + offset); - - } - - __syncthreads(); - - - if (threadIdx.x == 0) { - tbRunData.numLogEntries.store(0, memory_order::relaxed); - } -} - -Context::WaveInfo Context::computeWaveInfo() -{ - using namespace mwGPU; - - uint32_t active = __activemask(); - uint32_t num_active = __popc(active); - - uint32_t coalesced_idx = getNumLowerSetBits(active, lane_id_); - - return WaveInfo { - .activeMask = active, - .numActive = num_active, - .coalescedIDX = coalesced_idx, - }; -} - -void Context::stageChildJob(uint32_t func_id, uint32_t num_combined_jobs, - uint32_t bytes_per_job, void *containers) -{ - uint32_t offset = - tbScratch.waitQueue.numWaiting.fetch_add(1, memory_order::relaxed); - - tbScratch.waitingJobs[offset] = Job { - .data = (JobContainerBase *)containers_tmp, - .funcID = func_id, - .numCombinedJobs = num_combined_jobs, - .bytesPerJob = bytes_per_job, - }; -} - -JobID Context::waveSetupNewJob(uint32_t func_id, bool link_parent, - uint32_t num_invocations, uint32_t bytes_per_job, - void **thread_data_store) -{ - auto wave_info = computeWaveInfo(); - auto job_mgr = JobManager::get(); - auto job_trackers = getJobTrackers(job_mgr); - - JobID child_id; - char *tmp_data; - - uint32_t num_total_invocations = - __reduce_add_sync(wave_info.activeMask, num_invocations); - - if (wave_info.isLeader()) { - child_id = allocateJobTrackerSlot(job_mgr, job_trackers); - job_trackers[child_id.id].numOutstanding.store(1, memory_order::relaxed); - job_trackers[child_id.id].remainingInvocations.store( - num_total_invocations, memory_order::relaxed); - - uint32_t total_bytes = wave_info.numActive * bytes_per_job; - tmp_data = (char *)malloc(total_bytes); - - stageChildJob(func_id, tmp_data, bytes_per_job, wave_info.numActive); - } - - child_id = { - .gen = __shfl_sync(wave_info.activeMask, child_id.gen, - wave_info.leaderLane); - .id = __shfl_sync(wave_info.activeMask, child_id.id, - wave_info.leaderLane); - }; - - uint32_t parent_id; - if (link_parent) { - parent_id = job_id_.id; - - uint32_t parent_match = - __match_any_sync(wave_info.activeMask, parent_id); - - uint32_t top_thread = getHighestSetBit(parent_match); - uint32_t num_children = __popc(parent_match); - - if (lane_id_ == top_thread) { - trackers[parent_id].numOutstanding.fetch_add(num_children, - std::memory_order_relaxed); - } - - job_trackers[child_id.id].parents[wave_info.coalescedIDX] = parent_id; - } else { - parent_id = consts::jobTrackerTerm; - } - - tmp_data = (char *)__shfl_sync(wave_info.activeMask, (uintptr_t)tmp_data, - wave_info.leaderLane); - - *thread_data_store = tmp_data + wave_info.coalescedIDX * bytes_per_job; - - return child_id; -} - -JobID Context::getNewJobID(bool link_parent, uint32_t num_invocations) -{ - using namespace mwGPU; - - auto job_mgr = getJobManager(); - JobTracker *trackers = getJobTrackers(job_mgr); - - uint32_t parent_id; - if (link_parent) { - parent_id = job_id_.id; - trackers[parent_id].numOutstanding.fetch_add(1, - std::memory_order_release); - } else { - parent_id = consts::jobTrackerTerm; - } - - return allocateJobTrackerSlot(job_mgr, trackers, parent_id, num_invocations); -} - -// Allocates a shared block of memory for the active threads in wave, -// where lower threads are given a pointer to an early chunk of the block -JobContainerBase * Context::allocJob(uint32_t bytes_per_job, - WaveInfo wave_info) -{ - using namespace mwGPU; - - void *base_store; - if (lane_id_ == wave_info.leaderLane) { - base_store = allocateJobData(bytes_per_job * wave_info.numActive); - } - - // Sync store point & id with wave - base_store = (void *)__shfl_sync(wave_info.activeMask, - (uintptr_t)base_store, wave_info.leaderLane); - - return (JobContainerBase *)( - (char *)base_store + bytes_per_job * wave_info.coalescedIDX); -} - -void Context::logNewJob(uint32_t func_id, JobContainerBase *data, - uint32_t num_invocations, - uint32_t num_bytes_per_job, - uint32_t lane_id, - WaveInfo wave_info) -{ - using namespace mwGPU; - - uint32_t total_num_invocations = - warpSum(wave_info.activeMask, lane_id, num_invocations); - - if (lane_id_ == wave_info.leaderLane) { - queueMultiJobInWaitList(func_id, data, grid_id_, - wave_info.numActive, total_num_invocations, - num_bytes_per_job); - } -} - -void Context::markJobFinished() -{ - using namespace mwGPU; - auto job_mgr = JobManager::get(); - - JobTracker *trackers = getJobTracker(job_mgr); - -#pragma loop unroll - uint32_t job_idx = job_id_.id; - for (int i = 0; i < consts::numWarpThreads; i++) { - uint32_t other_job_idx = __shfl_sync(consts::allActive, job_idx); - - uint32_t job_match = __match_any_sync(wave_info.activeMask, job_idx); - - uint32_t top_thread = getHighestSetBit(job_match); - uint32_t num_invocations = __popc(job_match); - - if (lane_id_ == top_thread) { - uint32_t prev_invocations = - trackers[job_idx].remainingInvocations.fetch_sub(num_invocations, - memory_order::relaxed); - } - - if (prev_invocations == num_invocations) { - decrementJobTracker(trackers, job_id_); - } - } - } - - __syncthreads(); - - uint32_t num_waiting = - tbScratch.waitQueue.numWaiting.load(memory_order::relaxed); - - if (threadIdx.x == 0) { - uint32_t prev_invocations = metadata.remainingInvocations.fetch_sub( - num_invocations); - - if (prev_invocations == num_invocations) { - scratch = - metadata.numStaged.load(memory_order::relaxed); - } else { - scratch[0] = ~0_u32; - } - } - - __syncthreads(); - - uint32_t num_staged = scratch[0]; - if (num_staged == ~0_u32) { - return; - } - - // Merge children jobs - // FIXME: optimization - do a reduction at the thread block level first - // Or, launch a job with num_staged invocations that does this reduction - // globally - - StagedJob *staged_jobs = getStagedJobs(metadata); - - cuda::barrier cpy_barrier; - cuda::barrier::init(&cpy_barrier, consts::numMegakernelThreads); - - for (int i = 0; i < num_staged; i += consts::numMegakernelThreads) { - int offset = i + threadIdx.x; - bool inbounds = offset < num_staged; - offset = inbounds ? offset : num_staged - 1; - - StagedJob &staged_job = staged_jobs[offset]; - uint32_t func_id = staged_job.funcID; - - __syncthreads(); - } - - cpy_barrier.arrive_and_wait(); - - if (threadIdx.x == 0) { - decrementJobTracker(job_mgr, merged_job_id_); - } -} - -static inline uint32_t computeNumAllocatorChunks(uint64_t num_bytes) -{ - return utils::divideRoundUp(num_bytes, - (uint64_t)ChunkAllocator::chunkSize); -} - -} - -extern "C" __global__ void madronaMWGPUComputeConstants( - uint32_t num_worlds, - uint32_t num_world_data_bytes, - uint32_t world_data_alignment, - uint64_t num_allocator_bytes, - madrona::mwGPU::GPUImplConsts *out_constants, - size_t *job_system_buffer_size) -{ - using namespace madrona; - using namespace madrona::mwGPU; - - uint32_t max_num_jobs_per_grid = - madrona::mwGPU::computeMaxNumJobs(num_worlds); - - uint32_t max_num_jobs = consts::numJobGrids * max_num_jobs_per_grid; - - uint64_t total_bytes = sizeof(JobManager); - - uint64_t state_mgr_offset = utils::roundUp(total_bytes, - (uint64_t)alignof(StateManager)); - - total_bytes = state_mgr_offset + sizeof(StateManager); - - uint64_t chunk_allocator_offset = utils::roundUp(total_bytes, - (uint64_t)alignof(ChunkAllocator)); - - total_bytes = chunk_allocator_offset + sizeof(ChunkAllocator); - - uint64_t chunk_base_offset = utils::roundUp(total_bytes, - (uint64_t)alignof(ChunkAllocator::chunkSize)); - - uint64_t num_chunks = computeNumAllocatorChunks(num_allocator_bytes); - total_bytes += num_chunks * ChunkAllocator::chunkSize; - - uint64_t world_data_offset = - utils::roundUp(total_bytes, (uint64_t)world_data_alignment); - - total_bytes = - world_data_offset + (uint64_t)num_world_data_bytes * num_worlds; - - uint64_t grid_offset = utils::roundUp(total_bytes, - (uint64_t)alignof(JobGridInfo)); - - total_bytes = grid_offset + sizeof(JobGridInfo) * consts::numJobGrids; - - uint64_t wait_job_offset = madrona::utils::roundUp(total_bytes, - (uint64_t)alignof(Job)); - - uint64_t num_job_bytes = sizeof(Job) * max_num_jobs; - - total_bytes = wait_job_offset + num_job_bytes; - - uint64_t tracker_offset = madrona::utils::roundUp(total_bytes, - (uint64_t)alignof(JobTracker)); - - // FIXME: using max_num_jobs for this doesn't quite make sense, because - // there will be more outstanding trackers than waiting jobs due to - // parent trackers remaining alive until children finish - total_bytes = tracker_offset + sizeof(JobTracker) * max_num_jobs * 2; - - *out_constants = GPUImplConsts { - .jobSystemAddr = (void *)0ul, - .stateManagerAddr = (void *)state_mgr_offset, - .chunkAllocatorAddr = (void *)chunk_allocator_offset, - .chunkBaseAddr = (void *)chunk_base_offset, - .worldDataAddr = (void *)world_data_offset, - .numWorldDataBytes = num_world_data_bytes, - .numWorlds = num_worlds, - .jobGridsOffset = (uint32_t)grid_offset, - .jobListOffset = (uint32_t)wait_job_offset, - .maxJobsPerGrid = max_num_jobs_per_grid, - .jobTrackerOffset = (uint32_t)tracker_offset, - }; - - *job_system_buffer_size = total_bytes; -} - -extern "C" __global__ void madronaMWGPUInitialize( - uint64_t num_allocator_bytes) -{ - using namespace madrona; - using namespace madrona::mwGPU; - - auto job_mgr = getJobManager(); - - if (threadIdx.x == 0) { - new (job_mgr) JobManager(); - } - __syncthreads(); - - initializeJobSystem(job_mgr); - - if (threadIdx.x == 0) { - new (GPUImplConsts::get().stateManagerAddr) StateManager(1024); - } - - if (threadIdx.x == 0) { - new (GPUImplConsts::get().chunkAllocatorAddr) ChunkAllocator( - computeNumAllocatorChunks(num_allocator_bytes)); - } -} - -extern "C" __global__ void madronaTrainJobSystemKernel() -{ - madrona::mwGPU::jobLoop(); -} diff --git a/src/mw/device/megakernel_job_impl.inl b/src/mw/device/megakernel_job_impl.inl deleted file mode 100644 index ce18872c..00000000 --- a/src/mw/device/megakernel_job_impl.inl +++ /dev/null @@ -1,62 +0,0 @@ -#pragma once -#include -#include - -namespace madrona { -namespace mwGPU { - -#ifdef MADRONA_CLANG -#pragma clang diagnostic push -#pragma clang diagnostic ignored "-Wundefined-internal" -#endif -static inline __attribute__((always_inline)) void dispatch( - uint32_t func_id, - madrona::JobContainerBase *data, - uint32_t *data_indices, - uint32_t *invocation_offsets, - uint32_t num_launches, - uint32_t grid); -#ifdef MADRONA_CLANG -#pragma clang diagnostic pop -#endif - -static inline __attribute__((always_inline)) void megakernelImpl() -{ - static __shared__ RunnableJob threadblock_job; - static __shared__ bool job_found; - - while (true) { - { - JobManager *job_mgr = JobManager::get(); - if (job_mgr->numOutstandingInvocations.load( - std::memory_order_relaxed) == 0) { - break; - } - - if (threadIdx.x == 0) { - job_found = job_mgr->startBlockIter(&threadblock_job); - } - __syncthreads(); - - if (!job_found) { - break; - } - } - - RunnableJob runnable = threadblock_job; - - dispatch(func_id, data, data_indices, invocation_offsets, num_launches, - grid); - - // Iterate through submitted jobs, merge job IDs, etc - JobManager::get()->finishBlockIter(); - } -} - -} -} - -extern "C" void madronaMWGPUMegakernel() -{ - madrona::mwGPU::megakernelImpl(); -} diff --git a/src/mw/device/work_queue.hpp b/src/mw/device/work_queue.hpp deleted file mode 100644 index fc6af2f4..00000000 --- a/src/mw/device/work_queue.hpp +++ /dev/null @@ -1,42 +0,0 @@ -#pragma once - -#include -#include - -#include - -struct WorkItem { - uint32_t job_id; - void *data; -}; - -template -class WorkQueue { -public: - inline void enqueue(WorkItem work_item) - { - do { - uint32_t cur_count = count_.load_relaxed(); - if (cur_count >= queue_size) { - continue; - } - - uint32_t cur_head = head_.load_relaxed(); - uint32_t cur_tail = tail_.load_relaxed(); - } while(true); - while (!ensureEnqueue()) { - - - } - } - -private: - MADRONA_ALWAYS_INLINE bool ensureEnqueue() - { - } - - std::array tickets_; - AtomicU32 head_; - AtomicU32 tail_; - AtomicI32 count_; -}; diff --git a/src/mw/functionality.cpp b/src/mw/functionality.cpp deleted file mode 100644 index 4ad54f1c..00000000 --- a/src/mw/functionality.cpp +++ /dev/null @@ -1,15 +0,0 @@ -/* - * Copyright 2021-2022 Brennan Shacklett and contributors - * - * Use of this source code is governed by an MIT-style - * license that can be found in the LICENSE file or at - * https://opensource.org/licenses/MIT. - */ -#include - -#include "job.hpp" - -void set_val(float *data, uint32_t idx, float v) -{ - data[idx] = v; -} diff --git a/src/physics/CMakeLists.txt b/src/physics/CMakeLists.txt deleted file mode 100644 index cb7ac659..00000000 --- a/src/physics/CMakeLists.txt +++ /dev/null @@ -1,52 +0,0 @@ -set(INC_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../include/madrona") - -set(MADRONA_PHYSICS_SRCS - ${INC_DIR}/physics.hpp ${INC_DIR}/physics.inl physics.cpp - ${INC_DIR}/mesh_bvh.hpp ${INC_DIR}/mesh_bvh.inl - ${INC_DIR}/geo.hpp ${INC_DIR}/geo.inl geo.cpp - narrowphase.cpp broadphase.cpp - xpbd.hpp xpbd.cpp - tgs.hpp tgs.cpp -) - -add_library(madrona_physics STATIC - ${MADRONA_PHYSICS_SRCS} -) - -target_link_libraries(madrona_physics - PUBLIC - madrona_core -) - -add_library(madrona_mw_physics STATIC - ${MADRONA_PHYSICS_SRCS} -) - -target_link_libraries(madrona_mw_physics - PUBLIC - madrona_mw_core -) - -add_library(madrona_physics_assets STATIC - ${INC_DIR}/physics_assets.hpp physics_assets.cpp -) - -target_link_libraries(madrona_physics_assets PRIVATE - madrona_common -) - -add_library(madrona_physics_loader STATIC - ${INC_DIR}/physics_loader.hpp physics_loader.cpp -) - -target_link_libraries(madrona_physics_loader - PRIVATE - madrona_common - PUBLIC - madrona_physics_assets -) - -if (TARGET madrona_cuda) - target_link_libraries(madrona_physics_loader PRIVATE - madrona_cuda) -endif () diff --git a/src/physics/broadphase.cpp b/src/physics/broadphase.cpp deleted file mode 100644 index 77cb1fe3..00000000 --- a/src/physics/broadphase.cpp +++ /dev/null @@ -1,1054 +0,0 @@ -#include -#include - -#include - -#include "physics_impl.hpp" - -namespace madrona::phys::broadphase { - -using namespace base; -using namespace math; -using namespace geo; - -BVH::BVH(const ObjectManager *obj_mgr, - CountT max_leaves, - float leaf_velocity_expansion, - float leaf_accel_expansion) - : nodes_((Node *)rawAlloc(sizeof(Node) * - numInternalNodes(max_leaves))), - num_nodes_(0), - num_allocated_nodes_(numInternalNodes(max_leaves)), - leaf_entities_((Entity *)rawAlloc(sizeof(Entity) * max_leaves)), - obj_mgr_(obj_mgr), // FIXME, get rid of this - leaf_obj_ids_((ObjectID *) - rawAlloc(sizeof(ObjectID) * max_leaves)), - leaf_aabbs_((AABB *)rawAlloc(sizeof(AABB) * max_leaves)), - leaf_transforms_( - (LeafTransform *)rawAlloc(sizeof(LeafTransform) * max_leaves)), - leaf_parents_((uint32_t *)rawAlloc(sizeof(uint32_t) * max_leaves)), - sorted_leaves_((int32_t *)rawAlloc(sizeof(int32_t) * max_leaves)), - num_leaves_(0), - num_allocated_leaves_(max_leaves), - leaf_velocity_expansion_(leaf_velocity_expansion), - leaf_accel_expansion_(leaf_accel_expansion), - force_rebuild_(true) -{} - -CountT BVH::numInternalNodes(CountT num_leaves) const -{ - return std::max(utils::divideRoundUp(num_leaves - 1, CountT(3)), CountT(1)) + - num_leaves; // + num_leaves should not be necessary but the current - // top down build has an issue where leaves get - // unnecessarily split amongst internal nodes with only - // 1 or 2 children -} - -void BVH::rebuild() -{ - int32_t num_internal_nodes = numInternalNodes(num_leaves_.load_relaxed()); - num_nodes_ = num_internal_nodes; - assert(num_nodes_ <= num_allocated_nodes_); - - struct StackEntry { - int32_t nodeID; - int32_t parentID; - int32_t offset; - int32_t numObjs; - }; - - StackEntry stack[64]; - stack[0] = StackEntry { - sentinel_, - sentinel_, - 0, - int32_t(num_leaves_.load_relaxed()), - }; - - int32_t cur_node_offset = 0; - CountT stack_size = 1; - - while (stack_size > 0) { - StackEntry &entry = stack[stack_size - 1]; - int32_t node_id; - if (entry.numObjs <= 4) { - node_id = cur_node_offset++; - - Node &node = nodes_[node_id]; - node.parentID = entry.parentID; - - for (int i = 0; i < 4; i++) { - if (i < entry.numObjs) { - int32_t leaf_id = sorted_leaves_[entry.offset + i]; - - const auto &aabb = leaf_aabbs_[leaf_id]; - leaf_parents_[leaf_id] = - ((uint32_t)node_id << 2) | (uint32_t)i; - - node.setLeaf(i, leaf_id); - node.minX[i] = aabb.pMin.x; - node.minY[i] = aabb.pMin.y; - node.minZ[i] = aabb.pMin.z; - node.maxX[i] = aabb.pMax.x; - node.maxY[i] = aabb.pMax.y; - node.maxZ[i] = aabb.pMax.z; - } else { - node.clearChild(i); - node.minX[i] = FLT_MAX; - node.minY[i] = FLT_MAX; - node.minZ[i] = FLT_MAX; - node.maxX[i] = -FLT_MAX; - node.maxY[i] = -FLT_MAX; - node.maxZ[i] = -FLT_MAX; - } - } - } else if (entry.nodeID == sentinel_) { - node_id = cur_node_offset++; - // Record the node id in the stack entry for when this entry - // is reprocessed - entry.nodeID = node_id; - - Node &node = nodes_[node_id]; - for (CountT i = 0; i < 4; i++) { - node.clearChild(i); - } - node.parentID = entry.parentID; - - // midpoint sort items - auto midpoint_split = [this]( - int32_t base, int32_t num_elems) { - - auto get_center = [this, base](int32_t offset) { - AABB aabb = leaf_aabbs_[sorted_leaves_[base + offset]]; - - return (aabb.pMin + aabb.pMax) / 2.f; - }; - - Vector3 center_min { - FLT_MAX, - FLT_MAX, - FLT_MAX, - }; - - Vector3 center_max { - -FLT_MAX, - -FLT_MAX, - -FLT_MAX, - }; - - for (int i = 0; i < num_elems; i++) { - const Vector3 ¢er = get_center(i); - center_min = Vector3::min(center_min, center); - center_max = Vector3::max(center_max, center); - } - - auto split = [&](auto get_component) { - float split_val = 0.5f * (get_component(center_min) + - get_component(center_max)); - - int start = 0; - int end = num_elems; - - while (start < end) { - while (start < end && - get_component(get_center(start)) < split_val) { - ++start; - } - - while (start < end && get_component( - get_center(end - 1)) >= split_val) { - --end; - } - - if (start < end) { - std::swap(sorted_leaves_[base + start], - sorted_leaves_[base + end - 1]); - ++start; - --end; - } - } - - if (start > 0 && start < num_elems) { - return start; - } else { - return num_elems / 2; - } - }; - - Vector3 center_diff = center_max - center_min; - if (center_diff.x > center_diff.y && - center_diff.x > center_diff.z) { - return split([](Vector3 v) { - return v.x; - }); - } else if (center_diff.y > center_diff.x && - center_diff.y > center_diff.z) { - return split([](Vector3 v) { - return v.y; - }); - } else { - return split([](Vector3 v) { - return v.z; - }); - } - }; - - int32_t second_split = midpoint_split(entry.offset, entry.numObjs); - int32_t num_h1 = second_split; - int32_t num_h2 = entry.numObjs - second_split; - - int32_t first_split = midpoint_split(entry.offset, num_h1); - int32_t third_split = - midpoint_split(entry.offset + second_split, num_h2); - - // Setup stack to recurse into fourths. Put fourths on stack in - // reverse order to preserve left-right depth first ordering - - stack[stack_size++] = { - -1, - entry.nodeID, - entry.offset + num_h1 + third_split, - num_h2 - third_split, - }; - - stack[stack_size++] = { - -1, - entry.nodeID, - entry.offset + num_h1, - third_split, - }; - - stack[stack_size++] = { - -1, - entry.nodeID, - entry.offset + first_split, - num_h1 - first_split, - }; - - stack[stack_size++] = { - -1, - entry.nodeID, - entry.offset, - first_split, - }; - - // Don't finish processing this node until children are processed - continue; - } else { - // Revisiting this node after having processed children - node_id = entry.nodeID; - } - - // At this point, remove the current entry from the stack - stack_size -= 1; - - Node &node = nodes_[node_id]; - if (node.parentID == -1) { - continue; - } - - AABB combined_aabb = AABB::invalid(); - for (CountT i = 0; i < 4; i++) { - if (!node.hasChild(i)) { - break; - } - - combined_aabb = AABB::merge(combined_aabb, AABB { - /* .pMin = */ { - node.minX[i], - node.minY[i], - node.minZ[i], - }, - /* .pMax = */ { - node.maxX[i], - node.maxY[i], - node.maxZ[i], - }, - }); - } - - Node &parent = nodes_[node.parentID]; - CountT child_offset; - for (child_offset = 0; ; child_offset++) { - if (parent.children[child_offset] == sentinel_) { - break; - } - } - - parent.setInternal(child_offset, node_id); - parent.minX[child_offset] = combined_aabb.pMin.x; - parent.minY[child_offset] = combined_aabb.pMin.y; - parent.minZ[child_offset] = combined_aabb.pMin.z; - parent.maxX[child_offset] = combined_aabb.pMax.x; - parent.maxY[child_offset] = combined_aabb.pMax.y; - parent.maxZ[child_offset] = combined_aabb.pMax.z; - } - -#if 0 - { - // validate tree bottom up - int32_t num_leaves = num_leaves_.load_relaxed(); - for (int32_t i = 0; i < num_leaves; i++) { - const AABB &leaf_aabb = leaf_aabbs_[i]; - uint32_t leaf_parent = leaf_parents_[i]; - - int32_t node_idx = int32_t(leaf_parent >> 2_u32); - int32_t sub_idx = int32_t(leaf_parent & 3); - - Node *node = &nodes_[node_idx]; - while (true) { - auto invalid = [&]() { - printf("%d %d %d\n\t(%f %f %f) (%f %f %f)\n\t(%f %f %f) (%f %f %f)\n", - i, node_idx, sub_idx, leaf_aabb.pMin.x, leaf_aabb.pMin.y, leaf_aabb.pMin.z, - leaf_aabb.pMax.x, leaf_aabb.pMax.y, leaf_aabb.pMax.z, - node->minX[sub_idx], node->minY[sub_idx], node->minZ[sub_idx], - node->maxX[sub_idx], node->maxY[sub_idx], node->maxZ[sub_idx]); - assert(false); - }; - - if (leaf_aabb.pMin.x < node->minX[sub_idx]) { - invalid(); - } - if (leaf_aabb.pMin.y < node->minY[sub_idx]) { - invalid(); - } - if (leaf_aabb.pMin.z < node->minZ[sub_idx]) { - invalid(); - } - - if (leaf_aabb.pMax.x > node->maxX[sub_idx]) { - invalid(); - } - if (leaf_aabb.pMax.y > node->maxY[sub_idx]) { - invalid(); - } - if (leaf_aabb.pMax.z > node->maxZ[sub_idx]) { - invalid(); - } - - int child_idx = node_idx; - node_idx = node->parentID; - if (node_idx == sentinel_) { - break; - } - - node = &nodes_[node_idx]; - - int child_offset = -1; - for (int j = 0; j < 4; j++) { - if (node->children[j] == child_idx) { - child_offset = j; - break; - } - } - sub_idx = child_offset; - }; - } - } - - // Validate top down - { - int32_t stack[128]; - AABB aabb_stack[128]; - stack[0] = 0; - aabb_stack[0] = { - -FLT_MAX, - -FLT_MAX, - -FLT_MAX, - FLT_MAX, - FLT_MAX, - FLT_MAX, - }; - CountT stack_size = 1; - - while (stack_size > 0) { - int32_t stack_idx = --stack_size; - int32_t node_idx = stack[stack_idx]; - madrona::math::AABB parent_aabb = aabb_stack[stack_idx]; - - const Node &node = nodes_[node_idx]; - for (int i = 0; i < 4; i++) { - if (!node.hasChild(i)) { - continue; // Technically this could be break? - } - - AABB child_aabb { - /* .pMin = */ { - node.minX[i], - node.minY[i], - node.minZ[i], - }, - /* .pMax = */ { - node.maxX[i], - node.maxY[i], - node.maxZ[i], - }, - }; - - auto invalid = [&]() { - printf("Invalid top down %d %d (%f %f %f) (%f %f %f) (%f %f %f) (%f %f %f)\n", - node_idx, i, - parent_aabb.pMin.x, parent_aabb.pMin.y, parent_aabb.pMin.z, - parent_aabb.pMax.x, parent_aabb.pMax.y, parent_aabb.pMax.z, - child_aabb.pMin.x, child_aabb.pMin.y, child_aabb.pMin.z, - child_aabb.pMax.x, child_aabb.pMax.y, child_aabb.pMax.z); - assert(false); - }; - - if (child_aabb.pMin.x < parent_aabb.pMin.x) { - invalid(); - } - - if (child_aabb.pMin.y < parent_aabb.pMin.y) { - invalid(); - } - - if (child_aabb.pMin.z < parent_aabb.pMin.z) { - invalid(); - } - - if (child_aabb.pMax.x > parent_aabb.pMax.x) { - invalid(); - } - - if (child_aabb.pMax.y > parent_aabb.pMax.y) { - invalid(); - } - - if (child_aabb.pMax.z > parent_aabb.pMax.z) { - invalid(); - } - - if (node.isLeaf(i)) { - int32_t leaf_idx = node.leafIDX(i); - if (leaf_idx >= num_leaves_.load_relaxed()) { - printf("Out of bounds leaf %u %u %u\n", - i, leaf_idx, num_leaves_.load_relaxed()); - assert(false); - } - } else { - stack[stack_size] = node.children[i]; - aabb_stack[stack_size] = child_aabb; - stack_size += 1; - } - } - } - } -#endif -} - -static inline AABB expandAABBWithMotion( - AABB aabb, - const Vector3 &linear_velocity, - float velocity_expansion_factor, - float accel_expansion_factor) -{ - // FIXME include external velocity -#pragma unroll - for (int32_t i = 0; i < 3; i++) { - float pos_delta = - velocity_expansion_factor * linear_velocity[i]; - - float min_delta = pos_delta - accel_expansion_factor; - float max_delta = pos_delta + accel_expansion_factor; - - if (min_delta < 0.f) { - aabb.pMin[i] += min_delta; - } - if (max_delta > 0.f) { - aabb.pMax[i] += max_delta; - } - } - - return aabb; -} - -void BVH::updateLeafPosition(LeafID leaf_id, - const Vector3 &pos, - const Quat &rot, - const Diag3x3 &scale, - const Vector3 &linear_vel, - const AABB &obj_aabb) -{ - AABB world_aabb = obj_aabb.applyTRS(pos, rot, scale); - AABB expanded_aabb = expandAABBWithMotion(world_aabb, linear_vel, - leaf_velocity_expansion_, - leaf_accel_expansion_); - - leaf_aabbs_[leaf_id.id] = expanded_aabb; - leaf_transforms_[leaf_id.id] = { - pos, - rot, - scale, - }; - sorted_leaves_[leaf_id.id] = leaf_id.id; -} - -AABB BVH::expandLeaf(LeafID leaf_id, - const math::Vector3 &linear_vel) -{ - AABB aabb = leaf_aabbs_[leaf_id.id]; - AABB expanded_aabb = expandAABBWithMotion(aabb, linear_vel, - leaf_velocity_expansion_, leaf_accel_expansion_); - - leaf_aabbs_[leaf_id.id] = expanded_aabb; - - return expanded_aabb; -} - -MADRONA_ALWAYS_INLINE static inline float atomicMinF(float *addr, float value) -{ -#ifdef MADRONA_GPU_MODE - float old; - if (!signbit(value)) { - old = __int_as_float(atomicMin((int *)addr, __float_as_int(value))); - } else { - old = __uint_as_float( - atomicMax((unsigned int *)addr, __float_as_uint(value))); - } - - return old; -#else - AtomicFloatRef a(*addr); - float old = a.load(); - - while (old > value && - !a.compare_exchange_weak(old, value)) - {} - - return old; -#endif -} - -MADRONA_ALWAYS_INLINE static inline float atomicMaxF(float *addr, float value) -{ -#ifdef MADRONA_GPU_MODE - float old; - - // cuda::atomic::fetch_max does not seem to work properly (cuda 11.8) - if (!signbit(value)) { - old = __int_as_float( - atomicMax((int *)addr, __float_as_int(value))); - } else { - old = __uint_as_float( - atomicMin((unsigned int *)addr, __float_as_uint(value))); - } - - return old; -#else - AtomicFloatRef a(*addr); - float old = a.load(); - - while (old < value && - !a.compare_exchange_weak(old, value)) - {} - - return old; -#endif -} - -void BVH::refitLeaf(LeafID leaf_id, const AABB &leaf_aabb) -{ - uint32_t leaf_parent = leaf_parents_[leaf_id.id]; - - int32_t node_idx = int32_t(leaf_parent >> 2_u32); - int32_t sub_idx = int32_t(leaf_parent & 3); - - Node &leaf_node = nodes_[node_idx]; - - { - auto nonAtomicMinF = [](float *ptr, float v) { - AtomicFloatRef a(*ptr); - float old = a.load(); - if (v < old) { - a.store(v); - } - return old; - }; - - auto nonAtomicMaxF = [](float *ptr, float v) { - AtomicFloatRef a(*ptr); - float old = a.load(); - if (v > old) { - a.store(v); - } - return old; - }; - - float x_min_prev = - nonAtomicMinF(&leaf_node.minX[sub_idx], leaf_aabb.pMin.x); - float y_min_prev = - nonAtomicMinF(&leaf_node.minY[sub_idx], leaf_aabb.pMin.y); - float z_min_prev = - nonAtomicMinF(&leaf_node.minZ[sub_idx], leaf_aabb.pMin.z); - float x_max_prev = - nonAtomicMaxF(&leaf_node.maxX[sub_idx], leaf_aabb.pMax.x); - float y_max_prev = - nonAtomicMaxF(&leaf_node.maxY[sub_idx], leaf_aabb.pMax.y); - float z_max_prev = - nonAtomicMaxF(&leaf_node.maxZ[sub_idx], leaf_aabb.pMax.z); - - bool expanded = leaf_aabb.pMin.x < x_min_prev || - leaf_aabb.pMin.y < y_min_prev || - leaf_aabb.pMin.z < z_min_prev || - leaf_aabb.pMax.x > x_max_prev || - leaf_aabb.pMax.y > y_max_prev || - leaf_aabb.pMax.z > z_max_prev; - - if (!expanded) return; - } - - int32_t child_idx = node_idx; - node_idx = leaf_node.parentID; - - while (node_idx != sentinel_) { - Node &node = nodes_[node_idx]; - int child_offset = -1; - for (int j = 0; j < 4; j++) { - if (node.children[j] == child_idx) { - child_offset = j; - break; - } - } - assert(child_offset != -1); - - float x_min_prev = - atomicMinF(&node.minX[child_offset], leaf_aabb.pMin.x); - - float y_min_prev = - atomicMinF(&node.minY[child_offset], leaf_aabb.pMin.y); - - float z_min_prev = - atomicMinF(&node.minZ[child_offset], leaf_aabb.pMin.z); - - float x_max_prev = - atomicMaxF(&node.maxX[child_offset], leaf_aabb.pMax.x); - - float y_max_prev = - atomicMaxF(&node.maxY[child_offset], leaf_aabb.pMax.y); - - float z_max_prev = - atomicMaxF(&node.maxZ[child_offset], leaf_aabb.pMax.z); - - bool expanded = leaf_aabb.pMin.x < x_min_prev || - leaf_aabb.pMin.y < y_min_prev || - leaf_aabb.pMin.z < z_min_prev || - leaf_aabb.pMax.x > x_max_prev || - leaf_aabb.pMax.y > y_max_prev || - leaf_aabb.pMax.z > z_max_prev; - - if (!expanded) { - break; - } - - child_idx = node_idx; - node_idx = node.parentID; - } -} - -void BVH::updateTree() -{ - if (force_rebuild_) { - force_rebuild_ = false; - rebuild(); - } - //rebuild(); -} - -Entity BVH::traceRay(Vector3 o, - Vector3 d, - float *out_hit_t, - Vector3 *out_hit_normal, - float t_max) -{ - using namespace math; - - Diag3x3 inv_d = Diag3x3::fromVec(d).inv(); - - int32_t stack[32]; - stack[0] = 0; - CountT stack_size = 1; - - Entity closest_hit_entity = Entity::none(); - Vector3 closest_hit_normal; - - while (stack_size > 0) { - int32_t node_idx = stack[--stack_size]; - const Node &node = nodes_[node_idx]; - for (int i = 0; i < 4; i++) { - if (!node.hasChild(i)) { - continue; // Technically this could be break? - }; - - madrona::math::AABB child_aabb { - /* .pMin = */ { - node.minX[i], - node.minY[i], - node.minZ[i], - }, - /* .pMax = */ { - node.maxX[i], - node.maxY[i], - node.maxZ[i], - }, - }; - - if (child_aabb.rayIntersects(o, inv_d, 0.f, t_max)) { - if (node.isLeaf(i)) { - int32_t leaf_idx = node.leafIDX(i); - - float hit_t; - Vector3 leaf_hit_normal; - bool leaf_hit = traceRayIntoLeaf( - leaf_idx, o, d, 0.f, t_max, &hit_t, &leaf_hit_normal); - - if (leaf_hit) { - t_max = hit_t; - closest_hit_entity = leaf_entities_[leaf_idx]; - closest_hit_normal = leaf_hit_normal; - } - } else { - stack[stack_size++] = node.children[i]; - } - } - } - } - - if (closest_hit_entity == Entity::none()) { - return Entity::none(); - } - - *out_hit_t = t_max; - *out_hit_normal = closest_hit_normal; - return closest_hit_entity; -} - -static inline bool traceRayIntoPlane( - Vector3 ray_o, Vector3 ray_d, - float t_min, float t_max, - float *hit_t, - Vector3 *hit_normal) -{ - // ray_o and ray_d have already been transformed into the space of the - // plane. normal is (0, 0, 1), d is 0 - - float denom = ray_d.z; - - if (denom == 0) { - return false; - } - - float t = -ray_o.z / denom; - - if (t < t_min || t > t_max) { - return false; - } - - *hit_t = t; - *hit_normal = Vector3 { 0, 0, 1 }; - return true; -} - -// RTCD 5.3.8 (modified from segment to ray). Algorithm also in GPU Gems 2. -// Intersect ray r(t)=ray_o + , t_min <= t <=t_max against convex polyhedron -// specified by the n halfspaces defined by the planes p[]. On exit tfirst -// and tlast define the intersection, if any -static inline bool traceRayIntoConvexPolyhedron( - const HalfEdgeMesh &convex_mesh, - Vector3 ray_o, Vector3 ray_d, - float t_min, float t_max, - float *hit_t, - Vector3 *hit_normal) -{ - // Set initial interval based on t_min & t_max. For a ray, tlast should be - // set to +FLT_MAX. For a line, tfirst should also be set to –FLT_MAX - float tfirst = t_min; - float tlast = t_max; - - // Intersect segment against each plane - const CountT num_faces = convex_mesh.numFaces; - - // Our face normals point outside. RTCD uses plane normals pointing inside - // the polyhedron, so signs are flipped relative to the book - - Vector3 closest_normal = Vector3::zero(); - - for (CountT face_idx = 0; face_idx < num_faces; face_idx++) { - Plane plane = convex_mesh.facePlanes[face_idx]; - - float denom = dot(plane.normal, ray_d); - float neg_dist = plane.d - dot(plane.normal, ray_o); - - // Test if segment runs parallel to the plane - if (denom == 0.0f) { - // If so, return “no intersection” if segment lies outside plane - if (neg_dist < 0.0f) return false; - } else { - // Compute parameterized t value for intersection with current plane - float t = neg_dist / denom; - if (denom < 0.0f) { - // When entering halfspace, update tfirst if t is larger - if (t >= tfirst) { - tfirst = t; - closest_normal = plane.normal; - } - } else { - // When exiting halfspace, update tlast if t is smaller - if (t <= tlast) { - tlast = t; - } - } - // Exit with “no intersection” if intersection becomes empty - if (tfirst > tlast) return false; - } - } - - // Addition from RTCD algo: if ray only hits backfacing planes - // we don't set closest_normal and treat this as a miss - if (closest_normal.x == 0 && closest_normal.y == 0 && - closest_normal.z == 0) { - return false; - } - - *hit_t = tfirst; - *hit_normal = closest_normal; - - return true; -} - -bool BVH::traceRayIntoLeaf(int32_t leaf_idx, - math::Vector3 world_ray_o, - math::Vector3 world_ray_d, - float t_min, - float t_max, - float *hit_t, - math::Vector3 *hit_normal) -{ - ObjectID obj_id = leaf_obj_ids_[leaf_idx]; - LeafTransform leaf_txfm = leaf_transforms_[leaf_idx]; - - Quat rot_to_local = leaf_txfm.rot.inv(); - - Vector3 obj_ray_o = rot_to_local.rotateVec(world_ray_o - leaf_txfm.pos); - obj_ray_o.x /= leaf_txfm.scale.d0; - obj_ray_o.y /= leaf_txfm.scale.d1; - obj_ray_o.z /= leaf_txfm.scale.d2; - - Vector3 obj_ray_d = leaf_txfm.rot.inv().rotateVec(world_ray_d); - obj_ray_d.x /= leaf_txfm.scale.d0; - obj_ray_d.y /= leaf_txfm.scale.d1; - obj_ray_d.z /= leaf_txfm.scale.d2; - - auto inv_obj_ray_d = Diag3x3::fromVec(1.f / obj_ray_d); - - Vector3 obj_hit_normal; - - CountT prim_offset = obj_mgr_->rigidBodyPrimitiveOffsets[obj_id.idx]; - CountT num_prims = obj_mgr_->rigidBodyPrimitiveCounts[obj_id.idx]; - - bool hit_leaf = false; - for (CountT i = 0; i < (CountT)num_prims; i++) { - CountT prim_idx = prim_offset + i; - - AABB prim_aabb = obj_mgr_->primitiveAABBs[prim_idx]; - if (!prim_aabb.rayIntersects(obj_ray_o, inv_obj_ray_d, 0.f, t_max)) { - continue; - } - - bool hit_prim; - - const CollisionPrimitive *prim = - &obj_mgr_->collisionPrimitives[prim_idx]; - switch (prim->type) { - case CollisionPrimitive::Type::Hull: { - hit_prim = traceRayIntoConvexPolyhedron(prim->hull.halfEdgeMesh, - obj_ray_o, obj_ray_d, t_min, t_max, hit_t, &obj_hit_normal); - } break; - case CollisionPrimitive::Type::Plane: { - hit_prim = traceRayIntoPlane( - obj_ray_o, obj_ray_d, t_min, t_max, hit_t, &obj_hit_normal); - } break; - case CollisionPrimitive::Type::Sphere: { - assert(false); - } break; - default: MADRONA_UNREACHABLE(); - } - - if (hit_prim) { - hit_leaf = true; - t_max = *hit_t; - } - } - - if (hit_leaf) { - *hit_normal = leaf_txfm.rot.rotateVec(obj_hit_normal); - - return true; - } else { - return false; - } -} - -inline void updateLeafPositionsEntry( - Context &ctx, - const LeafID &leaf_id, - const Position &pos, - const Rotation &rot, - const Scale &scale, - const ObjectID &obj_id, - const Velocity &vel) -{ - BVH &bvh = ctx.singleton(); - ObjectManager &obj_mgr = *ctx.singleton().mgr; - AABB obj_aabb = obj_mgr.rigidBodyAABBs[obj_id.idx]; - - bvh.updateLeafPosition(leaf_id, pos, rot, scale, vel.linear, obj_aabb); -} - -// FIXME currently unused -inline void expandLeavesEntry( - Context &ctx, - const LeafID &leaf_id, - const Velocity &vel) -{ - BVH &bvh = ctx.singleton(); - AABB expanded = bvh.expandLeaf(leaf_id, vel.linear); - bvh.refitLeaf(leaf_id, expanded); -} - -inline void updateBVHEntry(Context &, BVH &bvh) -{ - bvh.updateTree(); -} - -inline void refitEntry(Context &ctx, LeafID leaf_id) -{ - BVH &bvh = ctx.singleton(); - bvh.refitLeaf(leaf_id, bvh.getLeafAABB(leaf_id)); -} - -inline void findIntersectingEntry( - Context &ctx, - const Entity &e, - LeafID leaf_id) -{ - BVH &bvh = ctx.singleton(); - ObjectManager &obj_mgr = *ctx.singleton().mgr; - - // FIXME: should have a flag for passing this - // directly into the system - Loc a_loc = ctx.loc(e); - bool a_is_static = - ctx.getDirect(RGDCols::ResponseType, a_loc) == - ResponseType::Static; - - ObjectID a_obj = ctx.getDirect(RGDCols::ObjectID, a_loc); - - CountT a_num_prims = obj_mgr.rigidBodyPrimitiveCounts[a_obj.idx]; - - bvh.findLeafIntersecting(leaf_id, [&](Entity intersecting_entity) { - if (e.id < intersecting_entity.id) { - Loc b_loc = ctx.loc(intersecting_entity); - - // FIXME: Change this so static objects are kept in a separate BVH - // and this check can be removed. - if (a_is_static && - ctx.getDirect(RGDCols::ResponseType, b_loc) == - ResponseType::Static) { - return; - } - - // We don't expand the primitive AABBs by movement (only object - // AABBs) so we just unconditionally emit narrowphase checks - // between each pair of primitives in the entity. Narrowphase - // will check transformed AABBs. - - ObjectID b_obj = ctx.getDirect(RGDCols::ObjectID, b_loc); - CountT b_num_prims = - obj_mgr.rigidBodyPrimitiveCounts[b_obj.idx]; - - // FIXME: would be nice to be able to make N temporaries all at - // once - - CountT total_narrowphase_checks = a_num_prims * b_num_prims; - - for (CountT prim_check_idx = 0; - prim_check_idx < total_narrowphase_checks; - prim_check_idx++) { - CountT a_prim_idx = prim_check_idx / b_num_prims; - CountT b_prim_idx = prim_check_idx % b_num_prims; - - Loc candidate_loc = ctx.makeTemporary(); - CandidateCollision &candidate = - ctx.getDirect( - RGDCols::CandidateCollision, candidate_loc); - - candidate.a = a_loc; - candidate.b = b_loc; - candidate.aPrim = a_prim_idx; - candidate.bPrim = b_prim_idx; - } - } - }); -} - -TaskGraphNodeID setupBVHTasks( - TaskGraphBuilder &builder, - Span deps) -{ - auto update_leaves = - builder.addToGraph>(deps); - - auto bvh_update = builder.addToGraph>({update_leaves}); - - // FIXME Unfortunately need to call refit here, because update - // won't necessarily do anything - auto refit = builder.addToGraph>({bvh_update}); - - return refit; -} - -TaskGraphNodeID setupPreIntegrationTasks( - TaskGraphBuilder &builder, - Span deps) -{ - auto find_intersects = builder.addToGraph>(deps); - - return find_intersects; -} - -TaskGraphNodeID setupPostIntegrationTasks( - TaskGraphBuilder &builder, - Span deps) -{ -#if 0 - auto expand_leaves = builder.addToGraph>({deps}); -#endif - - // FIXME: can we avoid doing a full tree refit here? - auto update_leaves = - builder.addToGraph>(deps); - - auto refit = builder.addToGraph>({update_leaves}); - - return refit; -} - -} diff --git a/src/physics/narrowphase.cpp b/src/physics/narrowphase.cpp deleted file mode 100644 index 4877673c..00000000 --- a/src/physics/narrowphase.cpp +++ /dev/null @@ -1,1964 +0,0 @@ -#include -#include -#include - -#include "physics_impl.hpp" - -#ifdef MADRONA_GPU_MODE -#include -#include -//#define COUNT_GPU_CLOCKS -#endif - -#ifdef COUNT_GPU_CLOCKS -#define MADRONA_COUNT_CLOCKS -extern "C" { -AtomicU64 narrowphaseAllClocks = 0; -AtomicU64 narrowphaseFetchWorldClocks = 0; -AtomicU64 narrowphaseSetupClocks = 0; -AtomicU64 narrowphasePrepClocks = 0; -AtomicU64 narrowphaseSwitchClocks = 0; -AtomicU64 narrowphaseSATFaceClocks = 0; -AtomicU64 narrowphaseSATEdgeClocks = 0; -AtomicU64 narrowphaseSATPlaneClocks = 0; -AtomicU64 narrowphaseSATContactClocks = 0; -AtomicU64 narrowphaseSATPlaneContactClocks = 0; -AtomicU64 narrowphaseSaveContactsClocks = 0; -AtomicU64 narrowphaseTxfmHullCtrs = 0; -AtomicU64 narrowphaseSATFinishClocks = 0; -} -#endif - -#ifdef MADRONA_COUNT_CLOCKS - -class ClockHelper { -public: - inline ClockHelper(AtomicU64 &counter) - : counter_(&counter) - { - cuda::atomic_thread_fence(cuda::memory_order_seq_cst, - cuda::thread_scope_thread); - start_ = timestamp(); - } - - inline void end() - { - cuda::atomic_thread_fence(cuda::memory_order_seq_cst, - cuda::thread_scope_thread); - auto end = timestamp(); - counter_->fetch_add_relaxed(end - start_); - counter_ = nullptr; - } - - inline ~ClockHelper() - { - if (counter_ != nullptr) { - end(); - } - } - -private: - inline uint64_t timestamp() const - { - uint64_t v; - asm volatile("mov.u64 %0, %%globaltimer;" - : "=l"(v)); - return v; - } - - AtomicU64 *counter_; - uint64_t start_; -}; - -#define PROF_START(name, counter) \ - ClockHelper name(counter) - -#define PROF_END(name) name.end() - -#endif - -#ifndef PROF_START -#define PROF_START(name, counter) -#define PROF_END(name) -#endif - -// Unconditionally disable GPU narrowphase version -#undef MADRONA_GPU_MODE -#undef MADRONA_GPU_COND -#define MADRONA_GPU_COND(...) - -namespace madrona::phys::narrowphase { - -using namespace base; -using namespace math; -using namespace geo; - -enum class NarrowphaseTest : uint32_t { - SphereSphere = 1, - HullHull = 2, - SphereHull = 3, - PlanePlane = 4, - SpherePlane = 5, - HullPlane = 6, -}; - -struct FaceQuery { - float separation; - CountT faceIdx; - Plane plane; -}; - -struct EdgeQuery { - float separation; - math::Vector3 normal; - int32_t edgeIdxA; - int32_t edgeIdxB; -}; - -struct HullState { - HalfEdgeMesh mesh; - Vector3 center; -}; - -struct Manifold { - math::Vector3 contactPoints[4]; - float penetrationDepths[4]; - int32_t numContactPoints; - math::Vector3 normal; -}; - -enum class ContactType { - None, - Sphere, - SATPlane, - SATFace, - SATEdge, -}; - -struct SphereContact { - Vector3 normal; - Vector3 pt; - float depth; -}; - -struct SATContact { - Vector3 normal; - float planeDOrSeparation; - uint32_t refFaceIdxOrEdgeIdxA; - uint32_t incidentFaceIdxOrEdgeIdxB; -}; - -static HullState makeHullState( - MADRONA_GPU_COND(const int32_t mwgpu_lane_id,) - const HalfEdgeMesh &mesh, - Vector3 translation, - Quat rotation, - Diag3x3 scale, - math::Vector3 *dst_vertices, - Plane *dst_planes) -{ - Mat3x3 unscaled_rot = Mat3x3::fromQuat(rotation); - Mat3x3 vertex_txfm = unscaled_rot * scale; - Mat3x3 normal_txfm = unscaled_rot * scale.inv(); - -#ifdef MADRONA_GPU_MODE - const CountT start_offset = mwgpu_lane_id; - constexpr CountT elems_per_iter = mwGPU::numWarpThreads; -#else - constexpr CountT start_offset = 0; - constexpr CountT elems_per_iter = 1; -#endif - - // FIXME: get rid of this center computation - store the offset from - // COM in each CollisionPrimitive and translate it - Vector3 center = Vector3::zero(); - - const CountT num_vertices = mesh.numVertices; - for (CountT i = start_offset; i < num_vertices; i += elems_per_iter) { - Vector3 world_pos = vertex_txfm * mesh.vertices[i] + translation; - dst_vertices[i] = world_pos; - center += world_pos; - } - - center /= num_vertices; - - // FIXME: could significantly optimize this with a uniform scale - // version - const CountT num_faces = mesh.numFaces; - for (CountT i = start_offset; i < num_faces; i += elems_per_iter) { - Plane obj_plane = mesh.facePlanes[i]; - Vector3 plane_origin = - vertex_txfm * (obj_plane.normal * obj_plane.d) + translation; - - Vector3 txfmed_normal = (normal_txfm * obj_plane.normal).normalize(); - float new_d = dot(txfmed_normal, plane_origin); - - dst_planes[i] = { - txfmed_normal, - new_d, - }; - - // Center should be behind each face plane (otherwise face normals - // are probably facing the wrong way) - this is too tight a loop - // to have this assert running all the time though - //assert(center.dot(txfmed_normal) - new_d < 0.0f); - } - - HalfEdgeMesh new_mesh { - .halfEdges = mesh.halfEdges, - .faceBaseHalfEdges = mesh.faceBaseHalfEdges, - .facePlanes = dst_planes, - .vertices = dst_vertices, - .numHalfEdges = mesh.numHalfEdges, - .numFaces = uint32_t(num_faces), - .numVertices = uint32_t(num_vertices), - }; - - return HullState { - new_mesh, - center, - }; - -} - -// Returns the signed distance -static inline float getDistanceFromPlane( - const Plane &plane, const Vector3 &a) -{ - float adotn = dot(a, plane.normal); - return adotn - plane.d; -} - -// Get intersection on plane of the line passing through 2 points -inline math::Vector3 planeIntersection(const Plane &plane, const math::Vector3 &p1, const math::Vector3 &p2) { - float distance = getDistanceFromPlane(plane, p1); - - return p1 + (p2 - p1) * (-distance / plane.normal.dot(p2 - p1)); -} - -#ifdef MADRONA_GPU_MODE - -MADRONA_ALWAYS_INLINE static inline std::pair -warpFloatMaxAndIdx(float val, int32_t idx) -{ -#pragma unroll - for (int32_t w = 16; w >= 1; w /= 2) { - float other_val = - __shfl_xor_sync(mwGPU::allActive, val, w); - int32_t other_idx = - __shfl_xor_sync(mwGPU::allActive, idx, w); - - if (other_val > val) { - val = other_val; - idx = other_idx; - } - } - - return { val, idx }; -} - -MADRONA_ALWAYS_INLINE static inline std::pair -warpFloatMinAndIdx(float val, int32_t idx) -{ -#pragma unroll - for (int32_t w = 16; w >= 1; w /= 2) { - float other_val = - __shfl_xor_sync(mwGPU::allActive, val, w); - int32_t other_idx = - __shfl_xor_sync(mwGPU::allActive, idx, w); - - if (other_val < val) { - val = other_val; - idx = other_idx; - } - } - - return { val, idx }; -} - -MADRONA_ALWAYS_INLINE static inline float warpFloatMin(float val) -{ -#pragma unroll - for (int32_t w = 16; w >= 1; w /= 2) { - float other_val = - __shfl_xor_sync(mwGPU::allActive, val, w); - - if (other_val < val) { - val = other_val; - } - } - - return val; -} - -#endif - -static float getHullDistanceFromPlane( - MADRONA_GPU_COND(const int32_t mwgpu_lane_id,) - const Plane &plane, const HullState &h) -{ -#ifdef MADRONA_GPU_MODE - constexpr CountT elems_per_iter = 32; -#else - constexpr CountT elems_per_iter = 1; -#endif - - float min_dot_n = FLT_MAX; - - auto computeVertexDotN = [&h, &plane](CountT vert_idx) { - Vector3 vertex = h.mesh.vertices[vert_idx]; - return dot(vertex, plane.normal); - }; - - const CountT num_verts = (CountT)h.mesh.numVertices; - for (int32_t offset = 0; offset < num_verts; offset += elems_per_iter) { -#ifdef MADRONA_GPU_MODE - int32_t vert_idx = offset + mwgpu_lane_id; - float cur_dot; - if (vert_idx < num_verts) { - cur_dot = computeVertexDotN(vert_idx); - } else { - cur_dot = FLT_MAX; - } -#else - float cur_dot = computeVertexDotN(offset); -#endif - - if (cur_dot < min_dot_n) { - min_dot_n = cur_dot; - } - } - -#ifdef MADRONA_GPU_MODE - min_dot_n = warpFloatMin(min_dot_n); -#endif - - return min_dot_n - plane.d; -} - -static FaceQuery queryFaceDirections( - MADRONA_GPU_COND(int32_t mwgpu_lane_id,) - const HullState &a, const HullState &b) -{ - Plane max_face_plane; - CountT max_dist_face = -1; - float max_dist = -FLT_MAX; - - const CountT num_a_faces = (CountT)a.mesh.numFaces; - for (CountT face_idx = 0; face_idx < num_a_faces; face_idx++) { - Plane plane = a.mesh.facePlanes[face_idx]; - float face_dist = getHullDistanceFromPlane( - MADRONA_GPU_COND(mwgpu_lane_id,) plane, b); - - if (face_dist > max_dist) { - max_dist = face_dist; - max_dist_face = face_idx; - max_face_plane = plane; - - if (max_dist > 0) { - break; - } - } - } - - return { max_dist, max_dist_face, max_face_plane }; -} - -static bool isMinkowskiFace( - const math::Vector3 &a, const math::Vector3 &b, - const math::Vector3 &c, const math::Vector3 &d) -{ - math::Vector3 bxa = b.cross(a); - math::Vector3 dxc = d.cross(c); - - float cba = c.dot(bxa); - float dba = d.dot(bxa); - float adc = a.dot(dxc); - float bdc = b.dot(dxc); - - return cba * dba < 0.0f && adc * bdc < 0.0f && cba * bdc > 0.0f; -} - -static inline std::pair getEdgeNormals( - const HalfEdgeMesh &mesh, HalfEdge cur_hedge, HalfEdge twin_hedge) -{ - Vector3 normal1 = mesh.facePlanes[cur_hedge.face].normal; - Vector3 normal2 = mesh.facePlanes[twin_hedge.face].normal; - - return { normal1, normal2 }; -} - -static inline Segment getEdgeSegment(const Vector3 *vertices, - const HalfEdge *hedges, - HalfEdge start) -{ - Vector3 a = vertices[start.rootVertex]; - - // FIXME: probably should put both vertex indices inline in the half edge - Vector3 b = vertices[hedges[start.next].rootVertex]; - - return { a, b }; -} - -static inline bool buildsMinkowskiFace( - const HalfEdgeMesh &a_mesh, const HalfEdgeMesh &b_mesh, - HalfEdge cur_hedge_a, HalfEdge twin_hedge_a, - HalfEdge cur_hedge_b, HalfEdge twin_hedge_b) -{ - auto [aNormal1, aNormal2] = - getEdgeNormals(a_mesh, cur_hedge_a, twin_hedge_a); - auto [bNormal1, bNormal2] = - getEdgeNormals(b_mesh, cur_hedge_b, twin_hedge_b); - - return isMinkowskiFace(aNormal1, aNormal2, -bNormal1, -bNormal2); -} - -struct EdgeTestResult { - Vector3 normal; - float separation; -}; - -static inline EdgeTestResult edgeDistance( - const HullState &a, const HullState &b, - HalfEdge hedge_a, HalfEdge hedge_b) -{ - Segment segment_a = - getEdgeSegment(a.mesh.vertices, a.mesh.halfEdges, hedge_a); - Segment segment_b = - getEdgeSegment(b.mesh.vertices, b.mesh.halfEdges, hedge_b); - - Vector3 dir_a = segment_a.p2 - segment_a.p1; - Vector3 dir_b = segment_b.p2 - segment_b.p1; - - Vector3 unnormalized_cross = dir_a.cross(dir_b); - float normal_len2 = unnormalized_cross.length2(); - - if (normal_len2 == 0) { - EdgeTestResult result; - result.separation = -FLT_MAX; - - return result; - } - - float inv_normal_len = -#ifdef MADRONA_GPU_MODE - rsqrtf(normal_len2); -#else - 1.f / sqrtf(normal_len2); -#endif - - math::Vector3 normal = unnormalized_cross * inv_normal_len; - - if (normal.dot(segment_a.p1 - a.center) < 0.0f) { - normal = -normal; - } - - float separation = normal.dot(segment_b.p1 - segment_a.p1); - - return { - normal, - separation, - }; -} - -static EdgeQuery queryEdgeDirections( - MADRONA_GPU_COND(int32_t mwgpu_lane_id,) - const HullState &a, const HullState &b) -{ - Vector3 normal {}; - int edgeAMaxDistance = 0; - int edgeBMaxDistance = 0; - float maxDistance = -FLT_MAX; - - auto testEdgeSeparation = [&a, &b](uint32_t hedge_idx_a, - uint32_t hedge_idx_b) { - HalfEdge cur_hedge_a = a.mesh.halfEdges[hedge_idx_a]; - HalfEdge twin_hedge_a = a.mesh.halfEdges[a.mesh.twinIDX(hedge_idx_a)]; - HalfEdge cur_hedge_b = b.mesh.halfEdges[hedge_idx_b]; - HalfEdge twin_hedge_b = b.mesh.halfEdges[b.mesh.twinIDX(hedge_idx_b)]; - - if (buildsMinkowskiFace(a.mesh, b.mesh, cur_hedge_a, twin_hedge_a, - cur_hedge_b, twin_hedge_b)) { - return edgeDistance(a, b, cur_hedge_a, cur_hedge_b); - } else { - EdgeTestResult result; - result.separation = -FLT_MAX; - return result; - } - }; - - const CountT a_num_edges = a.mesh.numEdges(); - const CountT b_num_edges = b.mesh.numEdges(); - -#ifdef MADRONA_GPU_MODE - const int32_t num_edge_tests = a_num_edges * b_num_edges; - for (int32_t edge_offset_linear = 0; edge_offset_linear < num_edge_tests; - edge_offset_linear += 32) { - int32_t edge_idx_linear = edge_offset_linear + mwgpu_lane_id; - - // FIXME: get rid of this level of indirection - int32_t edge_idx_a = edge_idx_linear / b_num_edges; - int32_t edge_idx_b = edge_idx_linear % b_num_edges; - - EdgeTestResult edge_cmp; - int32_t he_a_idx; - int32_t he_b_idx; - - if (edge_idx_linear >= num_edge_tests ) { - edge_cmp.separation = -FLT_MAX; - } else { - he_a_idx = a.mesh.edgeToHalfEdge(edge_idx_a); - he_b_idx = b.mesh.edgeToHalfEdge(edge_idx_b); - - edge_cmp = testEdgeSeparation(he_a_idx, he_b_idx); - } - - if (edge_cmp.separation > maxDistance) { - maxDistance = edge_cmp.separation; - normal = edge_cmp.normal; - edgeAMaxDistance = he_a_idx; - edgeBMaxDistance = he_b_idx; - } - - if (__ballot_sync(mwGPU::allActive, maxDistance > 0) != 0) { - break; - } - } - - int32_t max_lane_idx; - std::tie(maxDistance, max_lane_idx) = - warpFloatMaxAndIdx(maxDistance, mwgpu_lane_id); - - normal.x = __shfl_sync(mwGPU::allActive, normal.x, max_lane_idx); - normal.y = __shfl_sync(mwGPU::allActive, normal.y, max_lane_idx); - normal.z = __shfl_sync(mwGPU::allActive, normal.z, max_lane_idx); - - edgeAMaxDistance = __shfl_sync(mwGPU::allActive, - edgeAMaxDistance, max_lane_idx); - edgeBMaxDistance = __shfl_sync(mwGPU::allActive, - edgeBMaxDistance, max_lane_idx); - -#else - for (CountT edge_idx_a = 0; edge_idx_a < a_num_edges; edge_idx_a++) { - int32_t he_idx_a = a.mesh.edgeToHalfEdge(edge_idx_a); - for (CountT edge_idx_b = 0; edge_idx_b < b_num_edges; edge_idx_b++) { - int32_t he_idx_b = b.mesh.edgeToHalfEdge(edge_idx_b); - - EdgeTestResult edge_cmp = testEdgeSeparation(he_idx_a, he_idx_b); - - if (edge_cmp.separation > maxDistance) { - maxDistance = edge_cmp.separation; - normal = edge_cmp.normal; - edgeAMaxDistance = he_idx_a; - edgeBMaxDistance = he_idx_b; - - if (maxDistance > 0) { - // FIXME: this goto probably kills autovectorization - goto early_out; - } - } - } - } - - early_out: -#endif - - return { maxDistance, normal, edgeAMaxDistance, edgeBMaxDistance }; -} - -static CountT findIncidentFace(MADRONA_GPU_COND(int32_t mwgpu_lane_id,) - const HullState &h, Vector3 ref_normal) -{ -#ifdef MADRONA_GPU_MODE - constexpr CountT elems_per_iter = 32; -#else - constexpr CountT elems_per_iter = 1; -#endif - - auto computeFaceDotRef = [&h, ref_normal](CountT face_idx) { - Plane face_plane = h.mesh.facePlanes[face_idx]; - return dot(face_plane.normal, ref_normal); - }; - - float min_dot = FLT_MAX; - CountT minimizing_face = -1; - - const CountT num_faces = (CountT)h.mesh.numFaces; - for (CountT offset = 0; offset < num_faces; offset += elems_per_iter) { -#ifdef MADRONA_GPU_MODE - const CountT face_idx = offset + mwgpu_lane_id; - - float face_dot_ref; - if (face_idx < num_faces) { - face_dot_ref = computeFaceDotRef(face_idx); - } else{ - face_dot_ref = FLT_MAX; - } -#else - const CountT face_idx = offset; - float face_dot_ref = computeFaceDotRef(face_idx); -#endif - - if (face_dot_ref < min_dot) { - min_dot = face_dot_ref; - minimizing_face = face_idx; - } - } - -#ifdef MADRONA_GPU_MODE - std::tie(min_dot, minimizing_face) = - warpFloatMinAndIdx(min_dot, minimizing_face); -#endif - - assert(minimizing_face != -1); - return minimizing_face; -} - -static inline CountT clipPolygon(Vector3 *dst_vertices, - Plane clipping_plane, - const Vector3 *input_vertices, - CountT num_input_vertices) -{ - CountT num_new_vertices = 0; - - Vector3 v1 = input_vertices[num_input_vertices - 1]; - float d1 = getDistanceFromPlane(clipping_plane, v1); - - for (CountT i = 0; i < num_input_vertices; ++i) { - Vector3 v2 = input_vertices[i]; - float d2 = getDistanceFromPlane(clipping_plane, v2); - - if (d1 <= 0.0f && d2 <= 0.0f) { - // Both vertices are behind the plane, keep the second vertex - dst_vertices[num_new_vertices++] = v2; - } - else if (d1 <= 0.0f && d2 > 0.0f) { - // v1 is behind the plane, the other is in front (out) - Vector3 intersection = planeIntersection(clipping_plane, v1, v2); - dst_vertices[num_new_vertices++] = intersection; - } - else if (d2 <= 0.0f && d1 > 0.0f) { - math::Vector3 intersection = planeIntersection(clipping_plane, v1, v2); - dst_vertices[num_new_vertices++] = intersection; - dst_vertices[num_new_vertices++] = v2; - } - - // Now use v2 as the starting vertex - v1 = v2; - d1 = d2; - } - - return num_new_vertices; -} - -struct SATResult { - ContactType type; - SATContact contact; -}; - -static inline SATResult doSAT(MADRONA_GPU_COND(int32_t mwgpu_lane_id,) - const HullState &a, const HullState &b) -{ - PROF_START(sat_face_ctr, narrowphaseSATFaceClocks); - - FaceQuery faceQueryA = - queryFaceDirections(MADRONA_GPU_COND(mwgpu_lane_id,) a, b); - if (faceQueryA.separation > 0.0f) { - // There is a separating axis - no collision - SATResult result; - result.type = ContactType::None; - - return result; - } - - FaceQuery faceQueryB = - queryFaceDirections(MADRONA_GPU_COND(mwgpu_lane_id,) b, a); - if (faceQueryB.separation > 0.0f) { - // There is a separating axis - no collision - SATResult result; - result.type = ContactType::None; - - return result; - } - - PROF_END(sat_face_ctr); - PROF_START(sat_edge_ctr, narrowphaseSATEdgeClocks); - - EdgeQuery edgeQuery = - queryEdgeDirections(MADRONA_GPU_COND(mwgpu_lane_id,) a, b); - if (edgeQuery.separation > 0.0f) { - // There is a separating axis - no collision - SATResult result; - result.type = ContactType::None; - - return result; - } - - PROF_END(sat_edge_ctr); - - PROF_START(sat_finish_ctr, narrowphaseSATFinishClocks); - - bool bIsFaceContactA = faceQueryA.separation > edgeQuery.separation; - bool bIsFaceContactB = faceQueryB.separation > edgeQuery.separation; - - if (bIsFaceContactA || bIsFaceContactB) { - bool a_is_ref = faceQueryA.separation >= faceQueryB.separation; - - Plane ref_plane = a_is_ref ? faceQueryA.plane : faceQueryB.plane; - CountT ref_face_idx = - a_is_ref ? faceQueryA.faceIdx : faceQueryB.faceIdx; - const HullState &incident_hull = a_is_ref ? b : a; - - // Find incident face - CountT incident_face_idx = findIncidentFace( - MADRONA_GPU_COND(mwgpu_lane_id,) incident_hull, ref_plane.normal); - - SATResult result; - result.type = ContactType::SATFace, - result.contact.normal = ref_plane.normal; - result.contact.planeDOrSeparation = ref_plane.d; - uint32_t mask; - if (a_is_ref) { - mask = 0_u32; - } else { - mask = 1_u32 << 31_u32; - } - result.contact.refFaceIdxOrEdgeIdxA = uint32_t(ref_face_idx) | mask; - result.contact.incidentFaceIdxOrEdgeIdxB = uint32_t(incident_face_idx); - - return result; - } else { - SATResult result; - result.type = ContactType::SATEdge; - result.contact.normal = edgeQuery.normal; - result.contact.planeDOrSeparation = edgeQuery.separation; - result.contact.refFaceIdxOrEdgeIdxA = edgeQuery.edgeIdxA; - result.contact.incidentFaceIdxOrEdgeIdxB = edgeQuery.edgeIdxB; - return result; - } -} - -SATResult doSATPlane(MADRONA_GPU_COND(const int32_t mwgpu_lane_id,) - const Plane &plane, const HullState &h) -{ - PROF_START(sat_plane_ctr, narrowphaseSATPlaneClocks); - - float separation = getHullDistanceFromPlane( - MADRONA_GPU_COND(mwgpu_lane_id,) plane, h); - - if (separation > 0.0f) { - SATResult result; - result.type = ContactType::None; - - return result; - } - - PROF_START(sat_finish_ctr, narrowphaseSATFinishClocks); - - // Find incident face - CountT incident_face_idx = findIncidentFace( - MADRONA_GPU_COND(mwgpu_lane_id,) h, plane.normal); - - SATResult result; - result.type = ContactType::SATPlane; - result.contact.normal = plane.normal; - result.contact.planeDOrSeparation = plane.d; - result.contact.incidentFaceIdxOrEdgeIdxB = uint32_t(incident_face_idx); - - return result; -} - -static Manifold buildFaceContactManifold( - Vector3 contact_normal, - Vector3 *contacts, - float *penetration_depths, - CountT num_contacts, - Vector3 world_offset, - Quat to_world_frame) -{ - Manifold manifold; - if (num_contacts <= 4) { - manifold.numContactPoints = num_contacts; - for (CountT i = 0; i < num_contacts; i++) { - manifold.contactPoints[i] = contacts[i]; - manifold.penetrationDepths[i] = penetration_depths[i]; - } - } else { - // Going to select contact manifold comprised of points - // A B C and Q following Gregorious presentation. - - // Select point A as first point in contact list - manifold.numContactPoints = 4; - manifold.contactPoints[0] = contacts[0]; - manifold.penetrationDepths[0] = penetration_depths[0]; - - // Find point B furthest from point A - float max_dist_sq = 0.f; - for (CountT i = 1; i < num_contacts; i++) { - Vector3 cur_contact = contacts[i]; - float dist_sq = manifold.contactPoints[0].distance2(cur_contact); - if (dist_sq > max_dist_sq) { - max_dist_sq = dist_sq; - - manifold.contactPoints[1] = cur_contact; - manifold.penetrationDepths[1] = penetration_depths[i]; - } - } - - math::Vector3 ba = - manifold.contactPoints[1] - manifold.contactPoints[0]; - - // Find point C which maximizes area of triangle ABC - float max_tri_area = 0.0f; - bool max_tri_sign = 0.f; - for (CountT i = 1; i < num_contacts; i++) { - Vector3 cur_contact = contacts[i]; - math::Vector3 bc = cur_contact - manifold.contactPoints[1]; - float signed_area = contact_normal.dot(cross(ba, bc)); - float area = copysignf(signed_area, 1.f); - - if (area > max_tri_area) { - max_tri_area = area; - max_tri_sign = copysignf(1.f, signed_area); - - manifold.contactPoints[2] = cur_contact; - manifold.penetrationDepths[2] = penetration_depths[i]; - } - } - - // If we ultimately selected a triangle ABC with clockwise winding, - // flip around edge BA to make the triangle counterclockwise, so the - // next part only needs to search for negative area. - if (max_tri_sign == -1.f) { - ba = -ba; - std::swap(manifold.contactPoints[0], manifold.contactPoints[1]); - } - - // Select point Q that adds the most area to ABC - // Need to check ABQ (BA x AQ), BCQ (CB x QC), and CAQ (AC x QA) - - Vector3 cb = manifold.contactPoints[2] - manifold.contactPoints[1]; - Vector3 ac = manifold.contactPoints[0] - manifold.contactPoints[2]; - - float most_neg_area = 0.f; - for (CountT i = 1; i < num_contacts; i++) { - Vector3 cur_contact = contacts[i]; - - Vector3 aq = manifold.contactPoints[0] - cur_contact; - Vector3 qc = cur_contact - manifold.contactPoints[2]; - - float abq_area = contact_normal.dot(cross(ba, aq)); - float bcq_area = contact_normal.dot(cross(cb, qc)); - float caq_area = contact_normal.dot(cross(aq, ac)); - - float q_min_area = fminf(abq_area, fminf(bcq_area, caq_area)); - if (q_min_area < most_neg_area) { - most_neg_area = q_min_area; - - manifold.contactPoints[3] = cur_contact; - manifold.penetrationDepths[3] = penetration_depths[i]; - } - } - - if (max_dist_sq == 0.f || max_tri_area == 0.f || most_neg_area == 0.f) { - // FIXME: should not be possible - manifold.numContactPoints = 0; - manifold.normal = Vector3::zero(); - return manifold; - } - } - - for (CountT i = 0; i < (CountT)manifold.numContactPoints; i++) { - manifold.contactPoints[i] = - to_world_frame.rotateVec(manifold.contactPoints[i]) + world_offset; - } - - manifold.normal = to_world_frame.rotateVec(contact_normal); - - return manifold; -} - -MADRONA_ALWAYS_INLINE static inline Manifold createFaceContact( - Plane ref_plane, - int32_t ref_face_idx, - int32_t incident_face_idx, - const Vector3 *ref_vertices, - const Vector3 *other_vertices, - const HalfEdge *ref_hedges, - const HalfEdge *other_hedges, - const uint32_t *ref_face_hedges, - const uint32_t *other_face_hedges, - void *tmp_buf1, void *tmp_buf2, -#ifdef MADRONA_GPU_MODE - Mat3x4 ref_txfm, Mat3x4 other_txfm, -#endif - Vector3 world_offset, Quat to_world_frame) -{ - // Collect incident vertices: FIXME should have face indices - Vector3 *incident_vertices_tmp = (Vector3 *)tmp_buf1; - CountT num_incident_vertices = 0; - { - CountT hedge_idx = other_face_hedges[incident_face_idx]; - CountT start_hedge_idx = hedge_idx; - - do { - const auto &cur_hedge = other_hedges[hedge_idx]; - hedge_idx = cur_hedge.next; - - Vector3 cur_point = other_vertices[cur_hedge.rootVertex]; -#ifdef MADRONA_GPU_MODE - cur_point = other_txfm.txfmPoint(cur_point); -#endif - - incident_vertices_tmp[num_incident_vertices++] = cur_point; - } while (hedge_idx != start_hedge_idx); - } - - Vector3 *clipping_input = incident_vertices_tmp; - CountT num_clipped_vertices = num_incident_vertices; - - Vector3 *clipping_dst = (Vector3 *)tmp_buf2; - - // max output vertices is num_incident_vertices + num planes - // but we don't know num planes ahead of time without iterating - // through the reference face twice! Alternative would be to cache the - // side planes, or store max face size in each mesh. The worst case - // buffer sizes here is just the sum of the max face sizes - 1 - - // FIXME, this code assumes that clipping_input & clipping_dst have space - // to write incident_vertices + num_planes new vertices - // Loop over side planes - { - CountT hedge_idx = ref_face_hedges[ref_face_idx]; - CountT start_hedge_idx = hedge_idx; - - auto *cur_hedge = &ref_hedges[hedge_idx]; - Vector3 cur_point = ref_vertices[cur_hedge->rootVertex]; -#ifdef MADRONA_GPU_MODE - cur_point = ref_txfm.txfmPoint(cur_point); -#endif - do { - hedge_idx = cur_hedge->next; - cur_hedge = &ref_hedges[hedge_idx]; - Vector3 next_point = ref_vertices[cur_hedge->rootVertex]; -#ifdef MADRONA_GPU_MODE - next_point = ref_txfm.txfmPoint(next_point); -#endif - - Vector3 edge = next_point - cur_point; - Vector3 plane_normal = cross(edge, ref_plane.normal); - - float d = dot(plane_normal, cur_point); - cur_point = next_point; - - Plane side_plane { - plane_normal, - d, - }; - - num_clipped_vertices = clipPolygon(clipping_dst, side_plane, - clipping_input, num_clipped_vertices); - - std::swap(clipping_dst, clipping_input); - } while (hedge_idx != start_hedge_idx); - } - - // assert(num_clipped_vertices > 0); - - // clipping_input has the result due to the final swap - - // Filter clipping_input to ones below ref_plane and save penetration depth - float *penetration_depths = (float *)clipping_dst; - - CountT num_below_plane = 0; - for (CountT i = 0; i < num_clipped_vertices; ++i) { - Vector3 vertex = clipping_input[i]; - if (float d = getDistanceFromPlane(ref_plane, vertex); d <= 0.0f) { - // Project the point onto the reference plane - // (d guaranteed to be negative) - clipping_input[num_below_plane] = vertex - d * ref_plane.normal; - penetration_depths[num_below_plane] = -d; - - num_below_plane += 1; - } - } - - return buildFaceContactManifold(ref_plane.normal, clipping_input, - penetration_depths, num_below_plane, - world_offset, to_world_frame); -} - -static Manifold createFacePlaneContact(Plane plane, - int32_t incident_face_idx, - const Vector3 *vertices, - const HalfEdge *hedges, - const uint32_t *face_hedge_roots, - Vector3 *contacts_tmp, - float *penetration_depths_tmp, -#ifdef MADRONA_GPU_MODE - Mat3x4 hull_txfm, -#endif - Vector3 world_offset, - Quat to_world_frame) -{ - // Collect incident vertices: FIXME should have face indices - CountT num_incident_vertices = 0; - { - CountT hedge_idx = face_hedge_roots[incident_face_idx]; - CountT start_hedge_idx = hedge_idx; - - do { - const auto &cur_hedge = hedges[hedge_idx]; - hedge_idx = cur_hedge.next; - Vector3 vertex = vertices[cur_hedge.rootVertex]; - -#ifdef MADRONA_GPU_MODE - vertex = hull_txfm.txfmPoint(vertex); -#endif - - if (float d = getDistanceFromPlane(plane, vertex); d <= 0.0f) { - // Project the point onto the reference plane - // (d guaranteed to be negative) - contacts_tmp[num_incident_vertices] = - vertex - d * plane.normal; - penetration_depths_tmp[num_incident_vertices] = -d; - - num_incident_vertices += 1; - } - } while (hedge_idx != start_hedge_idx); - } - - return buildFaceContactManifold(plane.normal, contacts_tmp, - penetration_depths_tmp, num_incident_vertices, - world_offset, to_world_frame); -} - - -static Segment shortestSegmentBetween(const Segment &seg1, const Segment &seg2) -{ - math::Vector3 v1 = seg1.p2 - seg1.p1; - math::Vector3 v2 = seg2.p2 - seg2.p1; - - math::Vector3 v21 = seg2.p1 - seg1.p1; - - float dotv22 = v2.dot(v2); - float dotv11 = v1.dot(v1); - float dotv21 = v2.dot(v1); - float dotv211 = v21.dot(v1); - float dotv212 = v21.dot(v2); - - float denom = dotv21 * dotv21 - dotv22 * dotv11; - - float s, t; - - // FIXME: validate this epsilon - if (fabsf(denom) < 0.00001f) { - s = 0.0f; - t = (dotv11 * s - dotv211) / dotv21; - } - else { - s = (dotv212 * dotv21 - dotv22 * dotv211) / denom; - t = (-dotv211 * dotv21 + dotv11 * dotv212) / denom; - } - - s = fmaxf(fminf(s, 1.0f), 0.0f); - t = fmaxf(fminf(t, 1.0f), 0.0f); - - return { seg1.p1 + s * v1, seg2.p1 + t * v2 }; -} - -static Manifold createEdgeContact(Segment segA, - Segment segB, - Vector3 normal, - float separation, - Vector3 world_offset, - Quat to_world_frame) -{ -#if 0 - Segment s = shortestSegmentBetween(segA, segB); - Vector3 contact = 0.5f * (s.p1 + s.p2); - float depth = 0.5f * (s.p2 - s.p1).length(); -#endif - - // Deviation from Gregorius GDC 2015: - // Currently the solver expects the contact point to be ON object A. - // For Face-Face this means the point is on object A's face, which is - // the same as the presentation. In the edge-edge case, the presentation - // has the contact point between the edges. - // Our solver currently reconstructs the contact points on point B - // using the normal depth, so the presentation's method will - // reconstruct the wrong contact points. - // FIXME: revisit this after modifying solver to handle multi-point - // manifolds better - - // FIXME: don't need this full function call here. - Segment s = shortestSegmentBetween(segA, segB); - Vector3 contact = s.p1; - - Manifold manifold; - manifold.contactPoints[0] = - to_world_frame.rotateVec(contact) + world_offset, - manifold.penetrationDepths[0] = -separation; - manifold.numContactPoints = 1; - manifold.normal = to_world_frame.rotateVec(normal); - - return manifold; -} - -static Manifold createEdgeContact(Vector3 normal, - float separation, - int32_t hedge_idx_a, - int32_t hedge_idx_b, - const Vector3 *a_vertices, - const Vector3 *b_vertices, - const HalfEdge *a_hedges, - const HalfEdge *b_hedges, -#ifdef MADRONA_GPU_MODE - Vector3 a_pos, Quat a_rot, Diag3x3 a_scale, - Vector3 b_pos, Quat b_rot, Diag3x3 b_scale, -#endif - Vector3 world_offset, - Quat to_world_frame) -{ - Segment segA = getEdgeSegment(a_vertices, a_hedges, - a_hedges[hedge_idx_a]); - Segment segB = getEdgeSegment(b_vertices, b_hedges, - b_hedges[hedge_idx_b]); - -#ifdef MADRONA_GPU_MODE - segA.p1 = a_rot.rotateVec(a_scale * segA.p1) + a_pos; - segA.p2 = a_rot.rotateVec(a_scale * segA.p2) + a_pos; - segB.p1 = b_rot.rotateVec(b_scale * segB.p1) + b_pos; - segB.p2 = b_rot.rotateVec(b_scale * segB.p2) + b_pos; -#endif - - return createEdgeContact(segA, segB, - normal, separation, - world_offset, to_world_frame); -} - -static inline void addManifoldContacts( - Context &ctx, - Manifold manifold, - Loc ref_loc, Loc other_loc) -{ - PROF_START(save_contacts_ctr, narrowphaseSaveContactsClocks); - - const auto &physics_sys = ctx.singleton(); - - Loc c = ctx.makeTemporary(physics_sys.contactArchetypeID); - ctx.getDirect(RGDCols::ContactConstraint, c) = { - ref_loc, - other_loc, - { - Vector4::fromVec3W(manifold.contactPoints[0], - manifold.penetrationDepths[0]), - Vector4::fromVec3W(manifold.contactPoints[1], - manifold.penetrationDepths[1]), - Vector4::fromVec3W(manifold.contactPoints[2], - manifold.penetrationDepths[2]), - Vector4::fromVec3W(manifold.contactPoints[3], - manifold.penetrationDepths[3]), - }, - manifold.numContactPoints, - manifold.normal, - }; -} - -static inline void addSinglePointContact( - Context &ctx, - Vector3 point, - Vector3 normal, - float depth, - Loc ref_loc, - Loc other_loc) -{ - const auto &physics_sys = ctx.singleton(); - - Loc c = ctx.makeTemporary(physics_sys.contactArchetypeID); - - ctx.getDirect(RGDCols::ContactConstraint, c) = { - ref_loc, - other_loc, - { - Vector4::fromVec3W(point, depth), - Vector4::zero(), - Vector4::zero(), - Vector4::zero(), - }, - 1, - normal, - }; -} - -#ifdef MADRONA_GPU_MODE -namespace gpuImpl { -// FIXME: do something actually intelligent here -inline constexpr int32_t maxNumPlanes = 40; -inline constexpr int32_t numPlaneFloats = maxNumPlanes * 4; -} -#endif - -struct NarrowphaseResult { - ContactType type; - SphereContact sphere; - SATContact sat; - const Vector3 *aVertices; - const Vector3 *bVertices; - const HalfEdge *aHalfEdges; - const HalfEdge *bHalfEdges; - const uint32_t *aFaceHedgeRoots; - const uint32_t *bFaceHedgeRoots; -}; - -MADRONA_ALWAYS_INLINE static inline NarrowphaseResult narrowphaseDispatch( - MADRONA_GPU_COND(const int32_t mwgpu_lane_id,) - NarrowphaseTest test_type, - Vector3 a_pos, Vector3 b_pos, - Quat a_rot, Quat b_rot, - Diag3x3 a_scale, Diag3x3 b_scale, - const CollisionPrimitive *a_prim, const CollisionPrimitive *b_prim, - CountT max_num_tmp_vertices, - CountT max_num_tmp_faces, - Vector3 *txfm_vertex_buffer, - Plane *txfm_face_buffer) -{ - PROF_START(switch_body_ctr, narrowphaseSwitchClocks); - - switch (test_type) { - case NarrowphaseTest::SphereSphere: { - float a_radius, b_radius; - { - assert(a_scale.d0 == a_scale.d1 && a_scale.d0 == a_scale.d2); - assert(b_scale.d0 == b_scale.d1 && b_scale.d0 == b_scale.d2); - - a_radius = a_scale.d0 * a_prim->sphere.radius; - b_radius = b_scale.d0 * b_prim->sphere.radius; - } - - Vector3 to_b = b_pos - a_pos; - float dist = to_b.length(); - - if (dist > a_radius + b_radius) { - NarrowphaseResult result; - result.type = ContactType::None; - return result; - } - - Vector3 normal; - float penetration; - - if (dist > 0.f) { - normal = to_b / dist; - } else { - normal = math::up; - } - - penetration = a_radius + b_radius - dist; - - SphereContact contact { - .normal = normal, - .pt = a_pos + a_radius * normal, - .depth = penetration, - }; - - NarrowphaseResult result; - result.type = ContactType::Sphere; - result.sphere = contact; - result.aVertices = nullptr; - result.bVertices = nullptr; - result.aVertices = nullptr; - result.bVertices = nullptr; - result.aHalfEdges = nullptr; - result.bHalfEdges = nullptr; - result.aFaceHedgeRoots = nullptr; - result.bFaceHedgeRoots = nullptr; - return result; - } break; - case NarrowphaseTest::HullHull: { - // Get half edge mesh for hull A and hull B - const auto &a_he_mesh = a_prim->hull.halfEdgeMesh; - const auto &b_he_mesh = b_prim->hull.halfEdgeMesh; - - assert(a_he_mesh.numFaces + b_he_mesh.numFaces < - max_num_tmp_faces); - - assert(a_he_mesh.numVertices + b_he_mesh.numVertices < - max_num_tmp_vertices); - - PROF_START(txfm_hull_ctr, narrowphaseTxfmHullCtrs); - - HullState a_hull_state = makeHullState(MADRONA_GPU_COND(mwgpu_lane_id,) - a_he_mesh, a_pos, a_rot, a_scale, txfm_vertex_buffer, - txfm_face_buffer); - - txfm_vertex_buffer += a_hull_state.mesh.numVertices; - txfm_face_buffer += a_hull_state.mesh.numFaces; - - HullState b_hull_state = makeHullState(MADRONA_GPU_COND(mwgpu_lane_id,) - b_he_mesh, b_pos, b_rot, b_scale, txfm_vertex_buffer, - txfm_face_buffer); - - MADRONA_GPU_COND(__syncwarp(mwGPU::allActive)); - - PROF_END(txfm_hull_ctr); - - const SATResult sat = doSAT(MADRONA_GPU_COND(mwgpu_lane_id,) - a_hull_state, b_hull_state); - - NarrowphaseResult result; - result.type = sat.type; - result.sat = sat.contact; -#ifdef MADRONA_GPU_MODE - result.aVertices = a_he_mesh.vertices; - result.bVertices = b_he_mesh.vertices; -#else - result.aVertices = a_hull_state.mesh.vertices; - result.bVertices = b_hull_state.mesh.vertices; -#endif - result.aHalfEdges = a_hull_state.mesh.halfEdges; - result.bHalfEdges = b_hull_state.mesh.halfEdges; - result.aFaceHedgeRoots = a_hull_state.mesh.faceBaseHalfEdges; - result.bFaceHedgeRoots = b_hull_state.mesh.faceBaseHalfEdges; - - return result; - } break; - case NarrowphaseTest::SphereHull: { - float sphere_radius; - { - auto sphere = a_prim->sphere; - assert(a_scale.d0 == a_scale.d1 && a_scale.d0 == a_scale.d2); - sphere_radius = a_scale.d0 * sphere.radius; - } - - const auto &b_he_mesh = b_prim->hull.halfEdgeMesh; - assert(b_he_mesh.numFaces < max_num_tmp_faces); - assert(b_he_mesh.numVertices < max_num_tmp_vertices); - - PROF_START(txfm_hull_ctr, narrowphaseTxfmHullCtrs); - - Vector3 hull_origin = b_pos - a_pos; - - HullState b_hull_state = makeHullState(MADRONA_GPU_COND(mwgpu_lane_id,) - b_he_mesh, hull_origin, b_rot, b_scale, - txfm_vertex_buffer, txfm_face_buffer); - - MADRONA_GPU_COND(__syncwarp(mwGPU::allActive)); - - PROF_END(txfm_hull_ctr); - - Vector3 to_hull_closest_pt; - float hull_dist2 = hullClosestPointToOriginGJK( - b_hull_state.mesh, 1e-10f, &to_hull_closest_pt); - - if (hull_dist2 > sphere_radius * sphere_radius) { - NarrowphaseResult result; - result.type = ContactType::None; - return result; - } - - SphereContact sphere_contact; - - if (hull_dist2 == 0.f) { - // Need to do SAT - float max_sep = -FLT_MAX; - Vector3 sep_normal; - const CountT num_faces = b_hull_state.mesh.numFaces; - for (CountT i = 0; i < num_faces; i++) { - Plane plane = b_hull_state.mesh.facePlanes[i]; - // hull has already been moved so sphere is at origin - float face_dist = -plane.d; - - if (face_dist > max_sep) { - max_sep = face_dist; - sep_normal = plane.normal; - } - } - - // Discrepancy between SAT and GJK - if (max_sep > 0.f) { - assert(max_sep < 1e-5f); - NarrowphaseResult result; - result.type = ContactType::None; - return result; - } - sphere_contact.normal = sep_normal; - sphere_contact.pt = a_pos + sep_normal * sphere_radius; - sphere_contact.depth = -max_sep; - } else { - float to_hull_len = sqrtf(hull_dist2); - - float depth = sphere_radius - to_hull_len; - Vector3 normal = to_hull_closest_pt / to_hull_len; - - sphere_contact.normal = -normal; - sphere_contact.pt = a_pos + normal * sphere_radius; - sphere_contact.depth = depth; - } - - NarrowphaseResult result; - result.type = ContactType::Sphere; - result.sphere = sphere_contact; - result.aVertices = nullptr; - result.bVertices = nullptr; - result.aVertices = nullptr; - result.bVertices = nullptr; - result.aHalfEdges = nullptr; - result.bHalfEdges = nullptr; - result.aFaceHedgeRoots = nullptr; - result.bFaceHedgeRoots = nullptr; - return result; - } break; - case NarrowphaseTest::PlanePlane: { - // Planes must be static, this should never be called - assert(false); - MADRONA_UNREACHABLE(); - } break; - case NarrowphaseTest::SpherePlane: { - float sphere_radius; - { - auto sphere = a_prim->sphere; - assert(a_scale.d0 == a_scale.d1 && a_scale.d0 == a_scale.d2); - sphere_radius = a_scale.d0 * sphere.radius; - } - - constexpr Vector3 base_normal = { 0, 0, 1 }; - Vector3 plane_normal = b_rot.rotateVec(base_normal); - - float d = plane_normal.dot(b_pos); - float t = plane_normal.dot(a_pos) - d; - - float penetration = sphere_radius - t; - if (penetration < 0) { - NarrowphaseResult result; - result.type = ContactType::None; - return result; - } - - Vector3 contact_point = a_pos - t * plane_normal; - - SphereContact sphere_contact { - .normal = plane_normal, - .pt = contact_point, - .depth = penetration, - }; - - NarrowphaseResult result; - result.type = ContactType::Sphere; - result.sphere = sphere_contact; - result.aVertices = nullptr; - result.bVertices = nullptr; - result.aVertices = nullptr; - result.bVertices = nullptr; - result.aHalfEdges = nullptr; - result.bHalfEdges = nullptr; - result.aFaceHedgeRoots = nullptr; - result.bFaceHedgeRoots = nullptr; - return result; - } break; - case NarrowphaseTest::HullPlane: { - // Get half edge mesh for entity a (the hull) - const auto &a_he_mesh = a_prim->hull.halfEdgeMesh; - - assert(a_he_mesh.numFaces < max_num_tmp_faces); - assert(a_he_mesh.numVertices < max_num_tmp_vertices); - - PROF_START(txfm_hull_ctr, narrowphaseTxfmHullCtrs); - - HullState a_hull_state = makeHullState(MADRONA_GPU_COND(mwgpu_lane_id,) - a_he_mesh, a_pos, a_rot, a_scale, - txfm_vertex_buffer, txfm_face_buffer); - - MADRONA_GPU_COND(__syncwarp(mwGPU::allActive)); - - PROF_END(txfm_hull_ctr); - - constexpr Vector3 base_normal = { 0, 0, 1 }; -#if 0 - Quat inv_a_rot = a_rot.inv(); - Vector3 plane_origin_a_local = inv_a_rot.rotateVec(b_pos - a_pos); - Quat to_a_local_rot = (inv_a_rot * b_rot).normalize(); - - Vector3 plane_normal_a_local = - (to_a_local_rot.rotateVec(base_normal)).normalize(); -#endif - - Vector3 plane_normal = b_rot.rotateVec(base_normal); - - Plane plane { - plane_normal, - dot(plane_normal, b_pos), - }; - - const SATResult sat = doSATPlane( - MADRONA_GPU_COND(mwgpu_lane_id,) plane, a_hull_state); - - NarrowphaseResult result; - result.type = sat.type; - result.sat = sat.contact; -#ifdef MADRONA_GPU_MODE - result.aVertices = a_he_mesh.vertices; - result.bVertices = nullptr; -#else - result.aVertices = a_hull_state.mesh.vertices; - result.bVertices = nullptr; -#endif - result.aHalfEdges = a_hull_state.mesh.halfEdges; - result.bHalfEdges = nullptr; - result.aFaceHedgeRoots = a_hull_state.mesh.faceBaseHalfEdges; - result.bFaceHedgeRoots = nullptr; - return result; - } break; - default: MADRONA_UNREACHABLE(); - } -} - -MADRONA_ALWAYS_INLINE static inline void generateContacts( - Context &ctx, - NarrowphaseResult narrowphase_result, - Loc a_loc, Loc b_loc, -#ifdef MADRONA_GPU_MODE - Vector3 a_pos, Quat a_rot, Diag3x3 a_scale, - Vector3 b_pos, Quat b_rot, Diag3x3 b_scale, -#endif - void *thread_tmp_storage_a, void *thread_tmp_storage_b) -{ - switch (narrowphase_result.type) { - case ContactType::None: { - return; - } break; - case ContactType::Sphere: { - SphereContact sphere_contact = narrowphase_result.sphere; - - addSinglePointContact(ctx, sphere_contact.pt, sphere_contact.normal, - sphere_contact.depth, b_loc, a_loc); - } break; - case ContactType::SATPlane: { - // Plane is always b, always reference - Loc ref_loc = b_loc; - Loc other_loc = a_loc; - -#ifdef MADRONA_GPU_MODE - Mat3x4 hull_txfm = Mat3x4::fromTRS(a_pos, a_rot, a_scale); -#endif - - Plane plane { - narrowphase_result.sat.normal, - narrowphase_result.sat.planeDOrSeparation, - }; - - // Create plane contact - Manifold manifold = createFacePlaneContact( - plane, - int32_t(narrowphase_result.sat.incidentFaceIdxOrEdgeIdxB), - narrowphase_result.aVertices, - narrowphase_result.aHalfEdges, - narrowphase_result.aFaceHedgeRoots, - (Vector3 *)thread_tmp_storage_a, - (float *)thread_tmp_storage_b, -#ifdef MADRONA_GPU_MODE - hull_txfm, -#endif - { 0, 0, 0, }, - { 1, 0, 0, 0 }); - - // Sadly there are cases where two objects are just barely - // touching and post contact clipping all the clipped contacts - // are just barely separated due to FP32. For now just don't - // make a Contact in this situation. - if (manifold.numContactPoints > 0) { - addManifoldContacts(ctx, manifold, ref_loc, other_loc); - } - } break; - case ContactType::SATFace: { - const Vector3 *ref_vertices; - const Vector3 *other_vertices; - const HalfEdge *ref_hedges; - const HalfEdge *other_hedges; - const uint32_t *ref_face_hedges; - const uint32_t *other_face_hedges; - -#ifdef MADRONA_GPU_MODE - Mat3x4 ref_txfm; - Mat3x4 other_txfm; -#endif - - uint32_t ref_face_idx_and_ref_mask = - narrowphase_result.sat.refFaceIdxOrEdgeIdxA; - uint32_t incident_face_idx = - narrowphase_result.sat.incidentFaceIdxOrEdgeIdxB; - - uint32_t ref_face_idx = ref_face_idx_and_ref_mask & 0x7FFF'FFFF; - bool a_is_ref = ref_face_idx == ref_face_idx_and_ref_mask; - - Loc ref_loc, other_loc; - if (a_is_ref) { - ref_loc = a_loc; - other_loc = b_loc; - ref_vertices = narrowphase_result.aVertices; - other_vertices = narrowphase_result.bVertices; - ref_hedges = narrowphase_result.aHalfEdges; - other_hedges = narrowphase_result.bHalfEdges; - ref_face_hedges = narrowphase_result.aFaceHedgeRoots; - other_face_hedges = narrowphase_result.bFaceHedgeRoots; -#ifdef MADRONA_GPU_MODE - ref_txfm = Mat3x4::fromTRS(a_pos, a_rot, a_scale); - other_txfm = Mat3x4::fromTRS(b_pos, b_rot, b_scale); -#endif - } else { - ref_loc = b_loc; - other_loc = a_loc; - ref_vertices = narrowphase_result.bVertices; - other_vertices = narrowphase_result.aVertices; - ref_hedges = narrowphase_result.bHalfEdges; - other_hedges = narrowphase_result.aHalfEdges; - ref_face_hedges = narrowphase_result.bFaceHedgeRoots; - other_face_hedges = narrowphase_result.aFaceHedgeRoots; -#ifdef MADRONA_GPU_MODE - ref_txfm = Mat3x4::fromTRS(b_pos, b_rot, b_scale); - other_txfm = Mat3x4::fromTRS(a_pos, a_rot, a_scale); -#endif - } - - Plane ref_plane { - narrowphase_result.sat.normal, - narrowphase_result.sat.planeDOrSeparation, - }; - - // Create face contact - Manifold manifold = createFaceContact( - ref_plane, - int32_t(ref_face_idx), - int32_t(incident_face_idx), - ref_vertices, - other_vertices, - ref_hedges, - other_hedges, - ref_face_hedges, - other_face_hedges, - thread_tmp_storage_a, thread_tmp_storage_b, -#ifdef MADRONA_GPU_MODE - ref_txfm, - other_txfm, -#endif - { 0, 0, 0, }, - { 1, 0, 0, 0 }); - - // Sadly there are cases where two objects are just barely - // touching and post contact clipping all the clipped contacts - // are just barely separated due to FP32. For now just don't - // make a Contact in this situation. - if (manifold.numContactPoints > 0) { - addManifoldContacts(ctx, manifold, ref_loc, other_loc); - } - } break; - case ContactType::SATEdge: { - // A is always reference - Loc ref_loc = a_loc; - Loc other_loc = b_loc; - - // Create edge contact - Manifold manifold = createEdgeContact( - narrowphase_result.sat.normal, - narrowphase_result.sat.planeDOrSeparation, - int32_t(narrowphase_result.sat.refFaceIdxOrEdgeIdxA), - int32_t(narrowphase_result.sat.incidentFaceIdxOrEdgeIdxB), - narrowphase_result.aVertices, - narrowphase_result.bVertices, - narrowphase_result.aHalfEdges, - narrowphase_result.bHalfEdges, -#ifdef MADRONA_GPU_MODE - a_pos, a_rot, a_scale, - b_pos, b_rot, b_scale, -#endif - { 0, 0, 0 }, { 1, 0, 0, 0 }); - - addManifoldContacts(ctx, manifold, ref_loc, other_loc); - } break; - default: MADRONA_UNREACHABLE(); - } -} - -static inline void runNarrowphase( - Context &ctx, - const CandidateCollision &candidate_collision - MADRONA_GPU_COND(, - const int32_t mwgpu_warp_id, - const int32_t mwgpu_lane_id, - bool lane_active)) -{ - PROF_START(setup_ctr, narrowphaseSetupClocks); - -#ifdef MADRONA_GPU_MODE - const int32_t num_smem_bytes_per_warp = - mwGPU::SharedMemStorage::numBytesPerWarp(); - const int32_t num_smem_floats = num_smem_bytes_per_warp / sizeof(float); - int32_t num_vertex_floats = num_smem_floats - gpuImpl::numPlaneFloats; - int32_t max_num_vertices = num_vertex_floats / 3; - - constexpr int32_t max_num_tmp_faces = gpuImpl::maxNumPlanes; - int32_t max_num_tmp_vertices = max_num_vertices; - - Plane tmp_faces_buffer[max_num_tmp_faces]; - - Plane * smem_faces_buffer; - Vector3 * smem_vertices_buffer; - { - auto smem_buf = (char *)mwGPU::SharedMemStorage::buffer; - char *warp_smem_base = - smem_buf + num_smem_bytes_per_warp * mwgpu_warp_id; - - smem_faces_buffer = (Plane *)warp_smem_base; - smem_vertices_buffer = - (Vector3 *)(smem_faces_buffer + gpuImpl::maxNumPlanes); - } -#else - constexpr int32_t max_num_tmp_faces = 512; - constexpr int32_t max_num_tmp_vertices = 512; - - Plane tmp_faces_buffer[max_num_tmp_faces]; - Vector3 tmp_vertices_buffer[max_num_tmp_vertices]; -#endif - - PROF_END(setup_ctr); - - PROF_START(prep_ctr, narrowphasePrepClocks); - - Loc a_loc = candidate_collision.a; - Loc b_loc = candidate_collision.b; - - const ObjectManager &obj_mgr = *ctx.singleton().mgr; - - uint32_t a_prim_idx, b_prim_idx; - { - ObjectID a_obj = ctx.getDirect(RGDCols::ObjectID, a_loc); - ObjectID b_obj = ctx.getDirect(RGDCols::ObjectID, b_loc); - - const uint32_t a_prim_offset = - obj_mgr.rigidBodyPrimitiveOffsets[a_obj.idx]; - const uint32_t b_prim_offset = - obj_mgr.rigidBodyPrimitiveOffsets[b_obj.idx]; - - a_prim_idx = a_prim_offset + candidate_collision.aPrim; - b_prim_idx = b_prim_offset + candidate_collision.bPrim; - } - - const CollisionPrimitive *a_prim = - &obj_mgr.collisionPrimitives[a_prim_idx]; - const CollisionPrimitive *b_prim = - &obj_mgr.collisionPrimitives[b_prim_idx]; - - uint32_t raw_type_a = static_cast(a_prim->type); - uint32_t raw_type_b = static_cast(b_prim->type); - - // Swap a & b to be properly ordered based on object type - if (raw_type_a > raw_type_b) { - std::swap(a_loc, b_loc); - std::swap(a_prim, b_prim); - std::swap(a_prim_idx, b_prim_idx); - std::swap(raw_type_a, raw_type_b); - } - - const Vector3 a_pos = ctx.getDirect(RGDCols::Position, a_loc); - const Vector3 b_pos = ctx.getDirect(RGDCols::Position, b_loc); - const Quat a_rot = ctx.getDirect(RGDCols::Rotation, a_loc); - const Quat b_rot = ctx.getDirect(RGDCols::Rotation, b_loc); - const Diag3x3 a_scale(ctx.getDirect(RGDCols::Scale, a_loc)); - const Diag3x3 b_scale(ctx.getDirect(RGDCols::Scale, b_loc)); - - { - AABB a_obj_aabb = obj_mgr.primitiveAABBs[a_prim_idx]; - AABB b_obj_aabb = obj_mgr.primitiveAABBs[b_prim_idx]; - - AABB a_world_aabb = a_obj_aabb.applyTRS(a_pos, a_rot, a_scale); - AABB b_world_aabb = b_obj_aabb.applyTRS(b_pos, b_rot, b_scale); - - if (!a_world_aabb.intersects(b_world_aabb)) { -#ifdef MADRONA_GPU_MODE - lane_active = false; -#else - return; -#endif - } - } - -#ifdef MADRONA_GPU_MODE - const uint32_t active_mask = __ballot_sync(mwGPU::allActive, lane_active); - - if (active_mask == 0) { - return; - } -#endif - - const NarrowphaseTest test_type {raw_type_a | raw_type_b}; - - PROF_END(prep_ctr); - -#ifdef MADRONA_GPU_MODE - NarrowphaseResult thread_result; - -#if 0 - active_mask = __brev(active_mask); - int32_t leader_idx = __clz(active_mask); - active_mask <<= (leader_idx + 1); - do { -#endif -#pragma unroll - for (int32_t leader_idx = 0; leader_idx < 32; leader_idx++) { - if (!__shfl_sync(mwGPU::allActive, lane_active, leader_idx)) { - continue; - } - - auto warp_test_type = (NarrowphaseTest)__shfl_sync( - mwGPU::allActive, (uint32_t)test_type, leader_idx); - - Vector3 warp_a_pos { - __shfl_sync(mwGPU::allActive, a_pos.x, leader_idx), - __shfl_sync(mwGPU::allActive, a_pos.y, leader_idx), - __shfl_sync(mwGPU::allActive, a_pos.z, leader_idx), - }; - - Vector3 warp_b_pos { - __shfl_sync(mwGPU::allActive, b_pos.x, leader_idx), - __shfl_sync(mwGPU::allActive, b_pos.y, leader_idx), - __shfl_sync(mwGPU::allActive, b_pos.z, leader_idx), - }; - - Quat warp_a_rot { - __shfl_sync(mwGPU::allActive, a_rot.w, leader_idx), - __shfl_sync(mwGPU::allActive, a_rot.x, leader_idx), - __shfl_sync(mwGPU::allActive, a_rot.y, leader_idx), - __shfl_sync(mwGPU::allActive, a_rot.z, leader_idx), - }; - - Quat warp_b_rot { - __shfl_sync(mwGPU::allActive, b_rot.w, leader_idx), - __shfl_sync(mwGPU::allActive, b_rot.x, leader_idx), - __shfl_sync(mwGPU::allActive, b_rot.y, leader_idx), - __shfl_sync(mwGPU::allActive, b_rot.z, leader_idx), - }; - - Diag3x3 warp_a_scale { - __shfl_sync(mwGPU::allActive, a_scale.d0, leader_idx), - __shfl_sync(mwGPU::allActive, a_scale.d1, leader_idx), - __shfl_sync(mwGPU::allActive, a_scale.d2, leader_idx), - }; - - Diag3x3 warp_b_scale { - __shfl_sync(mwGPU::allActive, b_scale.d0, leader_idx), - __shfl_sync(mwGPU::allActive, b_scale.d1, leader_idx), - __shfl_sync(mwGPU::allActive, b_scale.d2, leader_idx), - }; - - auto warp_a_prim = (CollisionPrimitive *)__shfl_sync(mwGPU::allActive, - (uint64_t)a_prim, leader_idx); - - auto warp_b_prim = (CollisionPrimitive *)__shfl_sync(mwGPU::allActive, - (uint64_t)b_prim, leader_idx); - - NarrowphaseResult warp_result = narrowphaseDispatch( - mwgpu_lane_id, - warp_test_type, - warp_a_pos, warp_b_pos, - warp_a_rot, warp_b_rot, - warp_a_scale, warp_b_scale, - warp_a_prim, warp_b_prim, - max_num_tmp_vertices, max_num_tmp_faces, - smem_vertices_buffer, smem_faces_buffer); - - if (mwgpu_lane_id == leader_idx) { - thread_result = warp_result; - } - -#if 0 - uint32_t num_inactive = __clz(active_mask); - leader_idx += num_inactive + 1; - active_mask <<= (num_inactive + 1); - } while (leader_idx < 32); -#endif - } - - __syncwarp(mwGPU::allActive); - - if (lane_active) { - generateContacts(ctx, thread_result, - a_loc, b_loc, - a_pos, a_rot, a_scale, - b_pos, b_rot, b_scale, - tmp_faces_buffer, - tmp_faces_buffer + max_num_tmp_faces / 2); - - } -#else - NarrowphaseResult result = narrowphaseDispatch( - test_type, - a_pos, b_pos, - a_rot, b_rot, - a_scale, b_scale, - a_prim, b_prim, - max_num_tmp_vertices, max_num_tmp_faces, - tmp_vertices_buffer, tmp_faces_buffer); - - generateContacts(ctx, result, - a_loc, b_loc, - tmp_faces_buffer, - tmp_faces_buffer + max_num_tmp_faces / 2); -#endif -} - -inline void runNarrowphaseSystem( -#ifdef MADRONA_GPU_MODE - WorldID *world_ids, - const CandidateCollision *candidate_collisions, - int32_t num_candidates -#else - Context &ctx, - const CandidateCollision &candidate_collision -#endif - ) -{ - PROF_START(all_ctr, narrowphaseAllClocks); -#ifdef MADRONA_GPU_MODE - const int32_t mwgpu_warp_id = threadIdx.x / 32; - const int32_t mwgpu_lane_id = threadIdx.x % 32; - - PROF_START(world_get_ctr, narrowphaseFetchWorldClocks); - - const int32_t candidate_idx = min(mwgpu_lane_id, num_candidates - 1); - - bool lane_active = candidate_idx == mwgpu_lane_id; - - WorldID world_id = world_ids[candidate_idx]; - if (world_id.idx == -1) { - lane_active = false; - } - - Context ctx = TaskGraph::makeContext(world_id); - PROF_END(world_get_ctr); - - runNarrowphase(ctx, candidate_collisions[candidate_idx], - mwgpu_warp_id, mwgpu_lane_id, lane_active); - -#else - runNarrowphase(ctx, candidate_collision); -#endif -} - -TaskGraphNodeID setupTasks( - TaskGraphBuilder &builder, - Span deps) -{ -#ifdef MADRONA_GPU_MODE - auto narrowphase = builder.addToGraph>(deps); -#else - auto narrowphase = builder.addToGraph>(deps); -#endif - - auto finished = builder.addToGraph({narrowphase}); - - return finished; -} - -} diff --git a/src/physics/physics.cpp b/src/physics/physics.cpp deleted file mode 100644 index e12f6cf8..00000000 --- a/src/physics/physics.cpp +++ /dev/null @@ -1,410 +0,0 @@ -#include -#include - -#include "physics_impl.hpp" -#include "xpbd.hpp" -#include "tgs.hpp" - -namespace madrona::phys { - -using namespace base; -using namespace math; - -#ifdef MADRONA_GPU_MODE -//#define COUNT_GPU_CLOCKS -#endif - -#ifdef COUNT_GPU_CLOCKS -extern "C" { -extern AtomicU64 narrowphaseAllClocks; -extern AtomicU64 narrowphaseFetchWorldClocks; -extern AtomicU64 narrowphaseSetupClocks; -extern AtomicU64 narrowphasePrepClocks; -extern AtomicU64 narrowphaseSwitchClocks; -extern AtomicU64 narrowphaseSATFaceClocks; -extern AtomicU64 narrowphaseSATEdgeClocks; -extern AtomicU64 narrowphaseSATPlaneClocks; -extern AtomicU64 narrowphaseSATContactClocks; -extern AtomicU64 narrowphaseSATPlaneContactClocks; -extern AtomicU64 narrowphaseSaveContactsClocks; -extern AtomicU64 narrowphaseTxfmHullCtrs; -} - -inline void reportNarrowphaseClocks(Engine &ctx, - SolverData &) -{ - if (ctx.worldID().idx != 0) { - return; - } - - if (threadIdx.x == 0 && ctx.worldID().idx == 0) { - printf("[%lu, %lu, %lu, %lu, %lu, %lu, %lu, %lu, %lu, %lu, %lu, %lu]\n", - narrowphaseAllClocks.load(), - narrowphaseFetchWorldClocks.load(), - narrowphaseSetupClocks.load(), - narrowphasePrepClocks.load(), - narrowphaseSwitchClocks.load(), - narrowphaseSATFaceClocks.load(), - narrowphaseSATEdgeClocks.load(), - narrowphaseSATPlaneClocks.load(), - narrowphaseSATContactClocks.load(), - narrowphaseSATPlaneContactClocks.load(), - narrowphaseSaveContactsClocks.load(), - narrowphaseTxfmHullCtrs.load() - ); - - narrowphaseAllClocks.store(0); - narrowphaseFetchWorldClocks.store(0); - narrowphaseSetupClocks.store(0), - narrowphasePrepClocks.store(0); - narrowphaseSwitchClocks.store(0); - narrowphaseSATFaceClocks.store(0); - narrowphaseSATEdgeClocks.store(0); - narrowphaseSATPlaneClocks.store(0); - narrowphaseSATContactClocks.store(0); - narrowphaseSATPlaneContactClocks.store(0); - narrowphaseSaveContactsClocks.store(0); - narrowphaseTxfmHullCtrs.store(0); - } -} -#endif - - -static void initPhysicsState(Context &ctx, - float delta_t, - CountT num_substeps, - Vector3 gravity, - uint32_t contact_archetype_id, - uint32_t joint_archetype_id) -{ - float h = delta_t / (float)num_substeps; - float g_mag = gravity.length(); - - ctx.singleton() = { - .deltaT = delta_t, - .h = h, - .g = gravity, - .gMagnitude = g_mag, - .restitutionThreshold = 2.f * g_mag * h, - .contactArchetypeID = contact_archetype_id, - .jointArchetypeID = joint_archetype_id, - }; -} - -namespace PhysicsSystem { - -void init(Context &ctx, - ObjectManager *obj_mgr, - float delta_t, - CountT num_substeps, - math::Vector3 gravity, - CountT max_dynamic_objects, - Solver solver) -{ - broadphase::BVH &bvh = ctx.singleton(); - - // expansion factor is 2 * delta_t to give room - // for acceleration within the timestep - constexpr float max_inst_accel = 100.f; - new (&bvh) broadphase::BVH( - obj_mgr, max_dynamic_objects, 2.f * delta_t, - max_inst_accel * delta_t * delta_t); - - uint32_t contact_archetype_id, joint_archetype_id; - switch (solver) { - case Solver::XPBD: { - xpbd::getSolverArchetypeIDs(&contact_archetype_id, - &joint_archetype_id); - } break; - case Solver::TGS: { - tgs::getSolverArchetypeIDs(&contact_archetype_id, - &joint_archetype_id); - } break; - default: MADRONA_UNREACHABLE(); - } - - initPhysicsState( - ctx, delta_t, num_substeps, gravity, - contact_archetype_id, joint_archetype_id); - - switch (solver) { - case Solver::XPBD: { - xpbd::init(ctx); - } break; - case Solver::TGS: { - tgs::init(ctx); - } break; - default: MADRONA_UNREACHABLE(); - } - - ctx.singleton() = { obj_mgr }; -} - -void reset(Context &ctx) -{ - broadphase::BVH &bvh = ctx.singleton(); - bvh.rebuildOnUpdate(); - bvh.clearLeaves(); -} - -broadphase::LeafID registerEntity(Context &ctx, - Entity e, - ObjectID obj_id) -{ - auto &bvh = ctx.singleton(); - - return bvh.reserveLeaf(e, obj_id); -} - -bool checkEntityAABBOverlap( - Context &ctx, math::AABB aabb, Entity e) -{ - const ObjectManager &obj_mgr = *ctx.singleton().mgr; - - ObjectID e_obj_id = ctx.get(e); - Position e_pos = ctx.get(e); - Rotation e_rot = ctx.get(e); - Scale e_scale = ctx.get(e); - - uint32_t num_prims = obj_mgr.rigidBodyPrimitiveCounts[e_obj_id.idx]; - uint32_t base_prim_offset = obj_mgr.rigidBodyPrimitiveOffsets[e_obj_id.idx]; - - bool overlap = false; - for (uint32_t prim_offset = 0; prim_offset < num_prims; prim_offset++) { - uint32_t prim_idx = base_prim_offset + prim_offset; - - const CollisionPrimitive &prim = obj_mgr.collisionPrimitives[prim_idx]; - if (prim.type != CollisionPrimitive::Type::Hull) { - continue; - } - - AABB prim_aabb = obj_mgr.primitiveAABBs[prim_idx]; - AABB txfmed_aabb = prim_aabb.applyTRS(e_pos, e_rot, e_scale); - - if (!txfmed_aabb.overlaps(aabb)) { - continue; - } - - const Vector3 *vertices = prim.hull.halfEdgeMesh.vertices; - CountT num_verts = (CountT)prim.hull.halfEdgeMesh.numVertices; - - const std::array axes { - right, - fwd, - up, - }; - - std::array min_hull_projs { - FLT_MAX, - FLT_MAX, - FLT_MAX, - }; - std::array max_hull_projs { - -FLT_MAX, - -FLT_MAX, - -FLT_MAX, - }; - - for (CountT vert_idx = 0; vert_idx < num_verts; vert_idx++) { - Vector3 v = - e_rot.rotateVec(e_scale * vertices[vert_idx]) + e_pos; - -#pragma unroll - for (CountT i = 0; i < 3; i++) { - Vector3 axis = axes[i]; - - float proj = dot(v, axis); - if (proj < min_hull_projs[i]) { - min_hull_projs[i] = proj; - } - - if (proj > max_hull_projs[i]) { - max_hull_projs[i] = proj; - } - } - } - - bool axes_overlap = true; - -#pragma unroll - for (CountT i = 0; i < 3; i++) { - float min_aabb_proj = aabb.pMin[i]; - float max_aabb_proj = aabb.pMax[i]; - - float min_hull_proj = min_hull_projs[i]; - float max_hull_proj = max_hull_projs[i]; - - bool proj_overlap = max_hull_proj > min_aabb_proj && - max_aabb_proj > min_hull_proj; - - if (!proj_overlap) { - axes_overlap = false; - } - } - - if (axes_overlap) { - overlap = true; - break; - } - } - - return overlap; -} - - -Entity makeFixedJoint( - Context &ctx, - Entity e1, Entity e2, - math::Quat attach_rot1, math::Quat attach_rot2, - math::Vector3 r1, math::Vector3 r2, - float separation) -{ - const auto &physics_sys = ctx.singleton(); - Entity e = ctx.makeEntity(physics_sys.jointArchetypeID); - - ctx.get(e) = { - .e1 = e1, - .e2 = e2, - .type = JointConstraint::Type::Fixed, - .fixed = { - .attachRot1 = attach_rot1, - .attachRot2 = attach_rot2, - .separation = separation, - }, - .r1 = r1, - .r2 = r2, - }; - - return e; -} - -Entity makeHingeJoint( - Context &ctx, - Entity e1, Entity e2, - math::Vector3 a1_local, math::Vector3 a2_local, - math::Vector3 b1_local, math::Vector3 b2_local, - math::Vector3 r1, math::Vector3 r2) -{ - const auto &physics_sys = ctx.singleton(); - Entity e = ctx.makeEntity(physics_sys.jointArchetypeID); - - ctx.get(e) = { - .e1 = e1, - .e2 = e2, - .type = JointConstraint::Type::Hinge, - .hinge = { - .a1Local = a1_local, - .a2Local = a2_local, - .b1Local = b1_local, - .b2Local = b2_local, - }, - .r1 = r1, - .r2 = r2, - }; - - return e; -} - -void registerTypes(ECSRegistry ®istry, - Solver solver) -{ - registry.registerComponent(); - registry.registerComponent(); - registry.registerComponent(); - registry.registerComponent(); - registry.registerComponent(); - - registry.registerSingleton(); - - registry.registerComponent(); - registry.registerArchetype(); - - registry.registerComponent(); - registry.registerArchetype(); - - registry.registerComponent(); - registry.registerComponent(); - - registry.registerSingleton(); - registry.registerSingleton(); - - switch (solver) { - case Solver::XPBD: { - xpbd::registerTypes(registry); - } break; - case Solver::TGS: { - tgs::registerTypes(registry); - } break; - default: MADRONA_UNREACHABLE(); - } - - registry.registerBundle(); -} - -TaskGraphNodeID setupBroadphaseTasks( - TaskGraphBuilder &builder, - Span deps) -{ - return broadphase::setupBVHTasks(builder, deps); -} - -TaskGraphNodeID setupPhysicsStepTasks( - TaskGraphBuilder &builder, - Span deps, - CountT num_substeps, - Solver solver) -{ - auto broadphase_prep = - broadphase::setupPreIntegrationTasks(builder, deps); - - TaskGraphNodeID solver_finished; - switch (solver) { - case Solver::XPBD: { - solver_finished = xpbd::setupXPBDSolverTasks( - builder, broadphase_prep, num_substeps); - } break; - case Solver::TGS: { - solver_finished = tgs::setupTGSSolverTasks( - builder, broadphase_prep, num_substeps); - } break; - default: MADRONA_UNREACHABLE(); - } - - auto broadphase_post = - broadphase::setupPostIntegrationTasks(builder, {solver_finished}); - - auto physics_done = broadphase_post; - -#ifdef COUNT_GPU_CLOCKS - physics_done = builder.addToGraph>({physics_done}); -#endif - - return physics_done; -} - -TaskGraphNodeID setupCleanupTasks( - TaskGraphBuilder &builder, Span deps) -{ - return builder.addToGraph>(deps); -} - - -TaskGraphNodeID setupStandaloneBroadphaseOverlapTasks( - TaskGraphBuilder &builder, - Span deps) -{ - return broadphase::setupPreIntegrationTasks(builder, deps); -} - -TaskGraphNodeID setupStandaloneBroadphaseCleanupTasks( - TaskGraphBuilder &builder, - Span deps) -{ - return builder.addToGraph>(deps); -} - - -} - -} diff --git a/src/physics/physics_assets.cpp b/src/physics/physics_assets.cpp deleted file mode 100644 index 93fc787e..00000000 --- a/src/physics/physics_assets.cpp +++ /dev/null @@ -1,1408 +0,0 @@ -#include -#include - -#ifdef MADRONA_CUDA_SUPPORT -#include -#endif - -#include - -namespace madrona::phys { -using namespace geo; -using namespace math; - -namespace { - -struct EditMesh { - struct HEdge { - uint32_t next; - uint32_t prev; - uint32_t twin; - - uint32_t vert; - uint32_t face; - }; - - struct Face { - uint32_t hedge; - uint32_t next; - uint32_t prev; - - Plane plane; - }; - - struct Vert { - Vector3 pos; - uint32_t next; - uint32_t prev; - }; - - HEdge *hedges; - Face *faces; - Vert *verts; - - uint32_t numHedges; - uint32_t numFaces; - uint32_t numVerts; - - uint32_t hedgeFreeHead; - uint32_t faceFreeHead; - uint32_t vertFreeHead; -}; - -struct HullBuildData { - EditMesh mesh; - uint32_t *faceConflictLists; -}; - -struct MassProperties { - Diag3x3 inertiaTensor; - Vector3 centerOfMass; - Quat toDiagonal; -}; - -} - -static uint32_t allocMeshHedge(EditMesh &mesh) -{ - uint32_t hedge = mesh.hedgeFreeHead; - assert(hedge != 0); - mesh.hedgeFreeHead = mesh.hedges[hedge].next; - - return hedge; -} - -static void freeMeshHedge(EditMesh &mesh, uint32_t hedge) -{ - uint32_t old_head = mesh.hedgeFreeHead; - mesh.hedgeFreeHead = hedge; - mesh.hedges[hedge].next = old_head; -} - -static uint32_t createMeshFace(EditMesh &mesh) -{ - uint32_t face = mesh.faceFreeHead; - assert(face != 0); - mesh.faceFreeHead = mesh.faces[face].next; - - uint32_t prev_prev = mesh.faces[0].prev; - mesh.faces[0].prev = face; - mesh.faces[prev_prev].next = face; - - mesh.faces[face].next = 0; - mesh.faces[face].prev = prev_prev; - - mesh.numFaces += -1; - - return face; -} - -static void deleteMeshFace(EditMesh &mesh, uint32_t face) -{ - uint32_t next = mesh.faces[face].next; - uint32_t prev = mesh.faces[face].prev; - - mesh.faces[prev].next = next; - mesh.faces[next].prev = prev; - - uint32_t old_head = mesh.faceFreeHead; - mesh.faceFreeHead = face; - mesh.faces[face].next = old_head; -} - -static uint32_t allocMeshVert(EditMesh &mesh) -{ - uint32_t vert = mesh.vertFreeHead; - assert(vert != 0); - mesh.vertFreeHead = mesh.verts[vert].next; - - return vert; -} - -static void freeMeshVert(EditMesh &mesh, uint32_t vert) -{ - uint32_t old_head = mesh.vertFreeHead; - mesh.vertFreeHead = vert; - mesh.verts[vert].next = old_head; -} - -static uint32_t addVertToMesh(EditMesh &mesh, uint32_t vert) -{ - uint32_t prev_prev = mesh.verts[0].prev; - mesh.verts[0].prev = vert; - mesh.verts[prev_prev].next = vert; - - mesh.verts[vert].next = 0; - mesh.verts[vert].prev = prev_prev; - - mesh.numVerts += 1; - - return vert; -} - -static void removeVertFromMesh(EditMesh &mesh, uint32_t vert) -{ - uint32_t next = mesh.verts[vert].next; - uint32_t prev = mesh.verts[vert].prev; - - mesh.verts[prev].next = next; - mesh.verts[next].prev = prev; - - mesh.numVerts -= 1; -} - -static uint32_t addConflictVert(HullBuildData &hull_data, - uint32_t face, - Vector3 pos) -{ - auto &mesh = hull_data.mesh; - uint32_t vert = allocMeshVert(mesh); - - uint32_t next = hull_data.faceConflictLists[face]; - - hull_data.faceConflictLists[face] = vert; - mesh.verts[vert].next = next; - mesh.verts[vert].prev = 0; - - if (next != 0) { - mesh.verts[next].prev = vert; - } - - mesh.verts[vert].pos = pos; - - return vert; -} - -static void removeConflictVert(HullBuildData &hull_data, - uint32_t face, - uint32_t vert) -{ - auto &mesh = hull_data.mesh; - - uint32_t next = mesh.verts[vert].next; - uint32_t prev = mesh.verts[vert].prev; - - if (prev == 0) { - hull_data.faceConflictLists[face] = next; - } else { - mesh.verts[prev].next = next; - } - - if (next != 0) { - mesh.verts[next].prev = prev; - } -} - -// Gregorious, Implementing QuickHull, GDC 2014, Slide 77 -static float computePlaneEpsilon(Span verts) -{ - AABB aabb = AABB::invalid(); - - for (Vector3 v : verts) { - aabb.expand(v); - } - - Vector3 diff = aabb.pMax - aabb.pMin; - - return 3.f * (diff.x + diff.y + diff.z) * FLT_EPSILON; -} - -// RTCD 12.4.2 -template -static Plane computeNewellPlaneImpl(Fn &&iter_verts) -{ - Vector3 centroid { 0, 0, 0 }; - Vector3 n { 0, 0, 0 }; - - CountT num_verts = 0; - // Compute normal as being proportional to projected areas of polygon - // onto the yz, xz, and xy planes. Also compute centroid as - // representative point on the plane - iter_verts([¢roid, &n, &num_verts](Vector3 vi, Vector3 vj) { - n.x += (vi.y - vj.y) * (vi.z + vj.z); // projection on yz - n.y += (vi.z - vj.z) * (vi.x + vj.x); // projection on xz - n.z += (vi.x - vj.x) * (vi.y + vj.y); // projection on xy - - centroid += vj; - num_verts += 1; - }); - - assert(num_verts != 0); - - centroid /= num_verts; - - n = normalize(n); - return Plane { - .normal = n, - .d = dot(centroid, n), - }; -} - -static Plane computeNewellPlane(const Vector3 *verts, - Span indices) -{ - return computeNewellPlaneImpl([verts, indices](auto &&fn) { - for (CountT i = indices.size() - 1, j = 0; j < indices.size(); - i = j, j++) { - Vector3 vi = verts[indices[i]]; - Vector3 vj = verts[indices[j]]; - - fn(vi, vj); - } - }); -} - -static Plane computeNewellPlane(EditMesh &mesh, uint32_t face) -{ - return computeNewellPlaneImpl([mesh, face](auto &&fn) { - uint32_t start_hedge_idx = mesh.faces[face].hedge; - uint32_t cur_hedge_idx = start_hedge_idx; - - do { - const EditMesh::HEdge &cur_hedge = mesh.hedges[cur_hedge_idx]; - uint32_t next_hedge_idx = cur_hedge.next; - const EditMesh::HEdge &next_hedge = mesh.hedges[next_hedge_idx]; - - uint32_t i = cur_hedge.vert; - uint32_t j = next_hedge.vert; - - fn(mesh.verts[i].pos, mesh.verts[j].pos); - - cur_hedge_idx = next_hedge_idx; - } while (cur_hedge_idx != start_hedge_idx); - }); -} - -static float distToPlane(Plane plane, Vector3 v) -{ - return v.dot(plane.normal) - plane.d; -} - -static HullBuildData allocBuildData(StackAlloc &tmp_alloc, const CountT N) -{ - // + 1 for fake starting point for linked lists - const CountT max_num_verts = N + 1; - // Num edges = 3N - 6. Doubled for half edges, doubled for horizon - const CountT max_num_hedges = 4 * (3 * N - 6) + 1; - // Num edges = 2N - 4. Doubled for horizon - const CountT max_num_faces = 2 * (2 * N - 4) + 1; - - const auto buffer_sizes = std::to_array({ - int64_t(sizeof(EditMesh::HEdge) * max_num_hedges), // hedges - int64_t(sizeof(EditMesh::Face) * max_num_faces), // faces - int64_t(sizeof(EditMesh::Vert) * max_num_verts), // verts - int64_t(sizeof(uint32_t) * max_num_faces), // faceConflictLists - }); - - constexpr CountT sub_buffer_alignment = 128; - - int64_t buffer_offsets[buffer_sizes.size() - 1]; - int64_t total_bytes = utils::computeBufferOffsets( - buffer_sizes, buffer_offsets, sub_buffer_alignment); - - char *buf_base = - (char *)tmp_alloc.alloc(total_bytes, sub_buffer_alignment); - - EditMesh mesh { - .hedges = (EditMesh::HEdge *)(buf_base), - .faces = (EditMesh::Face *)(buf_base + buffer_offsets[0]), - .verts = (EditMesh::Vert *)(buf_base + buffer_offsets[1]), - .numHedges = 0, - .numFaces = 0, - .numVerts = 0, - .hedgeFreeHead = 1, - .faceFreeHead = 1, - .vertFreeHead = 1, - }; - - // Setup free lists - for (CountT i = 1; i < max_num_hedges; i++) { - mesh.hedges[i].next = uint32_t(i + 1); - } - mesh.hedges[max_num_hedges].next = 0; - - for (CountT i = 1; i < max_num_faces; i++) { - mesh.faces[i].next = uint32_t(i + 1); - } - mesh.faces[max_num_faces].next = 0; - - for (CountT i = 1; i < max_num_verts; i++) { - mesh.verts[i].next = uint32_t(i + 1); - } - mesh.verts[max_num_verts].next = 0; - - // Elem 0 is fake head / tail to avoid special cases - mesh.hedges[0].next = 0; - mesh.hedges[0].prev = 0; - - mesh.faces[0].next = 0; - mesh.faces[0].prev = 0; - - mesh.verts[0].next = 0; - mesh.verts[0].prev = 0; - - uint32_t *face_conflict_lists = (uint32_t *)(buf_base + buffer_offsets[2]); - for (CountT i = 0; i < max_num_faces; i++) { - face_conflict_lists[i] = 0; - } - - return HullBuildData { - .mesh = mesh, - .faceConflictLists = face_conflict_lists, - }; -} - -static bool initHullTetrahedron(EditMesh &mesh, - Span verts, - float epsilon, - uint32_t *tet_fids, - Plane *tet_face_planes) -{ - // Choose the initial 4 points for the hull - Vector3 v0 = verts[0]; - - Vector3 v1, e1; - float max_v1_dist = -FLT_MAX; - for (CountT i = 1; i < verts.size(); i++) { - Vector3 v = verts[i]; - Vector3 e = v - v0; - float e_len = e.length(); - if (e_len > max_v1_dist) { - v1 = v; - e1 = e; - max_v1_dist = e_len; - } - } - - if (max_v1_dist < epsilon) { - return false; - } - - Vector3 v2, e2; - float max_v2_area = -FLT_MAX; - for (CountT i = 1; i < verts.size(); i++) { - Vector3 v = verts[i]; - Vector3 e = v - v0; - - float area = cross(e, e1).length(); - - if (area > max_v2_area) { - v2 = v; - e2 = e; - max_v2_area = area; - } - } - - if (max_v2_area < epsilon) { - return false; - } - - Vector3 v3; - float max_v3_det = -FLT_MAX; - for (CountT i = 1; i < verts.size(); i++) { - Vector3 v = verts[i]; - Vector3 e = v - v0; - - Mat3x3 vol_mat {{ e1, e2, e }}; - float det = vol_mat.determinant(); - - if (det > max_v3_det) { - v3 = v; - max_v3_det = det; - } - } - - if (max_v3_det < epsilon) { - return false; - } - - // Setup initial halfedge mesh - uint32_t vids[4]; - vids[0] = allocMeshVert(mesh); - vids[1] = allocMeshVert(mesh); - vids[2] = allocMeshVert(mesh); - vids[3] = allocMeshVert(mesh); - addVertToMesh(mesh, vids[0]); - addVertToMesh(mesh, vids[1]); - addVertToMesh(mesh, vids[2]); - addVertToMesh(mesh, vids[3]); - mesh.verts[vids[0]].pos = v0; - mesh.verts[vids[1]].pos = v1; - mesh.verts[vids[2]].pos = v2; - mesh.verts[vids[3]].pos = v3; - - // Face 0: - // he0: 3 => 2, he1: 2 => 1, he2: 1 => 3, - // Face 1: - // he3: 2 => 3, he4: 3 => 0, he5: 0 => 2, - // Face 2: - // he6: 1 => 0, he7: 0 => 3, he8: 3 => 1, - // Face 3: - // he9: 0 => 1, he10: 1 => 2, he11: 2 => 0, - uint32_t eids[12]; - const uint32_t face_vert_indices[] = { - 3, 2, 1, - 2, 3, 0, - 1, 0, 3, - 0, 1, 2 - }; - - const uint32_t twin_hedge_indices[] = { - 3, 10, 8, - 0, 7, 11, - 9, 4, 2, - 6, 1, 5, - }; - - // Allocate half edges -#pragma unroll - for (CountT i = 0; i < 12; i++) { - eids[i] = allocMeshHedge(mesh); - } - - // Create faces and create halfedges - for (CountT i = 0; i < 4; i++) { - const uint32_t base_hedge_offset = i * 3; - uint32_t fid = tet_fids[i] = createMeshFace(mesh); - -#pragma unroll - for (CountT j = 0; j < 3; j++) { - const uint32_t cur_hedge_offset = base_hedge_offset + j; - const uint32_t next_hedge_offset = base_hedge_offset + ((j + 1) % 3); - const uint32_t prev_hedge_offset = base_hedge_offset + ((j + 2) % 3); - - uint32_t vid = vids[face_vert_indices[cur_hedge_offset]]; - uint32_t cur_eid = eids[cur_hedge_offset]; - - mesh.hedges[cur_eid].face = fid; - mesh.hedges[cur_eid].vert = vid; - - mesh.hedges[cur_eid].next = eids[next_hedge_offset]; - mesh.hedges[cur_eid].prev = eids[prev_hedge_offset]; - - mesh.hedges[cur_eid].twin = twin_hedge_indices[cur_hedge_offset]; - } - - mesh.faces[fid].hedge = eids[base_hedge_offset]; - - Plane face_plane = computeNewellPlane(mesh, fid); - mesh.faces[fid].plane = tet_face_planes[i] = face_plane; - } - - return true; -} - -static bool initHullBuild(Span verts, - StackAlloc &tmp_alloc, - HullBuildData *out) -{ - if (verts.size() < 4) { - return false; - } - - *out = allocBuildData(tmp_alloc, verts.size()); - EditMesh &mesh = out->mesh; - - float epsilon = computePlaneEpsilon(verts); - - uint32_t tet_face_ids[4]; - Plane tet_face_planes[4]; - // FIXME: choose proper epsilon not just plane epsilon - bool tet_success = initHullTetrahedron(mesh, verts, epsilon, tet_face_ids, - tet_face_planes); - if (!tet_success) { - return false; - } - - // Initial vertex binning - for (Vector3 pos : verts) { - float closest_plane_dist = FLT_MAX; - CountT closest_plane_idx = -1; - for (CountT i = 0; i < 4; i++) { - Plane cur_plane = tet_face_planes[i]; - float dist = distToPlane(cur_plane, pos); - - if (dist > epsilon) { - if (dist < closest_plane_dist) { - closest_plane_idx = i; - closest_plane_dist = dist; - } - } - } - - // This is an internal vertex - if (closest_plane_idx == -1) { - continue; - } - - addConflictVert(*out, tet_face_ids[closest_plane_idx], pos); - } - - return true; -} - -static void quickhullBuild(HullBuildData &build_data) -{ - auto &mesh = build_data.mesh; - // FIXME - (void)mesh; - (void)freeMeshHedge; - (void)deleteMeshFace; - (void)freeMeshVert; - (void)removeVertFromMesh; - (void)removeConflictVert; -} - -static HalfEdgeMesh editMeshToRuntimeMesh(StackAlloc &tmp_alloc, - EditMesh &edit_mesh) -{ - uint32_t *hedge_remap = tmp_alloc.allocN(edit_mesh.numHedges); - uint32_t *face_remap = tmp_alloc.allocN(edit_mesh.numFaces); - uint32_t *vert_remap = tmp_alloc.allocN(edit_mesh.numVerts); - - for (CountT i = 0; i < edit_mesh.numHedges; i++) { - hedge_remap[i] = 0xFFFF'FFFF; - } - - CountT num_new_hedges = 0; - for (uint32_t orig_eid = edit_mesh.hedges[0].next; - orig_eid != 0; orig_eid = edit_mesh.hedges[orig_eid].next) { - if (hedge_remap[orig_eid] != 0xFFFF'FFFF) { - continue; - } - - const EditMesh::HEdge &cur_hedge = edit_mesh.hedges[orig_eid]; - uint32_t twin_eid = cur_hedge.twin; - assert(hedge_remap[twin_eid] = 0xFFFF'FFFF); - - hedge_remap[orig_eid] = num_new_hedges; - hedge_remap[twin_eid] = num_new_hedges + 1; - num_new_hedges += 2; - } - - CountT num_new_verts = 0; - for (uint32_t orig_vid = edit_mesh.verts[0].next; - orig_vid != 0; orig_vid = edit_mesh.verts[orig_vid].next) { - vert_remap[orig_vid] = num_new_verts++; - } - - CountT num_new_faces = 0; - for (uint32_t orig_fid = edit_mesh.faces[0].next; - orig_fid != 0; orig_fid = edit_mesh.faces[orig_fid].next) { - face_remap[orig_fid] = num_new_faces++; - } - - auto hedges_out = tmp_alloc.allocN(num_new_hedges); - auto face_base_hedges_out = tmp_alloc.allocN(num_new_faces); - auto face_planes_out = tmp_alloc.allocN(num_new_faces); - auto positions_out = tmp_alloc.allocN(num_new_verts); - - for (uint32_t orig_eid = edit_mesh.hedges[0].next; - orig_eid != 0; orig_eid = edit_mesh.hedges[orig_eid].next) { - const EditMesh::HEdge &orig_hedge = edit_mesh.hedges[orig_eid]; - - hedges_out[hedge_remap[orig_eid]] = HalfEdge { - .next = hedge_remap[orig_hedge.next], - .rootVertex = vert_remap[orig_hedge.vert], - .face = face_remap[orig_hedge.face], - }; - } - - for (uint32_t orig_vid = edit_mesh.verts[0].next; - orig_vid != 0; orig_vid = edit_mesh.verts[orig_vid].next) { - const EditMesh::Vert &orig_vert = edit_mesh.verts[orig_vid]; - positions_out[vert_remap[orig_vid]] = orig_vert.pos; - } - - for (uint32_t orig_fid = edit_mesh.faces[0].next; - orig_fid != 0; orig_fid = edit_mesh.faces[orig_fid].next) { - const EditMesh::Face &orig_face = edit_mesh.faces[orig_fid]; - - uint32_t new_face_idx = face_remap[orig_fid]; - - face_base_hedges_out[new_face_idx] = hedge_remap[orig_face.hedge]; - face_planes_out[new_face_idx] = orig_face.plane; - } - - return HalfEdgeMesh { - .halfEdges = hedges_out, - .faceBaseHalfEdges = face_base_hedges_out, - .facePlanes = face_planes_out, - .vertices = positions_out, - .numHalfEdges = uint32_t(num_new_hedges), - .numFaces = uint32_t(num_new_faces), - .numVertices = uint32_t(num_new_verts), - }; -} - -static inline HalfEdgeMesh buildHalfEdgeMesh( - StackAlloc &tmp_alloc, - const imp::SourceMesh &src_mesh) -{ - auto numFaceVerts = [&src_mesh](CountT face_idx) { - if (src_mesh.faceCounts == nullptr) { - return 3_u32; - } else { - return src_mesh.faceCounts[face_idx]; - } - }; - - using namespace madrona::math; - - uint32_t num_hedges = 0; - for (CountT face_idx = 0; face_idx < (CountT)src_mesh.numFaces; - face_idx++) { - num_hedges += numFaceVerts(face_idx); - } - - assert(num_hedges % 2 == 0); - - // We already know how many polygons there are - auto hedges_out = tmp_alloc.allocN(num_hedges); - auto face_base_hedges_out = tmp_alloc.allocN(src_mesh.numFaces); - auto face_planes_out = tmp_alloc.allocN(src_mesh.numFaces); - - std::unordered_map edge_to_hedge; - - auto makeEdgeID = [](uint32_t a_idx, uint32_t b_idx) { - return ((uint64_t)a_idx << 32) | (uint64_t)b_idx; - }; - - CountT num_assigned_hedges = 0; - const uint32_t *cur_face_indices = src_mesh.indices; - for (CountT face_idx = 0; face_idx < (CountT)src_mesh.numFaces; - face_idx++) { - CountT num_face_vertices = numFaceVerts(face_idx); - - Plane face_plane = computeNewellPlane(src_mesh.positions, - Span(cur_face_indices, num_face_vertices)); - - face_planes_out[face_idx] = face_plane; - - for (CountT vert_offset = 0; vert_offset < num_face_vertices; - vert_offset++) { - uint32_t a_idx = cur_face_indices[vert_offset]; - uint32_t b_idx = cur_face_indices[ - (vert_offset + 1) % num_face_vertices]; - - uint64_t cur_edge_id = makeEdgeID(a_idx, b_idx); - - auto cur_edge_lookup = edge_to_hedge.find(cur_edge_id); - if (cur_edge_lookup == edge_to_hedge.end()) { - uint32_t cur_hedge_id = num_assigned_hedges; - uint32_t twin_hedge_id = num_assigned_hedges + 1; - - num_assigned_hedges += 2; - - uint64_t twin_edge_id = makeEdgeID(b_idx, a_idx); - - auto [new_edge_iter, cur_inserted] = - edge_to_hedge.emplace(cur_edge_id, cur_hedge_id); - assert(cur_inserted); - - auto [new_twin_iter, twin_inserted] = - edge_to_hedge.emplace(twin_edge_id, twin_hedge_id); - assert(twin_inserted); - - cur_edge_lookup = new_edge_iter; - } - - uint32_t hedge_idx = cur_edge_lookup->second; - if (vert_offset == 0) { - face_base_hedges_out[face_idx] = hedge_idx; - } - - uint32_t c_idx = cur_face_indices[ - (vert_offset + 2) % num_face_vertices]; - - auto next_edge_id = makeEdgeID(b_idx, c_idx); - auto next_edge_lookup = edge_to_hedge.find(next_edge_id); - - // If next doesn't exist yet, we can assume it will be the next - // allocated half edge - uint32_t next_hedge_idx = next_edge_lookup == edge_to_hedge.end() ? - num_assigned_hedges : next_edge_lookup->second; - - hedges_out[hedge_idx] = HalfEdge { - .next = next_hedge_idx, - .rootVertex = a_idx, - .face = uint32_t(face_idx), - }; - } - - cur_face_indices += num_face_vertices; - } - - assert(num_assigned_hedges == num_hedges); - - return HalfEdgeMesh { - .halfEdges = hedges_out, - .faceBaseHalfEdges = face_base_hedges_out, - .facePlanes = face_planes_out, - .vertices = src_mesh.positions, - .numHalfEdges = uint32_t(num_hedges), - .numFaces = src_mesh.numFaces, - .numVertices = src_mesh.numVertices, - }; -} - -static bool processConvexHull(const imp::SourceMesh &src_mesh, - bool build_hull, - StackAlloc &tmp_alloc, - HalfEdgeMesh *out_mesh) -{ - if (!build_hull) { - // Just assume the input geometry is a convex hull with coplanar faces - // merged - *out_mesh = buildHalfEdgeMesh(tmp_alloc, src_mesh); - } else { - HullBuildData hull_data; - bool valid_input = initHullBuild( - Span(src_mesh.positions, src_mesh.numVertices), tmp_alloc, - &hull_data); - - if (!valid_input) { - return false; - } - - quickhullBuild(hull_data); - - *out_mesh = editMeshToRuntimeMesh(tmp_alloc, hull_data.mesh); - } - - return true; -} - -static bool processConvexHulls( - Span in_meshes, - bool build_convex_hulls, - StackAlloc &tmp_alloc, - HalfEdgeMesh *out_meshes) -{ - for (CountT hull_idx = 0; hull_idx < in_meshes.size(); hull_idx++) { - const imp::SourceMesh &mesh = in_meshes[hull_idx]; - bool success = processConvexHull( - mesh, build_convex_hulls, tmp_alloc, &out_meshes[hull_idx]); - - if (!success) { - return false; - } - } - - return true; -} - -// Below functions diagonalize the inertia tensor and compute the necessary -// rotation for diagonalization as a quaternion. -// Source: Computing the Singular Value Decomposition of 3x3 matrices with -// minimal branching and elementary floating point operations. -// McAdams et al 2011 - -// McAdams Algorithm 2: -static std::pair approxGivensQuaternion(Symmetric3x3 m) -{ - - constexpr float gamma = 5.82842712474619f; - constexpr float c_star = 0.9238795325112867f; - constexpr float s_star = 0.3826834323650898f; - - float a11 = m.diag[0], a12 = m.off[0], a22 = m.diag[1]; - - float ch = 2.f * (a11 - a22); - float sh = a12; - - float sh2 = sh * sh; - - // This isn't in the paper, but basically want to make sure the quaternion - // performs an identity rotation for already diagonal matrices - if (sh2 < 1e-20f) { - return { 1.f, 0.f }; - } - - float ch2 = ch * ch; - - bool b = (gamma * sh2) < ch2; - - float omega = rsqrtApprox(ch2 + sh2); - - ch = b ? (omega * ch) : c_star; - sh = b ? (omega * sh) : s_star; - - return { ch, sh }; -} - -// Equation 12: approxGivensQuaternion returns an unscaled quaternion, -// need to rescale -static Symmetric3x3 jacobiIterConjugation(Symmetric3x3 m, float ch, float sh) -{ - float ch2 = ch * ch; - float sh2 = sh * sh; - float q_scale = ch2 + sh2; - - float q11 = (ch2 - sh2) / q_scale; - float q12 = (-2.f * sh * ch) / q_scale; - float q21 = (2.f * sh * ch) / q_scale; - float q22 = (ch2 - sh2) / q_scale; - - // Output = Q^T * m * Q. Given above values for Q, direct solution to - // compute output (given 0s for other terms) computed using SymPy - - auto [m11, m22, m33] = m.diag; - auto [m12, m13, m23] = m.off; - - float m11q11_m12q21 = m11 * q11 + m12 * q21; - float m11q12_m12q22 = m11 * q12 + m12 * q22; - - float m12q11_m22q21 = m12 * q11 + m22 * q21; - float m12q12_m22q22 = m12 * q12 + m22 * q22; - - return Symmetric3x3 { - .diag = { - q11 * m11q11_m12q21 + q21 * m12q11_m22q21, - q12 * m11q12_m12q22 + q22 * m12q12_m22q22, - m33, - }, - .off = { - q12 * m11q11_m12q21 + q22 * m12q11_m22q21, - m13 * q11 + m23 * q21, - m13 * q12 + m23 * q22, - }, - }; -} - -// Inertia tensor is symmetric positive semi definite, so we only need to -// perform the symmetric eigenanalysis part of the algorithm. -// -// Jacobi order: (p, q) = (1, 2), (1, 3), (2, 3), (1, 2), (1, 3) ... -// Pairs: (1, 2) = (a11, a22, a12); (1, 3) = (a11, a33, a13); -// (2, 3) = (a22, a33, a23) -static void diagonalizeInertiaTensor(const Symmetric3x3 &m, - Diag3x3 *out_diag, - Quat *out_rot) -{ - using namespace math; - - constexpr CountT num_jacobi_iters = 8; - - Symmetric3x3 cur_mat = m; - Quat accumulated_rot { 1, 0, 0, 0 }; - for (CountT i = 0; i < num_jacobi_iters; i++) { -#if 0 - printf("Cur:\n" - "%f %f %f\n" - "%f %f %f\n" - "%f %f %f\n", - cur_mat[0].x, cur_mat[1].x, cur_mat[2].x, - cur_mat[0].y, cur_mat[1].y, cur_mat[2].y, - cur_mat[0].z, cur_mat[1].z, cur_mat[2].z); -#endif - - auto [ch1, sh1] = approxGivensQuaternion(cur_mat); - cur_mat = jacobiIterConjugation(cur_mat, ch1, sh1); - - // Rearrange matrix so unrotated elements are in upper left corner - std::swap(cur_mat.diag[1], cur_mat.diag[2]); - std::swap(cur_mat.off[0], cur_mat.off[1]); - - auto [ch2, sh2] = approxGivensQuaternion(cur_mat); - cur_mat = jacobiIterConjugation(cur_mat, ch2, sh2); - - std::swap(cur_mat.diag[0], cur_mat.diag[2]); - std::swap(cur_mat.off[0], cur_mat.off[2]); - - auto [ch3, sh3] = approxGivensQuaternion(cur_mat); - cur_mat = jacobiIterConjugation(cur_mat, ch3, sh3); - - cur_mat = Symmetric3x3 { - .diag = { cur_mat.diag[2], cur_mat.diag[0], cur_mat.diag[1] }, - .off = { cur_mat.off[1], cur_mat.off[2], cur_mat.off[0] }, - }; - - // This could be optimized - accumulated_rot = Quat { ch1, 0, 0, sh1 } * Quat { ch2, 0, sh2, 0 } * - Quat { ch3, sh3, 0, 0 } * accumulated_rot; - } - - Quat final_rot = accumulated_rot.normalize(); - - // Compute the diagonal (all other terms should be ~0) - { - Mat3x3 q = Mat3x3::fromQuat(final_rot); - - auto [m11, m22, m33] = m.diag; - auto [m12, m13, m23] = m.off; - - auto [q11, q21, q31] = q[0]; - auto [q12, q22, q32] = q[1]; - auto [q13, q23, q33] = q[2]; - - out_diag->d0 = q11 * (m11 * q11 + m12 * q21 + m13 * q31) + - q21 * (m12 * q11 + m22 * q21 + m23 * q31) + - q31 * (m13 * q11 + m23 * q21 + m33 * q31); - - out_diag->d1 = q12 * (m11 * q12 + m12 * q22 + m13 * q32) + - q22 * (m12 * q12 + m22 * q22 + m23 * q32) + - q32 * (m13 * q12 + m23 * q22 + m33 * q32); - - out_diag->d2 = q13 * (m11 * q13 + m12 * q23 + m13 * q33) + - q23 * (m12 * q13 + m22 * q23 + m23 * q33) + - q33 * (m13 * q13 + m23 * q23 + m33 * q33); - } - - *out_rot = final_rot; -} - -// http://number-none.com/blow/inertia/ -static inline MassProperties computeMassProperties( - const HalfEdgeMesh *convex_hulls, - const SourceCollisionObject &src_obj) -{ - using namespace math; - const Symmetric3x3 C_canonical { - .diag = Vector3 { 1.f / 60.f, 1.f / 60.f, 1.f / 60.f }, - .off = Vector3 { 1.f / 120.f, 1.f / 120.f, 1.f / 120.f }, - }; - constexpr float density = 1.f; - - Symmetric3x3 C_total { - .diag = Vector3::zero(), - .off = Vector3::zero(), - }; - - float m_total = 0; - Vector3 x_total = Vector3::zero(); - - auto processTet = [&](Vector3 v1, Vector3 v2, Vector3 v3) { - // Reference point is (0, 0, 0) so tet edges are just the vertex - // positions - Vector3 e1 = v1; - Vector3 e2 = v2; - Vector3 e3 = v3; - - // Covariance matrix - Mat3x3 A {{ e1, e2, e3 }}; - float det_A = A.determinant(); - Symmetric3x3 C = det_A * Symmetric3x3::AXAT(A, C_canonical); - - // Mass - float volume = 1.f / 6.f * det_A; - float m = volume * density; - - Vector3 x = 0.25f * e1 + 0.25f * e2 + 0.25f * e3; - - // Accumulate tetrahedron properties - float old_m_total = m_total; - m_total += m; - x_total = (x * m + x_total * old_m_total) / m_total; - - C_total += C; - }; - - for (const SourceCollisionPrimitive &prim : src_obj.prims) { - if (prim.type == CollisionPrimitive::Type::Sphere) { - // FIXME: need to allow offset for primitives - m_total += 1.f; - - float r = prim.sphere.radius; - - // Note that we need the sphere's covariance matrix, - // not the inertia tensor (hence 1/2 standard formulas) - float v = 1.f / 5.f * r * r; - C_total += Symmetric3x3 { - .diag = Vector3 { v, v, v }, - .off = Vector3::zero(), - }; - continue; - } else if (prim.type == CollisionPrimitive::Type::Plane) { - // Plane has infinite mass / inertia. The rest of the - // object must as well - - return MassProperties { - Diag3x3::uniform(INFINITY), - Vector3::zero(), - Quat { 1, 0, 0, 0 }, - }; - } - - // Hull primitive - - const HalfEdgeMesh &convex_hull = convex_hulls[prim.hullInput.hullIDX]; - - for (CountT face_idx = 0; face_idx < (CountT)convex_hull.numFaces; - face_idx++) { - uint32_t root_hedge_idx = convex_hull.faceBaseHalfEdges[face_idx]; - HalfEdge root_hedge = convex_hull.halfEdges[root_hedge_idx]; - Vector3 v1 = convex_hull.vertices[root_hedge.rootVertex]; - uint32_t cur_hedge_idx = root_hedge.next; - - while (true) { - HalfEdge cur_hedge = convex_hull.halfEdges[cur_hedge_idx]; - uint32_t next_hedge_idx = cur_hedge.next; - if (next_hedge_idx == root_hedge_idx) { - break; - } - - HalfEdge next_hedge = convex_hull.halfEdges[next_hedge_idx]; - - Vector3 v2 = convex_hull.vertices[cur_hedge.rootVertex]; - Vector3 v3 = convex_hull.vertices[next_hedge.rootVertex]; - - processTet(v1, v2, v3); - - cur_hedge_idx = next_hedge_idx; - } - } - } - - auto translateCovariance = [](const Symmetric3x3 &C, - Vector3 x, // COM - float m, - Vector3 delta_x) { - Symmetric3x3 delta_xxT_plus_xdeltaxT { - .diag = 2.f * Vector3 { - x.x * delta_x.x, - x.y * delta_x.y, - x.z * delta_x.z, - }, - .off = Vector3 { - x.x * delta_x.y + x.y * delta_x.x, - x.x * delta_x.z + x.z * delta_x.x, - x.y * delta_x.z + x.z * delta_x.y, - }, - }; - - Symmetric3x3 delta_xdelta_xT = Symmetric3x3::vvT(delta_x); - return C + m * (delta_xxT_plus_xdeltaxT + delta_xdelta_xT); - }; - - // Move accumulated covariance matrix to center of mass - C_total = translateCovariance(C_total, x_total, m_total, -x_total); - - float tr_C = C_total[0][0] + C_total[1][1] + C_total[2][2]; - const Symmetric3x3 tr_C_diag { - .diag = Vector3 { tr_C, tr_C, tr_C }, - .off = Vector3::zero(), - }; - - // Compute inertia tensor - Symmetric3x3 inertia_tensor = tr_C_diag - C_total; - - // Rescale total mass of inertia tensor (unless infinity) - float inv_mass = 1.f / m_total; - inertia_tensor *= inv_mass; - -#if 0 - printf("Inertia Tensor:\n" - "%f %f %f\n" - "%f %f %f\n" - "%f %f %f\n" - "COM: (%f %f %f) mass: %f\n", - inertia_tensor[0].x, - inertia_tensor[1].x, - inertia_tensor[2].x, - inertia_tensor[0].y, - inertia_tensor[1].y, - inertia_tensor[2].y, - inertia_tensor[0].z, - inertia_tensor[1].z, - inertia_tensor[2].z, - x_total.x, - x_total.y, - x_total.z, - m_total - ); -#endif - - Diag3x3 diag_inertia; - Quat rot_to_diag; - diagonalizeInertiaTensor(inertia_tensor, &diag_inertia, &rot_to_diag); - -#if 0 - printf("Diag Inertia tensor: (%f %f %f) rot: (%f %f %f %f)\n\n", - diag_inertia.d0, diag_inertia.d1, diag_inertia.d2, - rot_to_diag.w, - rot_to_diag.x, - rot_to_diag.y, - rot_to_diag.z); -#endif - - return MassProperties { - diag_inertia, - x_total, - rot_to_diag, - }; -} - -static inline RigidBodyMassData toMassData(const MassProperties &mass_props, - float inv_m) -{ - Diag3x3 inv_inertia = inv_m / mass_props.inertiaTensor; - - return { - .invMass = inv_m, - .invInertiaTensor = Vector3 { // FIXME - inv_inertia.d0, - inv_inertia.d1, - inv_inertia.d2, - }, - .toCenterOfMass = mass_props.centerOfMass, - .toInteriaFrame = mass_props.toDiagonal, - }; -} - -static void computeRigidBodiesMetadata( - const HalfEdgeMesh *convex_hulls, - Span collision_objs, - RigidBodyMetadata *out_metadatas) -{ - for (CountT obj_idx = 0; obj_idx < collision_objs.size(); obj_idx++) { - const SourceCollisionObject &collision_obj = collision_objs[obj_idx]; - - MassProperties mass_props = computeMassProperties( - convex_hulls, collision_obj); - - out_metadatas[obj_idx] = RigidBodyMetadata { - .mass = toMassData(mass_props, collision_obj.invMass), - .friction = collision_obj.friction, - }; - } -} - -static void setupSpherePrimitive(const SourceCollisionPrimitive &src_prim, - CollisionPrimitive *out_prim, - AABB *out_aabb) -{ - out_prim->sphere = src_prim.sphere; - - const float r = src_prim.sphere.radius; - - *out_aabb = AABB { - .pMin = { -r, -r, -r }, - .pMax = { r, r, r }, - }; -} - -static void setupPlanePrimitive(const SourceCollisionPrimitive &, - CollisionPrimitive *out_prim, - AABB *out_aabb) -{ - out_prim->plane = CollisionPrimitive::Plane {}; - - *out_aabb = AABB { - .pMin = { -FLT_MAX, -FLT_MAX, -FLT_MAX }, - .pMax = { FLT_MAX, FLT_MAX, 0 }, - }; -} - -static void setupHullPrimitive(const SourceCollisionPrimitive &src_prim, - const HalfEdgeMesh *hull_meshes, - CollisionPrimitive *out_prim, - AABB *out_aabb) -{ - const HalfEdgeMesh &hull_mesh = hull_meshes[src_prim.hullInput.hullIDX]; - - AABB mesh_aabb = AABB::point(hull_mesh.vertices[0]); - for (CountT vert_idx = 1; vert_idx < (CountT)hull_mesh.numVertices; - vert_idx++) { - mesh_aabb.expand(hull_mesh.vertices[vert_idx]); - } - - out_prim->hull.halfEdgeMesh = hull_mesh; - *out_aabb = mesh_aabb; -} - -static void setupRigidBodyAABBsAndPrimitives( - HalfEdgeMesh *hull_meshes, - Span collision_objs, - CollisionPrimitive *out_prims, - AABB *out_prim_aabbs, - AABB *out_obj_aabbs, - uint32_t *out_prim_offsets, - uint32_t *out_prim_counts) -{ - using Type = CollisionPrimitive::Type; - - uint32_t cur_prim_offset = 0; - for (CountT obj_idx = 0; obj_idx < collision_objs.size(); obj_idx++) { - const SourceCollisionObject &collision_obj = collision_objs[obj_idx]; - - CountT num_prims = collision_obj.prims.size(); - CollisionPrimitive *obj_prims = out_prims + cur_prim_offset; - AABB *prim_aabbs = out_prim_aabbs + cur_prim_offset; - - auto obj_aabb = AABB::invalid(); - - for (CountT prim_idx = 0; prim_idx < num_prims; prim_idx++) { - const SourceCollisionPrimitive &src_prim = - collision_obj.prims[prim_idx]; - - CollisionPrimitive *out_prim = &obj_prims[prim_idx]; - out_prim->type = src_prim.type; - AABB prim_aabb; - - switch (src_prim.type) { - case Type::Sphere: { - setupSpherePrimitive(src_prim, out_prim, &prim_aabb); - } break; - case Type::Plane: { - setupPlanePrimitive(src_prim, out_prim, &prim_aabb); - } break; - case Type::Hull: { - setupHullPrimitive(src_prim, hull_meshes, - out_prim, &prim_aabb); - } break; - } - - prim_aabbs[prim_idx] = prim_aabb; - obj_aabb = AABB::merge(obj_aabb, prim_aabb); - } - - out_obj_aabbs[obj_idx] = obj_aabb; - out_prim_offsets[obj_idx] = cur_prim_offset; - out_prim_counts[obj_idx] = (uint32_t)num_prims; - - cur_prim_offset += (uint32_t)num_prims; - } -} - -void * RigidBodyAssets::processRigidBodyAssets( - Span convex_hull_meshes, - Span collision_objs, - bool build_convex_hulls, - StackAlloc &tmp_alloc, - RigidBodyAssets *out_assets, - CountT *out_num_bytes) -{ - auto tmp_frame = tmp_alloc.push(); - - HalfEdgeMesh *built_hulls = - tmp_alloc.allocN(convex_hull_meshes.size()); - - auto hull_build_frame = tmp_alloc.push(); - - bool hull_success = processConvexHulls(convex_hull_meshes, - build_convex_hulls, - tmp_alloc, - built_hulls); - - if (!hull_success) { - tmp_alloc.pop(hull_build_frame); - tmp_alloc.pop(tmp_frame); - return nullptr; - } - - CountT total_num_prims = 0; - for (CountT obj_idx = 0; obj_idx < collision_objs.size(); obj_idx++) { - const SourceCollisionObject &collision_obj = collision_objs[obj_idx]; - CountT cur_num_prims = collision_obj.prims.size(); - total_num_prims += cur_num_prims; - } - - CountT total_num_halfedges = 0; - CountT total_num_faces = 0; - CountT total_num_verts = 0; - for (CountT hull_idx = 0; hull_idx < convex_hull_meshes.size(); - hull_idx++) { - const HalfEdgeMesh &hull_mesh = built_hulls[hull_idx]; - - total_num_halfedges += hull_mesh.numHalfEdges; - total_num_faces += hull_mesh.numFaces; - total_num_verts += hull_mesh.numVertices; - } - - auto buffer_sizes = std::to_array({ - (int64_t)sizeof(HalfEdge) * total_num_halfedges, // halfEdges - (int64_t)sizeof(uint32_t) * total_num_faces, // faceBaseHalfEdges - (int64_t)sizeof(Plane) * total_num_faces, // facePlanes - (int64_t)sizeof(Vector3) * total_num_verts, // vertices - (int64_t)sizeof(CollisionPrimitive) * total_num_prims, // prims - (int64_t)sizeof(AABB) * total_num_prims, // primAABBs - (int64_t)sizeof(RigidBodyMetadata) * - collision_objs.size(), // metadatas - (int64_t)sizeof(AABB) * - collision_objs.size(), // obj_aabbs - (int64_t)sizeof(uint32_t) * - collision_objs.size(), // prim_offsets - (int64_t)sizeof(uint32_t) * - collision_objs.size(), // prim_counts - }); - - int64_t buffer_offsets[buffer_sizes.size() - 1]; - int64_t num_buffer_bytes = utils::computeBufferOffsets( - buffer_sizes, buffer_offsets, 64); - - char *buffer = (char *)malloc(num_buffer_bytes); - RigidBodyAssets assets { - .hullData = { - .halfEdges = (HalfEdge *)buffer, - .faceBaseHalfEdges = (uint32_t *)(buffer + buffer_offsets[0]), - .facePlanes = (Plane *)(buffer + buffer_offsets[1]), - .vertices = (Vector3 *)(buffer + buffer_offsets[2]), - .numHalfEdges = (uint32_t)total_num_halfedges, - .numFaces = (uint32_t)total_num_faces, - .numVerts = (uint32_t)total_num_verts, - }, - .primitives = (CollisionPrimitive *)(buffer + buffer_offsets[3]), - .primitiveAABBs = (AABB *)(buffer + buffer_offsets[4]), - .metadatas = (RigidBodyMetadata *)(buffer + buffer_offsets[5]), - .objAABBs = (AABB *)(buffer + buffer_offsets[6]), - .primOffsets = (uint32_t *)(buffer + buffer_offsets[7]), - .primCounts = (uint32_t *)(buffer + buffer_offsets[8]), - .numConvexHulls = (uint32_t)convex_hull_meshes.size(), - .totalNumPrimitives = (uint32_t)total_num_prims, - .numObjs = (uint32_t)collision_objs.size(), - }; - - CountT cur_halfedge_offset = 0; - CountT cur_face_offset = 0; - CountT cur_vert_offset = 0; - for (CountT hull_idx = 0; hull_idx < convex_hull_meshes.size(); - hull_idx++) { - HalfEdgeMesh &hull_mesh = built_hulls[hull_idx]; - - HalfEdge *he_out = &assets.hullData.halfEdges[cur_halfedge_offset]; - uint32_t *face_bases_out = - &assets.hullData.faceBaseHalfEdges[cur_face_offset]; - Plane *face_planes_out = &assets.hullData.facePlanes[cur_face_offset]; - Vector3 *verts_out = &assets.hullData.vertices[cur_vert_offset]; - - memcpy(he_out, hull_mesh.halfEdges, - sizeof(HalfEdge) * hull_mesh.numHalfEdges); - memcpy(face_bases_out, hull_mesh.faceBaseHalfEdges, - sizeof(uint32_t) * hull_mesh.numFaces); - memcpy(face_planes_out, hull_mesh.facePlanes, - sizeof(Plane) * hull_mesh.numFaces); - memcpy(verts_out, hull_mesh.vertices, - sizeof(Vector3) * hull_mesh.numVertices); - - hull_mesh.halfEdges = he_out; - hull_mesh.faceBaseHalfEdges = face_bases_out; - hull_mesh.facePlanes = face_planes_out; - hull_mesh.vertices = verts_out; - - cur_halfedge_offset += hull_mesh.numHalfEdges; - cur_face_offset += hull_mesh.numFaces; - cur_vert_offset += hull_mesh.numVertices; - } - - tmp_alloc.pop(hull_build_frame); - - setupRigidBodyAABBsAndPrimitives(built_hulls, - collision_objs, - assets.primitives, - assets.primitiveAABBs, - assets.objAABBs, - assets.primOffsets, - assets.primCounts); - - computeRigidBodiesMetadata( - built_hulls, collision_objs, assets.metadatas); - - tmp_alloc.pop(tmp_frame); - - *out_assets = assets; - *out_num_bytes = num_buffer_bytes; - return buffer; -} - -} diff --git a/src/physics/physics_impl.hpp b/src/physics/physics_impl.hpp deleted file mode 100644 index 1c5ee7ed..00000000 --- a/src/physics/physics_impl.hpp +++ /dev/null @@ -1,60 +0,0 @@ -#pragma once - -#include - -namespace madrona::phys { - -struct PhysicsSystemState { - float deltaT; - float h; - math::Vector3 g; - float gMagnitude; - float restitutionThreshold; - uint32_t contactArchetypeID; - uint32_t jointArchetypeID; -}; - -struct CandidateTemporary : Archetype {}; - -namespace broadphase { - -TaskGraphNodeID setupBVHTasks( - TaskGraphBuilder &builder, - Span deps); - -TaskGraphNodeID setupPreIntegrationTasks( - TaskGraphBuilder &builder, - Span deps); - -TaskGraphNodeID setupPostIntegrationTasks( - TaskGraphBuilder &builder, - Span deps); - -} - -namespace narrowphase { - -TaskGraphNodeID setupTasks( - TaskGraphBuilder &builder, - Span deps); - -} - -namespace RGDCols { - constexpr inline CountT Position = 2; - constexpr inline CountT Rotation = 3; - constexpr inline CountT Scale = 4; - constexpr inline CountT ObjectID = 5; - constexpr inline CountT ResponseType = 6; - constexpr inline CountT LeafID = 7; - constexpr inline CountT Velocity = 8; - constexpr inline CountT ExternalForce = 9; - constexpr inline CountT ExternalTorque = 10; - constexpr inline CountT SolverBase = 11; - - constexpr inline CountT CandidateCollision = 2; - constexpr inline CountT ContactConstraint = 2; - constexpr inline CountT JointConstraint = 2; -}; - -} diff --git a/src/physics/physics_loader.cpp b/src/physics/physics_loader.cpp deleted file mode 100644 index 7e65f155..00000000 --- a/src/physics/physics_loader.cpp +++ /dev/null @@ -1,341 +0,0 @@ -#include -#include - -#ifdef MADRONA_CUDA_SUPPORT -#include -#endif - -#include - -namespace madrona::phys { -using namespace geo; -using namespace math; - -#ifndef MADRONA_CUDA_SUPPORT -[[noreturn]] static void noCUDA() -{ - FATAL("PhysicsLoader: Not built with CUDA support"); -} -#endif - -struct PhysicsLoader::Impl { - CollisionPrimitive *primitives; - AABB *primAABBs; - - AABB *objAABBs; - uint32_t *rigidBodyPrimitiveOffsets; - uint32_t *rigidBodyPrimitiveCounts; - RigidBodyMetadata *metadatas; - - CountT curPrimOffset; - CountT curObjOffset; - - ObjectManager *mgr; - CountT maxPrims; - CountT maxObjs; - ExecMode execMode; - - static Impl * init(ExecMode exec_mode, CountT max_objects) - { - constexpr CountT max_prims_per_object = 20; - - size_t num_collision_prim_bytes = - sizeof(CollisionPrimitive) * max_objects * max_prims_per_object; - - size_t num_collision_aabb_bytes = - sizeof(AABB) * max_objects * max_prims_per_object; - - size_t num_obj_aabb_bytes = - sizeof(AABB) * max_objects; - - size_t num_offset_bytes = - sizeof(uint32_t) * max_objects; - - size_t num_count_bytes = - sizeof(uint32_t) * max_objects; - - size_t num_metadata_bytes = - sizeof(RigidBodyMetadata) * max_objects; - - CollisionPrimitive *primitives_ptr; - AABB *prim_aabb_ptr; - - AABB *obj_aabb_ptr; - uint32_t *offsets_ptr; - uint32_t *counts_ptr; - RigidBodyMetadata *metadata_ptr; - - ObjectManager *mgr; - - switch (exec_mode) { - case ExecMode::CPU: { - primitives_ptr = (CollisionPrimitive *)malloc( - num_collision_prim_bytes); - - prim_aabb_ptr = (AABB *)malloc( - num_collision_aabb_bytes); - - obj_aabb_ptr = (AABB *)malloc(num_obj_aabb_bytes); - - offsets_ptr = (uint32_t *)malloc(num_offset_bytes); - counts_ptr = (uint32_t *)malloc(num_count_bytes); - - metadata_ptr = - (RigidBodyMetadata *)malloc(num_metadata_bytes); - - mgr = new ObjectManager { - primitives_ptr, - prim_aabb_ptr, - obj_aabb_ptr, - offsets_ptr, - counts_ptr, - metadata_ptr, - }; - } break; - case ExecMode::CUDA: { -#ifndef MADRONA_CUDA_SUPPORT - noCUDA(); -#else - primitives_ptr = (CollisionPrimitive *)cu::allocGPU( - num_collision_prim_bytes); - - prim_aabb_ptr = (AABB *)cu::allocGPU( - num_collision_aabb_bytes); - - obj_aabb_ptr = (AABB *)cu::allocGPU(num_obj_aabb_bytes); - - offsets_ptr = (uint32_t *)cu::allocGPU(num_offset_bytes); - counts_ptr = (uint32_t *)cu::allocGPU(num_count_bytes); - - metadata_ptr = - (RigidBodyMetadata *)cu::allocGPU(num_metadata_bytes); - - mgr = (ObjectManager *)cu::allocGPU(sizeof(ObjectManager)); - - ObjectManager local { - primitives_ptr, - prim_aabb_ptr, - obj_aabb_ptr, - offsets_ptr, - counts_ptr, - metadata_ptr, - }; - - REQ_CUDA(cudaMemcpy(mgr, &local, sizeof(ObjectManager), - cudaMemcpyHostToDevice)); -#endif - } break; - default: MADRONA_UNREACHABLE(); - } - - return new Impl { - .primitives = primitives_ptr, - .primAABBs = prim_aabb_ptr, - .objAABBs = obj_aabb_ptr, - .rigidBodyPrimitiveOffsets = offsets_ptr, - .rigidBodyPrimitiveCounts = counts_ptr, - .metadatas = metadata_ptr, - .curPrimOffset = 0, - .curObjOffset = 0, - .mgr = mgr, - .maxPrims = max_objects * max_prims_per_object, - .maxObjs = max_objects, - .execMode = exec_mode, - }; - } -}; - -PhysicsLoader::PhysicsLoader(ExecMode exec_mode, CountT max_objects) - : impl_(Impl::init(exec_mode, max_objects)) -{} - -PhysicsLoader::~PhysicsLoader() -{ - if (impl_ == nullptr) { - return; - } - - switch (impl_->execMode) { - case ExecMode::CPU: { - delete impl_->mgr; - free(impl_->primitives); - free(impl_->primAABBs); - free(impl_->objAABBs); - free(impl_->rigidBodyPrimitiveOffsets); - free(impl_->rigidBodyPrimitiveCounts); - free(impl_->metadatas); - } break; - case ExecMode::CUDA: { -#ifndef MADRONA_CUDA_SUPPORT - noCUDA(); -#else - cu::deallocGPU(impl_->primitives); - cu::deallocGPU(impl_->primAABBs); - cu::deallocGPU(impl_->objAABBs); - cu::deallocGPU(impl_->rigidBodyPrimitiveOffsets); - cu::deallocGPU(impl_->rigidBodyPrimitiveCounts); - cu::deallocGPU(impl_->metadatas); -#endif - } break; - } -} - -PhysicsLoader::PhysicsLoader(PhysicsLoader &&o) = default; - -CountT PhysicsLoader::loadRigidBodies(const RigidBodyAssets &assets) -{ - CountT cur_obj_offset = impl_->curObjOffset; - impl_->curObjOffset += assets.numObjs; - CountT cur_prim_offset = impl_->curPrimOffset; - impl_->curPrimOffset += assets.totalNumPrimitives; - assert(impl_->curObjOffset <= impl_->maxObjs); - assert(impl_->curPrimOffset <= impl_->maxPrims); - - CollisionPrimitive *prims_dst = &impl_->primitives[cur_prim_offset]; - AABB *prim_aabbs_dst = &impl_->primAABBs[cur_prim_offset]; - - AABB *obj_aabbs_dst = &impl_->objAABBs[cur_obj_offset]; - uint32_t *offsets_dst = &impl_->rigidBodyPrimitiveOffsets[cur_obj_offset]; - uint32_t *counts_dst = &impl_->rigidBodyPrimitiveCounts[cur_obj_offset]; - RigidBodyMetadata *metadatas_dst = &impl_->metadatas[cur_obj_offset]; - - // FIXME: redo all this, leaks memory, slow, etc. Very non optimal on the - // CPU. - - uint32_t *offsets_tmp = (uint32_t *)malloc( - sizeof(uint32_t) * assets.numObjs); - for (CountT i = 0; i < (CountT)assets.numObjs; i++) { - offsets_tmp[i] = assets.primOffsets[i] + cur_prim_offset; - } - - HalfEdge *hull_halfedges; - uint32_t *hull_face_base_halfedges; - Plane *hull_face_planes; - Vector3 *hull_verts; - switch (impl_->execMode) { - case ExecMode::CPU: { - memcpy(prim_aabbs_dst, assets.primitiveAABBs, - sizeof(AABB) * assets.totalNumPrimitives); - - memcpy(obj_aabbs_dst, assets.objAABBs, - sizeof(AABB) * assets.numObjs); - memcpy(offsets_dst, offsets_tmp, - sizeof(uint32_t) * assets.numObjs); - memcpy(counts_dst, assets.primCounts, - sizeof(uint32_t) * assets.numObjs); - memcpy(metadatas_dst, assets.metadatas, - sizeof(RigidBodyMetadata) * assets.numObjs); - - hull_halfedges = (HalfEdge *)malloc( - sizeof(HalfEdge) * assets.hullData.numHalfEdges); - hull_face_base_halfedges = (uint32_t *)malloc( - sizeof(uint32_t) * assets.hullData.numFaces); - hull_face_planes = (Plane *)malloc( - sizeof(Plane) * assets.hullData.numFaces); - hull_verts = (Vector3 *)malloc( - sizeof(Vector3) * assets.hullData.numVerts); - - memcpy(hull_halfedges, assets.hullData.halfEdges, - sizeof(HalfEdge) * assets.hullData.numHalfEdges); - memcpy(hull_face_base_halfedges, assets.hullData.faceBaseHalfEdges, - sizeof(uint32_t) * assets.hullData.numFaces); - memcpy(hull_face_planes, assets.hullData.facePlanes, - sizeof(Plane) * assets.hullData.numFaces); - memcpy(hull_verts, assets.hullData.vertices, - sizeof(Vector3) * assets.hullData.numVerts); - } break; - case ExecMode::CUDA: { -#ifndef MADRONA_CUDA_SUPPORT - noCUDA(); -#else - cudaMemcpy(prim_aabbs_dst, assets.primitiveAABBs, - sizeof(AABB) * assets.totalNumPrimitives, - cudaMemcpyHostToDevice); - - cudaMemcpy(obj_aabbs_dst, assets.objAABBs, - sizeof(AABB) * assets.numObjs, - cudaMemcpyHostToDevice); - cudaMemcpy(offsets_dst, offsets_tmp, - sizeof(uint32_t) * assets.numObjs, - cudaMemcpyHostToDevice); - cudaMemcpy(counts_dst, assets.primCounts, - sizeof(uint32_t) * assets.numObjs, - cudaMemcpyHostToDevice); - cudaMemcpy(metadatas_dst, assets.metadatas, - sizeof(RigidBodyMetadata) * assets.numObjs, - cudaMemcpyHostToDevice); - - hull_halfedges = (HalfEdge *)cu::allocGPU( - sizeof(HalfEdge) * assets.hullData.numHalfEdges); - hull_face_base_halfedges = (uint32_t *)cu::allocGPU( - sizeof(uint32_t) * assets.hullData.numFaces); - hull_face_planes = (Plane *)cu::allocGPU( - sizeof(Plane) * assets.hullData.numFaces); - hull_verts = (Vector3 *)cu::allocGPU( - sizeof(Vector3) * assets.hullData.numVerts); - - cudaMemcpy(hull_halfedges, assets.hullData.halfEdges, - sizeof(HalfEdge) * assets.hullData.numHalfEdges, - cudaMemcpyHostToDevice); - cudaMemcpy(hull_face_base_halfedges, assets.hullData.faceBaseHalfEdges, - sizeof(uint32_t) * assets.hullData.numFaces, - cudaMemcpyHostToDevice); - cudaMemcpy(hull_face_planes, assets.hullData.facePlanes, - sizeof(Plane) * assets.hullData.numFaces, - cudaMemcpyHostToDevice); - cudaMemcpy(hull_verts, assets.hullData.vertices, - sizeof(Vector3) * assets.hullData.numVerts, - cudaMemcpyHostToDevice); -#endif - } - } - - auto primitives_tmp = (CollisionPrimitive *)malloc( - sizeof(CollisionPrimitive) * assets.totalNumPrimitives); - memcpy(primitives_tmp, assets.primitives, - sizeof(CollisionPrimitive) * assets.totalNumPrimitives); - - for (CountT i = 0; i < (CountT)assets.totalNumPrimitives; i++) { - CollisionPrimitive &cur_primitive = primitives_tmp[i]; - if (cur_primitive.type != CollisionPrimitive::Type::Hull) continue; - - HalfEdgeMesh &he_mesh = cur_primitive.hull.halfEdgeMesh; - - // FIXME: incoming HalfEdgeMeshes should have offsets or something - CountT hedge_offset = he_mesh.halfEdges - assets.hullData.halfEdges; - CountT face_offset = - he_mesh.facePlanes - assets.hullData.facePlanes; - CountT vert_offset = he_mesh.vertices - assets.hullData.vertices; - - he_mesh.halfEdges = hull_halfedges + hedge_offset; - he_mesh.faceBaseHalfEdges = hull_face_base_halfedges + face_offset; - he_mesh.facePlanes = hull_face_planes + face_offset; - he_mesh.vertices = hull_verts + vert_offset; - } - - switch (impl_->execMode) { - case ExecMode::CPU: { - memcpy(prims_dst, primitives_tmp, - sizeof(CollisionPrimitive) * assets.totalNumPrimitives); - } break; - case ExecMode::CUDA: { -#ifdef MADRONA_CUDA_SUPPORT - cudaMemcpy(prims_dst, primitives_tmp, - sizeof(CollisionPrimitive) * assets.totalNumPrimitives, - cudaMemcpyHostToDevice); -#endif - } break; - } - - free(primitives_tmp); - free(offsets_tmp); - - return cur_obj_offset; -} - -ObjectManager & PhysicsLoader::getObjectManager() -{ - return *impl_->mgr; -} - -} diff --git a/src/physics/tgs.cpp b/src/physics/tgs.cpp deleted file mode 100644 index 0218709a..00000000 --- a/src/physics/tgs.cpp +++ /dev/null @@ -1,304 +0,0 @@ -#include - -#include "physics_impl.hpp" -#include "tgs.hpp" - -// This implementation is inspired by the TGS-Soft solver from Solver2D: -// https://github.com/erincatto/solver2d/src/solve_tgs_soft.c -// Solver2D is MIT licensed, Copyright 2024 Erin Catto - -namespace madrona::phys::tgs { - -struct Contact : Archetype {}; -struct Joint : Archetype {}; - -// Any per-body solver state would go in components in this bundle -// (check XPBDRigidBodyState for example). -struct TGSRigidBodyState : Bundle< -> {}; - -struct SolverState { - Query jointQuery; - Query contactQuery; -}; - -using namespace base; -using namespace math; - -void registerTypes(ECSRegistry ®istry) -{ - registry.registerArchetype(); - registry.registerArchetype(); - - // Any components in the bundle specific to this solver must be - // registered first. - registry.registerBundle(); - - // This registers the solver's per-body state bundle for use in all - // rigid bodies in the system - registry.registerBundleAlias(); - - registry.registerSingleton(); -} - -void init(Context &ctx) -{ - new (&ctx.singleton()) SolverState { - .jointQuery = ctx.query(), - .contactQuery = ctx.query(), - }; -} - -void getSolverArchetypeIDs(uint32_t *contact_archetype_id, - uint32_t *joint_archetype_id) -{ - *contact_archetype_id = TypeTracker::typeID(); - *joint_archetype_id = TypeTracker::typeID(); -} - -static inline void solveJoints(Context &ctx, - SolverState &solver, - bool use_bias) -{ - (void)ctx; - (void)solver; - (void)use_bias; -} - -static inline void solveContacts(Context &ctx, - SolverState &solver, - bool use_bias) -{ - ctx.iterateQuery(solver.contactQuery, [&](ContactConstraint &contact) { - // Solve contact - (void)contact; - (void)use_bias; - }); -} - -inline void prepareContacts(Context &ctx, - ContactConstraint constraint) -{ - (void)ctx; - (void)constraint; -} - -inline void prepareJoints(Context &ctx, - JointConstraint constraint) -{ - (void)ctx; - (void)constraint; -} - -inline void integrateVelocities(Context &ctx, - Rotation q, - ResponseType response_type, - ExternalForce ext_force, - ExternalTorque ext_torque, - ObjectID obj_id, - Velocity &vel) -{ - if (response_type == ResponseType::Static) { - return; - } - - Vector3 v = vel.linear; - Vector3 omega = vel.angular; - - const auto &physics_sys = ctx.singleton(); - const ObjectManager &obj_mgr = *ctx.singleton().mgr; - const RigidBodyMetadata &metadata = obj_mgr.metadata[obj_id.idx]; - - float inv_m = metadata.mass.invMass; - Diag3x3 inv_I = Diag3x3::fromVec(metadata.mass.invInertiaTensor); - - float h = physics_sys.h; - - if (response_type == ResponseType::Dynamic) { - v += h * physics_sys.g; - } - - v += h * inv_m * ext_force; - - Diag3x3 I = { - (inv_I.d0 == 0) ? 0.0f : 1.0f / inv_I.d0, - (inv_I.d1 == 0) ? 0.0f : 1.0f / inv_I.d1, - (inv_I.d2 == 0) ? 0.0f : 1.0f / inv_I.d2, - }; - - Quat to_local = q.inv(); - - Vector3 tau_ext_local = to_local.rotateVec(ext_torque); - Vector3 omega_local = to_local.rotateVec(omega); - - // Integrate omega in local space - omega_local += - h * inv_I * (tau_ext_local - (cross(omega_local, I * omega))); - - omega = q.rotateVec(omega_local); - - // FIXME damping - - vel.linear = v; - vel.angular = omega; -} - -inline void warmStartContacts(Context &ctx, - ContactConstraint constraint) -{ - (void)ctx; - (void)constraint; -} - -inline void warmStartJoints(Context &ctx, - JointConstraint constraint) -{ - (void)ctx; - (void)constraint; -} - -inline void solveJointsBiased(Context &ctx, - SolverState &solver) -{ - solveJoints(ctx, solver, true); -} - -inline void solveContactsBiased(Context &ctx, - SolverState &solver) -{ - solveContacts(ctx, solver, true); -} - -inline void integratePositions(Context &ctx, - Position &pos, - Rotation &rot, - Velocity vel) -{ - Vector3 x = pos; - Quat q = rot; - - Vector3 v = vel.linear; - Vector3 omega = vel.angular; - - const auto &physics_sys = ctx.singleton(); - float h = physics_sys.h; - - x += h * v; - - Quat apply_omega = Quat::fromAngularVec(0.5f * h * omega); - - q += apply_omega * q; - q = q.normalize(); - - pos = x; - rot = q; -} - -inline void solveJointsUnbiased(Context &ctx, - SolverState &solver) -{ - solveJoints(ctx, solver, false); -} - -inline void solveContactsUnbiased(Context &ctx, - SolverState &solver) -{ - solveContacts(ctx, solver, false); -} - -TaskGraphNodeID setupTGSSolverTasks( - TaskGraphBuilder &builder, - TaskGraphNodeID broadphase, - CountT num_substeps) -{ - auto run_narrowphase = narrowphase::setupTasks(builder, {broadphase}); - auto clear_broadphase = builder.addToGraph< - ClearTmpNode>({run_narrowphase}); - - auto constraints_ready = clear_broadphase; -#ifdef MADRONA_GPU_MODE - // The GPU backend requires constraints to be sorted before iterateQuery - // can be called. This requires sorting both Contact and Joint entities. - constraints_ready = builder.addToGraph>( - {constraints_ready}); - - constraints_ready = - builder.addToGraph({constraints_ready}); - - constraints_ready = builder.addToGraph>( - {constraints_ready}); - - constraints_ready = - builder.addToGraph({constraints_ready}); -#endif - - auto cur_node = constraints_ready; - - cur_node = builder.addToGraph>({cur_node}); - - cur_node = builder.addToGraph>({cur_node}); - - for (CountT i = 0; i < num_substeps; i++) { - cur_node = builder.addToGraph>({cur_node}); - - cur_node = builder.addToGraph>({cur_node}); - - cur_node = builder.addToGraph>({cur_node}); - - cur_node = builder.addToGraph>({cur_node}); - - cur_node = builder.addToGraph>({cur_node}); - - cur_node = builder.addToGraph>({cur_node}); - - cur_node = builder.addToGraph>({cur_node}); - - cur_node = builder.addToGraph>({cur_node}); - } - - // For now, we don't have any persistent contacts support, so clear - // the contacts (free the memory). - auto clear_contacts = builder.addToGraph< - ClearTmpNode>({cur_node}); - - return clear_contacts; -} - -} diff --git a/src/physics/tgs.hpp b/src/physics/tgs.hpp deleted file mode 100644 index ce86f6c0..00000000 --- a/src/physics/tgs.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include "physics_impl.hpp" - -namespace madrona::phys::tgs { - -void registerTypes(ECSRegistry ®istry); - -void getSolverArchetypeIDs(uint32_t *contact_archetype_id, - uint32_t *joint_archetype_id); - -void init(Context &ctx); - -TaskGraphNodeID setupTGSSolverTasks( - TaskGraphBuilder &builder, - TaskGraphNodeID broadphase, - CountT num_substeps); - -} diff --git a/src/physics/xpbd.cpp b/src/physics/xpbd.cpp deleted file mode 100644 index 393129cc..00000000 --- a/src/physics/xpbd.cpp +++ /dev/null @@ -1,1146 +0,0 @@ -#include -#include - -#include "physics_impl.hpp" -#include "xpbd.hpp" - -namespace madrona::phys::xpbd { - -struct XPBDContactState { - float lambdaN[4]; -}; - -struct Contact : Archetype< - ContactConstraint, - XPBDContactState -> {}; - -struct Joint : Archetype {}; - -struct SolverState { - Query jointQuery; - Query contactQuery; -}; - -struct SubstepPrevState { - math::Vector3 prevPosition; - math::Quat prevRotation; -}; - -struct PreSolvePositional { - math::Vector3 x; - math::Quat q; -}; - -struct PreSolveVelocity { - math::Vector3 v; - math::Vector3 omega; -}; - -struct XPBDRigidBodyState : Bundle< - SubstepPrevState, - PreSolvePositional, - PreSolveVelocity -> {}; - -namespace XPBDCols { - constexpr inline CountT SubstepPrevState = RGDCols::SolverBase; - constexpr inline CountT PreSolvePositional = RGDCols::SolverBase + 1; - constexpr inline CountT PreSolveVelocity = RGDCols::SolverBase + 2; -}; - -using namespace base; -using namespace math; - -static inline bool hasNaN(Vector3 v) -{ - return isnan(v.x) || isnan(v.y) || isnan(v.z); -} - -static inline bool hasNaN(Quat q) -{ - return isnan(q.w) || isnan(q.x) || isnan(q.y) || isnan(q.z); -} - -static inline Vector3 multDiag(Vector3 diag, Vector3 v) -{ - return Vector3 { - diag.x * v.x, - diag.y * v.y, - diag.z * v.z, - }; -} - -[[maybe_unused]] static inline Vector3 computeEnergy( - float inv_m, Vector3 inv_I, Vector3 v, Vector3 omega, Quat q) -{ - if (inv_m == 0.f || inv_I.x == 0.f || inv_I.y == 0.f || inv_I.z == 0.f) { - return {0.f, 0.f, 0.f}; - } - - float m = 1.f / inv_m; - Vector3 I { - 1.f / inv_I.x, - 1.f / inv_I.y, - 1.f / inv_I.z, - }; - - float linear = m * v.length2(); - - Vector3 omega_local = q.inv().rotateVec(omega); - float rotational = dot(omega_local, multDiag(I, omega_local)); - - return { - 0.5f * (linear + rotational), - 0.5f * linear, - 0.5f * rotational, - }; -} - -inline void substepRigidBodies(Context &ctx, - Position &pos, - Rotation &rot, - const Velocity &vel, - const ObjectID &obj_id, - ResponseType response_type, - ExternalForce &ext_force, - ExternalTorque &ext_torque, - SubstepPrevState &prev_state, - PreSolvePositional &presolve_pos, - PreSolveVelocity &presolve_vel) -{ - Vector3 x = pos; - Quat q = rot; - - Vector3 v = vel.linear; - Vector3 omega = vel.angular; - - if (response_type == ResponseType::Static) { - // FIXME: currently presolve_pos and prev_state need to be set every - // frame even for static objects. A better solution would be on - // creation / making a non-static object static, these variables are - // set once. This would require a dedicated API to control this rather - // than just setting objects to ResponseType::Static - prev_state.prevPosition = x; - prev_state.prevRotation = q; - - presolve_pos.x = x; - presolve_pos.q = q; - presolve_vel.v = Vector3::zero(); - presolve_vel.omega = Vector3::zero(); - - return; - } - - prev_state.prevPosition = x; - prev_state.prevRotation = q; - - const auto &physics_sys = ctx.singleton(); - const ObjectManager &obj_mgr = *ctx.singleton().mgr; - const RigidBodyMetadata &metadata = obj_mgr.metadata[obj_id.idx]; - - float inv_m = metadata.mass.invMass; - Vector3 inv_I = metadata.mass.invInertiaTensor; - - float h = physics_sys.h; - - if (response_type == ResponseType::Dynamic) { - v += h * physics_sys.g; - } - - v += h * inv_m * ext_force; - - x += h * v; - - Vector3 I = { - (inv_I.x == 0) ? 0.0f : 1.0f / inv_I.x, - (inv_I.y == 0) ? 0.0f : 1.0f / inv_I.y, - (inv_I.z == 0) ? 0.0f : 1.0f / inv_I.z - }; - - Quat to_local = q.inv(); - - Vector3 tau_ext_local = to_local.rotateVec(ext_torque); - Vector3 omega_local = to_local.rotateVec(omega); - - Vector3 I_omega_local = multDiag(I, omega_local); - - omega_local += h * multDiag(inv_I, - tau_ext_local - (cross(omega_local, I_omega_local))); - - omega = q.rotateVec(omega_local); - - Quat apply_omega = Quat::fromAngularVec(0.5f * h * omega); - - q += apply_omega * q; - q = q.normalize(); - - pos = x; - rot = q; - - presolve_pos.x = x; - presolve_pos.q = q; - presolve_vel.v = v; - presolve_vel.omega = omega; -} - -[[maybe_unused]] inline void checkSubstep(Context &, - Entity, - const Position &pos, - const Rotation &rot, - const Velocity &vel, - const ObjectID &, - const ResponseType &, - const ExternalForce &, - const ExternalTorque &, - const SubstepPrevState &, - const PreSolvePositional &, - const PreSolveVelocity &) -{ - assert(!hasNaN(pos)); - assert(!hasNaN(rot)); - assert(!hasNaN(vel.linear)); - assert(!hasNaN(vel.angular)); -} - -static inline float generalizedInverseMass(Vector3 torque_axis, - Vector3 rot_axis, - float inv_m) -{ - return inv_m + dot(torque_axis, rot_axis); -} - -static float computePositionalLambda( - Vector3 torque_axis1, Vector3 torque_axis2, - Vector3 rot_axis1, Vector3 rot_axis2, - float inv_m1, float inv_m2, - float c, float alpha_tilde) -{ - float w1 = generalizedInverseMass(torque_axis1, rot_axis1, inv_m1); - float w2 = generalizedInverseMass(torque_axis2, rot_axis2, inv_m2); - - return -c / (w1 + w2 + alpha_tilde); -} - -MADRONA_ALWAYS_INLINE static inline void applyPositionalUpdate( - Vector3 &x1, Vector3 &x2, - Quat &q1, Quat &q2, - Vector3 rot_axis_local1, Vector3 rot_axis_local2, - float inv_m1, float inv_m2, - Vector3 n, - float delta_lambda) -{ - x1 += delta_lambda * inv_m1 * n; - x2 -= delta_lambda * inv_m2 * n; - - float half_lambda = 0.5f * delta_lambda; - - Vector3 q1_update_angular_local = half_lambda * rot_axis_local1; - Vector3 q1_update_angular = q1.rotateVec(q1_update_angular_local); - - Vector3 q2_update_angular_local = half_lambda * rot_axis_local2; - Vector3 q2_update_angular = q2.rotateVec(q2_update_angular_local); - - q1 += Quat::fromAngularVec(q1_update_angular) * q1; - q2 -= Quat::fromAngularVec(q2_update_angular) * q2; - - // Paper doesn't explicitly call for normalization but we immediately - // use q1 and q2 for the next constraint - q1 = q1.normalize(); - q2 = q2.normalize(); -} - -MADRONA_ALWAYS_INLINE static inline float applyPositionalUpdate( - Vector3 &x1, Vector3 &x2, - Quat &q1, Quat &q2, - Vector3 r1, Vector3 r2, - float inv_m1, float inv_m2, - Vector3 inv_I1, Vector3 inv_I2, - Vector3 n_world, - float c, float alpha_tilde) -{ - Vector3 n_local1 = q1.inv().rotateVec(n_world); - Vector3 n_local2 = q2.inv().rotateVec(n_world); - - Vector3 torque_axis_local1 = cross(r1, n_local1); - Vector3 torque_axis_local2 = cross(r2, n_local2); - - Vector3 rot_axis_local1 = multDiag(inv_I1, torque_axis_local1); - Vector3 rot_axis_local2 = multDiag(inv_I2, torque_axis_local2); - - float lambda = computePositionalLambda( - torque_axis_local1, torque_axis_local2, - rot_axis_local1, rot_axis_local2, - inv_m1, inv_m2, - c, alpha_tilde); - - applyPositionalUpdate( - x1, x2, - q1, q2, - rot_axis_local1, rot_axis_local2, - inv_m1, inv_m2, - n_world, lambda); - - return lambda; -} - -MADRONA_ALWAYS_INLINE static inline -std::pair computeAngularUpdate( - Quat q1, Quat q2, - Vector3 inv_I1, Vector3 inv_I2, - Vector3 n1, Vector3 n2, - float theta, - float alpha_tilde) -{ - Vector3 local_rot_axis1 = multDiag(inv_I1, n1); - Vector3 local_rot_axis2 = multDiag(inv_I2, n2); - - float w1 = dot(n1, local_rot_axis1); - float w2 = dot(n2, local_rot_axis2); - - float delta_lambda = -theta / (w1 + w2 + alpha_tilde); - - float half_lambda = 0.5f * delta_lambda; - Vector3 q1_update_angular_local = half_lambda * local_rot_axis1; - Vector3 q2_update_angular_local = half_lambda * local_rot_axis2; - - return { - Quat::fromAngularVec(q1.rotateVec(q1_update_angular_local)), - Quat::fromAngularVec(q2.rotateVec(q2_update_angular_local)), - }; -} - -static void applyAngularUpdate( - Quat &q1, Quat &q2, - Quat q1_update, Quat q2_update) -{ - q1 = (q1 + q1_update * q1).normalize(); - q2 = (q2 - q2_update * q2).normalize(); -} - -MADRONA_ALWAYS_INLINE static inline void handleContactConstraint( - Vector3 &x1, Vector3 &x2, - Quat &q1, Quat &q2, - SubstepPrevState prev1, SubstepPrevState prev2, - float inv_m1, float inv_m2, - Vector3 inv_I1, Vector3 inv_I2, - Vector3 r1, Vector3 r2, - Vector3 n_world, - float avg_mu_s, - float *lambda_n_out, - float *lambda_t_out) -{ - Vector3 p1 = q1.rotateVec(r1) + x1; - Vector3 p2 = q2.rotateVec(r2) + x2; - - float d = dot(p1 - p2, n_world); - - if (d <= 0) { - return; - } - - float lambda_n = applyPositionalUpdate( - x1, x2, - q1, q2, - r1, r2, - inv_m1, inv_m2, - inv_I1, inv_I2, - n_world, - d, 0); - *lambda_n_out = lambda_n; - - Vector3 x1_prev = prev1.prevPosition; - Quat q1_prev = prev1.prevRotation; - - Vector3 x2_prev = prev2.prevPosition; - Quat q2_prev = prev2.prevRotation; - - Vector3 p1_hat = q1_prev.rotateVec(r1) + x1_prev; - Vector3 p2_hat = q2_prev.rotateVec(r2) + x2_prev; - - // Update p1 and p2 so static friction covers any drift - // as a result of the positional correction along the normal - p1 = q1.rotateVec(r1) + x1; - p2 = q2.rotateVec(r2) + x2; - - Vector3 delta_p = (p1 - p1_hat) - (p2 - p2_hat); - Vector3 delta_p_t = delta_p - dot(delta_p, n_world) * n_world; - - float tangential_magnitude = delta_p_t.length(); - - if (tangential_magnitude > 0.f) { - Vector3 t_world = delta_p_t / tangential_magnitude; - Vector3 t_local1 = q1.inv().rotateVec(t_world); - Vector3 t_local2 = q2.inv().rotateVec(t_world); - - Vector3 friction_torque_axis_local1 = cross(r1, t_local1); - Vector3 friction_torque_axis_local2 = cross(r2, t_local2); - - Vector3 friction_rot_axis_local1 = - multDiag(inv_I1, friction_torque_axis_local1); - - Vector3 friction_rot_axis_local2 = - multDiag(inv_I2, friction_torque_axis_local2); - - float lambda_t = computePositionalLambda( - friction_torque_axis_local1, friction_torque_axis_local2, - friction_rot_axis_local1, friction_rot_axis_local2, - inv_m1, inv_m2, - tangential_magnitude, 0); - float lambda_threshold = lambda_n * avg_mu_s; - - if (lambda_t > lambda_threshold) { - *lambda_t_out = lambda_t; - - applyPositionalUpdate( - x1, x2, - q1, q2, - friction_rot_axis_local1, friction_rot_axis_local2, - inv_m1, inv_m2, - t_world, lambda_t); - } - } -} - -MADRONA_ALWAYS_INLINE static inline std::pair -getLocalSpaceContacts(const PreSolvePositional &presolve_pos1, - const PreSolvePositional &presolve_pos2, - Vector3 contact1, float penetration_depth, - Vector3 contact_normal) -{ - Vector3 contact2 = - contact1 - contact_normal * penetration_depth; - - // Transform the contact points into local space for a & b - Vector3 r1 = presolve_pos1.q.inv().rotateVec(contact1 - presolve_pos1.x); - Vector3 r2 = presolve_pos2.q.inv().rotateVec(contact2 - presolve_pos2.x); - - return { r1, r2 }; -} - -static bool getAvgContact(ContactConstraint contact, Vector3 *avg_out, float *penetration_out) -{ - Vector3 avg_contact = Vector3::zero(); - - float max_penetration = -FLT_MAX; - float penetration_sum = 0.f; - for (CountT i = 0; i < contact.numPoints; i++) { - Vector4 pt = contact.points[i]; - if (pt.w > max_penetration) { - max_penetration = pt.w; - } - - penetration_sum += pt.w; - } - - if (penetration_sum == 0.f) { - return true; - } - - for (CountT i = 0; i < contact.numPoints; i++) { - Vector4 pt = contact.points[i]; - avg_contact += pt.w / penetration_sum * pt.xyz(); - } - - *avg_out = avg_contact; - *penetration_out = max_penetration; - - return false; -} - -// For now, this function assumes both a & b are dynamic objects. -// FIXME: Need to add dynamic / static variant or handle missing the velocity -// component for static objects. -static inline void handleContact(Context &ctx, - ObjectManager &obj_mgr, - ContactConstraint contact, - float *lambdas) -{ - Position *x1_ptr = &ctx.getDirect(RGDCols::Position, contact.ref); - Position *x2_ptr = &ctx.getDirect(RGDCols::Position, contact.alt); - - Rotation *q1_ptr = &ctx.getDirect(RGDCols::Rotation, contact.ref); - Rotation *q2_ptr = &ctx.getDirect(RGDCols::Rotation, contact.alt); - - SubstepPrevState prev1 = ctx.getDirect( - XPBDCols::SubstepPrevState, contact.ref); - SubstepPrevState prev2 = ctx.getDirect( - XPBDCols::SubstepPrevState, contact.alt); - - PreSolvePositional presolve_pos1 = ctx.getDirect( - XPBDCols::PreSolvePositional, contact.ref); - PreSolvePositional presolve_pos2 = ctx.getDirect( - XPBDCols::PreSolvePositional, contact.alt); - - ObjectID obj_id1 = ctx.getDirect( - RGDCols::ObjectID, contact.ref); - ObjectID obj_id2 = ctx.getDirect(RGDCols::ObjectID, contact.alt); - - ResponseType resp_type1 = ctx.getDirect( - RGDCols::ResponseType, contact.ref); - ResponseType resp_type2 = ctx.getDirect( - RGDCols::ResponseType, contact.alt); - - RigidBodyMetadata metadata1 = obj_mgr.metadata[obj_id1.idx]; - RigidBodyMetadata metadata2 = obj_mgr.metadata[obj_id2.idx]; - - Vector3 x1 = *x1_ptr; - Vector3 x2 = *x2_ptr; - - Quat q1 = *q1_ptr; - Quat q2 = *q2_ptr; - - float inv_m1 = metadata1.mass.invMass; - float inv_m2 = metadata2.mass.invMass; - - Vector3 inv_I1 = metadata1.mass.invInertiaTensor; - Vector3 inv_I2 = metadata2.mass.invInertiaTensor; - - if (resp_type1 == ResponseType::Static) { - inv_m1 = 0.f; - inv_I1 = Vector3::zero(); - } - - if (resp_type2 == ResponseType::Static) { - inv_m2 = 0.f; - inv_I2 = Vector3::zero(); - } - - float mu_s1 = metadata1.friction.muS; - float mu_s2 = metadata2.friction.muS; - - float avg_mu_s = 0.5f * (mu_s1 + mu_s2); - - Vector3 avg_contact_pos; - float contact_pos_penetration; - bool zero_separation = getAvgContact(contact, &avg_contact_pos, &contact_pos_penetration); - if (zero_separation) { - return; - } - - { - CountT i = 0; - - auto [r1, r2] = getLocalSpaceContacts( - presolve_pos1, presolve_pos2, - avg_contact_pos, contact_pos_penetration, contact.normal); - - float lambda_n = 0.f; - float lambda_t = 0.f; - - handleContactConstraint(x1, x2, - q1, q2, - prev1, prev2, - inv_m1, inv_m2, - inv_I1, inv_I2, - r1, r2, - contact.normal, - avg_mu_s, - &lambda_n, - &lambda_t); - - lambdas[i] = lambda_n; - } - - *x1_ptr = x1; - *x2_ptr = x2; - - *q1_ptr = q1; - *q2_ptr = q2; -} - -static void applyJointOrientationConstraint( - Quat &q1, Quat &q2, - Quat attach_q1, Quat attach_q2, - Vector3 inv_I1, Vector3 inv_I2) -{ - Quat orientation1 = (q1 * attach_q1).normalize(); - Quat orientation2 = (q2 * attach_q2).normalize(); - - Quat diff = orientation1 * orientation2.inv(); - - Vector3 delta_q = - 2.f * math::Vector3 { diff.x, diff.y, diff.z }; - float delta_q_magnitude = delta_q.length(); - - if (delta_q_magnitude > 0) { - delta_q /= delta_q_magnitude; - Vector3 delta_q_local1 = q1.inv().rotateVec(delta_q); - Vector3 delta_q_local2 = q2.inv().rotateVec(delta_q); - - auto [update_q1, update_q2] = computeAngularUpdate( - q1, q2, - inv_I1, inv_I2, - delta_q_local1, delta_q_local2, - delta_q_magnitude, 0); - - applyAngularUpdate(q1, q2, update_q1, update_q2); - } -} - -static void applyJointAxisConstraint( - Quat &q1, Quat &q2, - Vector3 axis1_local, Vector3 axis2_local, - Vector3 inv_I1, Vector3 inv_I2) -{ - Vector3 axis1 = q1.rotateVec(axis1_local); - Vector3 axis2 = q2.rotateVec(axis2_local); - - Vector3 delta_q = cross(axis1, axis2); - float delta_q_magnitude = delta_q.length(); - - if (delta_q_magnitude > 0) { - delta_q /= delta_q_magnitude; - Vector3 delta_q_local1 = q1.inv().rotateVec(delta_q); - Vector3 delta_q_local2 = q2.inv().rotateVec(delta_q); - - auto [update_q1, update_q2] = computeAngularUpdate( - q1, q2, - inv_I1, inv_I2, - delta_q_local1, delta_q_local2, - delta_q_magnitude, 0); - - applyAngularUpdate(q1, q2, update_q1, update_q2); - } -} - -inline void handleJointConstraint(Context &ctx, - JointConstraint joint) -{ - Loc l1 = ctx.loc(joint.e1); - Loc l2 = ctx.loc(joint.e2); - - Vector3 *x1_ptr = &ctx.getDirect(RGDCols::Position, l1); - Vector3 *x2_ptr = &ctx.getDirect(RGDCols::Position, l2); - Quat *q1_ptr = &ctx.getDirect(RGDCols::Rotation, l1); - Quat *q2_ptr = &ctx.getDirect(RGDCols::Rotation, l2); - Vector3 x1 = *x1_ptr; - Vector3 x2 = *x2_ptr; - Quat q1 = *q1_ptr; - Quat q2 = *q2_ptr; - ResponseType resp_type1 = ctx.getDirect( - RGDCols::ResponseType, l1); - ResponseType resp_type2 = ctx.getDirect( - RGDCols::ResponseType, l2); - ObjectID obj_id1 = ctx.getDirect(RGDCols::ObjectID, l1); - ObjectID obj_id2 = ctx.getDirect(RGDCols::ObjectID, l2); - - ObjectManager &obj_mgr = *ctx.singleton().mgr; - RigidBodyMetadata metadata1 = obj_mgr.metadata[obj_id1.idx]; - RigidBodyMetadata metadata2 = obj_mgr.metadata[obj_id2.idx]; - - float inv_m1 = metadata1.mass.invMass; - Vector3 inv_I1 = metadata1.mass.invInertiaTensor; - - if (resp_type1 == ResponseType::Static) { - inv_m1 = 0.f; - inv_I1 = Vector3::zero(); - } - - float inv_m2 = metadata2.mass.invMass; - Vector3 inv_I2 = metadata2.mass.invInertiaTensor; - - if (resp_type2 == ResponseType::Static) { - inv_m2 = 0.f; - inv_I2 = Vector3::zero(); - } - - Vector3 pos_correction; - switch (joint.type) { - case JointConstraint::Type::Fixed: { - JointConstraint::Fixed fixed_data = joint.fixed; - - applyJointOrientationConstraint( - q1, q2, - fixed_data.attachRot1, fixed_data.attachRot2, - inv_I1, inv_I2); - - Vector3 r1_world = q1.rotateVec(joint.r1) + x1; - Vector3 r2_world = q2.rotateVec(joint.r2) + x2; - - Vector3 delta_r = r2_world - r1_world; - - Quat axes_rot = (q1 * fixed_data.attachRot1).normalize(); - - Vector3 a1 = axes_rot.rotateVec(math::fwd); - Vector3 b1 = axes_rot.rotateVec(math::right); - Vector3 c1 = cross(a1, b1); - - // This implements a fixed distance qlong the a1 axis and no distance - // along the other axes - - pos_correction = Vector3::zero(); - { - // Unlike paper, subtract from pos_correction because - // applyPositionalUpdate applies the negative magnitude to object 1 - float a_separation = dot(delta_r, a1); - pos_correction -= (a_separation - fixed_data.separation) * a1; - - float b_separation = dot(delta_r, b1); - pos_correction -= b_separation * b1; - - float c_separation = dot(delta_r, c1); - pos_correction -= c_separation * c1; - } - } break; - case JointConstraint::Type::Hinge: { - JointConstraint::Hinge hinge_data = joint.hinge; - - applyJointAxisConstraint(q1, q2, - hinge_data.a1Local, hinge_data.a2Local, - inv_I1, inv_I2); - - Vector3 r1_world = q1.rotateVec(joint.r1) + x1; - Vector3 r2_world = q2.rotateVec(joint.r2) + x2; - - pos_correction = r2_world - r1_world; - } break; - default: MADRONA_UNREACHABLE(); - } - - float pos_correction_magnitude = pos_correction.length(); - if (pos_correction_magnitude > 0.f) { - pos_correction /= pos_correction_magnitude; - - applyPositionalUpdate( - x1, x2, - q1, q2, - joint.r1, joint.r2, - inv_m1, inv_m2, - inv_I1, inv_I2, - pos_correction, pos_correction_magnitude, 0); - } - - *x1_ptr = x1; - *x2_ptr = x2; - *q1_ptr = q1; - *q2_ptr = q2; -} - -inline void solvePositions(Context &ctx, SolverState &solver_state) -{ - ObjectManager &obj_mgr = *ctx.singleton().mgr; - - ctx.iterateQuery(solver_state.contactQuery, - [&](ContactConstraint &contact, XPBDContactState &contact_solver_state) { - contact_solver_state.lambdaN[0] = 0.f; - contact_solver_state.lambdaN[1] = 0.f; - contact_solver_state.lambdaN[2] = 0.f; - contact_solver_state.lambdaN[3] = 0.f; - handleContact(ctx, obj_mgr, contact, contact_solver_state.lambdaN); - }); - - ctx.iterateQuery(solver_state.jointQuery, [&](JointConstraint joint) { - handleJointConstraint(ctx, joint); - }); -} - -inline void setVelocities(Context &ctx, - const Position &pos, - const Rotation &rot, - const SubstepPrevState &prev_state, - Velocity &vel) -{ - const auto &physics_sys = ctx.singleton(); - float h = physics_sys.h; - - Vector3 x = pos; - Quat q = rot; - - Vector3 x_prev = prev_state.prevPosition; - Quat q_prev = prev_state.prevRotation; - - // when cur and prev rotation are equal there should be 0 angular velocity - // Unfortunately, this computation introduces a small amount of FP error - // and in some rotations the object winds up with a small delta_q, so - // we do a check and if all components match, just set delta_q to the - // identity quaternion manually - Quat delta_q; - if (q.w != q_prev.w || q.x != q_prev.x || - q.y != q_prev.y || q.z != q_prev.z) { - delta_q = q * q_prev.inv(); - } else { - delta_q = { 1, 0, 0, 0 }; - } - - // FIXME: A noticeable amount of energy is lost for bodies that should - // be under a constant angular velocity with no resistance. The issue - // seems to be that delta_q here is consistently slightly less than - // apply_omega in substepRigidBodies, so the angular velocity reduces - // a little bit each frame. Investigate whether this is an FP - // precision issue or a consequence of linearized angular -> quaternion - // formulas and see if it can be mitigated. Dividing delta_q by delta_q.w - // seems to fix the issue - does that have unintended consequences? - Vector3 new_omega = - 2.f / h * Vector3 { delta_q.x, delta_q.y, delta_q.z }; - - vel.linear = (x - x_prev) / h; - vel.angular = delta_q.w > 0.f ? new_omega : -new_omega; -} - -static inline Vector3 computeRelativeVelocity( - Vector3 v1, Vector3 v2, - Vector3 omega1, Vector3 omega2, - Vector3 dir1, Vector3 dir2) -{ - return (v1 + cross(omega1, dir1)) - (v2 + cross(omega2, dir2)); -} - -static inline void applyFrictionVelocityUpdate( - Vector3 &v1, Vector3 &v2, - Vector3 &omega1, Vector3 &omega2, - Quat q1, Quat q2, - float inv_m1, float inv_m2, - Vector3 inv_I1, Vector3 inv_I2, - Vector3 n, - float mu_d, float h, - Vector3 r1_local, Vector3 r2_local, - Vector3 r1_world, Vector3 r2_world, - float lambda) -{ - Vector3 v = computeRelativeVelocity( - v1, v2, omega1, omega2, r1_world, r2_world); - - float vn = dot(n, v); - Vector3 vt = v - n * vn; - - float vt_len = vt.length(); - if (vt_len == 0.f) { - return; - } - - Vector3 delta_world = vt / vt_len; - - Vector3 delta_local1 = q1.inv().rotateVec(delta_world); - Vector3 delta_local2 = q2.inv().rotateVec(delta_world); - - Vector3 friction_torque_axis_local1 = - cross(r1_local, delta_local1); - Vector3 friction_torque_axis_local2 = - cross(r2_local, delta_local2); - - Vector3 friction_rot_axis_local1 = multDiag( - inv_I1, friction_torque_axis_local1); - Vector3 friction_rot_axis_local2 = multDiag( - inv_I2, friction_torque_axis_local2); - - float w1 = generalizedInverseMass( - friction_torque_axis_local1, friction_rot_axis_local1, inv_m1); - float w2 = generalizedInverseMass( - friction_torque_axis_local2, friction_rot_axis_local2, inv_m2); - - float inv_mass_scale = 1.f / (w1 + w2); - - // h * mu_d * |f_n| in paper. Note the paper is incorrect here - // (doesn't have w1 + w2 divisor). - float dynamic_friction_magnitude = - mu_d * fabsf(lambda) * inv_mass_scale / h; - - float corrected_magnitude = - -fminf(dynamic_friction_magnitude, vt_len); - - float impulse_magnitude = corrected_magnitude * inv_mass_scale; - - if (impulse_magnitude == 0.f) { - return; - } - - v1 += delta_world * impulse_magnitude * inv_m1; - v2 -= delta_world * impulse_magnitude * inv_m2; - - Vector3 omega1_update_local = - impulse_magnitude * friction_rot_axis_local1; - Vector3 omega2_update_local = - impulse_magnitude * friction_rot_axis_local2; - - omega1 += q1.rotateVec(omega1_update_local); - omega2 -= q2.rotateVec(omega2_update_local); -} - -static inline void applyRestitutionVelocityUpdate( - Vector3 &v1, Vector3 &v2, - Vector3 &omega1, Vector3 &omega2, - Quat q1, Quat q2, - float inv_m1, float inv_m2, - Vector3 inv_I1, Vector3 inv_I2, - Vector3 n, - float restitution_threshold, - Vector3 r1_world, - Vector3 r2_world, - Vector3 restitution_torque_axis_local1, - Vector3 restitution_torque_axis_local2, - float vn_bar) -{ - Vector3 v = computeRelativeVelocity( - v1, v2, omega1, omega2, r1_world, r2_world); - - float vn = dot(n, v); - - float e = 0.3f; // FIXME - if (fabsf(vn_bar) <= restitution_threshold) { - e = 0.f; - } - - float restitution_magnitude = fminf(-e * vn_bar, 0) - vn; - - Vector3 restitution_rot_axis_local1 = - multDiag(inv_I1, restitution_torque_axis_local1); - Vector3 restitution_rot_axis_local2 = - multDiag(inv_I2, restitution_torque_axis_local2); - - float w1 = generalizedInverseMass( - restitution_torque_axis_local1, restitution_rot_axis_local1, inv_m1); - float w2 = generalizedInverseMass( - restitution_torque_axis_local2, restitution_rot_axis_local2, inv_m2); - - float inv_mass_scale = 1.f / (w1 + w2); - - float impulse_magnitude = restitution_magnitude * inv_mass_scale; - - if (impulse_magnitude == 0.f) { - return; - } - - v1 += n * impulse_magnitude * inv_m1; - v2 -= n * impulse_magnitude * inv_m2; - - Vector3 omega1_update_local = - impulse_magnitude * restitution_rot_axis_local1; - Vector3 omega2_update_local = - impulse_magnitude * restitution_rot_axis_local2; - - omega1 += q1.rotateVec(omega1_update_local); - omega2 -= q2.rotateVec(omega2_update_local); -} - -static inline void solveVelocitiesForContact(Context &ctx, - ObjectManager &obj_mgr, - ContactConstraint contact, - float lambdaN[4], - float h, - float restitution_threshold) -{ - Velocity *v1_out = &ctx.getDirect(RGDCols::Velocity, contact.ref); - Velocity *v2_out = &ctx.getDirect(RGDCols::Velocity, contact.alt); - - Quat q1 = ctx.getDirect(RGDCols::Rotation, contact.ref); - Quat q2 = ctx.getDirect(RGDCols::Rotation, contact.alt); - - PreSolvePositional presolve_pos1 = ctx.getDirect( - XPBDCols::PreSolvePositional, contact.ref); - PreSolvePositional presolve_pos2 = ctx.getDirect( - XPBDCols::PreSolvePositional, contact.alt); - - PreSolveVelocity presolve_vel1 = - ctx.getDirect(XPBDCols::PreSolveVelocity, contact.ref); - PreSolveVelocity presolve_vel2 = - ctx.getDirect(XPBDCols::PreSolveVelocity, contact.alt); - - ObjectID obj_id1 = ctx.getDirect(RGDCols::ObjectID, contact.ref); - ObjectID obj_id2 = ctx.getDirect(RGDCols::ObjectID, contact.alt); - - ResponseType resp_type1 = ctx.getDirect( - RGDCols::ResponseType, contact.ref); - ResponseType resp_type2 = ctx.getDirect( - RGDCols::ResponseType, contact.alt); - - RigidBodyMetadata metadata1 = obj_mgr.metadata[obj_id1.idx]; - RigidBodyMetadata metadata2 = obj_mgr.metadata[obj_id2.idx]; - - auto [v1, omega1] = *v1_out; - auto [v2, omega2] = *v2_out; - - float inv_m1 = metadata1.mass.invMass; - float inv_m2 = metadata2.mass.invMass; - Vector3 inv_I1 = metadata1.mass.invInertiaTensor; - Vector3 inv_I2 = metadata2.mass.invInertiaTensor; - - if (resp_type1 == ResponseType::Static) { - inv_m1 = 0.f; - inv_I1 = Vector3::zero(); - } - - if (resp_type2 == ResponseType::Static) { - inv_m2 = 0.f; - inv_I2 = Vector3::zero(); - } - - float mu_d = 0.5f * (metadata1.friction.muD + metadata2.friction.muD); - - { - Vector3 avg_contact_pos; - float contact_pos_penetration; - bool zero_separation = getAvgContact(contact, &avg_contact_pos, &contact_pos_penetration); - if (zero_separation) { - return; - } - - auto [r1, r2] = getLocalSpaceContacts(presolve_pos1, presolve_pos2, - avg_contact_pos, contact_pos_penetration, contact.normal); - - Vector3 r1_presolve = presolve_pos1.q.rotateVec(r1); - Vector3 r2_presolve = presolve_pos2.q.rotateVec(r2); - - Vector3 v_bar = computeRelativeVelocity( - presolve_vel1.v, presolve_vel2.v, - presolve_vel1.omega, presolve_vel2.omega, - r1_presolve, r2_presolve); // FIXME r1_world or presolve? - - float vn_bar = dot(contact.normal, v_bar); - - Vector3 r1_world = q1.rotateVec(r1); - Vector3 r2_world = q2.rotateVec(r2); - - Vector3 restitution_torque_axis_local1 = - cross(r1, q1.inv().rotateVec(contact.normal)); - Vector3 restitution_torque_axis_local2 = - cross(r2, q2.inv().rotateVec(contact.normal)); - - applyRestitutionVelocityUpdate( - v1, v2, - omega1, omega2, - q1, q2, - inv_m1, inv_m2, - inv_I1, inv_I2, - contact.normal, - restitution_threshold, - r1_world, r2_world, - restitution_torque_axis_local1, restitution_torque_axis_local2, - vn_bar); - } - - float penetration_sum = 0.f; - for (CountT i = 0; i < contact.numPoints; i++) { - penetration_sum += contact.points[i].w; - } - - for (CountT i = 0; i < contact.numPoints; i++) { - auto [r1, r2] = getLocalSpaceContacts(presolve_pos1, presolve_pos2, - contact.points[i].xyz(), contact.points[i].w, contact.normal); - - Vector3 r1_world = q1.rotateVec(r1); - Vector3 r2_world = q2.rotateVec(r2); - - applyFrictionVelocityUpdate( - v1, v2, - omega1, omega2, - q1, q2, - inv_m1, inv_m2, - inv_I1, inv_I2, - contact.normal, - mu_d, h, - r1, r2, - r1_world, r2_world, - lambdaN[0] * (contact.points[i].w / penetration_sum)); - } - - *v1_out = Velocity { v1, omega1 }; - *v2_out = Velocity { v2, omega2 }; -} - -inline void solveVelocities(Context &ctx, SolverState &solver) -{ - ObjectManager &obj_mgr = *ctx.singleton().mgr; - PhysicsSystemState &physics_sys = ctx.singleton(); - - ctx.iterateQuery(solver.contactQuery, - [&](ContactConstraint &contact, XPBDContactState &contact_solver_state) { - solveVelocitiesForContact( - ctx, obj_mgr, contact, contact_solver_state.lambdaN, - physics_sys.h, physics_sys.restitutionThreshold); - }); -} - -void registerTypes(ECSRegistry ®istry) -{ - registry.registerComponent(); - registry.registerComponent(); - registry.registerComponent(); - registry.registerComponent(); - - registry.registerArchetype(); - registry.registerArchetype(); - - registry.registerSingleton(); - - registry.registerBundle(); - registry.registerBundleAlias(); -} - -void init(Context &ctx) -{ - new (&ctx.singleton()) SolverState { - .jointQuery = ctx.query(), - .contactQuery = ctx.query(), - }; -} - -void getSolverArchetypeIDs(uint32_t *contact_archetype_id, - uint32_t *joint_archetype_id) -{ - *contact_archetype_id = TypeTracker::typeID(); - *joint_archetype_id = TypeTracker::typeID(); -} - -TaskGraphNodeID setupXPBDSolverTasks( - TaskGraphBuilder &builder, - TaskGraphNodeID broadphase, - CountT num_substeps) -{ - auto cur_node = broadphase; - -#ifdef MADRONA_GPU_MODE - cur_node = - builder.addToGraph>({cur_node}); - cur_node = builder.addToGraph({cur_node}); -#endif - - for (CountT i = 0; i < num_substeps; i++) { - auto rgb_update = builder.addToGraph>({cur_node}); - - auto run_narrowphase = narrowphase::setupTasks(builder, {rgb_update}); - -#ifdef MADRONA_GPU_MODE - run_narrowphase = builder.addToGraph>( - {run_narrowphase}); - - run_narrowphase = builder.addToGraph( - {run_narrowphase}); -#endif - - auto solve_pos = builder.addToGraph>( - {run_narrowphase}); - - auto vel_set = builder.addToGraph>({solve_pos}); - - auto solve_vel = builder.addToGraph>({vel_set}); - - auto clear_contacts = builder.addToGraph< - ClearTmpNode>({solve_vel}); - - cur_node = builder.addToGraph({clear_contacts}); - -#if 0 - cur_node = builder.addToGraph>({cur_node}); -#endif - } - - auto clear_broadphase = builder.addToGraph< - ClearTmpNode>({cur_node}); - - return clear_broadphase; -} - -} diff --git a/src/physics/xpbd.hpp b/src/physics/xpbd.hpp deleted file mode 100644 index 8d51f71b..00000000 --- a/src/physics/xpbd.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include "physics_impl.hpp" - -namespace madrona::phys::xpbd { - -void registerTypes(ECSRegistry ®istry); - -void getSolverArchetypeIDs(uint32_t *contact_archetype_id, - uint32_t *joint_archetype_id); - -void init(Context &ctx); - -TaskGraphNodeID setupXPBDSolverTasks( - TaskGraphBuilder &builder, - TaskGraphNodeID broadphase, - CountT num_substeps); - -} diff --git a/src/python/bindings.cpp b/src/python/bindings.cpp index 049b389a..f6460c11 100644 --- a/src/python/bindings.cpp +++ b/src/python/bindings.cpp @@ -1,29 +1,12 @@ #include #include -#include - namespace nb = nanobind; namespace madrona::py { namespace { -struct JAXModule { - nb::object mod; - nb::object typeU8; - nb::object typeI8; - nb::object typeI16; - nb::object typeI32; - nb::object typeI64; - nb::object typeF16; - nb::object typeF32; - nb::object typeShapeDtypeStruct; - - static inline JAXModule imp(); - inline nb::object getDType(TensorElementType type) const; -}; - nb::dlpack::dtype toDLPackType(TensorElementType type) { switch (type) { @@ -71,312 +54,14 @@ nb::object tensor_to_pytorch(const Tensor &tensor) )); } -auto tensor_to_jax(const Tensor &tensor) -{ - nb::dlpack::dtype type = toDLPackType(tensor.type()); - - return nb::ndarray( - tensor.devicePtr(), - (size_t)tensor.numDims(), - (const size_t *)tensor.dims(), - nb::handle(), - nullptr, - type, - tensor.isOnGPU() ? - nb::device::cuda::value : - nb::device::cpu::value, - tensor.isOnGPU() ? tensor.gpuID() : 0 - ); -} - - -JAXModule JAXModule::imp() -{ - nb::object mod = nb::module_::import_("jax"); - nb::object jnp = mod.attr("numpy"); - - nb::object type_u8 = jnp.attr("uint8"); - nb::object type_i8 = jnp.attr("int8"); - nb::object type_i16 = jnp.attr("int16"); - nb::object type_i32 = jnp.attr("int32"); - nb::object type_i64 = jnp.attr("int64"); - nb::object type_f16 = jnp.attr("float16"); - nb::object type_f32 = jnp.attr("float32"); - - nb::object type_shapedtype = mod.attr("ShapeDtypeStruct"); - - return JAXModule { - .mod = mod, - .typeU8 = type_u8, - .typeI8 = type_i8, - .typeI16 = type_i16, - .typeI32 = type_i32, - .typeI64 = type_i64, - .typeF16 = type_f16, - .typeF32 = type_f32, - .typeShapeDtypeStruct = type_shapedtype, - }; -} - -nb::object JAXModule::getDType(TensorElementType type) const -{ - switch (type) { - case TensorElementType::UInt8: - return typeU8; - case TensorElementType::Int8: - return typeI8; - case TensorElementType::Int16: - return typeI16; - case TensorElementType::Int32: - return typeI32; - case TensorElementType::Int64: - return typeI64; - case TensorElementType::Float16: - return typeF16; - case TensorElementType::Float32: - return typeF32; - default: MADRONA_UNREACHABLE(); - } -} - -nb::object tensor_to_jax_shape_dtype( - const JAXModule &jax_mod, - const Tensor &t) -{ - nb::list shape; - - for (int64_t i = 0; i < t.numDims(); i++) { - int64_t d = t.dims()[i]; - shape.append(d); - } - - return jax_mod.typeShapeDtypeStruct( - nb::arg("dtype") = jax_mod.getDType(t.type()), - nb::arg("shape") = shape); -} - -nb::dict train_interface_inputs_to_pytree( - const JAXModule &jax_mod, - const TrainInterface &iface) -{ - const TrainStepInputInterface &inputs = iface.stepInputs(); - - nb::dict d; - - nb::dict actions; - for (const NamedTensor &t : inputs.actions) { - actions[t.name] = tensor_to_jax_shape_dtype(jax_mod, t.tensor); - } - - d["actions"] = actions; - - d["resets"] = tensor_to_jax_shape_dtype(jax_mod, inputs.resets); - d["sim_ctrl"] = tensor_to_jax_shape_dtype(jax_mod, inputs.simCtrl); - - nb::dict pbt; - for (const NamedTensor &t : inputs.pbt) { - pbt[t.name] = tensor_to_jax_shape_dtype(jax_mod, t.tensor); - } - - d["pbt"] = pbt; - - return d; -} - -nb::dict train_interface_outputs_to_pytree( - const JAXModule &jax_mod, - const TrainInterface &iface) -{ - const TrainStepOutputInterface &outputs = iface.stepOutputs(); - - nb::dict d; - - nb::dict obs; - for (const NamedTensor &t : outputs.observations) { - obs[t.name] = tensor_to_jax_shape_dtype(jax_mod, t.tensor); - } - - d["obs"] = obs; - - d["rewards"] = tensor_to_jax_shape_dtype(jax_mod, outputs.rewards); - d["dones"] = tensor_to_jax_shape_dtype(jax_mod, outputs.dones); - - nb::dict pbt; - for (const NamedTensor &t : outputs.pbt) { - pbt[t.name] = tensor_to_jax_shape_dtype(jax_mod, t.tensor); - } - - d["pbt"] = pbt; - - return d; -} - -nb::dict train_interface_checkpointing_to_pytree( - const JAXModule &jax_mod, - const TrainInterface &iface) -{ - nb::dict d; - - d["data"] = tensor_to_jax_shape_dtype( - jax_mod, iface.checkpointing()->checkpointData); - - return d; -} - -} - -nb::dict JAXInterface::setup(const TrainInterface &iface, - nb::object sim_obj, - void *sim_ptr, - void *init_fn, - void *step_fn, - void *save_ckpts_fn, - void *restore_ckpts_fn, - bool xla_gpu) -{ - JAXModule jax_mod = JAXModule::imp(); - - nb::capsule init_fn_capsule(init_fn, "xla._CUSTOM_CALL_TARGET"); - nb::capsule step_fn_capsule(step_fn, "xla._CUSTOM_CALL_TARGET"); - - auto sim_encode = nb::bytes((char *)&sim_ptr, sizeof(char *)); - - nb::dict input_iface = train_interface_inputs_to_pytree(jax_mod, iface); - nb::dict output_iface = train_interface_outputs_to_pytree(jax_mod, iface); - - nb::dict scope; - scope["sim_obj"] = sim_obj; - scope["sim_ptr"] = (uint64_t)sim_ptr; - scope["sim_encode"] = sim_encode; - scope["step_inputs_iface"] = input_iface; - scope["step_outputs_iface"] = output_iface; - scope["init_custom_call_capsule"] = init_fn_capsule; - scope["step_custom_call_capsule"] = step_fn_capsule; - scope["custom_call_platform"] = xla_gpu ? "gpu" : "cpu"; - - if (iface.checkpointing().has_value()) { - assert(save_ckpts_fn != nullptr && restore_ckpts_fn != nullptr); - - nb::dict checkpointing_iface = train_interface_checkpointing_to_pytree( - jax_mod, iface); - - nb::capsule save_ckpts_fn_capsule( - save_ckpts_fn, "xla._CUSTOM_CALL_TARGET"); - nb::capsule restore_ckpts_fn_capsule( - restore_ckpts_fn, "xla._CUSTOM_CALL_TARGET"); - - scope["ckpt_iface"] = checkpointing_iface; - scope["save_ckpts_custom_call_capsule"] = save_ckpts_fn_capsule; - scope["restore_ckpts_custom_call_capsule"] = restore_ckpts_fn_capsule; - } else { - scope["ckpt_iface"] = nb::none(); - } - - nb::exec( -#include "jax_register.py" - , scope); - - nb::dict fns; - fns["init"] = scope["init_func"]; - fns["step"] = scope["step_func"]; - - if (iface.checkpointing().has_value()) { - fns["save_ckpts"] = scope["save_ckpts_func"]; - fns["restore_ckpts"] = scope["restore_ckpts_func"]; - } - - return fns; -} - -static TensorElementType fromDLPackType(nb::dlpack::dtype dtype) -{ - using ET = TensorElementType; - - if (nb::dlpack::dtype_code(dtype.code) == nb::dlpack::dtype_code::Int) { - switch (dtype.bits) { - case 8: - return ET::Int8; - case 16: - return ET::Int16; - case 32: - return ET::Int32; - case 64: - return ET::Int64; - default: - break; - } - } else if (nb::dlpack::dtype_code(dtype.code) == - nb::dlpack::dtype_code::UInt && dtype.bits == 8) { - return ET::UInt8; - } else if (nb::dlpack::dtype_code(dtype.code) == - nb::dlpack::dtype_code::Float) { - if (dtype.bits == 16) { - return ET::Float16; - } - if (dtype.bits == 32) { - return ET::Float32; - } - } - - FATAL("Tensor: Invalid tensor dtype"); } void setupMadronaSubmodule(nb::module_ parent_mod) { auto m = parent_mod.def_submodule("madrona"); - nb::class_(m, "ExecMode") - .def_prop_ro_static("CPU", [](nb::handle) { - return madrona::py::PyExecMode(madrona::ExecMode::CPU); - }) - .def_prop_ro_static("CUDA", [](nb::handle) { - return madrona::py::PyExecMode(madrona::ExecMode::CUDA); - }); - nb::class_(m, "Tensor") - .def("__init__", [](Tensor *dst, nb::ndarray<> torch_tensor) { - Optional gpu_id = Optional::none(); - if (torch_tensor.device_type() == nb::device::cuda::value) { - gpu_id = torch_tensor.device_id(); - } else if (torch_tensor.device_type() != nb::device::cpu::value) { - FATAL("madrona::Tensor: failed to import unknown tensor type"); - } - - std::array dims; - - if (torch_tensor.ndim() > Tensor::maxDimensions) { - FATAL("Cannot construct Tensor with more than %ld dimensions\n", - Tensor::maxDimensions); - } - - for (size_t i = 0; i < torch_tensor.ndim(); i++) { - dims[i] = torch_tensor.shape(i); - } - - new (dst) Tensor(torch_tensor.data(), - fromDLPackType(torch_tensor.dtype()), - { dims.data(), (CountT)torch_tensor.ndim() }, - gpu_id); - }) .def("to_torch", tensor_to_pytorch, nb::rv_policy::automatic_reference) - .def("to_jax", tensor_to_jax, nb::rv_policy::automatic_reference) - ; - -#ifdef MADRONA_CUDA_SUPPORT - nb::class_(m, "CudaSync") - .def("wait", &CudaSync::wait) - ; -#endif - - nb::class_(m, "TrainInterface") - .def("step_inputs", [](const TrainInterface &iface) { - return train_interface_inputs_to_pytree( - JAXModule::imp(), iface); - }) - .def("step_outputs", [](const TrainInterface &iface) { - return train_interface_outputs_to_pytree( - JAXModule::imp(), iface); - }) ; } diff --git a/src/python/jax_register.py b/src/python/jax_register.py deleted file mode 100644 index f50d6a27..00000000 --- a/src/python/jax_register.py +++ /dev/null @@ -1,350 +0,0 @@ -R"===(#" - -from functools import partial - -import jax -from jax import numpy as jnp -from jax import core, dtypes -from jax.core import ShapedArray -from jax.lib import xla_client -from jax.interpreters import xla -from jax.interpreters import mlir -from jax.interpreters.mlir import ir, dtype_to_ir_type -from jaxlib.hlo_helpers import custom_call -import builtins as __builtins__ -import numpy as np - -from jax._src import effects -from jax._src.lib.mlir.dialects import hlo - -custom_call_prefix = f"{type(sim_obj).__name__}_{id(sim_obj)}" -init_custom_call_name = f"{custom_call_prefix}_init" -step_custom_call_name = f"{custom_call_prefix}_step" -del sim_obj - -xla_client.register_custom_call_target(init_custom_call_name, init_custom_call_capsule, platform=custom_call_platform) - -xla_client.register_custom_call_target(step_custom_call_name, step_custom_call_capsule, platform=custom_call_platform) - - -def _row_major_layout(shape): - return tuple(range(len(shape) - 1, -1, -1)) - - -def _shape_dtype_to_abstract_vals(vs): - return tuple(ShapedArray(v.shape, v.dtype) for v in vs) - - -def _lower_shape_dtypes(shape_dtypes): - return [ir.RankedTensorType.get(i.shape, dtype_to_ir_type(i.dtype)) for i in shape_dtypes] - - -# Below code uses ordered effects, which is internal logic taken from -# jax io_callback code and emit_python_callback code. The idea is a -# token is threaded through the custom_call, which preserves ordering and -# prevents sim_step calls from being elided if their outputs aren't used. -# This code deviates slightly from the jax convention which is to put -# the token in the first input / output on GPU. Instead, we put the token -# in the first input and the *last* output, which means we can just skip -# the first buffer passed to the custom call target (the input token) -# and write to the rest of the buffers normally, leaving the final token -# output buffer untouched. - - -def _fake_token_type(): - return ir.RankedTensorType.get((0,), dtype_to_ir_type(np.dtype("float32"))) - - -def _prepend_token_to_inputs(types, layouts): - return [_fake_token_type(), *types], [(0,), *layouts] - - -def _append_token_to_results(types, layouts): - return [*types, _fake_token_type()], [*layouts, (0,)] - - -def _init_lowering(ctx): - token = mlir.ir_constant(np.empty((0,), np.float32)) - - result_types = _lower_shape_dtypes(step_outputs_iface["obs"].values()) - result_layouts = [_row_major_layout(t.shape) for t in result_types] - - result_types, result_layouts = _append_token_to_results(result_types, result_layouts) - - results = custom_call( - init_custom_call_name, - backend_config=sim_encode, - operands=[mlir.ir_constant(sim_ptr), token], - operand_layouts=[(), (0,)], - result_types=result_types, - result_layouts=result_layouts, - has_side_effect=True, - ).results - - *results, token = results - return token, *results - - -def _init_abstract(): - return ( - ShapedArray((0,), jnp.float32), - *_shape_dtype_to_abstract_vals(step_outputs_iface["obs"].values()), - ) - - -def _flatten_step_output_shape_dtypes(): - result_shape_dtypes = list(step_outputs_iface["obs"].values()) - - result_shape_dtypes.append(step_outputs_iface["rewards"]) - result_shape_dtypes.append(step_outputs_iface["dones"]) - - result_shape_dtypes += step_outputs_iface["pbt"].values() - - return result_shape_dtypes - - -def _step_lowering(ctx, *flattened_inputs): - token, *flattened_inputs = flattened_inputs - - input_types = [ir.RankedTensorType(i.type) for i in flattened_inputs] - input_layouts = [_row_major_layout(t.shape) for t in input_types] - input_types, input_layouts = _prepend_token_to_inputs(input_types, input_layouts) - - result_types = _lower_shape_dtypes(_flatten_step_output_shape_dtypes()) - result_layouts = [_row_major_layout(t.shape) for t in result_types] - result_types, result_layouts = _append_token_to_results(result_types, result_layouts) - - inputs = [token, *flattened_inputs] - - results = custom_call( - step_custom_call_name, - backend_config=sim_encode, - operands=[mlir.ir_constant(sim_ptr), *inputs], - operand_layouts=[(), *input_layouts], - result_types=result_types, - result_layouts=result_layouts, - has_side_effect=True, - ).results - - *results, token = results - return token, *results - - -def _step_abstract(*inputs): - return ( - ShapedArray((0,), jnp.float32), - *_shape_dtype_to_abstract_vals(_flatten_step_output_shape_dtypes()), - ) - - -_init_primitive = core.Primitive(init_custom_call_name) -_init_primitive.multiple_results = True -_init_primitive.def_impl(partial(xla.apply_primitive, _init_primitive)) -_init_primitive.def_abstract_eval(_init_abstract) - -mlir.register_lowering( - _init_primitive, - _init_lowering, - platform=custom_call_platform, -) - -_step_primitive = core.Primitive(step_custom_call_name) -_step_primitive.multiple_results = True -_step_primitive.def_impl(partial(xla.apply_primitive, _step_primitive)) -_step_primitive.def_abstract_eval(_step_abstract) - -mlir.register_lowering( - _step_primitive, - _step_lowering, - platform=custom_call_platform, -) - - -def init_func(): - sim_state, *flattened_out = _init_primitive.bind() - return { - "state": sim_state, - "obs": {k: o for k, o in zip(step_outputs_iface["obs"].keys(), flattened_out)}, - } - - -def step_func(step_inputs): - flattened_in = [step_inputs["state"]] - - for k in step_inputs_iface["actions"].keys(): - print(step_inputs["actions"]) - flattened_in.append(step_inputs["actions"][k]) - - flattened_in.append(step_inputs["resets"]) - flattened_in.append(step_inputs["sim_ctrl"]) - - for k in step_inputs_iface["pbt"].keys(): - flattened_in.append(step_inputs["pbt"][k]) - - sim_state, *flattened_out = _step_primitive.bind(*flattened_in) - - out = {} - - cur_idx = 0 - - def next_out(): - nonlocal cur_idx - o = flattened_out[cur_idx] - cur_idx += 1 - return o - - out["state"] = sim_state - out["obs"] = {} - for k in step_outputs_iface["obs"].keys(): - out["obs"][k] = next_out() - - out["rewards"] = next_out() - out["dones"] = next_out() - - out["pbt"] = {} - for k in step_outputs_iface["pbt"].keys(): - out["pbt"][k] = next_out() - - return out - - -init_func = jax.jit(init_func) -step_func = jax.jit(step_func) - -if ckpt_iface != None: - save_ckpts_custom_call_name = f"{custom_call_prefix}_save_ckpts" - restore_ckpts_custom_call_name = f"{custom_call_prefix}_restore_ckpts" - - xla_client.register_custom_call_target( - save_ckpts_custom_call_name, - save_ckpts_custom_call_capsule, - platform=custom_call_platform, - ) - - xla_client.register_custom_call_target( - restore_ckpts_custom_call_name, - restore_ckpts_custom_call_capsule, - platform=custom_call_platform, - ) - - def _flatten_save_ckpts_output_shape_dtypes(): - result_shape_dtypes = [ckpt_iface["data"]] - return result_shape_dtypes - - def _save_ckpts_lowering(ctx, *flattened_inputs): - token, *flattened_inputs = flattened_inputs - - input_types = [ir.RankedTensorType(i.type) for i in flattened_inputs] - input_layouts = [_row_major_layout(t.shape) for t in input_types] - input_types, input_layouts = _prepend_token_to_inputs(input_types, input_layouts) - - result_types = _lower_shape_dtypes(_flatten_save_ckpts_output_shape_dtypes()) - result_layouts = [_row_major_layout(t.shape) for t in result_types] - result_types, result_layouts = _append_token_to_results(result_types, result_layouts) - - inputs = [token, *flattened_inputs] - - results = custom_call( - save_ckpts_custom_call_name, - backend_config=sim_encode, - operands=[mlir.ir_constant(sim_ptr), *inputs], - operand_layouts=[(), *input_layouts], - result_types=result_types, - result_layouts=result_layouts, - has_side_effect=True, - ).results - - *results, token = results - return token, *results - - def _save_ckpts_abstract(*inputs): - return ( - core.abstract_token, - *_shape_dtype_to_abstract_vals(_flatten_save_ckpts_output_shape_dtypes()), - ) - - _save_ckpts_primitive = core.Primitive(save_ckpts_custom_call_name) - _save_ckpts_primitive.multiple_results = True - _save_ckpts_primitive.def_impl(partial(xla.apply_primitive, _save_ckpts_primitive)) - _save_ckpts_primitive.def_abstract_eval(_save_ckpts_abstract) - - mlir.register_lowering( - _save_ckpts_primitive, - _save_ckpts_lowering, - platform=custom_call_platform, - ) - - def _flatten_restore_ckpts_output_shape_dtypes(): - result_shape_dtypes = list(step_outputs_iface["obs"].values()) - return result_shape_dtypes - - def _restore_ckpts_lowering(ctx, *flattened_inputs): - token, *flattened_inputs = flattened_inputs - - input_types = [ir.RankedTensorType(i.type) for i in flattened_inputs] - input_layouts = [_row_major_layout(t.shape) for t in input_types] - input_types, input_layouts = _prepend_token_to_inputs(input_types, input_layouts) - - result_types = _lower_shape_dtypes(_flatten_restore_ckpts_output_shape_dtypes()) - result_layouts = [_row_major_layout(t.shape) for t in result_types] - result_types, result_layouts = _append_token_to_results(result_types, result_layouts) - - inputs = [token, *flattened_inputs] - - results = custom_call( - restore_ckpts_custom_call_name, - backend_config=sim_encode, - operands=[mlir.ir_constant(sim_ptr), *inputs], - operand_layouts=[(), *input_layouts], - result_types=result_types, - result_layouts=result_layouts, - has_side_effect=True, - ).results - - *results, token = results - return token, *results - - def _restore_ckpts_abstract(*inputs): - return ( - core.abstract_token, - *_shape_dtype_to_abstract_vals(_flatten_restore_ckpts_output_shape_dtypes()), - ) - - _restore_ckpts_primitive = core.Primitive(restore_ckpts_custom_call_name) - _restore_ckpts_primitive.multiple_results = True - _restore_ckpts_primitive.def_impl(partial(xla.apply_primitive, _restore_ckpts_primitive)) - _restore_ckpts_primitive.def_abstract_eval(_restore_ckpts_abstract) - - mlir.register_lowering( - _restore_ckpts_primitive, - _restore_ckpts_lowering, - platform=custom_call_platform, - ) - - def save_ckpts_func(save_inputs): - flattened_in = [save_inputs["state"]] - flattened_in.append(save_inputs["should_save"]) - - sim_state, *flattened_out = _save_ckpts_primitive.bind(*flattened_in) - - return { - "state": sim_state, - "ckpts": flattened_out[0], - } - - def restore_ckpts_func(restore_inputs): - flattened_in = [restore_inputs["state"]] - flattened_in.append(restore_inputs["should_restore"]) - flattened_in.append(restore_inputs["ckpt_data"]) - - sim_state, *flattened_out = _restore_ckpts_primitive.bind(*flattened_in) - - return { - "state": sim_state, - "obs": {k: o for k, o in zip(step_outputs_iface["obs"].keys(), flattened_out)}, - } - - save_ckpts_func = jax.jit(save_ckpts_func) - restore_ckpts_func = jax.jit(restore_ckpts_func) - -# )===" diff --git a/src/python/utils.cpp b/src/python/utils.cpp index a8409465..a105b501 100644 --- a/src/python/utils.cpp +++ b/src/python/utils.cpp @@ -12,169 +12,6 @@ namespace madrona::py { -#ifdef MADRONA_CUDA_SUPPORT -CudaSync::CudaSync(cudaExternalSemaphore_t sema) - : sema_(sema) -{} - -void CudaSync::wait(uint64_t strm) -{ - // Get the current CUDA stream from pytorch and force it to wait - // on an external semaphore to finish - cudaStream_t cuda_strm = (cudaStream_t)strm; - cudaExternalSemaphoreWaitParams params {}; - REQ_CUDA(cudaWaitExternalSemaphoresAsync(&sema_, ¶ms, 1, cuda_strm)); -} - -#ifdef MADRONA_LINUX -void CudaSync::key_() {} -#endif -#endif - -struct TrainInterface::Impl { - HeapArray nameBuffer; - HeapArray namedTensors; - - TrainStepInputInterface inputs; - TrainStepOutputInterface outputs; - Optional checkpointing; - - static inline Impl * init( - TrainStepInputInterface inputs, - TrainStepOutputInterface outputs, - Optional checkpointing); -}; - -TrainInterface::Impl * TrainInterface::Impl::init( - TrainStepInputInterface inputs, - TrainStepOutputInterface outputs, - Optional checkpointing) -{ - CountT num_total_name_chars = 0; - CountT num_total_named_tensors = 0; - - auto sumStorageRequirements = [ - &num_total_name_chars, &num_total_named_tensors - ](Span tensors) - { - for (NamedTensor named_tensor : tensors) { - num_total_name_chars += strlen(named_tensor.name) + 1; - } - - num_total_named_tensors += tensors.size(); - }; - - sumStorageRequirements(inputs.actions); - sumStorageRequirements(inputs.pbt); - - sumStorageRequirements(outputs.observations); - sumStorageRequirements(outputs.stats); - sumStorageRequirements(outputs.pbt); - - HeapArray name_buffer(num_total_name_chars); - HeapArray named_tensors(num_total_named_tensors); - - char *cur_name_ptr = name_buffer.data(); - NamedTensor *cur_named_tensor_ptr = named_tensors.data(); - - auto makeOwnedName = [ - &cur_name_ptr - ](const char *in_name) - { - char *out_name = cur_name_ptr; - - size_t name_len = strlen(in_name) + 1; - memcpy(out_name, in_name, name_len); - cur_name_ptr += name_len; - - return out_name; - }; - - auto makeOwnedNamedTensors = [ - &makeOwnedName, &cur_named_tensor_ptr - ](Span inputs) - { - NamedTensor *out_start = cur_named_tensor_ptr; - - for (const NamedTensor &named_in : inputs) { - const char *owned_name = makeOwnedName(named_in.name); - *(cur_named_tensor_ptr++) = NamedTensor { - .name = owned_name, - .tensor = named_in.tensor, - }; - } - - return Span(out_start, inputs.size()); - }; - - TrainStepInputInterface owned_inputs { - .actions = makeOwnedNamedTensors(inputs.actions), - .resets = inputs.resets, - .simCtrl = inputs.simCtrl, - .pbt = makeOwnedNamedTensors(inputs.pbt), - }; - - TrainStepOutputInterface owned_outputs { - .observations = makeOwnedNamedTensors(outputs.observations), - .rewards = outputs.rewards, - .dones = outputs.dones, - .stats = makeOwnedNamedTensors(outputs.stats), - .pbt = makeOwnedNamedTensors(outputs.pbt), - }; - - Optional owned_checkpointing = - Optional::none(); - - if (checkpointing.has_value()) { - owned_checkpointing = TrainCheckpointingInterface { - .checkpointData = checkpointing->checkpointData, - }; - } - - assert(cur_name_ptr == name_buffer.data() + name_buffer.size()); - assert(cur_named_tensor_ptr == - named_tensors.data() + named_tensors.size()); - - return new Impl { - .nameBuffer = std::move(name_buffer), - .namedTensors = std::move(named_tensors), - .inputs = owned_inputs, - .outputs = owned_outputs, - .checkpointing = std::move(owned_checkpointing), - }; -} - -TrainInterface::TrainInterface() - : impl_(nullptr) -{} - -TrainInterface::TrainInterface( - TrainStepInputInterface step_inputs, - TrainStepOutputInterface step_outputs, - Optional checkpointing) - : impl_(Impl::init(step_inputs, step_outputs, std::move(checkpointing))) -{} - -TrainInterface::TrainInterface(TrainInterface &&) = default; -TrainInterface::~TrainInterface() = default; - -TrainInterface & TrainInterface::operator=(TrainInterface &&o) = default; - -TrainStepInputInterface TrainInterface::stepInputs() const -{ - return impl_->inputs; -} - -TrainStepOutputInterface TrainInterface::stepOutputs() const -{ - return impl_->outputs; -} - -Optional TrainInterface::checkpointing() const -{ - return impl_->checkpointing; -} - Tensor::Printer::Printer(Printer &&o) : dev_ptr_(o.dev_ptr_), print_ptr_(o.print_ptr_) @@ -397,176 +234,7 @@ TensorInterface Tensor::interface() const } #ifdef MADRONA_LINUX -void PyExecMode::key_() {} void Tensor::key_() {} -void TrainInterface::key_() {} -#endif - -[[maybe_unused]] static inline uint64_t numTensorBytes(const Tensor &t) -{ - uint64_t num_items = 1; - uint64_t num_dims = t.numDims(); - for (uint64_t i = 0; i < num_dims; i++) { - num_items *= t.dims()[i]; - } - - return num_items * (uint64_t)t.numBytesPerItem(); -} - -void TrainInterface::cpuCopyStepInputs(void **buffers) -{ - auto copyToSim = [](const Tensor &dst, void *src) { - uint64_t num_bytes = numTensorBytes(dst); - - if (dst.isOnGPU()) { -#ifdef MADRONA_CUDA_SUPPORT - REQ_CUDA(cudaMemcpy(dst.devicePtr(), src, num_bytes, cudaMemcpyHostToDevice)); -#else - assert(false); -#endif - } else { - memcpy(dst.devicePtr(), src, num_bytes); - } - }; - - TrainStepInputInterface &inputs = impl_->inputs; - - for (const NamedTensor &t : inputs.actions) { - copyToSim(t.tensor, *buffers++); - } - - copyToSim(inputs.resets, *buffers++); - copyToSim(inputs.simCtrl, *buffers++); - - for (const NamedTensor &t : inputs.pbt) { - copyToSim(t.tensor, *buffers++); - } -} - -void TrainInterface::cpuCopyObservations(void **buffers) -{ - auto copyFromSim = [](void *dst, const Tensor &src) { - uint64_t num_bytes = numTensorBytes(src); - - if (src.isOnGPU()) { -#ifdef MADRONA_CUDA_SUPPORT - REQ_CUDA(cudaMemcpy(dst, src.devicePtr(), num_bytes, cudaMemcpyHostToDevice)); -#else - assert(false); -#endif - } else { - memcpy(dst, src.devicePtr(), num_bytes); - } - }; - - for (const NamedTensor &t : impl_->outputs.observations) { - copyFromSim(*buffers++, t.tensor); - } -} - -void TrainInterface::cpuCopyStepOutputs(void **buffers) -{ - auto copyFromSim = [](void *dst, const Tensor &src) { - uint64_t num_bytes = numTensorBytes(src); - - if (src.isOnGPU()) { -#ifdef MADRONA_CUDA_SUPPORT - REQ_CUDA(cudaMemcpy(dst, src.devicePtr(), num_bytes, cudaMemcpyHostToDevice)); -#else - assert(false); -#endif - } else { - memcpy(dst, src.devicePtr(), num_bytes); - } - }; - - TrainStepOutputInterface &outputs = impl_->outputs; - - for (const NamedTensor &t : outputs.observations) { - copyFromSim(*buffers++, t.tensor); - } - - copyFromSim(*buffers++, outputs.rewards); - copyFromSim(*buffers++, outputs.dones); - - for (const NamedTensor &t : outputs.stats) { - copyFromSim(*buffers++, t.tensor); - } - - for (const NamedTensor &t : outputs.pbt) { - copyFromSim(*buffers++, t.tensor); - } -} - - -#ifdef MADRONA_CUDA_SUPPORT -void ** TrainInterface::cudaCopyStepInputs(cudaStream_t strm, void **buffers) -{ - auto copyToSim = [&strm](const Tensor &dst, void *src) { - uint64_t num_bytes = numTensorBytes(dst); - - REQ_CUDA(cudaMemcpyAsync(dst.devicePtr(), src, num_bytes, - dst.isOnGPU() ? cudaMemcpyDeviceToDevice : cudaMemcpyDeviceToHost, - strm)); - }; - - TrainStepInputInterface &inputs = impl_->inputs; - - for (const NamedTensor &t : inputs.actions) { - copyToSim(t.tensor, *buffers++); - } - - copyToSim(inputs.resets, *buffers++); - copyToSim(inputs.simCtrl, *buffers++); - - for (const NamedTensor &t : inputs.pbt) { - copyToSim(t.tensor, *buffers++); - } - - return buffers; -} - -void TrainInterface::cudaCopyObservations(cudaStream_t strm, void **buffers) -{ - auto copyFromSim = [&strm](void *dst, const Tensor &src) { - uint64_t num_bytes = numTensorBytes(src); - - REQ_CUDA(cudaMemcpyAsync(dst, src.devicePtr(), num_bytes, - src.isOnGPU() ? cudaMemcpyDeviceToDevice : cudaMemcpyHostToDevice, strm)); - }; - - for (const NamedTensor &t : impl_->outputs.observations) { - copyFromSim(*buffers++, t.tensor); - } -} - -void TrainInterface::cudaCopyStepOutputs(cudaStream_t strm, void **buffers) -{ - auto copyFromSim = [&strm](void *dst, const Tensor &src) { - uint64_t num_bytes = numTensorBytes(src); - - REQ_CUDA(cudaMemcpyAsync(dst, src.devicePtr(), num_bytes, - src.isOnGPU() ? cudaMemcpyDeviceToDevice : cudaMemcpyHostToDevice, strm)); - }; - - TrainStepOutputInterface &outputs = impl_->outputs; - - for (const NamedTensor &t : outputs.observations) { - copyFromSim(*buffers++, t.tensor); - } - - copyFromSim(*buffers++, outputs.rewards); - copyFromSim(*buffers++, outputs.dones); - - for (const NamedTensor &t : outputs.stats) { - copyFromSim(*buffers++, t.tensor); - } - - for (const NamedTensor &t : outputs.pbt) { - copyFromSim(*buffers++, t.tensor); - } -} - #endif } diff --git a/src/render/batch_renderer.cpp b/src/render/batch_renderer.cpp index cbdc0b87..41d4a503 100644 --- a/src/render/batch_renderer.cpp +++ b/src/render/batch_renderer.cpp @@ -1557,17 +1557,14 @@ struct BatchRenderer::Impl { // Required whether we do batch rendering or not PipelineMP<1> prepareViews; PipelineMP<1> batchDraw; - PipelineMP<1> createVisualization; PipelineMP<1> lighting; PipelineMP<1> shadowGen; PipelineMP<1> shadowDraw; - Optional> postProcess; // Add post-processing pipeline //One frame is on simulation frame HeapArray batchFrames; VkDescriptorSet assetSetPrepare; - VkDescriptorSet assetSetDraw; VkDescriptorSet assetSetTextureMat; VkDescriptorSet assetSetLighting; @@ -1606,11 +1603,6 @@ BatchRenderer::Impl::Impl(const Config &cfg, RenderContext &rctx): dev, rctx.pipelineCache, VK_NULL_HANDLE, consts::numDrawCmdBuffers * cfg.numFrames, 5, rctx.repeatSampler, cfg.maxTextures)), - createVisualization( - makeComputePipeline( - dev, rctx.pipelineCache, 1, sizeof(uint32_t) * 2, - consts::numDrawCmdBuffers * cfg.numFrames, rctx.repeatSampler, - "visualize_tris.hlsl", "visualize", makeShaders)), lighting( makeComputePipeline( dev, rctx.pipelineCache, 3, sizeof(shader::DeferredLightingPushConstBR), @@ -1625,17 +1617,8 @@ BatchRenderer::Impl::Impl(const Config &cfg, RenderContext &rctx): makeShadowDrawPipeline( dev, rctx.pipelineCache, VK_NULL_HANDLE, consts::numDrawCmdBuffers * cfg.numFrames, 3)), - postProcess( - cfg.enableBatchRenderer ? - makeComputePipeline( - dev, rctx.pipelineCache, 1, sizeof(uint32_t) * 4, // push constants for width, height, view count, etc. - consts::numDrawCmdBuffers * cfg.numFrames, rctx.repeatSampler, - "post_process.hlsl", "main", makeShaders) : - Optional>::none() - ), batchFrames(cfg.numFrames), assetSetPrepare(rctx.asset_set_cull_), - assetSetDraw(rctx.asset_set_draw_), assetSetTextureMat(rctx.asset_set_mat_tex_), assetSetLighting(rctx.asset_batch_lighting_set_), renderExtent { cfg.renderWidth, cfg.renderHeight }, @@ -1693,9 +1676,6 @@ BatchRenderer::~BatchRenderer() impl->dev.dt.destroyPipeline(impl->dev.hdl, impl->batchDraw.hdls[0], nullptr); impl->dev.dt.destroyPipelineLayout(impl->dev.hdl, impl->batchDraw.layout, nullptr); - impl->dev.dt.destroyPipeline(impl->dev.hdl, impl->createVisualization.hdls[0], nullptr); - impl->dev.dt.destroyPipelineLayout(impl->dev.hdl, impl->createVisualization.layout, nullptr); - impl->dev.dt.destroyPipeline(impl->dev.hdl, impl->lighting.hdls[0], nullptr); impl->dev.dt.destroyPipelineLayout(impl->dev.hdl, impl->lighting.layout, nullptr); @@ -1705,11 +1685,6 @@ BatchRenderer::~BatchRenderer() impl->dev.dt.destroyPipeline(impl->dev.hdl, impl->shadowDraw.hdls[0], nullptr); impl->dev.dt.destroyPipelineLayout(impl->dev.hdl, impl->shadowDraw.layout, nullptr); - if (impl->postProcess.has_value()) { - impl->dev.dt.destroyPipeline(impl->dev.hdl, impl->postProcess->hdls[0], nullptr); - impl->dev.dt.destroyPipelineLayout(impl->dev.hdl, impl->postProcess->layout, nullptr); - } - for (CountT i = 0; i < impl->batchFrames.size(); i++) { impl->dev.dt.destroyCommandPool(impl->dev.hdl, impl->batchFrames[i].prepareCmdPool, nullptr); impl->dev.dt.destroyCommandPool(impl->dev.hdl, impl->batchFrames[i].renderCmdPool, nullptr); @@ -2001,11 +1976,6 @@ void BatchRenderer::prepareForRendering(BatchRenderInfo info, interop->aabbCPU->flush(impl->dev); } - if (interop->voxelInputCPU.has_value()) { - // Need to flush engine input state before copy - interop->voxelInputCPU->flush(impl->dev); - } - if (interop->lightsCPU.has_value()) { *interop->bridge.totalNumLights = interop->bridge.totalNumLightsCPUInc->load_acquire(); diff --git a/src/render/ecs_interop.hpp b/src/render/ecs_interop.hpp index fc2fcd1c..f334169e 100644 --- a/src/render/ecs_interop.hpp +++ b/src/render/ecs_interop.hpp @@ -27,7 +27,6 @@ struct RenderECSBridge { uint64_t *lightWorldIDs; int32_t renderWidth; int32_t renderHeight; - uint32_t *voxels; uint32_t maxViewsPerworld; uint32_t maxInstancesPerWorld; diff --git a/src/render/ecs_system.cpp b/src/render/ecs_system.cpp index 4b9e6df1..ee243c03 100644 --- a/src/render/ecs_system.cpp +++ b/src/render/ecs_system.cpp @@ -28,7 +28,6 @@ struct RenderingSystemState { uint32_t *totalNumViews; uint32_t *totalNumInstances; uint32_t *totalNumLights; - uint32_t *voxels; float aspectRatio; // This is used if on the CPU backend @@ -655,7 +654,6 @@ void init(Context &ctx, const RenderECSBridge *bridge) #endif system_state.aspectRatio = (float)bridge->renderWidth / (float)bridge->renderHeight; - system_state.voxels = bridge->voxels; } #if 0 diff --git a/src/render/font.ttf b/src/render/font.ttf deleted file mode 100644 index d017d72d..00000000 Binary files a/src/render/font.ttf and /dev/null differ diff --git a/src/render/render_common.hpp b/src/render/render_common.hpp index 707c7cee..a7bec041 100644 --- a/src/render/render_common.hpp +++ b/src/render/render_common.hpp @@ -172,14 +172,6 @@ struct EngineInterop { uint32_t maxInstancesPerWorld; uint32_t maxLightsPerWorld; - Optional voxelInputCPU; -#ifdef MADRONA_VK_CUDA_SUPPORT - Optional voxelInputGPU; - Optional voxelInputCUDA; -#endif - - VkBuffer voxelHdl; - uint32_t *iotaArrayInstancesCPU; uint32_t *iotaArrayViewsCPU; uint32_t *iotaArrayLightOffsetsCPU; diff --git a/src/render/render_ctx.cpp b/src/render/render_ctx.cpp index bbde3ef1..ed0d74a3 100644 --- a/src/render/render_ctx.cpp +++ b/src/render/render_ctx.cpp @@ -57,8 +57,6 @@ using PackedVertex = render::shader::PackedVertex; using MeshData = render::shader::MeshData; using MaterialDataShader = render::shader::MaterialData; using ObjectData = render::shader::ObjectData; -using DrawPushConst = render::shader::DrawPushConst; -using CullPushConst = render::shader::CullPushConst; using DeferredLightingPushConst = render::shader::DeferredLightingPushConst; using DrawCmd = render::shader::DrawCmd; using DrawData = render::shader::DrawData; @@ -277,7 +275,9 @@ static PipelineShaders makeDrawShaders( std::filesystem::path root_dir = py_root_env ? (std::string(py_root_env) + "/src/render") : STRINGIFY(MADRONA_RENDER_DATA_DIR); std::filesystem::path shader_dir = std::filesystem::weakly_canonical(root_dir / "shaders"); - auto shader_path = (shader_dir / "viewer_draw.hlsl").string(); + // Compiled only to reflect the material-texture descriptor-set layout + // (set 3) that the live batch-draw path binds; the pipeline is never built. + auto shader_path = (shader_dir / "batch_draw_rgb.hlsl").string(); ShaderCompiler compiler; SPIRVShader vert_spirv = compiler.compileHLSLFileToSPV( @@ -306,14 +306,21 @@ static PipelineShaders makeDrawShaders( return PipelineShaders(dev, tmp_alloc, shaders, Span({ BindingOverride { - 2, + 3, 0, VK_NULL_HANDLE, max_textures, VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT, }, BindingOverride { - 2, + 4, + 0, + VK_NULL_HANDLE, + max_textures, + VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT, + }, + BindingOverride { + 3, 1, repeat_sampler, 1, @@ -328,232 +335,18 @@ static PipelineShaders makeCullShader(const Device &dev) std::filesystem::path root_dir = py_root_env ? (std::string(py_root_env) + "/src/render") : STRINGIFY(MADRONA_RENDER_DATA_DIR); std::filesystem::path shader_dir = std::filesystem::weakly_canonical(root_dir / "shaders"); + // Compiled only to reflect the instance-cull descriptor-set layout (set 2) + // that the live prepare-views path binds; the pipeline is never built. ShaderCompiler compiler; SPIRVShader spirv = compiler.compileHLSLFileToSPV( - (shader_dir / "viewer_cull.hlsl").string().c_str(), {}, - {}, { "instanceCull", ShaderStage::Compute }); + (shader_dir / "prepare_views.hlsl").string().c_str(), {}, + {}, { "main", ShaderStage::Compute }); StackAlloc tmp_alloc; return PipelineShaders(dev, tmp_alloc, Span(&spirv, 1), {}); } -static Pipeline<1> makeDrawPipeline(const Device &dev, - VkPipelineCache pipeline_cache, - VkRenderPass render_pass, - VkSampler repeat_sampler, - VkSampler clamp_sampler, - uint32_t num_frames, - uint32_t max_textures) -{ - auto shaders = makeDrawShaders( - dev, repeat_sampler, clamp_sampler, max_textures); - VkPipelineVertexInputStateCreateInfo vert_info {}; - VkPipelineInputAssemblyStateCreateInfo input_assembly_info {}; - VkPipelineViewportStateCreateInfo viewport_info {}; - VkPipelineMultisampleStateCreateInfo multisample_info {}; - VkPipelineRasterizationStateCreateInfo raster_info {}; - - initCommonDrawPipelineInfo(vert_info, input_assembly_info, - viewport_info, multisample_info, raster_info); - - // Depth/Stencil - VkPipelineDepthStencilStateCreateInfo depth_info {}; - depth_info.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO; - depth_info.depthTestEnable = VK_TRUE; - depth_info.depthWriteEnable = VK_TRUE; - depth_info.depthCompareOp = VK_COMPARE_OP_GREATER_OR_EQUAL; - depth_info.depthBoundsTestEnable = VK_FALSE; - depth_info.stencilTestEnable = VK_FALSE; - depth_info.back.compareOp = VK_COMPARE_OP_ALWAYS; - - // Blend - VkPipelineColorBlendAttachmentState blend_attach {}; - blend_attach.blendEnable = VK_FALSE; - blend_attach.colorWriteMask = - VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | - VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT; - - array blend_attachments {{ - blend_attach, - blend_attach, - blend_attach - }}; - - VkPipelineColorBlendStateCreateInfo blend_info {}; - blend_info.sType = - VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO; - blend_info.logicOpEnable = VK_FALSE; - blend_info.attachmentCount = - static_cast(blend_attachments.size()); - blend_info.pAttachments = blend_attachments.data(); - - // Dynamic - array dyn_enable { - VK_DYNAMIC_STATE_VIEWPORT, - VK_DYNAMIC_STATE_SCISSOR, - }; - - VkPipelineDynamicStateCreateInfo dyn_info {}; - dyn_info.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO; - dyn_info.dynamicStateCount = dyn_enable.size(); - dyn_info.pDynamicStates = dyn_enable.data(); - - // Push constant - VkPushConstantRange push_const { - VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT, - 0, - sizeof(DrawPushConst), - }; - - // Layout configuration - - array draw_desc_layouts {{ - shaders.getLayout(0), - shaders.getLayout(1), - shaders.getLayout(2), - }}; - - VkPipelineLayoutCreateInfo gfx_layout_info; - gfx_layout_info.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO; - gfx_layout_info.pNext = nullptr; - gfx_layout_info.flags = 0; - gfx_layout_info.setLayoutCount = static_cast(draw_desc_layouts.size()); - gfx_layout_info.pSetLayouts = draw_desc_layouts.data(); - gfx_layout_info.pushConstantRangeCount = 1; - gfx_layout_info.pPushConstantRanges = &push_const; - - VkPipelineLayout draw_layout; - REQ_VK(dev.dt.createPipelineLayout(dev.hdl, &gfx_layout_info, nullptr, &draw_layout)); - array gfx_stages {{ - { - VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO, - nullptr, - 0, - VK_SHADER_STAGE_VERTEX_BIT, - shaders.getShader(0), - "vert", - nullptr, - }, - { - VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO, - nullptr, - 0, - VK_SHADER_STAGE_FRAGMENT_BIT, - shaders.getShader(1), - "frag", - nullptr, - }, - }}; - - VkGraphicsPipelineCreateInfo gfx_info; - gfx_info.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO; - gfx_info.pNext = nullptr; - gfx_info.flags = 0; - gfx_info.stageCount = gfx_stages.size(); - gfx_info.pStages = gfx_stages.data(); - gfx_info.pVertexInputState = &vert_info; - gfx_info.pInputAssemblyState = &input_assembly_info; - gfx_info.pTessellationState = nullptr; - gfx_info.pViewportState = &viewport_info; - gfx_info.pRasterizationState = &raster_info; - gfx_info.pMultisampleState = &multisample_info; - gfx_info.pDepthStencilState = &depth_info; - gfx_info.pColorBlendState = &blend_info; - gfx_info.pDynamicState = &dyn_info; - gfx_info.layout = draw_layout; - gfx_info.renderPass = render_pass; - gfx_info.subpass = 0; - gfx_info.basePipelineHandle = VK_NULL_HANDLE; - gfx_info.basePipelineIndex = -1; - - VkPipeline draw_pipeline; - REQ_VK(dev.dt.createGraphicsPipelines(dev.hdl, pipeline_cache, 1, - &gfx_info, nullptr, &draw_pipeline)); - - FixedDescriptorPool desc_pool(dev, shaders, 0, num_frames); - - return { - std::move(shaders), - draw_layout, - { draw_pipeline }, - std::move(desc_pool), - }; -} - -static Pipeline<1> makeCullPipeline(const Device &dev, - VkPipelineCache pipeline_cache, - CountT num_frames) -{ - PipelineShaders shader = makeCullShader(dev); - - // Push constant - VkPushConstantRange push_const { - VK_SHADER_STAGE_COMPUTE_BIT, - 0, - sizeof(CullPushConst), - }; - - // Layout configuration - std::array desc_layouts { - shader.getLayout(0), - shader.getLayout(1), - }; - - VkPipelineLayoutCreateInfo cull_layout_info; - cull_layout_info.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO; - cull_layout_info.pNext = nullptr; - cull_layout_info.flags = 0; - cull_layout_info.setLayoutCount = - static_cast(desc_layouts.size()); - cull_layout_info.pSetLayouts = desc_layouts.data(); - cull_layout_info.pushConstantRangeCount = 1; - cull_layout_info.pPushConstantRanges = &push_const; - - VkPipelineLayout cull_layout; - REQ_VK(dev.dt.createPipelineLayout(dev.hdl, &cull_layout_info, nullptr, - &cull_layout)); - - std::array compute_infos; -#if 0 - VkPipelineShaderStageRequiredSubgroupSizeCreateInfoEXT subgroup_size; - subgroup_size.sType = - VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_REQUIRED_SUBGROUP_SIZE_CREATE_INFO_EXT; - subgroup_size.pNext = nullptr; - subgroup_size.requiredSubgroupSize = 32; -#endif - - compute_infos[0].sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO; - compute_infos[0].pNext = nullptr; - compute_infos[0].flags = 0; - compute_infos[0].stage = { - VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO, - nullptr, //&subgroup_size, - VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT_EXT, - VK_SHADER_STAGE_COMPUTE_BIT, - shader.getShader(0), - "instanceCull", - nullptr, - }; - compute_infos[0].layout = cull_layout; - compute_infos[0].basePipelineHandle = VK_NULL_HANDLE; - compute_infos[0].basePipelineIndex = -1; - - std::array pipelines; - REQ_VK(dev.dt.createComputePipelines(dev.hdl, pipeline_cache, - compute_infos.size(), - compute_infos.data(), nullptr, - pipelines.data())); - - FixedDescriptorPool desc_pool(dev, shader, 0, num_frames); - - return Pipeline<1> { - std::move(shader), - cull_layout, - pipelines, - std::move(desc_pool), - }; -} @@ -565,8 +358,7 @@ static EngineInterop setupEngineInterop(Device &dev, uint32_t max_instances_per_world, uint32_t max_lights_per_world, uint32_t render_width, - uint32_t render_height, - VoxelConfig voxel_config) + uint32_t render_height) { (void)dev; @@ -753,31 +545,6 @@ static EngineInterop setupEngineInterop(Device &dev, } } - const uint32_t num_voxels = voxel_config.xLength * voxel_config.yLength * voxel_config.zLength; - const uint32_t staging_size = num_voxels > 0 ? num_voxels * sizeof(int32_t) : 4; - - auto voxel_cpu = Optional::none(); - VkBuffer voxel_buffer_hdl = VK_NULL_HANDLE; - uint32_t *voxel_buffer_ptr = nullptr; - -#ifdef MADRONA_VK_CUDA_SUPPORT - auto voxel_gpu = Optional::none(); - auto voxel_cuda = Optional::none(); -#endif - - if (!gpu_input) { - voxel_cpu = alloc.makeStagingBuffer(staging_size); - voxel_buffer_ptr = num_voxels ? (uint32_t *)voxel_cpu->ptr : nullptr; - voxel_buffer_hdl = voxel_cpu->buffer; - } else { -#ifdef MADRONA_VK_CUDA_SUPPORT - voxel_gpu = alloc.makeDedicatedBuffer(staging_size, false, true); - voxel_cuda.emplace(dev, voxel_gpu->mem, staging_size); - voxel_buffer_hdl = voxel_gpu->buf.buffer; - voxel_buffer_ptr = num_voxels ? (uint32_t *)voxel_cuda->getDevicePointer() : nullptr; -#endif - } - uint32_t *total_num_views_readback = nullptr; uint32_t *total_num_instances_readback = nullptr; uint32_t *total_num_lights_readback = nullptr; @@ -825,7 +592,6 @@ static EngineInterop setupEngineInterop(Device &dev, .lightWorldIDs = (uint64_t *)world_ids_lights_base, .renderWidth = (int32_t)render_width, .renderHeight = (int32_t)render_height, - .voxels = voxel_buffer_ptr, .maxViewsPerworld = max_views_per_world, .maxInstancesPerWorld = max_instances_per_world, .maxLightsPerWorld = max_lights_per_world, @@ -900,12 +666,6 @@ static EngineInterop setupEngineInterop(Device &dev, max_views_per_world, max_instances_per_world, max_lights_per_world, - std::move(voxel_cpu), -#ifdef MADRONA_VK_CUDA_SUPPORT - std::move(voxel_gpu), - std::move(voxel_cuda), -#endif - voxel_buffer_hdl, iota_array_instances, iota_array_views, iota_array_lights, @@ -1310,15 +1070,11 @@ RenderContext::RenderContext( InternalConfig::gbufferFormat, InternalConfig::depthFormat)), shadowPass(makeShadowRenderPass( dev, InternalConfig::varianceFormat, InternalConfig::depthFormat)), - instanceCull(makeCullPipeline(dev, pipelineCache, InternalConfig::numFrames)), - objectDraw(makeDrawPipeline( - dev, pipelineCache, renderPass, repeatSampler, clampSampler, - InternalConfig::numFrames, max_textures_)), - asset_desc_pool_cull_(dev, instanceCull.shaders, 1, 1), - asset_desc_pool_draw_(dev, objectDraw.shaders, 1, 1), - asset_desc_pool_mat_tx_(dev, objectDraw.shaders, 2, 1), + instanceCull(makeCullShader(dev)), + objectDraw(makeDrawShaders(dev, repeatSampler, clampSampler, max_textures_)), + asset_desc_pool_cull_(dev, instanceCull, 2, 1), + asset_desc_pool_mat_tx_(dev, objectDraw, 3, 1), asset_set_cull_(asset_desc_pool_cull_.makeSet()), - asset_set_draw_(asset_desc_pool_draw_.makeSet()), asset_set_mat_tex_(asset_desc_pool_mat_tx_.makeSet()), load_cmd_pool_(makeCmdPool(dev, dev.gfxQF)), load_cmd_(makeCmdBuffer(dev, load_cmd_pool_)), @@ -1327,12 +1083,11 @@ RenderContext::RenderContext( dev, alloc, cfg.execMode == ExecMode::CUDA, cfg.numWorlds, cfg.maxViewsPerWorld, cfg.maxInstancesPerWorld, cfg.maxLightsPerWorld, - br_width_, br_height_, cfg.voxelCfg)), + br_width_, br_height_)), lights_(InternalConfig::maxLights), loaded_assets_(0), sky_(loadSky(dev, alloc, renderQueue)), material_textures_(0), - voxel_config_(cfg.voxelCfg), num_worlds_(cfg.numWorlds), gpu_input_(cfg.execMode == ExecMode::CUDA) { @@ -1583,11 +1338,6 @@ RenderContext::~RenderContext() dev.dt.destroyDescriptorPool(dev.hdl, asset_pool_, nullptr); - dev.dt.destroyPipeline(dev.hdl, objectDraw.hdls[0], nullptr); - dev.dt.destroyPipelineLayout(dev.hdl, objectDraw.layout, nullptr); - - dev.dt.destroyPipeline(dev.hdl, instanceCull.hdls[0], nullptr); - dev.dt.destroyPipelineLayout(dev.hdl, instanceCull.layout, nullptr); dev.dt.destroyRenderPass(dev.hdl, renderPass, nullptr); dev.dt.destroyRenderPass(dev.hdl, shadowPass, nullptr); @@ -2176,14 +1926,12 @@ CountT RenderContext::loadObjects(Span src_objs, vert_info.buffer = asset_buffer.buffer; vert_info.offset = buffer_offsets[1]; vert_info.range = buffer_sizes[2]; - DescHelper::storage(desc_updates.emplace_back(), asset_set_draw_, &vert_info, 0); DescHelper::storage(desc_updates.emplace_back(), asset_batch_lighting_set_, &vert_info, 0); VkDescriptorBufferInfo mat_info; mat_info.buffer = asset_buffer.buffer; mat_info.offset = buffer_offsets[3]; mat_info.range = buffer_sizes[4]; - DescHelper::storage(desc_updates.emplace_back(), asset_set_draw_, &mat_info, 1); DescHelper::storage(desc_updates.emplace_back(), asset_batch_lighting_set_, &mat_info, 2); VkDescriptorBufferInfo index_set_info; diff --git a/src/render/render_ctx.hpp b/src/render/render_ctx.hpp index 8b6ec624..582971f6 100644 --- a/src/render/render_ctx.hpp +++ b/src/render/render_ctx.hpp @@ -39,15 +39,16 @@ struct RenderContext { VkRenderPass renderPass; VkRenderPass shadowPass; - Pipeline<1> instanceCull; - Pipeline<1> objectDraw; + // Compiled batch-path shaders kept only to source the descriptor-set + // layouts below (instanceCull <- prepare_views set 2, objectDraw <- + // batch_draw_rgb set 3); no pipelines are built from them. + render::vk::PipelineShaders instanceCull; + render::vk::PipelineShaders objectDraw; render::vk::FixedDescriptorPool asset_desc_pool_cull_; - render::vk::FixedDescriptorPool asset_desc_pool_draw_; render::vk::FixedDescriptorPool asset_desc_pool_mat_tx_; VkDescriptorSet asset_set_cull_; - VkDescriptorSet asset_set_draw_; VkDescriptorSet asset_set_mat_tex_; VkCommandPool load_cmd_pool_; @@ -64,7 +65,6 @@ struct RenderContext { Sky sky_; DynArray material_textures_; - VoxelConfig voxel_config_; uint32_t num_worlds_; std::unique_ptr batchRenderer; diff --git a/src/render/shaders/batch_draw_depth.hlsl b/src/render/shaders/batch_draw_depth.hlsl deleted file mode 100644 index 83644e49..00000000 --- a/src/render/shaders/batch_draw_depth.hlsl +++ /dev/null @@ -1,102 +0,0 @@ -#include "shader_utils.hlsl" - -[[vk::push_constant]] -BatchDrawPushConst pushConst; - -// Instances and views -[[vk::binding(0, 0)]] -StructuredBuffer viewDataBuffer; - -[[vk::binding(1, 0)]] -StructuredBuffer engineInstanceBuffer; - -[[vk::binding(2, 0)]] -StructuredBuffer instanceOffsets; - -// Draw information -[[vk::binding(0, 1)]] -RWStructuredBuffer drawCount; - -[[vk::binding(1, 1)]] -RWStructuredBuffer drawCommandBuffer; - -[[vk::binding(2, 1)]] -RWStructuredBuffer drawDataBuffer; - -// Asset descriptor bindings -[[vk::binding(0, 2)]] -StructuredBuffer vertexDataBuffer; - -[[vk::binding(1, 2)]] -StructuredBuffer meshDataBuffer; - -[[vk::binding(2, 2)]] -StructuredBuffer materialBuffer; - -struct V2F { - [[vk::location(0)]] float3 vsCoord : TEXCOORD0; -}; - -[shader("vertex")] -float4 vert(in uint vid : SV_VertexID, - in uint draw_id : SV_InstanceID, - out V2F v2f) : SV_Position -{ - DrawDataBR draw_data = drawDataBuffer[draw_id + pushConst.drawDataOffset]; - - Vertex vert = unpackVertex(vertexDataBuffer[vid]); - uint instance_id = draw_data.instanceID; - - PerspectiveCameraData view_data = unpackViewData(viewDataBuffer[draw_data.viewID]); - EngineInstanceData instance_data = unpackEngineInstanceData(engineInstanceBuffer[instance_id]); - - float3 to_view_translation; - float4 to_view_rotation; - computeCompositeTransform(instance_data.position, instance_data.rotation, - view_data.pos, view_data.rot, - to_view_translation, to_view_rotation); - - float3 view_pos = - rotateVec(to_view_rotation, instance_data.scale * vert.position) + - to_view_translation; - - float depth = length(view_pos); - float z_far = view_data.zFar; - float z_near = view_data.zNear; - - float clip_z = z_far / (z_far - z_near) * view_pos.y + - (z_far * z_near) / (z_near - z_far); - float4 clip_pos = projectToClip( - view_data, - view_pos, - clip_z); - -#if 1 - uint something = min(0, instanceOffsets[0]) + - min(0, drawCount[0]) + - min(0, drawCommandBuffer[0].vertexOffset) + - min(0, int(ceil(meshDataBuffer[0].vertexOffset))); - - // v2f.meshID = draw_data.meshID; -#endif - - clip_pos.x += min(0.0, abs(float(draw_data.meshID))) + - min(0.0, abs(float(draw_data.instanceID))) + - something; - - v2f.vsCoord = view_pos; - - return clip_pos; -} - -struct PixelOutput { - float depthOut : SV_Target0; -}; - -[shader("pixel")] -PixelOutput frag(in V2F v2f, in uint prim_id : SV_PrimitiveID) -{ - PixelOutput output; - output.depthOut = length(v2f.vsCoord) + min(0.0, abs(materialBuffer[0].color.x)); - return output; -} diff --git a/src/render/shaders/blur.hlsl b/src/render/shaders/blur.hlsl deleted file mode 100644 index 680758f3..00000000 --- a/src/render/shaders/blur.hlsl +++ /dev/null @@ -1,58 +0,0 @@ -#include "shader_utils.hlsl" - -[[vk::push_constant]] -BlurPushConst pushConst; - -[[vk::binding(0, 0)]] -RWTexture2D attachment; - -[[vk::binding(1, 0)]] -RWTexture2D intermediate; - -#define WEIGHT_COUNT 5 - -[numThreads(32, 32, 1)] -[shader("compute")] -void blur(uint3 idx : SV_DispatchThreadID) -{ - uint2 targetDim; - attachment.GetDimensions(targetDim.x, targetDim.y); - - if (idx.x < targetDim.x && idx.y < targetDim.y) { - uint2 targetPixel = idx.xy; - - const float weights[WEIGHT_COUNT] = { -#if 0 - 20.0 / 64.0, - 15.0 / 64.0, - 6.0 / 64.0, - 1.0 / 64.0 -#endif - 0.227027, 0.1945946, 0.1216216, 0.054054, 0.016216 - // 1.0 - }; - - if (pushConst.isVertical) { - // Read from intermediate, write to attachment - float2 output = intermediate[targetPixel].rg * weights[0]; - - for (int i = 1; i < WEIGHT_COUNT; ++i) { - output += intermediate[targetPixel + uint2(0, i)].rg * weights[i]; - output += intermediate[targetPixel - uint2(0, i)].rg * weights[i]; - } - - attachment[targetPixel] = output; - } - else { - // Read from attachment, write to intermediate - float2 output = attachment[targetPixel].rg * weights[0]; - - for (int i = 1; i < WEIGHT_COUNT; ++i) { - output += attachment[targetPixel + uint2(i, 0)].rg * weights[i]; - output += attachment[targetPixel - uint2(i, 0)].rg * weights[i]; - } - - intermediate[targetPixel] = output; - } - } -} diff --git a/src/render/shaders/draw_deferred_depth.hlsl b/src/render/shaders/draw_deferred_depth.hlsl deleted file mode 100644 index d9d42e05..00000000 --- a/src/render/shaders/draw_deferred_depth.hlsl +++ /dev/null @@ -1,160 +0,0 @@ -#include "shader_utils.hlsl" - -// GBuffer descriptor bindings - -[[vk::push_constant]] -DeferredLightingPushConstBR pushConst; - -// This is an array of all the textures -[[vk::binding(0, 0)]] -RWTexture2DArray rgbInBuffer[]; - -[[vk::binding(1, 0)]] -RWStructuredBuffer rgbOutputBuffer; - -[[vk::binding(2, 0)]] -RWStructuredBuffer depthOutputBuffer; - -[[vk::binding(3, 0)]] -Texture2D depthInBuffer[]; - -[[vk::binding(4, 0)]] -SamplerState linearSampler; - -[[vk::binding(0, 1)]] -StructuredBuffer indexBuffer; - -// Instances and views -[[vk::binding(0, 2)]] -StructuredBuffer viewDataBuffer; - -[[vk::binding(1, 2)]] -StructuredBuffer engineInstanceBuffer; - -[[vk::binding(2, 2)]] -StructuredBuffer instanceOffsets; - - -// Lighting -[[vk::binding(0, 3)]] -StructuredBuffer lights; - -[[vk::binding(1, 3)]] -Texture2D transmittanceLUT; - -[[vk::binding(2, 3)]] -Texture2D irradianceLUT; - -[[vk::binding(3, 3)]] -Texture3D scatteringLUT; - -[[vk::binding(4, 3)]] -StructuredBuffer skyBuffer; - - -#include "lighting.h" - -uint zeroDummy() -{ - uint zero_dummy = min(asuint(viewDataBuffer[0].data[2].w), 0) + - min(asuint(engineInstanceBuffer[0].data[0].x), 0) + - min(indexBuffer[0], 0) + - min(instanceOffsets[0], 0) + - min(0.0, abs(transmittanceLUT.SampleLevel( - linearSampler, float2(0.0, 0.0f), 0).x)) + - min(0.0, abs(irradianceLUT.SampleLevel( - linearSampler, float2(0.0, 0.0f), 0).x)) + - min(0.0, abs(scatteringLUT.SampleLevel( - linearSampler, float3(0.0, 0.0f, 0.0f), 0).x)) + - min(0.0, abs(skyBuffer[0].solarIrradiance.x)) + - min(0.0, abs(float(rgbInBuffer[0][uint3(0,0,0)].x))) + - min(0.0, abs(viewDataBuffer[0].data[0].x)) + - min(0.0, abs(engineInstanceBuffer[0].data[0].x)) + - min(0.0, abs(float(indexBuffer[0]))) + - min(0.0, abs(depthInBuffer[0].SampleLevel(linearSampler, float2(0,0), 0).x)); - - - return zero_dummy; -} - -float linearToSRGB(float v) -{ - if (v <= 0.00031308f) { - return 12.92f * v; - } else { - return 1.055f*pow(v,(1.f / 2.4f)) - 0.055f; - } -} - -uint32_t linearToSRGB8(float3 rgb) -{ - float3 srgb = float3( - linearToSRGB(rgb.x), - linearToSRGB(rgb.y), - linearToSRGB(rgb.z)); - - uint3 quant = (uint3)(255 * clamp(srgb, 0.f, 1.f)); - - return quant.r | (quant.g << 8) | (quant.b << 16) | ((uint32_t)255 << 24); -} - -// idx.x is the x coordinate of the image -// idx.y is the y coordinate of the image -// idx.z is the global view index -[numThreads(32, 32, 1)] -[shader("compute")] -void lighting(uint3 idx : SV_DispatchThreadID) -{ - uint view_idx = idx.z; - - uint num_views_per_image = pushConst.maxImagesXPerTarget * - pushConst.maxImagesYPerTarget; - - // Figure out which image to render to - uint target_idx = view_idx / num_views_per_image; - - // View index within that target - uint target_view_idx = view_idx % num_views_per_image; - - uint target_view_idx_x = target_view_idx % - pushConst.maxImagesXPerTarget; - uint target_view_idx_y = target_view_idx / - pushConst.maxImagesXPerTarget; - - float x_pixel_offset = target_view_idx_x * pushConst.viewWidth; - float y_pixel_offset = target_view_idx_y * pushConst.viewHeight; - - if (idx.x >= pushConst.viewWidth || idx.y >= pushConst.viewHeight) { - return; - } - - uint3 vbuffer_pixel = uint3(idx.x, idx.y, 0); - - float2 vbuffer_pixel_clip = - float2(float(vbuffer_pixel.x) + 0.5f, float(vbuffer_pixel.y) + 0.5f) / - float2(pushConst.viewWidth, pushConst.viewHeight); - - vbuffer_pixel_clip = vbuffer_pixel_clip * 2.0f - float2(1.0f, 1.0f); - vbuffer_pixel_clip.y *= -1.0; - - uint2 sample_uv_u32 = vbuffer_pixel.xy + uint2(x_pixel_offset, y_pixel_offset); - - float2 total_res = float2(pushConst.viewWidth * pushConst.maxImagesXPerTarget, pushConst.viewHeight * pushConst.maxImagesYPerTarget); - - float2 sample_uv = float2(sample_uv_u32) / total_res; - sample_uv.y = 1.0 - sample_uv.y; - - float depth = rgbInBuffer[target_idx][vbuffer_pixel + - uint3(x_pixel_offset, y_pixel_offset, 0)]; - - float3 out_color = float3(depth, depth, depth); - - out_color.x += zeroDummy(); - - uint32_t out_pixel_idx = - view_idx * pushConst.viewWidth * pushConst.viewHeight + - idx.y * pushConst.viewWidth + idx.x; - - rgbOutputBuffer[out_pixel_idx] = linearToSRGB8(float3(0 + zeroDummy(), 0, 0)); - depthOutputBuffer[out_pixel_idx] = depth; -} diff --git a/src/render/shaders/grid_draw.hlsl b/src/render/shaders/grid_draw.hlsl deleted file mode 100644 index 32e6a3e3..00000000 --- a/src/render/shaders/grid_draw.hlsl +++ /dev/null @@ -1,122 +0,0 @@ -#include "shader_utils.hlsl" - -[[vk::push_constant]] -GridDrawPushConst pushConst; - -[[vk::binding(0, 0)]] -RWTexture2D gridOutput; - -[[vk::binding(1, 0)]] -StructuredBuffer batchRenderRGBOut; - -[[vk::binding(2, 0)]] -StructuredBuffer batchRenderDepthOut; - -float vizDepth(float v) -{ - return log(v + 0.9f) / log(10000.f); -} - -float srgbToLinear(float srgb) -{ - if (srgb <= 0.04045f) { - return srgb / 12.92f; - } - - return pow((srgb + 0.055f) / 1.055f, 2.4f); -} - -float4 rgb8ToFloat(uint r, uint g, uint b) -{ - return float4( - srgbToLinear((float)r / 255.f), - srgbToLinear((float)g / 255.f), - srgbToLinear((float)b / 255.f), - 1.f); -} - -float4 fetchBatchRenderPixel(uint3 pixel_idx) -{ - pixel_idx.x = clamp(pixel_idx.x, 0, pushConst.viewWidth - 1); - pixel_idx.y = clamp(pixel_idx.y, 0, pushConst.viewHeight - 1); - - uint linear_pixel_idx = - pixel_idx.z * pushConst.viewHeight * pushConst.viewWidth + - pixel_idx.y * pushConst.viewWidth + pixel_idx.x; - - if (pushConst.showDepth == 1) { - float depth = vizDepth(batchRenderDepthOut[linear_pixel_idx]); - return float4(float3(depth, depth, depth), 1); - } else { - uint packed = batchRenderRGBOut[linear_pixel_idx]; - - uint r = packed & 0xFF; - uint g = (packed >> 8) & 0xFF; - uint b = (packed >> 16) & 0xFF; - - return rgb8ToFloat(r, g, b); - } -} - -float4 sampleBatchRenderOutput(float2 uv, uint view_idx) -{ - float2 coords = - uv * float2(pushConst.viewWidth, pushConst.viewHeight); - - float2 base = floor(coords); - float2 diff = coords - base; - uint2 c00 = (uint2)base; - uint2 c01 = c00 + uint2(1, 0); - uint2 c10 = c00 + uint2(0, 1); - uint2 c11 = c00 + uint2(1, 1); - - float4 p00 = fetchBatchRenderPixel(uint3(c00, view_idx)); - float4 p01 = fetchBatchRenderPixel(uint3(c01, view_idx)); - float4 p10 = fetchBatchRenderPixel(uint3(c10, view_idx)); - float4 p11 = fetchBatchRenderPixel(uint3(c11, view_idx)); - - float4 a = p00 + diff.x * (p01 - p00); - float4 b = p10 + diff.x * (p11 - p10); - - return a + diff.y * (b - a); -} - -[numThreads(32, 32, 1)] -[shader("compute")] -void gridDraw(uint3 idx : SV_DispatchThreadID) -{ - uint2 target_dim; - gridOutput.GetDimensions(target_dim.x, target_dim.y); - - if (idx.x >= target_dim.x || idx.y >= target_dim.y) { - return; - } - - // Get the view index that this pixel is going to sample from - float global_pixel_x = pushConst.offsetX + float(idx.x); - float global_pixel_y = pushConst.offsetY + float(idx.y); - - float ratio_x = global_pixel_x / float(pushConst.gridViewSize); - float ratio_y = global_pixel_y / float(pushConst.gridViewSize); - - int view_idx_x = floor(ratio_x); - int view_idx_y = floor(ratio_y); - - float uv_x = ratio_x - float(view_idx_x); - float uv_y = ratio_y - float(view_idx_y); - - if (view_idx_x < 0 || view_idx_y < 0 || view_idx_x >= pushConst.gridWidth) { - gridOutput[idx.xy] = float4(0, 0, 0, 0); - return; - } - - // Get linear view index - int view_idx = view_idx_x + view_idx_y * pushConst.gridWidth; - - if (view_idx >= pushConst.numViews) { - gridOutput[idx.xy] = float4(0, 0, 0, 0); - return; - } - - gridOutput[idx.xy] = sampleBatchRenderOutput(float2(uv_x, uv_y), view_idx); -} diff --git a/src/render/shaders/post_process.hlsl b/src/render/shaders/post_process.hlsl deleted file mode 100644 index 65b4c46c..00000000 --- a/src/render/shaders/post_process.hlsl +++ /dev/null @@ -1,394 +0,0 @@ -#include "shader_utils.hlsl" - -// Post-processing push constants -struct PostProcessPushConst { - uint32_t viewWidth; - uint32_t viewHeight; - uint32_t totalViews; - uint32_t blurRadius; // blur kernel radius -}; - -[[vk::push_constant]] -PostProcessPushConst pushConst; - -// Input/Output buffers -[[vk::binding(0, 0)]] -RWStructuredBuffer rgbBuffer; // Input and output RGB buffer - -// ------------------------------------------------------------------------------------------------ -// Sample RGB from buffer with bounds checking -uint32_t sampleRGB(uint32_t view_idx, int2 coord) { - // Clamp coordinates to view bounds - coord = clamp(coord, int2(0, 0), int2(pushConst.viewWidth - 1, pushConst.viewHeight - 1)); - - uint32_t pixel_idx = view_idx * pushConst.viewWidth * pushConst.viewHeight + - coord.y * pushConst.viewWidth + coord.x; - - return rgbBuffer[pixel_idx]; -} - -// ------------------------------------------------------------------------------------------------ -float3 unpackRGB8(uint32_t packed) { - uint3 rgb; - rgb.r = (packed) & 0xFF; - rgb.g = (packed >> 8) & 0xFF; - rgb.b = (packed >> 16) & 0xFF; - - return float3(rgb) / 255.0f; -} - -// ------------------------------------------------------------------------------------------------ -uint32_t packRGB8(float3 rgb) { - uint3 quant = (uint3)(255 * clamp(rgb, 0.f, 1.f)); - return quant.r | (quant.g << 8) | (quant.b << 16) | ((uint32_t)255 << 24); -} - -// ------------------------------------------------------------------------------------------------ -// Sample RGB as float3 with bounds checking -float3 sampleRGBFloat(uint32_t view_idx, int2 coord) { - return unpackRGB8(sampleRGB(view_idx, coord)); -} - -// ------------------------------------------------------------------------------------------------ -// Bilinear interpolation helper -float3 bilerp(float3 a, float3 b, float3 c, float3 d, float2 t) { - float3 top = lerp(a, b, t.x); // Interpolate between top-left and top-right - float3 bottom = lerp(c, d, t.x); // Interpolate between bottom-left and bottom-right - return lerp(top, bottom, t.y); // Interpolate between top and bottom -} - -// ------------------------------------------------------------------------------------------------ -// Texture sampling with bilinear filtering -// uvs are normalized coordinates [0,1] -// offset is in pixels -// Equivalent to GLSL textureOffset() function -float3 sampleBilinearColor(uint32_t view_idx, float2 uv, float2 offset = float2(0.f, 0.f)) { - // Convert pixel offset to normalized coordinates - float2 offsetUV = float2(offset) / float2(pushConst.viewWidth, pushConst.viewHeight); - - // Apply offset to normalized coordinates - uv += offsetUV; - - // Convert normalized coordinates to pixel coordinates - float2 texelCoord = uv * float2(pushConst.viewWidth, pushConst.viewHeight); - - // Get the integer coordinates of the four surrounding texels - int2 coord00 = int2(floor(texelCoord)); - int2 coord10 = coord00 + int2(1, 0); - int2 coord01 = coord00 + int2(0, 1); - int2 coord11 = coord00 + int2(1, 1); - - // Calculate fractional part for interpolation - float2 fract = texelCoord - float2(coord00); - - // Sample the four surrounding texels - float3 sample00 = sampleRGBFloat(view_idx, coord00); // Top-left - float3 sample10 = sampleRGBFloat(view_idx, coord10); // Top-right - float3 sample01 = sampleRGBFloat(view_idx, coord01); // Bottom-left - float3 sample11 = sampleRGBFloat(view_idx, coord11); // Bottom-right - - // Perform bilinear interpolation - return bilerp(sample00, sample10, sample01, sample11, fract); -} - -// ------------------------------------------------------------------------------------------------ -// Texture sampling with guassian filtering -float3 sampleGaussianColor(uint view_idx, float2 uv, float2 offset = float2(0.f, 0.f)) { - float2 resolution = float2(pushConst.viewWidth, pushConst.viewHeight); - - // Convert pixel offset to normalized coordinates - float2 offsetUV = float2(offset) / resolution; - // Apply offset to normalized coordinates - uv += offsetUV; - - float2 texelCoord = uv * resolution + 0.5; - int2 baseCoord = int2(floor(texelCoord)); - - // Fractional part is ignored here — this is a centered filter. - float3 result = float3(0.0, 0.0, 0.0); - - float weightSum = 0; - - const int2 offsets[9] = { - int2(-1, -1), int2(0, -1), int2(1, -1), - int2(-1, 0), int2(0, 0), int2(1, 0), - int2(-1, 1), int2(0, 1), int2(1, 1) - }; - - const float weights[9] = { - 0.5, 0.7, 0.5, - 0.7, 2.5, 0.7, - 0.5, 0.7, 0.5 - }; - - for (int i = 0; i < 9; ++i) { - int2 coord = baseCoord + offsets[i]; - float3 color = sampleRGBFloat(view_idx, coord); - result += weights[i] * color; - weightSum += weights[i]; - } - - return result / float(weightSum); -} - -// ------------------------------------------------------------------------------------------------ -float3 sampleColor(uint32_t view_idx, float2 uv, float2 offset = float2(0.f, 0.f)) { - return sampleBilinearColor(view_idx, uv, offset); -} - -// ------------------------------------------------------------------------------------------------ -// Similar to textureOffset but return luminance directly -float sampleLuminance(uint32_t view_idx, float2 uv, float2 offset = float2(0.f, 0.f)) { - return rgbToLuminance(sampleColor(view_idx, uv, offset)); -} - -// ------------------------------------------------------------------------------------------------ -// Edge detection result -struct EdgeResult { - bool isEdge; - float lumaRange; -}; - -// given center and neighbor luminances, is anti-aliasing needed? -EdgeResult detectEdge(float lumaM, float lumaN, float lumaS, float lumaE, float lumaW) { - // Find the minimum and maximum luma around the current pixel - float lumaMin = min(lumaM, min(min(lumaN, lumaS), min(lumaE, lumaW))); - float lumaMax = max(lumaM, max(max(lumaN, lumaS), max(lumaE, lumaW))); - - EdgeResult result; - - // Compute the delta - result.lumaRange = lumaMax - lumaMin; - - // If the luma variation is lower than a threshold, we are not on an edge - const float fxaaQualityEdgeThreshold = 0.02; - const float fxaaQualityEdgeMax = 0.125; - - result.isEdge = (result.lumaRange >= max(fxaaQualityEdgeThreshold, lumaMax * fxaaQualityEdgeMax)); - - return result; -} - -// ------------------------------------------------------------------------------------------------ -// is horizontal or vertical edge? -bool isHorizontal(float lumaM, float lumaN, float lumaS, float lumaE, float lumaW, - float lumaNE, float lumaNW, float lumaSE, float lumaSW) { - // Combine the four edge lumas - float lumaNS = lumaN + lumaS; - float lumaWE = lumaW + lumaE; - - // Combine the four corner lumas - float lumaW_ = lumaNW + lumaSW; - float lumaE_ = lumaNE + lumaSE; - float lumaN_ = lumaNW + lumaNE; - float lumaS_ = lumaSW + lumaSE; - - // Compute an estimation of the gradient along the horizontal and vertical axis. - float edgeHoriz = abs(-2.0 * lumaW + lumaW_) + abs(-2.0 * lumaM + lumaNS) * 2.0 + abs(-2.0 * lumaE + lumaE_); - float edgeVert = abs(-2.0 * lumaS + lumaS_) + abs(-2.0 * lumaM + lumaWE) * 2.0 + abs(-2.0 * lumaN + lumaN_); - - // Is the local edge horizontal or vertical? - return edgeHoriz >= edgeVert; -} - -// ------------------------------------------------------------------------------------------------ -float3 applyFXAA(uint32_t view_idx, int2 coord) { - // Sample the center pixel - float2 texelSize = float2(1.0 / float(pushConst.viewWidth), 1.0 / float(pushConst.viewHeight)); - float2 coordf = float2(coord) * texelSize;// + 0.5 * texelSize; - float3 rgbM = sampleColor(view_idx, coordf); - float lumaM = rgbToLuminance(rgbM); - - // 1.Detecting where to apply AA - - // Sample neighbors - float lumaS = sampleLuminance(view_idx, coordf, float2( 0.f, -1.f)); // South - float lumaN = sampleLuminance(view_idx, coordf, float2( 0.f, 1.f)); // North - float lumaE = sampleLuminance(view_idx, coordf, float2(-1.f, 0.f)); // East - float lumaW = sampleLuminance(view_idx, coordf, float2( 1.f, 0.f)); // West - - EdgeResult edgeDetection = detectEdge(lumaM, lumaN, lumaS, lumaE, lumaW); - if (!edgeDetection.isEdge) { - return sampleRGBFloat(view_idx, coord); // No anti-aliasing needed - } - float lumaRange = edgeDetection.lumaRange; - - // 2. Estimating gradient and choosing edge direction - - // Sample the corners - float lumaNW = sampleLuminance(view_idx, coordf, float2( 1.f, 1.f)); - float lumaNE = sampleLuminance(view_idx, coordf, float2(-1.f, 1.f)); - float lumaSW = sampleLuminance(view_idx, coordf, float2( 1.f, -1.f)); - float lumaSE = sampleLuminance(view_idx, coordf, float2(-1.f, -1.f)); - - // Is the local edge horizontal or vertical? - bool horzSpan = isHorizontal(lumaM, lumaN, lumaS, lumaE, lumaW, lumaNE, lumaNW, lumaSE, lumaSW); - - // Select the two neighboring texels lumas in the opposite direction to the local edge - float luma1 = horzSpan ? lumaS : lumaE; - float luma2 = horzSpan ? lumaN : lumaW; - - // Compute gradients in this direction - float gradient1 = luma1 - lumaM; - float gradient2 = luma2 - lumaM; - - // Which direction is the steepest? - bool is1Steepest = abs(gradient1) >= abs(gradient2); - - // Gradient in the corresponding direction, normalized - float gradientScaled = 0.25 * max(abs(gradient1), abs(gradient2)); - - // Choose the step size (how far to go on each iteration) according to the edge direction - float stepLength = horzSpan ? texelSize.y : texelSize.x; - - // Average luma in the correct direction - float lumaLocalAverage = 0.0; - if (is1Steepest) { - // Switch the direction - stepLength = -stepLength; - lumaLocalAverage = 0.5 * (luma1 + lumaM); - } else { - lumaLocalAverage = 0.5 * (luma2 + lumaM); - } - - // Shift UV in the correct direction - float2 currentUv = float2(coordf); - if (horzSpan) { - currentUv.y += stepLength * 0.5; - } else { - currentUv.x += stepLength * 0.5; - } - - // Compute offset (for each iteration step) in the right direction - float2 offset = horzSpan ? float2(texelSize.x, 0.0) : float2(0.0, texelSize.y); - - // Compute UVs to explore on each side of the edge, orthogonally - float2 uv1 = currentUv - offset; - float2 uv2 = currentUv + offset; - - // Read the lumas at both current extremities of the exploration segment - float lumaEnd1 = sampleLuminance(view_idx, uv1); - float lumaEnd2 = sampleLuminance(view_idx, uv2); - - lumaEnd1 -= lumaLocalAverage; - lumaEnd2 -= lumaLocalAverage; - - // If the luma deltas at the current extremities are larger than the local gradient, we have reached the side of the edge - bool reached1 = abs(lumaEnd1) >= gradientScaled; - bool reached2 = abs(lumaEnd2) >= gradientScaled; - bool reachedBoth = reached1 && reached2; - - // If the side is not reached, we continue to explore in this direction - if (!reached1) { - uv1 -= offset; - } - if (!reached2) { - uv2 += offset; - } - - // If both sides have not been reached, continue to explore - #define ITERATIONS 10 - const float QUALITY[ITERATIONS] = {1.5f, 2.0f, 2.0f, 2.0f, 2.0f, 4.0f, 8.0f, 16.0f, 32.0f, 64.0f}; - if (!reachedBoth) { - for (int i = 0; i < ITERATIONS; i++) { // Maximum 12 iterations - // If needed, read luma in 1st direction, compute delta - if (!reached1) { - lumaEnd1 = sampleLuminance(view_idx, uv1); - lumaEnd1 = lumaEnd1 - lumaLocalAverage; - } - // If needed, read luma in opposite direction, compute delta - if (!reached2) { - lumaEnd2 = sampleLuminance(view_idx, uv2); - lumaEnd2 = lumaEnd2 - lumaLocalAverage; - } - // If the luma deltas at the current extremities is larger than the local gradient, we have reached the side of the edge - reached1 = abs(lumaEnd1) >= gradientScaled; - reached2 = abs(lumaEnd2) >= gradientScaled; - reachedBoth = reached1 && reached2; - - // If the side is not reached, we continue to explore in this direction, with a variable quality - if (!reached1) { - uv1 -= offset * QUALITY[i]; - } - if (!reached2) { - uv2 += offset * QUALITY[i]; - } - - // If both sides have been reached, stop the exploration - if (reachedBoth) { - break; - } - } - } - - // Compute the distances to each extremity of the edge - float distance1 = horzSpan ? (coordf.x - uv1.x) : (coordf.y - uv1.y); - float distance2 = horzSpan ? (uv2.x - coordf.x) : (uv2.y - coordf.y); - - // In which direction is the extremity of the edge closer? - bool isDirection1 = distance1 < distance2; - float distanceFinal = min(distance1, distance2); - - // Length of the edge - float edgeLength = (distance1 + distance2); - - // UV offset: read in the direction of the closest side of the edge - float pixelOffset = -distanceFinal / edgeLength + 0.5; - - // Is the luma at center smaller than the local average? - bool isLumaMLowerThanAvg = lumaM < lumaLocalAverage; - - // If the luma at center is smaller than at its neighbor, the delta luma at each end should be positive (same variation) - bool correctVariation = ((isDirection1 ? lumaEnd1 : lumaEnd2) < 0.0) != isLumaMLowerThanAvg; - - // If the luma variation is incorrect, do not offset - float finalOffset = correctVariation ? pixelOffset : 0.0; - - // Sub-pixel shifting - // Full weighted average of the luma over the 3x3 neighborhood - float lumaAverage = (1.0/12.0) * (2.0 * (lumaN + lumaE + lumaS + lumaW) + lumaNE + lumaNW + lumaSE + lumaSW); - - // Ratio of the delta between the global average and the center luma, over the luma range in the 3x3 neighborhood - float subPixelOffset1 = clamp(abs(lumaAverage - lumaM) / lumaRange, 0.0, 1.0); - float subPixelOffset2 = (-2.0 * subPixelOffset1 + 3.0) * subPixelOffset1 * subPixelOffset1; - - // Compute a sub-pixel offset based on this delta - const float fxaaQualitySubpixel = 0.75; - float subPixelOffsetFinal = subPixelOffset2 * subPixelOffset2 * fxaaQualitySubpixel; // Sub-pixel quality - - // Pick the biggest of the two offsets - finalOffset = max(finalOffset, subPixelOffsetFinal); - - // Compute the final UV coordinates - float2 finalUv = float2(coordf); - if (horzSpan) { - finalUv.y += finalOffset * stepLength; - } else { - finalUv.x += finalOffset * stepLength; - } - - // Read the color at the new UV coordinates, and return it - return sampleColor(view_idx, finalUv); -} - -// ------------------------------------------------------------------------------------------------ -[numThreads(32, 32, 1)] -[shader("compute")] -void main(uint3 idx : SV_DispatchThreadID) -{ - if (idx.x >= pushConst.viewWidth || idx.y >= pushConst.viewHeight || idx.z >= pushConst.totalViews) { - return; - } - - uint view_idx = idx.z; - int2 coord = int2(idx.x, idx.y); - - float3 final_color; - - final_color = applyFXAA(view_idx, coord); - - uint32_t out_pixel_idx = view_idx * pushConst.viewWidth * pushConst.viewHeight + - idx.y * pushConst.viewWidth + idx.x; - - rgbBuffer[out_pixel_idx] = packRGB8(final_color); -} diff --git a/src/render/shaders/shader_common.h b/src/render/shaders/shader_common.h index d89c32b7..43214a12 100644 --- a/src/render/shaders/shader_common.h +++ b/src/render/shaders/shader_common.h @@ -131,14 +131,6 @@ struct ShadowGenPushConst { uint32_t worldIdx; }; -struct VoxelGenPushConst { - uint32_t worldX; - uint32_t worldY; - uint32_t worldZ; - float blockWidth; - uint32_t numBlocks; -}; - struct Vertex { float3 position; float3 normal; diff --git a/src/render/shaders/shadow_gen.hlsl b/src/render/shaders/shadow_gen.hlsl deleted file mode 100644 index 0f9688f7..00000000 --- a/src/render/shaders/shadow_gen.hlsl +++ /dev/null @@ -1,164 +0,0 @@ -#include "shader_utils.hlsl" - -[[vk::push_constant]] -ShadowGenPushConst pushConst; - -[[vk::binding(0, 0)]] -RWStructuredBuffer shadowViewDataBuffer; - -[[vk::binding(1, 0)]] -StructuredBuffer flycamBuffer; - -[[vk::binding(2, 0)]] -StructuredBuffer lights; - -[[vk::binding(3, 0)]] -StructuredBuffer viewDataBuffer; - -[[vk::binding(4, 0)]] -StructuredBuffer viewOffsetsBuffer; - -float4 invQuat(float4 rot) -{ - return float4(-rot.x, -rot.y, -rot.z, rot.w); -} - -PerspectiveCameraData getCameraData() -{ - if (pushConst.viewIdx == 0) { - return unpackViewData(flycamBuffer[0]); - } else { - int view_idx = (pushConst.viewIdx - 1) + viewOffsetsBuffer[pushConst.worldIdx]; - return unpackViewData(viewDataBuffer[view_idx]); - } -} - -[numThreads(32, 1, 1)] -[shader("compute")] -void shadowGen(uint3 idx : SV_DispatchThreadID) -{ - /* Assume that the sun is from lights[0] */ - if (idx.x != 0) - return; - - PerspectiveCameraData unpackedView = getCameraData(); - - float3 cam_pos = unpackedView.pos; - float4 cam_rot = invQuat(unpackedView.rot); - - float3 cam_fwd = rotateVec(cam_rot, float3(0.0f, 1.0f, 0.0f)); - float3 cam_up = rotateVec(cam_rot, float3(0.0f, 0.0f, 1.0f)); - float3 cam_right = rotateVec(cam_rot, float3(1.0f, 0.0f, 0.0f)); - - // Construct orthonormal basis - ShaderLightData light = unpackLightData(lights[0]); - float3 light_fwd = normalize(light.direction.xyz); - float3 light_up = (light_fwd.x < 0.9999f) ? - normalize(cross(float3(1.f, 0.f, 0.f), light_fwd)) : - float3(0.f, 0.f, 1.f); - float3 light_right = cross(light_fwd, light_up); - - // Note that we use the basis vectors as the *rows* of the to_light - // transform matrix, because we want the inverse of the light to world - // matrix (which is just the transpose for rotation matrices). - float3x3 to_light = float3x3( - light_right.x, light_right.y, light_right.z, - light_fwd.x, light_fwd.y, light_fwd.z, - light_up.x, light_up.y, light_up.z - ); - - float far_width, near_width, far_height, near_height; - - float tan_half_fov = -1.0f / unpackedView.yScale; - float aspect = -unpackedView.yScale / unpackedView.xScale; - float near = 1.0f; - float far = 80.0f; - - far_height = 2.0f * far * tan_half_fov; - near_height = 2.0f * near * tan_half_fov; - far_width = far_height * aspect; - near_width = near_height * aspect; - - float3 center_near = cam_pos + cam_fwd * near; - float3 center_far = cam_pos + cam_fwd * far; - - float far_width_half = far_width / 2.0f; - float near_width_half = near_width / 2.0f; - float far_height_half = far_height / 2.0f; - float near_height_half = near_height / 2.0f; - - // f = far, n = near, l = left, r = right, t = top, b = bottom - enum OrthoCorner { - flt, flb, - frt, frb, - nlt, nlb, - nrt, nrb - }; - - float3 ls_corners[8]; - -#if 0 - ls_corners[flt] = mul(view, float4(cam_pos + ws_direction * far - right_view_ax * far_width_half + up_view_ax * far_height_half, 1.0f)); - ls_corners[flb] = mul(view, float4(ws_position + ws_direction * far - right_view_ax * far_width_half - up_view_ax * far_height_half, 1.0f)); - ls_corners[frt] = mul(view, float4(ws_position + ws_direction * far + right_view_ax * far_width_half + up_view_ax * far_height_half, 1.0f)); - ls_corners[frb] = mul(view, float4(ws_position + ws_direction * far + right_view_ax * far_width_half - up_view_ax * far_height_half, 1.0f)); - ls_corners[nlt] = mul(view, float4(ws_position + ws_direction * near - right_view_ax * near_width_half + up_view_ax * near_height_half, 1.0f)); - ls_corners[nlb] = mul(view, float4(ws_position + ws_direction * near - right_view_ax * near_width_half - up_view_ax * near_height_half, 1.0f)); - ls_corners[nrt] = mul(view, float4(ws_position + ws_direction * near + right_view_ax * near_width_half + up_view_ax * near_height_half, 1.0f)); - ls_corners[nrb] = mul(view, float4(ws_position + ws_direction * near + right_view_ax * near_width_half - up_view_ax * near_height_half, 1.0f)); -#endif - - ls_corners[flt] = mul(to_light, center_far - cam_right * far_width_half + cam_up * far_height_half); - ls_corners[flb] = mul(to_light, center_far - cam_right * far_width_half - cam_up * far_height_half); - ls_corners[frt] = mul(to_light, center_far + cam_right * far_width_half + cam_up * far_height_half); - ls_corners[frb] = mul(to_light, center_far + cam_right * far_width_half - cam_up * far_height_half); - ls_corners[nlt] = mul(to_light, center_near - cam_right * near_width_half + cam_up * near_height_half); - ls_corners[nlb] = mul(to_light, center_near - cam_right * near_width_half - cam_up * near_height_half); - ls_corners[nrt] = mul(to_light, center_near + cam_right * near_width_half + cam_up * near_height_half); - ls_corners[nrb] = mul(to_light, center_near + cam_right * near_width_half - cam_up * near_height_half); - - float x_min, x_max, y_min, y_max, z_min, z_max; - - x_min = x_max = ls_corners[0].x; - y_min = y_max = ls_corners[0].y; - z_min = z_max = ls_corners[0].z; - - for (uint32_t i = 1; i < 8; ++i) { - if (x_min > ls_corners[i].x) x_min = ls_corners[i].x; - if (x_max < ls_corners[i].x) x_max = ls_corners[i].x; - - if (y_min > ls_corners[i].y) y_min = ls_corners[i].y; - if (y_max < ls_corners[i].y) y_max = ls_corners[i].y; - - if (z_min > ls_corners[i].z) z_min = ls_corners[i].z; - if (z_max < ls_corners[i].z) z_max = ls_corners[i].z; - } - - { - float tmp = y_max; - y_max = y_min; - y_min = tmp; - } - - - float4x4 projection =(float4x4( - float4(2.0f / (x_max - x_min), 0.0f, 0.0f, -(x_max + x_min) / (x_max - x_min)), - float4(0.0f, 0.0f, -2.0f / (z_max - z_min), (z_max+z_min) / (z_max - z_min)), - float4(0.0f, 1.0f / (y_max - y_min), 0.0f, -(y_min) / (y_max - y_min)), - float4(0.0f, 0.0f, 0.0f, 1.0f))); - - shadowViewDataBuffer[pushConst.viewIdx].viewProjectionMatrix = mul( - projection, float4x4( - float4(to_light[0].xyz, 0.f), - float4(to_light[1].xyz, 0.f), - float4(to_light[2].xyz, 0.f), - float4(0.f, 0.f, 0.f, 1.f) - ) - ); - - { - shadowViewDataBuffer[pushConst.viewIdx].cameraRight = float4(cam_right, 1.f); - shadowViewDataBuffer[pushConst.viewIdx].cameraUp = float4(cam_up, 1.f); - shadowViewDataBuffer[pushConst.viewIdx].cameraForward = float4(cam_fwd, 1.f); - } -} diff --git a/src/render/shaders/textured_quad.hlsl b/src/render/shaders/textured_quad.hlsl deleted file mode 100644 index 6ff5afec..00000000 --- a/src/render/shaders/textured_quad.hlsl +++ /dev/null @@ -1,64 +0,0 @@ -#include "shader_utils.hlsl" - -[[vk::push_constant]] -TexturedQuadPushConst pushConst; - -[[vk::binding(0, 0)]] -Texture2D toDisplay; - -[[vk::binding(1, 0)]] -SamplerState samplerState; - -#if 0 -[numThreads(32, 32, 1)] -[shader("compute")] -void main(uint3 idx : SV_DispatchThreadID) -{ - if (idx.x >= pushConst.extentPixels.x || idx.y >= pushConst.extentPixels.y) { - return; - } - - uint2 dst_pixel = pushConst.startPixels + idx.xy; - uint2 src_pixel = idx.xy; - - // Source pixel UV coordinate for sampling - float2 src_pixel_uv = float2(idx.xy) / float2(pushConst.extentPixels); - float4 src_pixel_value = toDisplay.SampleLevel(samplerState, src_pixel_uv, 0); - - outputImage[dst_pixel] = src_pixel_value; -} -#endif - -struct V2F { - [[vk::location(0)]] float2 uv : TEXCOORD0; -}; - -[shader("vertex")] -float4 vert(in uint vid : SV_VertexID, - out V2F v2f) : SV_Position -{ - float2 vertices[4] = { - pushConst.startPixels, - pushConst.startPixels + float2(pushConst.extentPixels.x, 0), - pushConst.startPixels + float2(0, pushConst.extentPixels.y), - pushConst.startPixels + pushConst.extentPixels, - }; - - v2f.uv = (vertices[vid] - pushConst.startPixels) / pushConst.extentPixels; - - float4 pos = float4(2.0f * (vertices[vid] / pushConst.targetExtent) - float2(1.0, 1.0), 0.0, 1.0); - - return pos; -} - -struct PixelOutput { - float4 color : SV_Target0; -}; - -[shader("pixel")] -PixelOutput frag(in V2F v2f) -{ - PixelOutput output; - output.color = toDisplay.SampleLevel(samplerState, v2f.uv, 0); - return output; -} diff --git a/src/render/shaders/viewer_cull.hlsl b/src/render/shaders/viewer_cull.hlsl deleted file mode 100644 index f4e0da7f..00000000 --- a/src/render/shaders/viewer_cull.hlsl +++ /dev/null @@ -1,112 +0,0 @@ -#include "shader_utils.hlsl" - -[[vk::push_constant]] -CullPushConst pushConst; - -// Contains the view just for the fly cam -[[vk::binding(0, 0)]] -StructuredBuffer viewDataBuffer; - -// Contains the instances for all the worlds -[[vk::binding(1, 0)]] -StructuredBuffer engineInstanceBuffer; - -[[vk::binding(2, 0)]] -RWStructuredBuffer drawCount; - -[[vk::binding(3, 0)]] -RWStructuredBuffer drawCommandBuffer; - -[[vk::binding(4, 0)]] -RWStructuredBuffer drawDataBuffer; - -[[vk::binding(5, 0)]] -RWStructuredBuffer instanceOffsets; - -// Asset descriptor bindings - -[[vk::binding(0, 1)]] -StructuredBuffer objectDataBuffer; - -[[vk::binding(1, 1)]] -StructuredBuffer meshDataBuffer; - -uint getNumInstancesForWorld(uint world_idx) -{ - if (world_idx == pushConst.numWorlds - 1) { - return pushConst.numInstances - instanceOffsets[world_idx]; - } else { - return instanceOffsets[world_idx+1] - instanceOffsets[world_idx]; - } -} - -uint getInstanceOffsetsForWorld(uint world_idx) -{ - return instanceOffsets[world_idx]; -} - -struct SharedData { - uint numInstances; - uint numInstancesPerThread; - uint instancesOffset; -}; - -groupshared SharedData sm; - -// No actual culling performed yet -[numThreads(32, 1, 1)] -[shader("compute")] -void instanceCull(uint3 tid : SV_DispatchThreadID, - uint3 tid_local : SV_GroupThreadID, - uint3 gid : SV_GroupID) -{ - if (tid_local.x == 0) { - sm.numInstances = getNumInstancesForWorld(pushConst.worldID); - sm.numInstancesPerThread = (sm.numInstances + pushConst.numThreads-1) / - pushConst.numThreads; - sm.instancesOffset = getInstanceOffsetsForWorld(pushConst.worldID); - // printf("%d\n", sm.numInstances); - } - - GroupMemoryBarrierWithGroupSync(); - - for (int i = 0; i < sm.numInstancesPerThread; ++i) { - uint local_idx = i * pushConst.numThreads + tid.x; - - if (local_idx >= sm.numInstances) { - return; - } - - uint current_instance_idx = sm.instancesOffset + local_idx; - - EngineInstanceData instance_data = unpackEngineInstanceData(engineInstanceBuffer[current_instance_idx]); - ObjectData obj = objectDataBuffer[instance_data.objectID]; - - uint draw_offset; - InterlockedAdd(drawCount[0], obj.numMeshes, draw_offset); - - for (int32_t i = 0; i < obj.numMeshes; i++) { - MeshData mesh = meshDataBuffer[obj.meshOffset + i]; - - uint draw_id = draw_offset + i; - DrawCmd draw_cmd; - draw_cmd.indexCount = mesh.numIndices; - draw_cmd.instanceCount = 1; - draw_cmd.firstIndex = mesh.indexOffset; - draw_cmd.vertexOffset = mesh.vertexOffset; - draw_cmd.firstInstance = draw_id; - - DrawData draw_data; - - if (instance_data.matID == -1) { - draw_data.materialID = mesh.materialIndex; - } else { - draw_data.materialID = instance_data.matID; - } - draw_data.instanceID = current_instance_idx; - - drawCommandBuffer[draw_id] = draw_cmd; - drawDataBuffer[draw_id] = draw_data; - } - } -} diff --git a/src/render/shaders/viewer_deferred_lighting.hlsl b/src/render/shaders/viewer_deferred_lighting.hlsl deleted file mode 100644 index 07fdac98..00000000 --- a/src/render/shaders/viewer_deferred_lighting.hlsl +++ /dev/null @@ -1,334 +0,0 @@ -#include "shader_utils.hlsl" - -// GBuffer descriptor bindings - -[[vk::push_constant]] -DeferredLightingPushConst pushConst; - -[[vk::binding(0, 0)]] -RWTexture2D gbufferAlbedo; - -[[vk::binding(1, 0)]] -RWTexture2D gbufferNormal; - -[[vk::binding(2, 0)]] -RWTexture2D gbufferPosition; - -// Assume stuff is Y-UP from here -[[vk::binding(3, 0)]] -StructuredBuffer lights; - -// Atmosphere -[[vk::binding(4, 0)]] -Texture2D transmittanceLUT; - -[[vk::binding(5, 0)]] -Texture2D irradianceLUT; - -[[vk::binding(6, 0)]] -Texture3D scatteringLUT; - -// Shadows -[[vk::binding(7, 0)]] -Texture2D shadowMap; - -// Assume stuff is Y-UP from here -[[vk::binding(8, 0)]] -StructuredBuffer shadowViewDataBuffer; - -// Sampler -[[vk::binding(9, 0)]] -SamplerState linearSampler; - -// Assume stuff from here is Y-UP -[[vk::binding(10, 0)]] -StructuredBuffer skyBuffer; - -#include "lighting.h" - -#define SHADOW_BIAS 0.002f - -float linear_step(float low, float high, float v) { - return clamp((v - low) / (high - low), 0, 1); -} - -/* Shadowing is done using variance shadow mapping. */ -float shadowFactorVSM(float3 world_pos, uint2 target_pixel) -{ - uint2 shadow_map_dim; - shadowMap.GetDimensions(shadow_map_dim.x, shadow_map_dim.y); - - float2 texel_size = float2(1.f, 1.f) / float2(shadow_map_dim); - - float4 world_pos_v4 = float4(world_pos.xyz, 1.f); - - /* Light space position */ - float4 ls_pos = mul(shadowViewDataBuffer[pushConst.viewIdx].viewProjectionMatrix, - world_pos_v4); - - ls_pos.xyz /= ls_pos.w; - ls_pos.z += SHADOW_BIAS; - - /* UV to use when sampling in the shadow map. */ - float2 uv = ls_pos.xy * 0.5 + float2(0.5, 0.5); - - /* Only deal with points which are within the shadow map. */ - if (uv.x > 1.0 || uv.x < 0.0 || uv.y > 1.0 || uv.y < 0.0 || - ls_pos.z > 1.0 || ls_pos.z < 0.0) - return 1.0; - - float2 moment = shadowMap.SampleLevel(linearSampler, uv, 0); - - float occlusion = 0.0f; - - float pcf_count = 1; - - for (int x = int(-pcf_count); x <= int(pcf_count); ++x) { - for (int y = int(-pcf_count); y <= int(pcf_count); ++y) { - float2 moment = shadowMap.SampleLevel(linearSampler, - uv + float2(x, y) * texel_size, 0).rg; - - // Chebychev's inequality - float p = (ls_pos.z > moment.x); - float sigma = max(moment.y - moment.x * moment.x, 0.0); - - float dist_from_mean = (ls_pos.z - moment.x); - - float pmax = linear_step(0.9, 1.0, sigma / (sigma + dist_from_mean * dist_from_mean)); - float occ = min(1.0f, max(pmax, p)); - - occlusion += occ; - } - } - - occlusion /= (pcf_count * 2.0f + 1.0f) * (pcf_count * 2.0f + 1.0f); - - return occlusion; -} - -/* Coordinates which are prefixed with w are in world space. */ -struct GBufferData { - float3 wPosition; - float3 wNormal; - float3 albedo; - float3 wCameraPos; -}; - -/* BRDF calculations. */ -float distributionGGX(float ndoth, float roughness) { - float a2 = roughness * roughness; - float f = (ndoth * a2 - ndoth) * ndoth + 1.0f; - return a2 / (M_PI * f * f); -} - -float smithGGX(float ndotv, float ndotl, float roughness) { - float r = roughness + 1.0; - float k = (r * r) / 8.0; - float ggx1 = ndotv / (ndotv * (1.0f - k) + k); - float ggx2 = ndotl / (ndotl * (1.0f - k) + k); - return ggx1 * ggx2; -} - -float3 fresnel(float hdotv, float3 base) { - return base + (1.0f - base) * pow(1.0f - clamp(hdotv, 0.0f, 1.0f), 5.0f); -} - -float3 fresnelRoughness(float ndotv, float3 base, float roughness) { - float one_minus_rough = 1.0f - roughness; - return base + (max(float3(one_minus_rough, one_minus_rough, one_minus_rough), base) - base) * - pow(1.0f - ndotv, 5.0f); -} - -/* Gets the BRDF contribution on a pixel which lies on a surface for any generic - * incoming radiance value. */ -float3 directionalRadianceBRDF(in GBufferData gbuffer, - in float3 base_reflectivity, - in float roughness, - in float metalness, - in float3 view_direction, - in float3 incoming_radiance, - in float3 light_direction) { - float3 view_dir = -view_direction; - - float3 halfway = normalize(light_direction + view_dir); - - float ndotv = max(dot(gbuffer.wNormal.xyz, view_dir), 0.000001f); - float ndotl = max(dot(gbuffer.wNormal.xyz, light_direction), 0.000001f); - float hdotv = max(dot(halfway, view_dir), 0.000001f); - float ndoth = max(dot(gbuffer.wNormal.xyz, halfway), 0.000001f); - - float distribution_term = distributionGGX(ndoth, roughness); - float smith_term = smithGGX(ndotv, ndotl, roughness); - float3 fresnel_term = fresnel(hdotv, base_reflectivity); - - float3 specular = smith_term * distribution_term * fresnel_term; - specular /= 4.0 * ndotv * ndotl; - - float3 kd = float3(1.0, 1.0, 1.0) - fresnel_term; - - kd *= 1.0f - metalness; - - return (kd * gbuffer.albedo.rgb / M_PI + specular) * incoming_radiance * ndotl; -} - -/* Gets the radiance on a pixel which lies on a surface taking into account BRDF as - * well as the sky transmittance and radiance incoming from the sun. Also takes into - * account shadows using variance shadow maps. */ -float3 accumulateSunRadianceBRDF(in GBufferData gbuffer, - float roughness, - float metal, - float r, - float mu_sun, - uint2 target_pixel) -{ - float3 ret = float3(0.0, 0.0, 0.0); - - /* We get the radiance incoming from the sun to this point by taking into account - * the transmittance of the sky. */ - float3 radiance_from_sun = skyBuffer[0].solarIrradiance.xyz * - getTransmittanceToSun(skyBuffer[0], transmittanceLUT, r, mu_sun); - - ShaderLightData light = unpackLightData(lights[0]); - ret += directionalRadianceBRDF(gbuffer, - lerp(float3(0.04, 0.04, 0.04), gbuffer.albedo.rgb, metal), - roughness, - metal, - normalize(gbuffer.wPosition.xyz - pushConst.viewPos.xyz), - radiance_from_sun, - normalize(-light.direction.xyz)); - - float shadow_factor = shadowFactorVSM(gbuffer.wPosition, target_pixel); - - return ret * shadow_factor; -} - -/* Caculates the radiance which actually arrives to the camera. This will require calling - * accumulateSunRadianceBRDF to get the radiance hitting the surface, and then using - * getSkyRadianceToPoint to calculate the amount of in/out scatter happening on the ray - * from the surface point to the camera. */ -float4 getPointRadianceBRDF(float roughness, float metal, in GBufferData gbuffer, uint2 target_pixel) -{ - float3 sky_irradiance, sun_irradiance, point_radiance; - ShaderLightData light = unpackLightData(lights[0]); - - { /* Calculate sun and sky irradiance which will contribute to the final BRDF. */ - float3 p = gbuffer.wPosition / 1000.0 - skyBuffer[0].wPlanetCenter.xyz; - float3 normal = gbuffer.wNormal; - - float3 sun_direction = -normalize(light.direction.xyz); - - float3 view_direction = normalize(gbuffer.wPosition - pushConst.viewPos.xyz); - - float r = length(p); - float mu_sun = dot(p, sun_direction) / r; - - sky_irradiance = getIrradiance(skyBuffer[0], irradianceLUT, r, mu_sun) * - (1.0 + dot(normal, p) / r) * 0.5; - - float3 accumulated_radiance = accumulateSunRadianceBRDF(gbuffer, roughness, metal, - r, mu_sun, target_pixel); - - point_radiance = accumulated_radiance + gbuffer.albedo.rgb * (1.0 / M_PI) * sky_irradiance; - } - - /* How much is scattered towards us. */ - float3 transmittance; - float3 in_scatter = getSkyRadianceToPoint(skyBuffer[0], transmittanceLUT, - scatteringLUT, scatteringLUT, - pushConst.viewPos.xyz / 1000.0 - skyBuffer[0].wPlanetCenter.xyz, - gbuffer.wPosition / 1000.0 - skyBuffer[0].wPlanetCenter.xyz, 0.0, - -normalize(light.direction.xyz), - transmittance); - - point_radiance = point_radiance * transmittance + in_scatter; - - return float4(point_radiance, 1.0); -} - -/* Calculates the outgoing camera view ray for a given pixel. */ -float3 getOutgoingRay(float2 target_pixel, float2 target_dim) -{ - float aspect_ratio = target_dim.x / target_dim.y; - float tan_fov = tan(pushConst.fovy / 2.0f); - - float right_scale = aspect_ratio * tan_fov; - float up_scale = tan_fov; - - float2 raster = float2(target_pixel.x + 0.5, target_pixel.y + 0.5); - float2 screen = float2((2.0f * raster.x) / target_dim.x - 1.0f, - (2.0f * raster.y) / target_dim.y - 1.0f); - - float3 cam_right = shadowViewDataBuffer[pushConst.viewIdx].cameraRight.xyz; - float3 cam_up = shadowViewDataBuffer[pushConst.viewIdx].cameraUp.xyz; - float3 cam_forward = shadowViewDataBuffer[pushConst.viewIdx].cameraForward.xyz; - - float3 dir = screen.x * cam_right * right_scale - screen.y * cam_up * up_scale + cam_forward; - - return normalize(dir); -} - -[numThreads(32, 32, 1)] -[shader("compute")] -void lighting(uint3 idx : SV_DispatchThreadID) -{ - uint2 target_dim; - gbufferAlbedo.GetDimensions(target_dim.x, target_dim.y); - - if (idx.x < target_dim.x && idx.y < target_dim.y) - { - uint2 target_pixel = uint2(idx.x, idx.y); - - float3 outgoing_ray = getOutgoingRay((float2)target_pixel, (float2)target_dim); - - float4 color = gbufferAlbedo[target_pixel]; - - float4 normal = gbufferNormal[target_pixel]; - float4 position = gbufferPosition[target_pixel]; - - /* If normal.w is more than 0, this object was rasterized. */ - float point_alpha = normal.w; - - GBufferData gbuffer_data; - gbuffer_data.wPosition = position.xyz; - gbuffer_data.wNormal = normal.xyz; - gbuffer_data.albedo = color.rgb; - gbuffer_data.wCameraPos = pushConst.viewPos.xyz; - - float roughness = color.a; - float metalness = position.a; - - /* Radiance from the rasterized pixel. */ - float4 point_radiance = getPointRadianceBRDF(roughness, metalness, - gbuffer_data, target_pixel); - - ShaderLightData light = unpackLightData(lights[0]); - float3 sun_direction = normalize(-light.direction.xyz); - - /* Incoming radiance from the sky: */ - float3 transmittance; - float3 radiance = getSkyRadiance(skyBuffer[0], transmittanceLUT, - scatteringLUT, scatteringLUT, - (gbuffer_data.wCameraPos / 1000.0 - skyBuffer[0].wPlanetCenter.xyz), - outgoing_ray, 0.0, sun_direction, - transmittance); - - if (dot(outgoing_ray, sun_direction) > - skyBuffer[0].sunSize.y * 0.99999) { - radiance = radiance + transmittance * getSolarRadiance(skyBuffer[0]) * 0.06; - } - - radiance = lerp(radiance, point_radiance.xyz, point_alpha); - - /* Tone Mapping. */ - float3 one = float3(1.0, 1.0, 1.0); - float3 exp_value = exp(-radiance / float3(2.0f, 2.0f, 2.0f) * pushConst.exposure); - - float3 diff = one - exp_value; - float3 out_color = diff; - - //float viz = abs(dot(normal.xyz, outgoing_ray)); - //out_color = float3(viz, viz, viz) * color.xyz + 1e-6f * out_color; - - gbufferAlbedo[target_pixel] = float4(out_color, 1.0); - } -} diff --git a/src/render/shaders/viewer_draw.hlsl b/src/render/shaders/viewer_draw.hlsl deleted file mode 100644 index 24e6e97e..00000000 --- a/src/render/shaders/viewer_draw.hlsl +++ /dev/null @@ -1,173 +0,0 @@ -#include "shader_utils.hlsl" - -[[vk::push_constant]] -DrawPushConst push_const; - -[[vk::binding(0, 0)]] -StructuredBuffer flycamBuffer; - -[[vk::binding(1, 0)]] -StructuredBuffer engineInstanceBuffer; - -[[vk::binding(2, 0)]] -StructuredBuffer drawDataBuffer; - -[[vk::binding(3, 0)]] -StructuredBuffer shadowViewDataBuffer; - -[[vk::binding(4, 0)]] -StructuredBuffer viewDataBuffer; - -[[vk::binding(5, 0)]] -StructuredBuffer viewOffsetsBuffer; - -// Asset descriptor bindings -[[vk::binding(0, 1)]] -StructuredBuffer vertexDataBuffer; - -[[vk::binding(1, 1)]] -StructuredBuffer materialBuffer; - -// Texture descriptor bindings -[[vk::binding(0, 2)]] -Texture2D materialTexturesArray[]; - -[[vk::binding(1, 2)]] -SamplerState linearSampler; - -[[vk::binding(2, 2)]] -StructuredBuffer materialTexturesIndices; - -struct V2F { - [[vk::location(0)]] float3 normal : TEXCOORD0; - [[vk::location(1)]] float3 position : TEXCOORD1; - [[vk::location(2)]] float dummy : TEXCOORD3; - [[vk::location(3)]] float2 uv : TEXCOORD4; - [[vk::location(4)]] int materialIdx : TEXCOORD5; - [[vk::location(5)]] float worldIdx : TEXCOORD6; -}; - -PerspectiveCameraData getCameraData() -{ - PerspectiveCameraData camera_data; - - if (push_const.viewIdx == 0) { - camera_data = unpackViewData(flycamBuffer[0]); - } else { - PerspectiveCameraData fly_cam = unpackViewData(flycamBuffer[0]); - - int view_idx = (push_const.viewIdx - 1) + viewOffsetsBuffer[push_const.worldIdx]; - camera_data = unpackViewData(viewDataBuffer[view_idx]); - - // We want to inherit the aspect ratio from the flycam camera - camera_data.xScale = fly_cam.xScale; - camera_data.yScale = fly_cam.yScale; - } - - return camera_data; -} - -[shader("vertex")] -float4 vert(in uint vid : SV_VertexID, - in uint draw_id : SV_InstanceID, - out V2F v2f) : SV_Position -{ - DrawData draw_data = drawDataBuffer[draw_id]; - Vertex vert = unpackVertex(vertexDataBuffer[vid]); - - uint instance_id = draw_data.instanceID; - EngineInstanceData instance_data = unpackEngineInstanceData(engineInstanceBuffer[instance_id]); - PerspectiveCameraData view_data = getCameraData(); - - float3 to_view_translation; - float4 to_view_rotation; - computeCompositeTransform( - instance_data.position, instance_data.rotation, - view_data.pos, view_data.rot, - to_view_translation, to_view_rotation - ); - - float3 view_pos = rotateVec(to_view_rotation, instance_data.scale * vert.position) + to_view_translation; - -#if 0 - float4 clip_pos = float4( - view_data.xScale * view_pos.x, - view_data.yScale * view_pos.z, - view_data.zNear, - view_pos.y); -#endif - - float4 clip_pos; - - if (push_const.isOrtho == 1) { - float x_max = push_const.xMax * view_data.xScale; - float x_min = push_const.xMin * view_data.xScale; - - float y_max = push_const.yMax; - float y_min = push_const.yMin; - - float z_max = push_const.zMax * (1.0f / -view_data.yScale); - float z_min = push_const.zMin * (1.0f / -view_data.yScale); - - float4x4 m1 = float4x4( - float4(2.0f / (x_max - x_min), 0.0f, 0.0f, -(x_max + x_min) / (x_max - x_min)), - float4(0.0f, 0.0f, -2.0f / (z_max - z_min), -(z_max+z_min) / (z_max - z_min)), - float4(0.0f, 1.0f / (y_max - y_min), 0.0f, -(y_min) / (y_max - y_min)), - float4(0.0f, 0.0f, 0.0f, 1.0f)); - - clip_pos = mul(m1, float4(view_pos.x, view_pos.y, view_pos.z, 1.0f)); - clip_pos.z = 1.0 - clip_pos.z; - } - else { -#if 0 - clip_pos = float4( view_data.xScale * view_pos.x, - view_data.yScale * view_pos.z, - view_data.zNear, - 1.0); -#endif - - clip_pos = projectToClip(view_data, view_pos, view_data.zNear); - } - - v2f.normal = normalize(rotateVec(instance_data.rotation, (vert.normal / instance_data.scale))); - v2f.uv = float2(vert.uv.x, 1.0f - vert.uv.y); - v2f.position = rotateVec(instance_data.rotation, instance_data.scale * vert.position) + instance_data.position; - v2f.dummy = shadowViewDataBuffer[0].viewProjectionMatrix[0][0]; - v2f.materialIdx = draw_data.materialID; - v2f.worldIdx = instance_data.worldID; - - return clip_pos; -} - -struct PixelOutput { - float4 color : SV_Target0; - float4 normal : SV_Target1; - float4 position : SV_Target2; -}; - -[shader("pixel")] -PixelOutput frag(in V2F v2f) -{ - PixelOutput output; - - MaterialData mat_data = materialBuffer[v2f.materialIdx]; - float metalness = mat_data.metalness; - float roughness = mat_data.roughness; - float4 color = mat_data.color; - - int texture_idx = -1; - uint texture_count = mat_data.numTextures; - if (texture_count > 0) { - uint texture_start = mat_data.textureOffset; - texture_idx = materialTexturesIndices[texture_start + v2f.worldIdx % texture_count]; - } - if (texture_idx != -1) { - color *= materialTexturesArray[texture_idx].SampleLevel(linearSampler, v2f.uv, 0); - } - - output.color = color; - output.color.a = roughness; - output.normal = float4(normalize(v2f.normal), 1.f); - output.position = float4(v2f.position, v2f.dummy * 0.0000001f + metalness); - return output; -} diff --git a/src/render/shaders/viewer_shadow_draw.hlsl b/src/render/shaders/viewer_shadow_draw.hlsl deleted file mode 100644 index 2bcbb50e..00000000 --- a/src/render/shaders/viewer_shadow_draw.hlsl +++ /dev/null @@ -1,76 +0,0 @@ -#include "shader_utils.hlsl" - -[[vk::push_constant]] -DrawPushConst push_const; - -[[vk::binding(0, 0)]] -StructuredBuffer flycamBuffer; - -[[vk::binding(1, 0)]] -StructuredBuffer engineInstanceBuffer; - -[[vk::binding(2, 0)]] -StructuredBuffer drawDataBuffer; - -[[vk::binding(3, 0)]] -StructuredBuffer shadowViewDataBuffer; - -[[vk::binding(4, 0)]] -StructuredBuffer viewDataBuffer; - -[[vk::binding(5, 0)]] -StructuredBuffer viewOffsetsBuffer; - -// Asset descriptor bindings - -[[vk::binding(0, 1)]] -StructuredBuffer vertexDataBuffer; - -[[vk::binding(1, 1)]] -StructuredBuffer materialBuffer; - -#if 0 -struct V2F { - [[vk::location(0)]] float depth : TEXCOORD0; -}; -#endif - - -[shader("vertex")] -float4 vert(in uint vid : SV_VertexID, - in uint draw_id : SV_InstanceID) : SV_Position -{ - Vertex vert = unpackVertex(vertexDataBuffer[vid]); - DrawData draw_data = drawDataBuffer[draw_id]; - float4 color = materialBuffer[vert.materialIdx].color; - uint instance_id = draw_data.instanceID; - - float4x4 shadow_matrix = shadowViewDataBuffer[push_const.viewIdx].viewProjectionMatrix; - - EngineInstanceData instance_data = unpackEngineInstanceData( - engineInstanceBuffer[instance_id]); - - float dummy = 0.00000000001f * float(drawDataBuffer[0].materialID + - flycamBuffer[0].data[0].w + materialBuffer[0].color.w + - viewDataBuffer[0].data[0].w + float(viewOffsetsBuffer[0])); - - float4 world_space_pos = float4( - instance_data.position + mul(toMat(instance_data.rotation), (instance_data.scale * vert.position)), - 1.f + dummy); - - float4 clip_pos = mul(shadow_matrix, world_space_pos); - - return clip_pos; -} - -[shader("pixel")] -float2 frag(in float4 position : SV_Position) : SV_Target0 -{ - float depth = position.z; - - float dx = ddx(depth); - float dy = ddy(depth); - float sigma = depth * depth + 0.25 * (dx * dx + dy * dy); - - return float2(depth, sigma); -} diff --git a/src/render/shaders/visualize_tris.hlsl b/src/render/shaders/visualize_tris.hlsl deleted file mode 100644 index d31df5b1..00000000 --- a/src/render/shaders/visualize_tris.hlsl +++ /dev/null @@ -1,26 +0,0 @@ -#include "shader_utils.hlsl" - -[[vk::binding(0, 0)]] -RWTexture2D vizBuffer; - -[[vk::binding(1, 0)]] -RWTexture2D colorBuffer; - -[numThreads(32, 32, 1)] -[shader("compute")] -void visualize(uint3 idx : SV_DispatchThreadID) -{ - uint2 target_dim; - vizBuffer.GetDimensions(target_dim.x, target_dim.y); - - if (idx.x >= target_dim.x || idx.y >= target_dim.y) { - return; - } - - uint2 target_pixel = uint2(idx.x, idx.y); - - uint2 in_value = vizBuffer[target_pixel]; - float4 generated_color = intToColor(hash(in_value.x + in_value.y)); - - colorBuffer[target_pixel] = generated_color; -} diff --git a/src/render/shaders/voxel_draw.hlsl b/src/render/shaders/voxel_draw.hlsl deleted file mode 100644 index d3eb4aa5..00000000 --- a/src/render/shaders/voxel_draw.hlsl +++ /dev/null @@ -1,114 +0,0 @@ -#include "shader_utils.hlsl" - -[[vk::push_constant]] -DrawPushConst push_const; - -[[vk::binding(0, 0)]] -StructuredBuffer flycamBuffer; - -// Asset descriptor bindings -[[vk::binding(1, 0)]] -StructuredBuffer vertexDataBuffer; - -[[vk::binding(2, 0)]] -StructuredBuffer viewDataBuffer; - -[[vk::binding(3, 0)]] -StructuredBuffer viewOffsetsBuffer; - -// Texture descriptor bindings -[[vk::binding(0, 1)]] -Texture2D materialTexturesArray[]; - -[[vk::binding(1, 1)]] -SamplerState linearSampler; - -[[vk::binding(2, 1)]] -StructuredBuffer materialTexturesIndices; - -struct V2F { - [[vk::location(0)]] float3 normal : TEXCOORD0; - [[vk::location(1)]] float3 position : TEXCOORD1; - [[vk::location(2)]] float4 color : TEXCOORD2; - [[vk::location(3)]] float dummy : TEXCOORD3; - [[vk::location(4)]] float2 uv : TEXCOORD4; - [[vk::location(5)]] int texIdx : TEXCOORD5; - [[vk::location(6)]] float roughness : TEXCOORD6; - [[vk::location(7)]] float metalness : TEXCOORD7; -}; - -[shader("vertex")] -float4 vert(in uint vid : SV_VertexID, - in uint draw_id : SV_InstanceID, - out V2F v2f) : SV_Position -{ - Vertex vert = unpackVertex(vertexDataBuffer[vid]); - float4 color = float4(1,1,1,1); - - PerspectiveCameraData view_data = getCameraData(); - //unpackViewData(viewDataBuffer[push_const.viewIdx]); - - float3 to_view_translation; - float4 to_view_rotation; - - float3 objectScale = float3(1,1,1); - float3 objectPos = float3(0,0,0); - float4 objectRotation = float4(0,0,0,1); - - computeCompositeTransform(objectPos, objectRotation, - view_data.pos, view_data.rot, - to_view_translation, to_view_rotation); - - float3 view_pos = - rotateVec(to_view_rotation, objectScale * vert.position) + - to_view_translation; - - float4 clip_pos = projectToClip(view_data, view_pos, view_data.zNear); - // v2f.viewPos = view_pos; -#if 0 - v2f.normal = normalize( - rotateVec(to_view_rotation, (vert.normal / objectScale))); -#endif - v2f.normal = normalize( - rotateVec(objectRotation, (vert.normal / objectScale))); - v2f.uv = vert.uv; - v2f.color = color; - v2f.position = rotateVec(objectRotation, - objectScale * vert.position) + objectPos; - v2f.dummy = 1; - v2f.texIdx = 2; - v2f.roughness = 0; - v2f.metalness = 0; - - return clip_pos; -} - -struct PixelOutput { - float4 color : SV_Target0; - float4 normal : SV_Target1; - float4 position : SV_Target2; -}; - -[shader("pixel")] -PixelOutput frag(in V2F v2f) -{ - PixelOutput output; - output.color = v2f.color; - output.color.a = v2f.roughness; - output.normal = float4(normalize(v2f.normal), 1.f); - output.position = float4(v2f.position, v2f.dummy * 0.0000001f); - output.position.a += v2f.metalness; - - // output.color.rgb = v2f.normal.xyz; - - if ( v2f.texIdx != -1) { - output.color *= materialTexturesArray[v2f.texIdx].SampleLevel( - linearSampler, float2(v2f.uv.x, 1.f - v2f.uv.y), 0); - } - - //output.color = max(float4(1,1,1,1),output.color); - //output.color = max(float4(1,1,1,1),output.color); - //output.color = output.color*float4((output.normal.xyz+float3(1,1,1))*0.5,1); - - return output; -} diff --git a/src/render/shaders/voxel_gen.hlsl b/src/render/shaders/voxel_gen.hlsl deleted file mode 100644 index bec853c4..00000000 --- a/src/render/shaders/voxel_gen.hlsl +++ /dev/null @@ -1,500 +0,0 @@ -#include "shader_utils.hlsl" - -[[vk::push_constant]] -VoxelGenPushConst pushConst; - -[[vk::binding(0, 0)]] -RWStructuredBuffer vbo; - -[[vk::binding(1, 0)]] -RWStructuredBuffer ibo; - -[[vk::binding(2, 0)]] -StructuredBuffer voxels; - - -int coord(int x,int y,int z,int blockX,int blockY,int blockZ) -{ - return x * blockY * blockZ + y * blockZ + z; -} - -#define xDim 64 -#define yDim 1 -#define zDim 1 - -[numThreads(xDim, yDim, zDim)] -[shader("compute")] -void voxelGen(uint3 idx : SV_DispatchThreadID) -{ - - int indexVal = 0; - float halfBlockSize = pushConst.blockWidth/2; - float texAtlasStep = 1.0 / pushConst.numBlocks; - - int workPerThread = ceil(pushConst.worldX / (float)xDim); - - for(int i2=0;i2= pushConst.worldX){ - return; - } - - for (int j = 0; j < pushConst.worldY; j++) { - for (int k = 0; k < pushConst.worldZ; k++) { - uint data = voxels[coord(i,j,k,pushConst.worldX,pushConst.worldY,pushConst.worldZ)]; - if (data != 0) { - float blockCenterX = halfBlockSize + pushConst.blockWidth * i; - float blockCenterY = halfBlockSize + pushConst.blockWidth * j; - float blockCenterZ = halfBlockSize + pushConst.blockWidth * k; - - int index = 32*6*coord(i,j,k,pushConst.worldX,pushConst.worldY,pushConst.worldZ); - int indexindex = 6*6*coord(i,j,k,pushConst.worldX,pushConst.worldY,pushConst.worldZ); - int vertexIndex = coord(i,j,k,pushConst.worldX,pushConst.worldY,pushConst.worldZ)*6*4; - - int blockID = data - 1; - - uint leftN = (i > 0) ? voxels[coord(i-1,j,k,pushConst.worldX,pushConst.worldY,pushConst.worldZ)]: 0; - if (!leftN) { - //position - vbo[index + 0] = blockCenterX - halfBlockSize; - vbo[index + 1] = blockCenterY - halfBlockSize; - vbo[index + 2] = blockCenterZ - halfBlockSize; - - //tex coord - vbo[index + 3] = texAtlasStep * (blockID); - vbo[index + 4] = -0.333; - - //normals - vbo[index+5] = -1; - vbo[index+6] = 0; - vbo[index+7] = 0; - - - //position - vbo[index + 8] = blockCenterX - halfBlockSize; - vbo[index + 9] = blockCenterY + halfBlockSize; - vbo[index + 10] = blockCenterZ - halfBlockSize; - - //tex coord - vbo[index + 11] = texAtlasStep * (blockID + 1); - vbo[index + 12] = -0.333; - - //normals - vbo[index+13] = -1; - vbo[index+14] = 0; - vbo[index+15] = 0; - - - //position - vbo[index + 16] = blockCenterX - halfBlockSize; - vbo[index + 17] = blockCenterY - halfBlockSize; - vbo[index + 18] = blockCenterZ + halfBlockSize; - - //tex coord - vbo[index + 19] = texAtlasStep * (blockID); - vbo[index + 20] = 0; - - //normals - vbo[index+21] = -1; - vbo[index+22] = 0; - vbo[index+23] = 0; - - - //position - vbo[index + 24] = blockCenterX - halfBlockSize; - vbo[index + 25] = blockCenterY + halfBlockSize; - vbo[index + 26] = blockCenterZ + halfBlockSize; - - //tex coord - vbo[index + 27] = texAtlasStep * (blockID + 1); - vbo[index + 28] = 0; - - //normals - vbo[index+29] = -1; - vbo[index+30] = 0; - vbo[index+31] = 0; - } - - ibo[indexindex+0] = vertexIndex + 2; - ibo[indexindex+1] = vertexIndex + 1; - ibo[indexindex+2] = vertexIndex + 0; - ibo[indexindex+3] = vertexIndex + 3; - ibo[indexindex+4] = vertexIndex + 1; - ibo[indexindex+5] = vertexIndex + 2; - - - index+=32; - indexindex+=6; - vertexIndex+=4; - - uint rightN = (i < pushConst.worldX - 1) ? voxels[coord(i+1,j,k,pushConst.worldX,pushConst.worldY,pushConst.worldZ)] : 0; - if (!rightN) { - //position - vbo[index + 0] = blockCenterX + halfBlockSize; - vbo[index + 1] = blockCenterY - halfBlockSize; - vbo[index + 2] = blockCenterZ - halfBlockSize; - - //tex coord - vbo[index + 3] = texAtlasStep * (blockID); - vbo[index + 4] = -0.333; - - //normals - vbo[index+5] = 1; - vbo[index+6] = 0; - vbo[index+7] = 0; - - - //position - vbo[index + 8] = blockCenterX + halfBlockSize; - vbo[index + 9] = blockCenterY + halfBlockSize; - vbo[index + 10] = blockCenterZ - halfBlockSize; - - //tex coord - vbo[index + 11] = texAtlasStep * (blockID + 1); - vbo[index + 12] = -0.333; - - //normals - vbo[index+13] = 1; - vbo[index+14] = 0; - vbo[index+15] = 0; - - - //position - vbo[index + 16] = blockCenterX + halfBlockSize; - vbo[index + 17] = blockCenterY - halfBlockSize; - vbo[index + 18] = blockCenterZ + halfBlockSize; - - //tex coord - vbo[index + 19] = texAtlasStep * (blockID); - vbo[index + 20] = 0; - - //normals - vbo[index+21] = 1; - vbo[index+22] = 0; - vbo[index+23] = 0; - - - //position - vbo[index + 24] = blockCenterX + halfBlockSize; - vbo[index + 25] = blockCenterY + halfBlockSize; - vbo[index + 26] = blockCenterZ + halfBlockSize; - - //tex coord - vbo[index + 27] = texAtlasStep * (blockID + 1); - vbo[index + 28] = 0; - - //normals - vbo[index+29] = 1; - vbo[index+30] = 0; - vbo[index+31] = 0; - } - - ibo[indexindex+0] = vertexIndex + 0; - ibo[indexindex+1] = vertexIndex + 1; - ibo[indexindex+2] = vertexIndex + 2; - ibo[indexindex+3] = vertexIndex + 2; - ibo[indexindex+4] = vertexIndex + 1; - ibo[indexindex+5] = vertexIndex + 3; - - index+=32; - indexindex+=6; - vertexIndex+=4; - - uint topN = (k < pushConst.worldZ - 1) ? voxels[coord(i,j,k+1,pushConst.worldX,pushConst.worldY,pushConst.worldZ)] : 0; - if (!topN) { - //position - vbo[index + 0] = (blockCenterX - halfBlockSize); - vbo[index + 1] = blockCenterY - halfBlockSize; - vbo[index + 2] = blockCenterZ + halfBlockSize; - - //tex coord - vbo[index + 3] = texAtlasStep * (blockID); - vbo[index + 4] = 0.333; - - //normals - vbo[index+5] = 0; - vbo[index+6] = 0; - vbo[index+7] = 1; - - - //position - vbo[index + 8] = blockCenterX + halfBlockSize; - vbo[index + 9] = blockCenterY - halfBlockSize; - vbo[index + 10] = blockCenterZ + halfBlockSize; - - //tex coord - vbo[index + 11] = texAtlasStep * (blockID); - vbo[index + 12] = 0.667; - - //normals - vbo[index+13] = 0; - vbo[index+14] = 0; - vbo[index+15] = 1; - - - //position - vbo[index + 16] = blockCenterX - halfBlockSize; - vbo[index + 17] = blockCenterY + halfBlockSize; - vbo[index + 18] = blockCenterZ + halfBlockSize; - - //tex coord - vbo[index + 19] = texAtlasStep * (blockID+1); - vbo[index + 20] = 0.333; - - //normals - vbo[index+21] = 0; - vbo[index+22] = 0; - vbo[index+23] = 1; - - - //position - vbo[index + 24] = blockCenterX + halfBlockSize; - vbo[index + 25] = blockCenterY + halfBlockSize; - vbo[index + 26] = blockCenterZ + halfBlockSize; - - //tex coord - vbo[index + 27] = texAtlasStep * (blockID + 1); - vbo[index + 28] = 0.667; - - //normals - vbo[index+29] = 0; - vbo[index+30] = 0; - vbo[index+31] = 1; - } - - ibo[indexindex+0] = vertexIndex + 0; - ibo[indexindex+1] = vertexIndex + 1; - ibo[indexindex+2] = vertexIndex + 2; - ibo[indexindex+3] = vertexIndex + 2; - ibo[indexindex+4] = vertexIndex + 1; - ibo[indexindex+5] = vertexIndex + 3; - - index+=32; - indexindex+=6; - vertexIndex+=4; - - uint bottomN = (k > 0) ? voxels[coord(i,j,k-1,pushConst.worldX,pushConst.worldY,pushConst.worldZ)] : 0; - if (!bottomN) { - //position - vbo[index + 0] = blockCenterX - halfBlockSize; - vbo[index + 1] = blockCenterY - halfBlockSize; - vbo[index + 2] = blockCenterZ - halfBlockSize; - - //tex coord - vbo[index + 3] = texAtlasStep * (blockID); - vbo[index + 4] = 0.667; - - //normals - vbo[index+5] = 0; - vbo[index+6] = 0; - vbo[index+7] = -1; - - - //position - vbo[index + 8] = blockCenterX + halfBlockSize; - vbo[index + 9] = blockCenterY - halfBlockSize; - vbo[index + 10] = blockCenterZ - halfBlockSize; - - //tex coord - vbo[index + 11] = texAtlasStep * (blockID); - vbo[index + 12] = 1; - - //normals - vbo[index+13] = 0; - vbo[index+14] = 0; - vbo[index+15] = -1; - - - //position - vbo[index + 16] = blockCenterX - halfBlockSize; - vbo[index + 17] = blockCenterY + halfBlockSize; - vbo[index + 18] = blockCenterZ - halfBlockSize; - - //tex coord - vbo[index + 19] = texAtlasStep * (blockID+1); - vbo[index + 20] = 0.667; - - //normals - vbo[index+21] = 0; - vbo[index+22] = 0; - vbo[index+23] = -1; - - - //position - vbo[index + 24] = blockCenterX + halfBlockSize; - vbo[index + 25] = blockCenterY + halfBlockSize; - vbo[index + 26] = blockCenterZ - halfBlockSize; - - //tex coord - vbo[index + 27] = texAtlasStep * (blockID + 1); - vbo[index + 28] = 1; - - //normals - vbo[index+29] = 0; - vbo[index+30] = 0; - vbo[index+31] = -1; - } - - ibo[indexindex+0] = vertexIndex + 2; - ibo[indexindex+1] = vertexIndex + 1; - ibo[indexindex+2] = vertexIndex + 0; - ibo[indexindex+3] = vertexIndex + 3; - ibo[indexindex+4] = vertexIndex + 1; - ibo[indexindex+5] = vertexIndex + 2; - - index+=32; - indexindex+=6; - vertexIndex+=4; - - uint frontN = (j < pushConst.worldY - 1) ? voxels[coord(i,j+1,k,pushConst.worldX,pushConst.worldY,pushConst.worldZ)] : 0; - if (!frontN) { - - - //position - vbo[index + 0] = blockCenterX - halfBlockSize; - vbo[index + 1] = blockCenterY + halfBlockSize; - vbo[index + 2] = blockCenterZ - halfBlockSize; - - //tex coord - vbo[index + 3] = texAtlasStep * (blockID); - vbo[index + 4] = -0.333; - - //normals - vbo[index+5] = 0; - vbo[index+6] = 1; - vbo[index+7] = 0; - - - //position - vbo[index + 8] = blockCenterX + halfBlockSize; - vbo[index + 9] = blockCenterY + halfBlockSize; - vbo[index + 10] = blockCenterZ - halfBlockSize; - - //tex coord - vbo[index + 11] = texAtlasStep * (blockID+1); - vbo[index + 12] = -0.333; - - //normals - vbo[index+13] = 0; - vbo[index+14] = 1; - vbo[index+15] = 0; - - - //position - vbo[index + 16] = blockCenterX - halfBlockSize; - vbo[index + 17] = blockCenterY + halfBlockSize; - vbo[index + 18] = blockCenterZ + halfBlockSize; - - //tex coord - vbo[index + 19] = texAtlasStep * (blockID); - vbo[index + 20] = 0; - - //normals - vbo[index+21] = 0; - vbo[index+22] = 1; - vbo[index+23] = 0; - - - //position - vbo[index + 24] = blockCenterX + halfBlockSize; - vbo[index + 25] = blockCenterY + halfBlockSize; - vbo[index + 26] = blockCenterZ + halfBlockSize; - - //tex coord - vbo[index + 27] = texAtlasStep * (blockID + 1); - vbo[index + 28] = 0; - - //normals - vbo[index+29] = 0; - vbo[index+30] = 1; - vbo[index+31] = 0; - } - - ibo[indexindex+0] = vertexIndex + 2; - ibo[indexindex+1] = vertexIndex + 1; - ibo[indexindex+2] = vertexIndex + 0; - ibo[indexindex+3] = vertexIndex + 3; - ibo[indexindex+4] = vertexIndex + 1; - ibo[indexindex+5] = vertexIndex + 2; - - index+=32; - indexindex+=6; - vertexIndex+=4; - - uint backN = (j > 0) ? voxels[coord(i,j-1,k,pushConst.worldX,pushConst.worldY,pushConst.worldZ)] : 0; - if (!backN) { - - - //position - vbo[index + 0] = blockCenterX - halfBlockSize; - vbo[index + 1] = blockCenterY - halfBlockSize; - vbo[index + 2] = blockCenterZ - halfBlockSize; - - //tex coord - vbo[index + 3] = texAtlasStep * (blockID); - vbo[index + 4] = -0.333; - - //normals - vbo[index+5] = 0; - vbo[index+6] = -1; - vbo[index+7] = 0; - - - //position - vbo[index + 8] = blockCenterX + halfBlockSize; - vbo[index + 9] = blockCenterY - halfBlockSize; - vbo[index + 10] = blockCenterZ - halfBlockSize; - - //tex coord - vbo[index + 11] = texAtlasStep * (blockID+1); - vbo[index + 12] = -0.333; - - //normals - vbo[index+13] = 0; - vbo[index+14] = -1; - vbo[index+15] = 0; - - - //position - vbo[index + 16] = blockCenterX - halfBlockSize; - vbo[index + 17] = blockCenterY - halfBlockSize; - vbo[index + 18] = blockCenterZ + halfBlockSize; - - //tex coord - vbo[index + 19] = texAtlasStep * (blockID); - vbo[index + 20] = 0; - - //normals - vbo[index+21] = 0; - vbo[index+22] = -1; - vbo[index+23] = 0; - - - //position - vbo[index + 24] = blockCenterX + halfBlockSize; - vbo[index + 25] = blockCenterY - halfBlockSize; - vbo[index + 26] = blockCenterZ + halfBlockSize; - - //tex coord - vbo[index + 27] = texAtlasStep * (blockID + 1); - vbo[index + 28] = 0; - - //normals - vbo[index+29] = 0; - vbo[index+30] = -1; - vbo[index+31] = 0; - - } - - ibo[indexindex + 0] = vertexIndex + 0; - ibo[indexindex + 1] = vertexIndex + 1; - ibo[indexindex + 2] = vertexIndex + 2; - ibo[indexindex + 3] = vertexIndex + 2; - ibo[indexindex + 4] = vertexIndex + 1; - ibo[indexindex + 5] = vertexIndex + 3; - - } - } - } - } -}