Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: polyline preprocess #14

Merged
merged 4 commits into from
Apr 18, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
117 changes: 62 additions & 55 deletions lib/include/preprocess/polyline_preprocess_kernel.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -16,67 +16,71 @@
#define PREPROCESS__POLYLINE_PREPROCESS_KERNEL_CUH_

/**
* @brief Transform to target agent coordinates system.
* @brief Transform polylines to target agent coordinate system and extend its feature with previous
* x, y`(D->D+2)`.
*
* Points which all elements are 0.0 except of typeID are filled by 0.0 and the corresponding
* mask value is 0.0 too.
*
* @param K The number of polylines.
* @param P The number of points contained in each polyline.
* @param PointDim The number of point state dimensions.
* @param in_polyline Source polylines, in shape [K*P*PointDim].
* @param PointDim The number of point dimensions, expecting (x, y, z, dx, dy, dz, typeID).
* @param inPolyline Source polyline, in shape [K*P*D].
* @param B The number of target agents.
* @param AgentDim The number of agent state dimensions.
* @param target_state Source target state at the latest timestamp, in shape [B*AgentDim].
* @param out_polyline Output polylines, in shape [B*K*P*(PointDim+2)].
* @param out_polyline_mask Output polyline mask, in shape [B*K*P].
* @param AgentDim The number of agent state dimensions, expecting (x, y, z, length, width, height,
* yaw, vx, vy, ax, ay).
* @param targetState Source target agent states, in shape [B*AgentDim].
* @param outPolyline Output polyline, in shape [B*K*P*(PointDim+2)].
* @param outPolylineMask Output polyline mask, in shape [B*K*P].
*/
__global__ void transformPolylineKernel(
const int K, const int P, const int PointDim, const float * in_polyline, const int B,
const int AgentDim, const float * target_state, float * out_polyline, bool * out_polyline_mask);
const int K, const int P, const int PointDim, const float * inPolyline, const int B,
const int AgentDim, const float * targetState, float * outPolyline, bool * outPolylineMask);

/**
* @brief Set the previous xy position at the end of element.
* @brief Set the Previous Position Kernel object
*
* @param B The number of target agents.
* @param K The number of polylines.
* @param P The number of points contained in each polyline.
* @param D The number of point dimensions.
* @param mask The polyline mask, in shape [B*K*P].
* @param polyline The container of polylines, in shape [B*K*P*D]
* @param D The number of point dimensions, expecting (x, y, ..., preX, preY).
* @param polyline Source polyline, in shape [B*K*P*D].
*/
__global__ void setPreviousPositionKernel(
const int B, const int K, const int P, const int D, const bool * mask, float * polyline);
const int B, const int K, const int P, const int D, float * polyline);

/**
* @brief Calculate distances from targets to polylines.
* @brief Calculate center distance from target agent to each polyline.
*
* @note `inPolyline` must have been transformed to target coordinates system.
* @note Polyline must have been transformed to target agent coordinates system.
*
* @param B The number of target agents.
* @param L The number of polylines.
* @param K The number of polylines.
* @param P The number of points contained in each polyline.
* @param AgentDim The number agent state dimensions.
* @param targetState Source target agent state, in shape [B*AgentDim].
* @param PointDim The number of point dimensions.
* @param inPolyline Source polyline, in shape [B*L*P*PointDim].
* @param inPolylineMask Source polyline mask, in shape [B*K*P].
* @param outDistance Output distance, in shape [B*L].
* @param D The number of point dimensions, expecting [x, y, ...].
* @param polyline Source polyline, in shape [B*K*P*D].
* @param polylineMask Source polyline mask, in shape [B*K*P].
* @param distance Output calculated distances, in shape [B*K].
*/
__global__ void calculateCenterDistanceKernel(
const int B, const int L, const int P, const int AgentDim, const float * targetState,
const int PointDim, const float * inPolyline, const bool * inPolylineMask, float * outDistance);
const int B, const int K, const int P, const int D, const float * polyline,
const bool * polylineMask, float * distance);

/**
* @brief Extract K polylines with the smallest distances.
*
* @note This kernel is implemented using shared memory.
* @note Because this kernel supposes to allocate shared memory dynamically it is necessary to
* specify `sizeof(float) * L` in the kernel execution configuration.
*
* @param K The number elements to be extracted (K <= L).
* @param K The number of polylines to be extracted.
* @param B The number of target agents.
* @param L The number of polylines.
* @param L The number of original polylines.
* @param P The number of points contained in each polyline.
* @param D The number of point dimensions.
* @param inPolyline Source polyline, in shape [B*L*P*D].
* @param inPolylineMask Source polyline mask, in shape [B*L*P].
* @param inDistance Source polyline distances, in shape [B*L].
* @param inDistance Source distances from target agents to the centers of each polyline, in shape
* [B*L].
* @param outPolyline Output polyline, in shape [B*K*P*D].
* @param outPolylineMask Output polyline mask, in shape [B*K*P].
*/
Expand All @@ -86,64 +90,67 @@ __global__ void extractTopKPolylineKernel(
bool * outPolylineMask);

/**
* @brief Calculate the magnitudes of polylines.
* @brief Calculate center positions of each polyline with respect to target agent coordinates
* system.
*
* @note Polyline must have been transformed to target agent coordinates system.
*
* @param B The number of target agents.
* @param K The number of polylines.
* @param P The number of points contained in each polyline.
* @param D The number of point dimensions.
* @param polyline Source polylines, in shape [B*K*P*(D + 2)].
* @param mask Source polyline masks, in shape [B*K*P].
* @param center Output magnitudes of polylines, in shape [B*K*3].
* @param D The number of point dimensions, expecting (x, y, z, ...).
* @param polyline Source polyline, in shape [B*K*P*D].
* @param polylineMask Source polyline mask, in shape [B*K*P].
* @param center Output centers, in shape [B*K*3].
*/
__global__ void calculatePolylineCenterKernel(
const int B, const int K, const int P, const int PointDim, const float * polyline,
const bool * mask, float * center);
const int B, const int K, const int P, const int D, const float * polyline,
const bool * polylineMask, float * center);

/**
* @brief In cases of the number of batch polylines (L) is greater than K,
* extracts the topK elements.
*
* @param L The number of source polylines.
* @param K The number of polylines expected as the model input.
* @param K The number of polylines to be extracted.
* @param L The number of original polylines.
* @param P The number of points contained in each polyline.
* @param PointDim The number of point state dimensions.
* @param AgentDim The number of agent state dimensions.
* @param in_polyline Source polylines, in shape [L*P*PointDim].
* @param inPolyline Source polylines, in shape [L*P*PointDim].
* @param B The number of target agents.
* @param target_state Target agent state at the latest timestamp, in shape [B, AgentDim].
* @param out_polyline Output polylines, in shape [B*K*P*(PointDim+2)].
* @param out_polyline_mask Output polyline masks, in shape [B*K*P].
* @param out_polyline_center Output magnitudes of each polyline with respect to target coords,
* @param AgentDim The number of agent state dimensions.
* @param targetState Target agent state at the latest timestamp, in shape [B*AgentDim].
* @param outPolyline Output polylines, in shape [B*K*P*(PointDim+2)].
* @param outPolylineMask Output polyline masks, in shape [B*K*P].
* @param outPolylineCenter Output magnitudes of each polyline with respect to target coords,
* in shape [B*K*3].
* @param stream CUDA stream.
* @return cudaError_t
*/
cudaError_t polylinePreprocessWithTopkLauncher(
const int L, const int K, const int P, const int PointDim, const float * in_polyline, const int B,
const int AgentDim, const float * target_state, float * out_polyline, bool * out_polyline_mask,
float * out_polyline_center, cudaStream_t stream);
const int K, const int L, const int P, const int PointDim, const float * inPolyline, const int B,
const int AgentDim, const float * targetState, float * outPolyline, bool * outPolylineMask,
float * outPolylineCenter, cudaStream_t stream);

/**
* @brief Do preprocess for polyline if the number of batched polylines is K.
*
* @param K The number of polylines.
* @param P The number of points contained in each polyline.
* @param PointDim The number of point state dimensions.
* @param in_polyline Source polylines, in shape [K*P*PointDim].
* @param inPolyline Source polylines, in shape [K*P*PointDim].
* @param B The number of target agents.
* @param AgentDim The number of agent state dimensions.
* @param target_state Target agent state at the latest timestamp, in shape [B, AgentDim].
* @param out_polyline Output polylines, in shape [B*K*P*(PointDim + 2)].
* @param out_polyline_mask Output polyline masks, in shape [B*K*P].
* @param out_polyline_center Output magnitudes of each polyline with respect to target coords,
* @param targetState Target agent state at the latest timestamp, in shape [B*AgentDim].
* @param outPolyline Output polylines, in shape [B*K*P*(PointDim+2)].
* @param outPolylineMask Output polyline masks, in shape [B*K*P].
* @param outPolylineCenter Output magnitudes of each polyline with respect to target coords,
* in shape [B*K*3].
* @param stream CUDA stream.
* @return cudaError_t
*/
cudaError_t polylinePreprocessLauncher(
const int K, const int P, const int PointDim, const float * in_polyline, const int B,
const int AgentDim, const float * target_state, float * out_polyline, bool * out_polyline_mask,
float * out_polyline_center, cudaStream_t stream);
const int K, const int P, const int PointDim, const float * inPolyline, const int B,
const int AgentDim, const float * targetState, float * outPolyline, bool * outPolylineMask,
float * outPolylineCenter, cudaStream_t stream);

#endif // PREPROCESS__POLYLINE_PREPROCESS_KERNEL_CUH_
Loading
Loading