Skip to content

Commit 5e3a671

Browse files
ndgrigorianndgrigorian
committed
Added functions for querying reduction atomic support per type per function
1 parent 89a10cc commit 5e3a671

File tree

6 files changed

+245
-104
lines changed

6 files changed

+245
-104
lines changed

‎dpctl/tensor/libtensor/source/reductions/max.cpp‎

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,11 @@
3030
#include <vector>
3131

3232
#include "kernels/reductions.hpp"
33-
#include "reduction_over_axis.hpp"
3433
#include "utils/type_dispatch.hpp"
3534

35+
#include "reduction_atomic_support.hpp"
36+
#include "reduction_over_axis.hpp"
37+
3638
namespace py = pybind11;
3739

3840
namespace dpctl
@@ -71,8 +73,6 @@ static reduction_contig_impl_fn_ptr
7173

7274
void populate_max_over_axis_dispatch_tables(void)
7375
{
74-
using dpctl::tensor::kernels::reduction_contig_impl_fn_ptr;
75-
using dpctl::tensor::kernels::reduction_strided_impl_fn_ptr;
7676
using td_ns::DispatchTableBuilder;
7777

7878
using dpctl::tensor::kernels::MaxOverAxisAtomicStridedFactory;
@@ -112,6 +112,20 @@ void populate_max_over_axis_dispatch_tables(void)
112112
dtb6.populate_dispatch_table(max_over_axis0_contig_temps_dispatch_table);
113113
}
114114

115+
using atomic_support::atomic_support_fn_ptr_t;
116+
static atomic_support_fn_ptr_t max_atomic_support_vector[td_ns::num_types];
117+
118+
void populate_max_atomic_support_dispatch_vector(void)
119+
{
120+
using td_ns::DispatchVectorBuilder;
121+
122+
using atomic_support::MaxAtomicSupportFactory;
123+
DispatchVectorBuilder<atomic_support_fn_ptr_t, MaxAtomicSupportFactory,
124+
td_ns::num_types>
125+
dvb;
126+
dvb.populate_dispatch_vector(max_atomic_support_vector);
127+
}
128+
115129
} // namespace impl
116130

117131
void init_max(py::module_ m)
@@ -128,11 +142,9 @@ void init_max(py::module_ m)
128142
using impl::max_over_axis_strided_atomic_dispatch_table;
129143
using impl::max_over_axis_strided_temps_dispatch_table;
130144

131-
using dpctl::tensor::py_internal::check_atomic_support;
132-
const auto &check_atomic_support_size4 =
133-
check_atomic_support</*require_atomic64*/ false>;
134-
const auto &check_atomic_support_size8 =
135-
check_atomic_support</*require_atomic64*/ true>;
145+
using impl::populate_max_atomic_support_dispatch_vector;
146+
populate_max_atomic_support_dispatch_vector();
147+
using impl::max_atomic_support_vector;
136148

137149
auto max_pyapi = [&](const arrayT &src, int trailing_dims_to_reduce,
138150
const arrayT &dst, sycl::queue &exec_q,
@@ -146,7 +158,7 @@ void init_max(py::module_ m)
146158
max_over_axis_strided_temps_dispatch_table,
147159
max_over_axis0_contig_temps_dispatch_table,
148160
max_over_axis1_contig_temps_dispatch_table,
149-
check_atomic_support_size4, check_atomic_support_size8);
161+
max_atomic_support_vector);
150162
};
151163
m.def("_max_over_axis", max_pyapi, "", py::arg("src"),
152164
py::arg("trailing_dims_to_reduce"), py::arg("dst"),

‎dpctl/tensor/libtensor/source/reductions/min.cpp‎

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,11 @@
3030
#include <vector>
3131

3232
#include "kernels/reductions.hpp"
33-
#include "reduction_over_axis.hpp"
3433
#include "utils/type_dispatch.hpp"
3534

35+
#include "reduction_atomic_support.hpp"
36+
#include "reduction_over_axis.hpp"
37+
3638
namespace py = pybind11;
3739

3840
namespace dpctl
@@ -112,6 +114,20 @@ void populate_min_over_axis_dispatch_tables(void)
112114
dtb6.populate_dispatch_table(min_over_axis0_contig_temps_dispatch_table);
113115
}
114116

117+
using atomic_support::atomic_support_fn_ptr_t;
118+
static atomic_support_fn_ptr_t min_atomic_support_vector[td_ns::num_types];
119+
120+
void populate_min_atomic_support_dispatch_vector(void)
121+
{
122+
using td_ns::DispatchVectorBuilder;
123+
124+
using atomic_support::MinAtomicSupportFactory;
125+
DispatchVectorBuilder<atomic_support_fn_ptr_t, MinAtomicSupportFactory,
126+
td_ns::num_types>
127+
dvb;
128+
dvb.populate_dispatch_vector(min_atomic_support_vector);
129+
}
130+
115131
} // namespace impl
116132

117133
void init_min(py::module_ m)
@@ -128,11 +144,9 @@ void init_min(py::module_ m)
128144
using impl::min_over_axis_strided_atomic_dispatch_table;
129145
using impl::min_over_axis_strided_temps_dispatch_table;
130146

131-
using dpctl::tensor::py_internal::check_atomic_support;
132-
const auto &check_atomic_support_size4 =
133-
check_atomic_support</*require_atomic64*/ false>;
134-
const auto &check_atomic_support_size8 =
135-
check_atomic_support</*require_atomic64*/ true>;
147+
using impl::populate_min_atomic_support_dispatch_vector;
148+
populate_min_atomic_support_dispatch_vector();
149+
using impl::min_atomic_support_vector;
136150

137151
auto min_pyapi = [&](const arrayT &src, int trailing_dims_to_reduce,
138152
const arrayT &dst, sycl::queue &exec_q,
@@ -146,7 +160,7 @@ void init_min(py::module_ m)
146160
min_over_axis_strided_temps_dispatch_table,
147161
min_over_axis0_contig_temps_dispatch_table,
148162
min_over_axis1_contig_temps_dispatch_table,
149-
check_atomic_support_size4, check_atomic_support_size8);
163+
min_atomic_support_vector);
150164
};
151165
m.def("_min_over_axis", min_pyapi, "", py::arg("src"),
152166
py::arg("trailing_dims_to_reduce"), py::arg("dst"),

‎dpctl/tensor/libtensor/source/reductions/prod.cpp‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,11 @@
3030
#include <vector>
3131

3232
#include "kernels/reductions.hpp"
33-
#include "reduction_over_axis.hpp"
3433
#include "utils/type_dispatch.hpp"
3534

35+
#include "reduction_atomic_support.hpp"
36+
#include "reduction_over_axis.hpp"
37+
3638
namespace py = pybind11;
3739

3840
namespace dpctl
@@ -112,6 +114,20 @@ void populate_prod_over_axis_dispatch_tables(void)
112114
dtb6.populate_dispatch_table(prod_over_axis0_contig_temps_dispatch_table);
113115
}
114116

117+
using atomic_support::atomic_support_fn_ptr_t;
118+
static atomic_support_fn_ptr_t prod_atomic_support_vector[td_ns::num_types];
119+
120+
void populate_prod_atomic_support_dispatch_vector(void)
121+
{
122+
using td_ns::DispatchVectorBuilder;
123+
124+
using atomic_support::ProductAtomicSupportFactory;
125+
DispatchVectorBuilder<atomic_support_fn_ptr_t, ProductAtomicSupportFactory,
126+
td_ns::num_types>
127+
dvb;
128+
dvb.populate_dispatch_vector(prod_atomic_support_vector);
129+
}
130+
115131
} // namespace impl
116132

117133
void init_prod(py::module_ m)
@@ -128,11 +144,9 @@ void init_prod(py::module_ m)
128144
using impl::prod_over_axis_strided_atomic_dispatch_table;
129145
using impl::prod_over_axis_strided_temps_dispatch_table;
130146

131-
using dpctl::tensor::py_internal::check_atomic_support;
132-
const auto &check_atomic_support_size4 =
133-
check_atomic_support</*require_atomic64*/ false>;
134-
const auto &check_atomic_support_size8 =
135-
check_atomic_support</*require_atomic64*/ true>;
147+
using impl::populate_prod_atomic_support_dispatch_vector;
148+
populate_prod_atomic_support_dispatch_vector();
149+
using impl::prod_atomic_support_vector;
136150

137151
auto prod_pyapi = [&](const arrayT &src, int trailing_dims_to_reduce,
138152
const arrayT &dst, sycl::queue &exec_q,
@@ -146,7 +160,7 @@ void init_prod(py::module_ m)
146160
prod_over_axis_strided_temps_dispatch_table,
147161
prod_over_axis0_contig_temps_dispatch_table,
148162
prod_over_axis1_contig_temps_dispatch_table,
149-
check_atomic_support_size4, check_atomic_support_size8);
163+
prod_atomic_support_vector);
150164
};
151165
m.def("_prod_over_axis", prod_pyapi, "", py::arg("src"),
152166
py::arg("trailing_dims_to_reduce"), py::arg("dst"),
@@ -160,7 +174,7 @@ void init_prod(py::module_ m)
160174
input_dtype, output_dtype, dst_usm_type, q,
161175
prod_over_axis_strided_atomic_dispatch_table,
162176
prod_over_axis_strided_temps_dispatch_table,
163-
check_atomic_support_size4, check_atomic_support_size8);
177+
prod_atomic_support_vector);
164178
};
165179
m.def("_prod_over_axis_dtype_supported", prod_dtype_supported, "",
166180
py::arg("arg_dtype"), py::arg("out_dtype"),

‎dpctl/tensor/libtensor/source/reductions/reduction_atomic_support.hpp‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
//===-- ------------ Implementation of _tensor_impl module ----*-C++-*-/===//
2+
//
3+
// Data Parallel Control (dpctl)
4+
//
5+
// Copyright 2020-2023 Intel Corporation
6+
//
7+
// Licensed under the Apache License, Version 2.0 (the "License");
8+
// you may not use this file except in compliance with the License.
9+
// You may obtain a copy of the License at
10+
//
11+
// http://www.apache.org/licenses/LICENSE-2.0
12+
//
13+
// Unless required by applicable law or agreed to in writing, software
14+
// distributed under the License is distributed on an "AS IS" BASIS,
15+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16+
// See the License for the specific language governing permissions and
17+
// limitations under the License.
18+
//
19+
//===--------------------------------------------------------------------===//
20+
///
21+
/// \file
22+
/// This file defines functions of dpctl.tensor._tensor_impl extensions
23+
//===--------------------------------------------------------------------===//
24+
25+
#pragma once
26+
#include <CL/sycl.hpp>
27+
#include <complex>
28+
#include <type_traits>
29+
30+
#include "utils/type_utils.hpp"
31+
32+
namespace dpctl
33+
{
34+
namespace tensor
35+
{
36+
namespace py_internal
37+
{
38+
namespace atomic_support
39+
{
40+
41+
typedef bool (*atomic_support_fn_ptr_t)(const sycl::queue &, sycl::usm::alloc);
42+
43+
/*! @brief Function which returns a constant value for atomic support */
44+
template <bool return_value>
45+
bool fixed_decision(const sycl::queue &, sycl::usm::alloc)
46+
{
47+
return return_value;
48+
}
49+
50+
/*! @brief Template for querying atomic support for a type on a device */
51+
template <typename T>
52+
bool check_atomic_support(const sycl::queue &exec_q,
53+
sycl::usm::alloc usm_alloc_type)
54+
{
55+
constexpr bool atomic32 = (sizeof(T) == 4);
56+
constexpr bool atomic64 = (sizeof(T) == 8);
57+
using dpctl::tensor::type_utils::is_complex;
58+
if constexpr ((!atomic32 && !atomic64) || is_complex<T>::value) {
59+
return fixed_decision<false>(exec_q, usm_alloc_type);
60+
}
61+
else {
62+
bool supports_atomics = false;
63+
const sycl::device &dev = exec_q.get_device();
64+
if constexpr (atomic64) {
65+
if (!dev.has(sycl::aspect::atomic64)) {
66+
return false;
67+
}
68+
}
69+
switch (usm_alloc_type) {
70+
case sycl::usm::alloc::shared:
71+
supports_atomics =
72+
dev.has(sycl::aspect::usm_atomic_shared_allocations);
73+
break;
74+
case sycl::usm::alloc::host:
75+
supports_atomics =
76+
dev.has(sycl::aspect::usm_atomic_host_allocations);
77+
break;
78+
case sycl::usm::alloc::device:
79+
supports_atomics = true;
80+
break;
81+
default:
82+
supports_atomics = false;
83+
}
84+
return supports_atomics;
85+
}
86+
}
87+
88+
template <typename fnT, typename T> struct MaxAtomicSupportFactory
89+
{
90+
fnT get()
91+
{
92+
if constexpr (std::is_floating_point_v<T>) {
93+
return fixed_decision<false>;
94+
}
95+
else {
96+
return check_atomic_support<T>;
97+
}
98+
}
99+
};
100+
101+
template <typename fnT, typename T> struct MinAtomicSupportFactory
102+
{
103+
fnT get()
104+
{
105+
if constexpr (std::is_floating_point_v<T>) {
106+
return fixed_decision<false>;
107+
}
108+
else {
109+
return check_atomic_support<T>;
110+
}
111+
}
112+
};
113+
114+
template <typename fnT, typename T> struct SumAtomicSupportFactory
115+
{
116+
fnT get()
117+
{
118+
if constexpr (std::is_floating_point_v<T>) {
119+
return fixed_decision<false>;
120+
}
121+
else {
122+
return check_atomic_support<T>;
123+
}
124+
}
125+
};
126+
127+
template <typename fnT, typename T> struct ProductAtomicSupportFactory
128+
{
129+
fnT get()
130+
{
131+
if constexpr (std::is_floating_point_v<T>) {
132+
return fixed_decision<false>;
133+
}
134+
else {
135+
return check_atomic_support<T>;
136+
}
137+
}
138+
};
139+
140+
} // namespace atomic_support
141+
} // namespace py_internal
142+
} // namespace tensor
143+
} // namespace dpctl

0 commit comments

Comments
 (0)