warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl

warp/native/tile_radix_sort.h

@@ -921,7 +921,7 @@ template <typename K, typename V, typename KeyToUint>
 void radix_sort_pairs_cpu_core(K* keys, K* aux_keys, V* values, V* aux_values, int n)
 {
     KeyToUint converter;
-	static unsigned int tables[2][1 << 16];
+	unsigned int tables[2][1 << 16];
 	memset(tables, 0, sizeof(tables));
 	
 	// build histograms

warp/native/tile_reduce.h

@@ -19,6 +19,12 @@
 
 #include "tile.h"
 
+#ifdef __clang__
+// disable warnings related to C++17 extensions on CPU JIT builds
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wc++17-extensions"
+#endif // __clang__
+
 #define WP_TILE_WARP_SIZE 32
 
 namespace wp
@@ -76,7 +82,7 @@ inline CUDA_CALLABLE T warp_shuffle_down(T val, int offset, int mask)
   return output;
 }
 
-// Vector overload
+// vector overload
 template <unsigned Length, typename T>
 inline CUDA_CALLABLE wp::vec_t<Length, T> warp_shuffle_down(wp::vec_t<Length, T> val, int offset, int mask)
 {
@@ -88,7 +94,7 @@ inline CUDA_CALLABLE wp::vec_t<Length, T> warp_shuffle_down(wp::vec_t<Length, T>
     return result;
 }
 
-// Matrix overload
+// matrix overload
 template <unsigned Rows, unsigned Cols, typename T>
 inline CUDA_CALLABLE wp::mat_t<Rows, Cols, T> warp_shuffle_down(wp::mat_t<Rows, Cols, T> val, int offset, int mask)
 {
@@ -117,7 +123,7 @@ inline CUDA_CALLABLE T warp_reduce(T val, Op f, unsigned int mask)
     }
     else
     {
-        // handle partial warp case
+        // handle partial warp case - works for contiguous masks
         for (int offset=WP_TILE_WARP_SIZE/2; offset > 0; offset /= 2)
         {            
             T shfl_val = warp_shuffle_down(sum, offset, mask);
@@ -175,6 +181,51 @@ inline CUDA_CALLABLE ValueAndIndex<T> warp_reduce_tracked(T val, int idx, Op f,
     return result;
 }
 
+// combines per-thread reduction results across warps and the entire block
+// assumes each thread has already reduced its local data to thread_sum
+// returns the block-wide reduced value (only valid in thread 0)
+template <typename T, typename Op>
+inline CUDA_CALLABLE T block_combine_thread_results(T thread_sum, bool thread_has_data, Op f, 
+                                                     T* partials, int& active_warps)
+{
+    constexpr int warp_count = (WP_TILE_BLOCK_DIM + WP_TILE_WARP_SIZE - 1) / WP_TILE_WARP_SIZE;
+    const int warp_index = threadIdx.x / WP_TILE_WARP_SIZE;
+    const int lane_index = threadIdx.x % WP_TILE_WARP_SIZE;
+    
+    // determine which threads have data
+    unsigned int mask = __ballot_sync(0xFFFFFFFF, thread_has_data);
+    bool warp_is_active = mask != 0;
+    
+    // warp reduction
+    T warp_sum;
+    if (thread_has_data)
+        warp_sum = warp_reduce(thread_sum, f, mask);
+    
+    // lane 0 of each active warp writes to shared memory and increments counter
+    if (lane_index == 0 && warp_is_active)
+    {
+        partials[warp_index] = warp_sum;
+        atomicAdd(&active_warps, 1);
+    }
+    
+    // sync to ensure all warps have written their partials
+    WP_TILE_SYNC();
+    
+    // thread 0 performs final reduction across active warps
+    T block_sum;
+    if (threadIdx.x == 0)
+    {
+        block_sum = partials[0];
+        
+        for (int w = 1; w < active_warps; ++w)
+        {
+            block_sum = f(block_sum, partials[w]);
+        }
+    }
+    
+    return block_sum;
+}
+
 // non-axis version which computes sum 
 // across the entire tile using the whole block
 template <typename Tile, typename Op>
@@ -185,15 +236,14 @@ auto tile_reduce_impl(Op f, Tile& t)
     auto input = t.copy_to_register();
     auto output = tile_register_t<T, tile_layout_register_t<tile_shape_t<1>>>();
 
-    const int warp_count = (WP_TILE_BLOCK_DIM + WP_TILE_WARP_SIZE - 1)/WP_TILE_WARP_SIZE;
-    const int warp_index = threadIdx.x/WP_TILE_WARP_SIZE;
-    const int lane_index = threadIdx.x%WP_TILE_WARP_SIZE;
-
-    T thread_sum = input.data[0];
+    constexpr int warp_count = (WP_TILE_BLOCK_DIM + WP_TILE_WARP_SIZE - 1) / WP_TILE_WARP_SIZE;
 
     using Layout = typename decltype(input)::Layout;
 
-    // thread reduction
+    // step 1: each thread reduces its own registers locally
+    T thread_sum = input.data[0];
+    bool thread_has_data = Layout::valid(Layout::linear_from_register(0));
+
     WP_PRAGMA_UNROLL
     for (int i=1; i < Layout::NumRegs; ++i)
     {
@@ -204,48 +254,190 @@ auto tile_reduce_impl(Op f, Tile& t)
         thread_sum = f(thread_sum, input.data[i]);
     }
 
-    // ensure that only threads with at least one valid item participate in the reduction
-    unsigned int mask = __ballot_sync(__activemask(), Layout::valid(Layout::linear_from_register(0)));
-    bool warp_is_active = mask != 0;
-
-    // warp reduction
-    T warp_sum = warp_reduce(thread_sum, f, mask);
-
-    // fixed size scratch pad for partial results in shared memory
-    WP_TILE_SHARED T partials[warp_count];
-
-    // count of active warps
-    WP_TILE_SHARED int active_warps;
+    // shared memory for cross-warp reduction
+    __shared__ T partials[warp_count];
+    __shared__ int active_warps;
+    
     if (threadIdx.x == 0)
         active_warps = 0;
     
-    // ensure active_warps is initialized
     WP_TILE_SYNC();
 
-    if (lane_index == 0 && warp_is_active)
-    {
-        partials[warp_index] = warp_sum;
-        atomicAdd(&active_warps, 1);
-    }
+    // step 2-3: combine thread results across warps and block
+    T block_sum = block_combine_thread_results(thread_sum, thread_has_data, f, partials, active_warps);
 
-    // ensure partials are ready
-    WP_TILE_SYNC();
-
-    // reduce across block, todo: use warp_reduce() here
     if (threadIdx.x == 0)
-    {
-        T block_sum = partials[0];
-        
-        WP_PRAGMA_UNROLL
-        for (int i=1; i < active_warps; ++i)
-            block_sum = f(block_sum, partials[i]);
-
         output.data[0] = block_sum;
-    }
 
     return output;
 }
 
+template <int Axis, typename Op, typename Tile>
+auto tile_reduce_axis_impl(Op f, Tile& t)
+{
+    using T = typename Tile::Type;
+    using InputShape = typename Tile::Layout::Shape;
+    using OutputShape = typename tile_shape_remove_dim<Axis, InputShape>::type;
+    
+    constexpr int reduce_dim_size = InputShape::dim(Axis);
+    constexpr int output_size = OutputShape::size();
+    
+    // special case: 1D input delegates to block-wide tile_reduce_impl for optimal performance
+    if constexpr (InputShape::N == 1)
+    {
+        return tile_reduce_impl(f, t);
+    }
+    
+    // shared memory buffer for the output (used by all tiers)
+    __shared__ T output_buffer[output_size];
+    
+    // create output layout for coordinate conversion (used by all tiers)
+    using OutputLayout = tile_layout_strided_t<OutputShape>;
+    
+    if constexpr (reduce_dim_size <= 32)
+    {
+        // Tier 1: Single thread per output element (optimal for small reductions)
+        
+        // each thread processes output elements, performing reduction along the axis
+        for (int out_idx = WP_TILE_THREAD_IDX; out_idx < output_size; out_idx += WP_TILE_BLOCK_DIM)
+        {
+            // convert output linear index to output coordinates
+            auto out_coord = OutputLayout::coord_from_linear(out_idx);
+            
+            // initialize accumulator with first element along the reduction axis
+            T accumulator = t.data(tile_coord_insert_axis<Axis>(out_coord, 0));
+            
+            // reduce across the axis
+            for (int i = 1; i < reduce_dim_size; ++i)
+            {
+                accumulator = f(accumulator, t.data(tile_coord_insert_axis<Axis>(out_coord, i)));
+            }
+            
+            // store to output buffer
+            output_buffer[out_idx] = accumulator;
+        }
+        
+        // sync before reading output
+        WP_TILE_SYNC();
+    }
+    else if constexpr (reduce_dim_size <= 256)
+    {
+        // Tier 2: Warp-based reduction (one warp per output element)
+        constexpr int warp_count = (WP_TILE_BLOCK_DIM + WP_TILE_WARP_SIZE - 1) / WP_TILE_WARP_SIZE;
+        const int warp_index = threadIdx.x / WP_TILE_WARP_SIZE;
+        const int lane_index = threadIdx.x % WP_TILE_WARP_SIZE;
+        
+        constexpr int chunks_per_slice = (reduce_dim_size + WP_TILE_WARP_SIZE - 1) / WP_TILE_WARP_SIZE;
+        
+        // shared memory: one accumulator per warp
+        __shared__ T warp_partials[warp_count];
+        
+        // each warp processes output slices
+        for (int out_idx = warp_index; out_idx < output_size; out_idx += warp_count)
+        {
+            auto out_coord = OutputLayout::coord_from_linear(out_idx);
+            
+            // process the reduction axis in chunks of 32
+            for (int chunk = 0; chunk < chunks_per_slice; ++chunk)
+            {
+                int axis_idx = chunk * WP_TILE_WARP_SIZE + lane_index;
+                bool valid = axis_idx < reduce_dim_size;
+                
+                T val;
+                if (valid)
+                {
+                    auto in_coord = tile_coord_insert_axis<Axis>(out_coord, axis_idx);
+                    val = t.data(in_coord);
+                }
+                
+                // warp reduce this chunk (only valid lanes participate)
+                unsigned int mask = __ballot_sync(0xFFFFFFFF, valid);
+                T chunk_result = warp_reduce(val, f, mask);
+                
+                // lane 0 accumulates the chunk result
+                if (lane_index == 0)
+                {
+                    if (chunk == 0)
+                        warp_partials[warp_index] = chunk_result;
+                    else
+                        warp_partials[warp_index] = f(warp_partials[warp_index], chunk_result);
+                }
+            }
+            
+            // lane 0 writes final result for this output element
+            if (lane_index == 0)
+                output_buffer[out_idx] = warp_partials[warp_index];
+        }
+        
+        // sync before reading output
+        WP_TILE_SYNC();
+    }
+    else
+    {
+        // Tier 3: Block-level reduction (entire block collaborates on each output element)
+        constexpr int warp_count = (WP_TILE_BLOCK_DIM + WP_TILE_WARP_SIZE - 1) / WP_TILE_WARP_SIZE;
+        
+        // shared memory for cross-warp reduction
+        __shared__ T partials[warp_count];
+        __shared__ int active_warps;
+        
+        // process each output element sequentially with full block cooperation
+        for (int out_idx = 0; out_idx < output_size; ++out_idx)
+        {
+            auto out_coord = OutputLayout::coord_from_linear(out_idx);
+            
+            // step 1: each thread reduces its strided subset of the slice locally
+            bool thread_has_data = threadIdx.x < reduce_dim_size;
+            T thread_sum;
+            
+            if (thread_has_data)
+            {
+                // initialize with first element
+                auto in_coord = tile_coord_insert_axis<Axis>(out_coord, threadIdx.x);
+                thread_sum = t.data(in_coord);
+                
+                // reduce remaining elements with stride
+                for (int i = threadIdx.x + WP_TILE_BLOCK_DIM; i < reduce_dim_size; i += WP_TILE_BLOCK_DIM)
+                {
+                    auto in_coord = tile_coord_insert_axis<Axis>(out_coord, i);
+                    T val = t.data(in_coord);
+                    thread_sum = f(thread_sum, val);
+                }
+            }
+            
+            // initialize active warp counter
+            if (threadIdx.x == 0)
+                active_warps = 0;
+            
+            WP_TILE_SYNC();
+            
+            // step 2-3: combine thread results across warps and block
+            T block_sum = block_combine_thread_results(thread_sum, thread_has_data, f, partials, active_warps);
+            
+            if (threadIdx.x == 0)
+                output_buffer[out_idx] = block_sum;
+            
+            // sync before next output element
+            WP_TILE_SYNC();
+        }
+    }
+    
+    // copy from shared memory buffer to register tile (common to all tiers)
+    auto output = tile_register_t<T, tile_layout_register_t<OutputShape>>();
+    using OutputRegLayout = typename decltype(output)::Layout;
+    
+    WP_PRAGMA_UNROLL
+    for (int i = 0; i < OutputRegLayout::NumRegs; ++i)
+    {
+        int linear = OutputRegLayout::linear_from_register(i);
+        if (!OutputRegLayout::valid(linear))
+            break;
+        
+        output.data[i] = output_buffer[linear];
+    }
+    
+    return output;
+}
 
 // non-axis version which computes sum 
 // across the entire tile using the whole block
@@ -286,11 +478,11 @@ auto tile_arg_reduce_impl(Op f, OpTrack track, Tile& t)
     ValueAndIndex<T> warp_sum = warp_reduce_tracked(thread_sum, champion_index, f, track, mask);
 
     // fixed size scratch pad for partial results in shared memory
-    WP_TILE_SHARED T partials[warp_count];
-    WP_TILE_SHARED int partials_idx[warp_count];
+    __shared__ T partials[warp_count];
+    __shared__ int partials_idx[warp_count];
 
     // count of active warps
-    WP_TILE_SHARED int active_warps;
+    __shared__ int active_warps;
     if (threadIdx.x == 0)
         active_warps = 0;
     
@@ -356,6 +548,65 @@ auto tile_reduce_impl(Op f, Tile& t)
     return output;
 }
 
+template <int Axis, typename Op, typename Tile>
+auto tile_reduce_axis_impl(Op f, Tile& t)
+{
+    using T = typename Tile::Type;
+    using InputShape = typename Tile::Layout::Shape;
+    using OutputShape = typename tile_shape_remove_dim<Axis, InputShape>::type;
+    
+    constexpr int reduce_dim_size = InputShape::dim(Axis);
+    
+    // CPU version - work directly with register tiles, no thread coordination needed
+    auto input = t.copy_to_register();
+    auto output = tile_register_t<T, tile_layout_register_t<OutputShape>>();
+    using OutputLayout = typename decltype(output)::Layout;
+    
+    // iterate through each output element and reduce along the axis
+    constexpr int output_size = OutputShape::size();
+    for (int out_idx = 0; out_idx < output_size; ++out_idx)
+    {
+        T accumulator;
+        
+        // special case for 1D input (reduces to single value)
+        if constexpr (InputShape::N == 1)
+        {
+            accumulator = input.data[0];
+            for (int i = 1; i < reduce_dim_size; ++i)
+            {
+                // input is in registers, linear access
+                accumulator = f(accumulator, input.data[i]);
+            }
+        }
+        else
+        {
+            // multi-dimensional case
+            auto out_coord = OutputLayout::coord_from_linear(out_idx);
+            
+            // get input coordinates by inserting axis values
+            auto coord_0 = tile_coord_insert_axis<Axis>(out_coord, 0);
+            int input_linear_0 = tile_layout_register_t<InputShape>::linear_from_coord(coord_0);
+            int input_reg_0 = tile_layout_register_t<InputShape>::register_from_linear(input_linear_0);
+            accumulator = input.data[input_reg_0];
+            
+            // reduce across the axis
+            for (int i = 1; i < reduce_dim_size; ++i)
+            {
+                auto coord_i = tile_coord_insert_axis<Axis>(out_coord, i);
+                int input_linear_i = tile_layout_register_t<InputShape>::linear_from_coord(coord_i);
+                int input_reg_i = tile_layout_register_t<InputShape>::register_from_linear(input_linear_i);
+                accumulator = f(accumulator, input.data[input_reg_i]);
+            }
+        }
+        
+        // store to output register
+        int output_reg = OutputLayout::register_from_linear(out_idx);
+        output.data[output_reg] = accumulator;
+    }
+    
+    return output;
+}
+
 template <typename Tile, typename Op, typename OpTrack>
 auto tile_arg_reduce_impl(Op f, OpTrack track, Tile& t)
 {
@@ -391,15 +642,25 @@ inline void adj_tile_reduce_impl()
     // todo: general purpose reduction gradients not implemented 
 }
 
+inline void adj_tile_reduce_axis_impl() 
+{
+    // todo: axis-specific reduction gradients not implemented 
+}
+
 // entry point for Python code-gen, wraps op in a lambda to perform overload resolution
 #define tile_reduce(op, t) tile_reduce_impl([](auto x, auto y) { return op(x, y);}, t)
-#define adj_tile_reduce(op, a, adj_op, adj_a, adj_ret) adj_tile_reduce_impl()
+#define adj_tile_reduce(op, t, adj_op, adj_t, adj_ret) adj_tile_reduce_impl()
 
 #define tile_arg_reduce(op, opTrack, t) tile_arg_reduce_impl([](auto x, auto y) { return op(x, y);}, [](auto a, auto b, auto c, auto d) { return opTrack(a, b, c, d); }, t)
-#define adj_tile_arg_reduce(op, a, adj_op, adj_a, adj_ret) adj_tile_arg_reduce_impl()
+#define adj_tile_arg_reduce(op, t, adj_op, adj_t, adj_ret) adj_tile_arg_reduce_impl()
+
+// axis-specific reduction entry points
+#define tile_reduce_axis(op, t, axis) tile_reduce_axis_impl<axis>([](auto x, auto y) { return op(x, y);}, t)
+#define adj_tile_reduce_axis(op, t, axis, adj_op, adj_t, adj_axis, adj_ret) adj_tile_reduce_axis_impl()
 
 // convenience methods for specific reductions
 
+// whole-tile sum
 template <typename Tile>
 auto tile_sum(Tile& t)
 {
@@ -418,7 +679,7 @@ void adj_tile_sum(Tile& t, Tile& adj_t, AdjTile& adj_ret)
     T scratch = adj_reg.data[0];
 #else
     // broadcast incoming adjoint to block
-    WP_TILE_SHARED T scratch;
+    __shared__ T scratch;
     if (WP_TILE_THREAD_IDX == 0)
         scratch = adj_reg.data[0];
 
@@ -434,6 +695,90 @@ void adj_tile_sum(Tile& t, Tile& adj_t, AdjTile& adj_ret)
     adj_t.grad_add(adj_ret_reg);
 }
 
+// axis-specific sum
+template <int Axis, typename Tile>
+auto tile_sum(Tile& t)
+{
+    return tile_reduce_axis_impl<Axis>([](auto x, auto y) { return add(x, y); }, t);
+}
+
+// special case adjoint for axis-specific summation
+template<int Axis, typename Tile, typename AdjTile>
+void adj_tile_sum(Tile& t, Tile& adj_t, AdjTile& adj_ret)
+{
+    using InputShape = typename Tile::Layout::Shape;
+    
+    if constexpr (InputShape::N == 1)
+    {
+        // 1D -> scalar case: broadcast scalar to 1D
+        auto broadcasted = tile_broadcast<InputShape::dim(0), 0>(adj_ret);
+        tile_add_inplace(adj_t, broadcasted);
+    }
+    else if constexpr (InputShape::N == 2)
+    {
+        if constexpr (Axis == 0)
+        {
+            // broadcast from (D1,) to (D0, D1) with strides (0, 1)
+            auto broadcasted = tile_broadcast<InputShape::dim(0), InputShape::dim(1), 0, 1>(adj_ret);
+            tile_add_inplace(adj_t, broadcasted);
+        }
+        else // Axis == 1
+        {
+            // broadcast from (D0,) to (D0, D1) with strides (1, 0)  
+            auto broadcasted = tile_broadcast<InputShape::dim(0), InputShape::dim(1), 1, 0>(adj_ret);
+            tile_add_inplace(adj_t, broadcasted);
+        }
+    }
+    else if constexpr (InputShape::N == 3)
+    {
+        if constexpr (Axis == 0)
+        {
+            // broadcast from (D1, D2) to (D0, D1, D2) with strides (0, D2, 1)
+            auto broadcasted = tile_broadcast<InputShape::dim(0), InputShape::dim(1), InputShape::dim(2), 0, InputShape::dim(2), 1>(adj_ret);
+            tile_add_inplace(adj_t, broadcasted);
+        }
+        else if constexpr (Axis == 1)
+        {
+            // broadcast from (D0, D2) to (D0, D1, D2) with strides (D2, 0, 1)
+            auto broadcasted = tile_broadcast<InputShape::dim(0), InputShape::dim(1), InputShape::dim(2), InputShape::dim(2), 0, 1>(adj_ret);
+            tile_add_inplace(adj_t, broadcasted);
+        }
+        else // Axis == 2
+        {
+            // broadcast from (D0, D1) to (D0, D1, D2) with strides (D1, 1, 0)
+            auto broadcasted = tile_broadcast<InputShape::dim(0), InputShape::dim(1), InputShape::dim(2), InputShape::dim(1), 1, 0>(adj_ret);
+            tile_add_inplace(adj_t, broadcasted);
+        }
+    }
+    else if constexpr (InputShape::N == 4)
+    {
+        if constexpr (Axis == 0)
+        {
+            // broadcast from (D1, D2, D3) to (D0, D1, D2, D3) with strides (0, D2*D3, D3, 1)
+            auto broadcasted = tile_broadcast<InputShape::dim(0), InputShape::dim(1), InputShape::dim(2), InputShape::dim(3), 0, InputShape::dim(2)*InputShape::dim(3), InputShape::dim(3), 1>(adj_ret);
+            tile_add_inplace(adj_t, broadcasted);
+        }
+        else if constexpr (Axis == 1)
+        {
+            // broadcast from (D0, D2, D3) to (D0, D1, D2, D3) with strides (D2*D3, 0, D3, 1)
+            auto broadcasted = tile_broadcast<InputShape::dim(0), InputShape::dim(1), InputShape::dim(2), InputShape::dim(3), InputShape::dim(2)*InputShape::dim(3), 0, InputShape::dim(3), 1>(adj_ret);
+            tile_add_inplace(adj_t, broadcasted);
+        }
+        else if constexpr (Axis == 2)
+        {
+            // broadcast from (D0, D1, D3) to (D0, D1, D2, D3) with strides (D1*D3, D3, 0, 1)
+            auto broadcasted = tile_broadcast<InputShape::dim(0), InputShape::dim(1), InputShape::dim(2), InputShape::dim(3), InputShape::dim(1)*InputShape::dim(3), InputShape::dim(3), 0, 1>(adj_ret);
+            tile_add_inplace(adj_t, broadcasted);
+        }
+        else // Axis == 3
+        {
+            // broadcast from (D0, D1, D2) to (D0, D1, D2, D3) with strides (D1*D2, D2, 1, 0)
+            auto broadcasted = tile_broadcast<InputShape::dim(0), InputShape::dim(1), InputShape::dim(2), InputShape::dim(3), InputShape::dim(1)*InputShape::dim(2), InputShape::dim(2), 1, 0>(adj_ret);
+            tile_add_inplace(adj_t, broadcasted);
+        }
+    }
+}
+
 template <typename Tile>
 auto tile_max(Tile& t)
 {
@@ -485,6 +830,9 @@ void adj_tile_argmin(Tile& t, Tile& adj_t, AdjTile& adj_ret)
 }
 
 
+} // namespace wp
 
 
-} // namespace wp
+#ifdef __clang__
+#pragma clang diagnostic pop
+#endif

warp/native/tile_scan.h

@@ -50,7 +50,7 @@ inline CUDA_CALLABLE T scan_warp_inclusive(int lane, T value)
 template<typename T>
 inline CUDA_CALLABLE T thread_block_scan_inclusive(int lane, int warp_index, int num_warps, T value)
 {
-    WP_TILE_SHARED T sums[1024 / WP_TILE_WARP_SIZE]; // 1024 is the maximum number of threads per block
+    __shared__ T sums[1024 / WP_TILE_WARP_SIZE]; // 1024 is the maximum number of threads per block
 
     value = scan_warp_inclusive(lane, value);
 
@@ -85,7 +85,7 @@ inline CUDA_CALLABLE void thread_block_scan(T* values, int num_elements)
     const int num_threads_in_block = blockDim.x;
     const int num_iterations = (num_elements + num_threads_in_block - 1) / num_threads_in_block;
 
-    WP_TILE_SHARED T offset;
+    __shared__ T offset;
     if (threadIdx.x == 0)
         offset = T(0);
 
@@ -124,7 +124,7 @@ inline CUDA_CALLABLE auto tile_scan_inclusive_impl(Tile& t)
     constexpr int num_elements_to_scan = Tile::Layout::Shape::size(); 
 
     // create a temporary shared tile to hold the input values
-    WP_TILE_SHARED T smem[num_elements_to_scan];
+    __shared__ T smem[num_elements_to_scan];
     tile_shared_t<T, tile_layout_strided_t<typename Tile::Layout::Shape>, false> scratch(smem, nullptr);
 
     // copy input values to scratch space
@@ -147,7 +147,7 @@ inline CUDA_CALLABLE auto tile_scan_exclusive_impl(Tile& t)
     constexpr int num_elements_to_scan = Tile::Layout::Shape::size(); 
 
     // create a temporary shared tile to hold the input values
-    WP_TILE_SHARED T smem[num_elements_to_scan];
+    __shared__ T smem[num_elements_to_scan];
     tile_shared_t<T, tile_layout_strided_t<typename Tile::Layout::Shape>, false> scratch(smem, nullptr);
 
     // copy input values to scratch space