Skip to content

Commit 4694a25

Browse files
committed
Indexer output type changes in boolean reducton kernels
1 parent f8ef141 commit 4694a25

File tree

1 file changed

+10
-9
lines changed

1 file changed

+10
-9
lines changed

dpctl/tensor/libtensor/include/kernels/boolean_reductions.hpp

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -187,15 +187,16 @@ struct SequentialBooleanReduction
187187
{
188188

189189
auto inp_out_iter_offsets_ = inp_out_iter_indexer_(id[0]);
190-
const size_t &inp_iter_offset =
190+
const py::ssize_t &inp_iter_offset =
191191
inp_out_iter_offsets_.get_first_offset();
192-
const size_t &out_iter_offset =
192+
const py::ssize_t &out_iter_offset =
193193
inp_out_iter_offsets_.get_second_offset();
194194

195195
outT red_val(identity_);
196196
for (size_t m = 0; m < reduction_max_gid_; ++m) {
197-
auto inp_reduction_offset = inp_reduced_dims_indexer_(m);
198-
auto inp_offset = inp_iter_offset + inp_reduction_offset;
197+
py::ssize_t inp_reduction_offset =
198+
static_cast<py::ssize_t>(inp_reduced_dims_indexer_(m));
199+
py::ssize_t inp_offset = inp_iter_offset + inp_reduction_offset;
199200

200201
// must convert to boolean first to handle nans
201202
using dpctl::tensor::type_utils::convert_impl;
@@ -435,9 +436,9 @@ struct StridedBooleanReduction
435436
size_t wg_size = it.get_local_range(1);
436437

437438
auto inp_out_iter_offsets_ = inp_out_iter_indexer_(reduction_id);
438-
const size_t &inp_iter_offset =
439+
const py::ssize_t &inp_iter_offset =
439440
inp_out_iter_offsets_.get_first_offset();
440-
const size_t &out_iter_offset =
441+
const py::ssize_t &out_iter_offset =
441442
inp_out_iter_offsets_.get_second_offset();
442443

443444
outT local_red_val(identity_);
@@ -447,9 +448,9 @@ struct StridedBooleanReduction
447448
size_t arg_reduce_gid = arg_reduce_gid0 + m * wg_size;
448449

449450
if (arg_reduce_gid < reduction_max_gid_) {
450-
auto inp_reduction_offset =
451-
inp_reduced_dims_indexer_(arg_reduce_gid);
452-
auto inp_offset = inp_iter_offset + inp_reduction_offset;
451+
py::ssize_t inp_reduction_offset = static_cast<py::ssize_t>(
452+
inp_reduced_dims_indexer_(arg_reduce_gid));
453+
py::ssize_t inp_offset = inp_iter_offset + inp_reduction_offset;
453454

454455
// must convert to boolean first to handle nans
455456
using dpctl::tensor::type_utils::convert_impl;

0 commit comments

Comments
 (0)