@@ -188,21 +188,26 @@ rocblas_status rocblas_trtri_trsm_template(rocblas_handle handle,
188188 static constexpr size_t sub_blockSize = 128 ;
189189 size_t tri_elements_to_zero = rocblas_num_non_tri_elements (NB ) * sub_blocks;
190190 size_t num_sub_blocks = (tri_elements_to_zero + sub_blockSize - 1 ) / sub_blockSize;
191- ROCBLAS_LAUNCH_KERNEL ((rocblas_trtri_fill<sub_blockSize, T>),
192- dim3 (num_sub_blocks, batch_count),
193- dim3 (sub_blockSize),
194- 0 ,
195- handle->get_stream (),
196- handle,
197- uplo == rocblas_fill_lower ? rocblas_fill_upper : rocblas_fill_lower,
198- NB ,
199- rocblas_num_non_tri_elements (NB ),
200- NB ,
201- NB * NB ,
202- invA,
203- offset_invAin,
204- stride_invA,
205- sub_blocks);
191+
192+ dim3 grid_fill (num_sub_blocks, batch_count);
193+ dim3 threads_fill (sub_blockSize);
194+ ROCBLAS_LAUNCH_KERNEL_GRID (grid_fill,
195+ (rocblas_trtri_fill<sub_blockSize, T>),
196+ grid_fill,
197+ threads_fill,
198+ 0 ,
199+ handle->get_stream (),
200+ handle,
201+ uplo == rocblas_fill_lower ? rocblas_fill_upper
202+ : rocblas_fill_lower,
203+ NB ,
204+ rocblas_num_non_tri_elements (NB ),
205+ NB ,
206+ NB * NB ,
207+ invA,
208+ offset_invAin,
209+ stride_invA,
210+ sub_blocks);
206211
207212 constexpr rocblas_int JB = IB * 4 ;
208213 rocblas_stride sub_stride_A = NB * size_t (lda) + NB ;
@@ -309,21 +314,25 @@ rocblas_status rocblas_trtri_trsm_template(rocblas_handle handle,
309314 size_t tri_elements_to_zero = rocblas_num_non_tri_elements (rem);
310315 size_t num_sub_blocks = (tri_elements_to_zero + sub_blockSize - 1 ) / sub_blockSize;
311316
312- ROCBLAS_LAUNCH_KERNEL ((rocblas_trtri_fill<sub_blockSize, T>),
313- dim3 (num_sub_blocks, batch_count),
314- dim3 (sub_blockSize),
315- 0 ,
316- handle->get_stream (),
317- handle,
318- uplo == rocblas_fill_lower ? rocblas_fill_upper : rocblas_fill_lower,
319- rem,
320- rocblas_num_non_tri_elements (rem),
321- NB ,
322- 0 ,
323- invA,
324- sub_blocks * NB * size_t (NB ) + offset_invAin,
325- stride_invA,
326- 1 );
317+ dim3 grid_fill (num_sub_blocks, batch_count);
318+ dim3 threads_fill (sub_blockSize);
319+ ROCBLAS_LAUNCH_KERNEL_GRID (grid_fill,
320+ (rocblas_trtri_fill<sub_blockSize, T>),
321+ grid_fill,
322+ threads_fill,
323+ 0 ,
324+ handle->get_stream (),
325+ handle,
326+ uplo == rocblas_fill_lower ? rocblas_fill_upper
327+ : rocblas_fill_lower,
328+ rem,
329+ rocblas_num_non_tri_elements (rem),
330+ NB ,
331+ 0 ,
332+ invA,
333+ sub_blocks * NB * size_t (NB ) + offset_invAin,
334+ stride_invA,
335+ 1 );
327336
328337 if constexpr (BATCHED )
329338 status = rocblas_internal_trtri_batched_template (
0 commit comments