Skip to content

Commit 35e917c

Browse files
nathanielsimardnathanielsimard
authored
Refactor: Device locking (#959)
1 parent 3b4d937 commit 35e917c

File tree

176 files changed

+1632
-2139
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

176 files changed

+1632
-2139
lines changed

‎crates/cubecl-attention/src/base.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ pub enum Strategy {
2424
#[allow(clippy::result_large_err)]
2525
pub fn launch<R: Runtime, AP: AttentionPrecision>(
2626
strategy: &Strategy,
27-
client: &ComputeClient<R::Server, R::Channel>,
27+
client: &ComputeClient<R::Server>,
2828
query: TensorHandle<R, QG<AP>>,
2929
key: TensorHandle<R, KG<AP>>,
3030
value: TensorHandle<R, VG<AP>>,
@@ -43,7 +43,7 @@ pub fn launch<R: Runtime, AP: AttentionPrecision>(
4343
#[allow(clippy::result_large_err)]
4444
pub fn launch_ref<R: Runtime, AP: AttentionPrecision>(
4545
strategy: &Strategy,
46-
client: &ComputeClient<R::Server, R::Channel>,
46+
client: &ComputeClient<R::Server>,
4747
query: &TensorHandleRef<R>,
4848
key: &TensorHandleRef<R>,
4949
value: &TensorHandleRef<R>,
@@ -55,7 +55,7 @@ pub fn launch_ref<R: Runtime, AP: AttentionPrecision>(
5555
}
5656

5757
pub fn launch_tmp<R: Runtime, AP: AttentionPrecision>(
58-
client: &ComputeClient<R::Server, R::Channel>,
58+
client: &ComputeClient<R::Server>,
5959
query: &TensorHandleRef<R>,
6060
key: &TensorHandleRef<R>,
6161
value: &TensorHandleRef<R>,

‎crates/cubecl-attention/src/components/batch/base.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ pub trait BatchAttentionFamily: Send + Sync + 'static {
2525
///
2626
/// Out-of-bounds can happen
2727
unsafe fn launch_unchecked<'a, MS: AttentionSpec, R: Runtime>(
28-
client: &ComputeClient<<R as Runtime>::Server, <R as Runtime>::Channel>,
28+
client: &ComputeClient<<R as Runtime>::Server>,
2929
cube_dim: CubeDim,
3030
cube_count: CubeCount,
3131
input: InputRuntimeArg<'a, MS, R>,
@@ -38,7 +38,7 @@ pub trait BatchAttentionFamily: Send + Sync + 'static {
3838
///
3939
/// This function may return an error if the configuration cannot be supported on the current runtime.
4040
fn setup<AP: AttentionPrecision, R: Runtime>(
41-
client: &ComputeClient<R::Server, R::Channel>,
41+
client: &ComputeClient<R::Server>,
4242
problem: &AttentionProblem,
4343
selection: &AttentionSelection,
4444
line_sizes: &AttentionLineSizes,

‎crates/cubecl-attention/src/components/batch/dummy/setup.rs‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ impl<GA: GlobalAttentionFamily> BatchAttentionFamily for DummyBatchAttentionFami
2222
type Config = DummyBatchConfig<GA::Config>;
2323

2424
fn setup<AP: crate::components::AttentionPrecision, R: cubecl_core::Runtime>(
25-
client: &ComputeClient<R::Server, R::Channel>,
25+
client: &ComputeClient<R::Server>,
2626
problem: &AttentionProblem,
2727
selection: &AttentionSelection,
2828
line_sizes: &AttentionLineSizes,
@@ -44,10 +44,7 @@ impl<GA: GlobalAttentionFamily> BatchAttentionFamily for DummyBatchAttentionFami
4444
AS: crate::components::AttentionSpec,
4545
R: cubecl_core::Runtime,
4646
>(
47-
client: &cubecl_core::prelude::ComputeClient<
48-
<R as cubecl_core::Runtime>::Server,
49-
<R as cubecl_core::Runtime>::Channel,
50-
>,
47+
client: &cubecl_core::prelude::ComputeClient<<R as cubecl_core::Runtime>::Server>,
5148
cube_dim: cubecl_core::CubeDim,
5249
cube_count: cubecl_core::CubeCount,
5350
input: crate::components::InputRuntimeArg<'a, AS, R>,

‎crates/cubecl-attention/src/components/global/base.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ pub trait GlobalAttentionFamily: Send + Sync + 'static {
2222
///
2323
/// This function may return an error if the configuration cannot be supported on the current runtime.
2424
fn setup<AP: AttentionPrecision, R: Runtime>(
25-
client: &ComputeClient<R::Server, R::Channel>,
25+
client: &ComputeClient<R::Server>,
2626
problem: &AttentionProblem,
2727
selection: &AttentionSelection,
2828
line_sizes: &AttentionLineSizes,

‎crates/cubecl-attention/src/components/global/dummy/setup.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ impl<
3030
type Config = DummyGlobalConfig<SA::Config>;
3131

3232
fn setup<AP: crate::components::AttentionPrecision, R: cubecl_core::Runtime>(
33-
client: &ComputeClient<R::Server, R::Channel>,
33+
client: &ComputeClient<R::Server>,
3434
problem: &AttentionProblem,
3535
selection: &AttentionSelection,
3636
line_sizes: &AttentionLineSizes,

‎crates/cubecl-attention/src/components/stage/base.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ pub trait StageAttentionFamily: Send + Sync + 'static {
3939
///
4040
/// This function may return an error if the configuration cannot be supported on the current runtime.
4141
fn setup<AP: AttentionPrecision, R: Runtime>(
42-
client: &ComputeClient<R::Server, R::Channel>,
42+
client: &ComputeClient<R::Server>,
4343
problem: &AttentionProblem,
4444
selection: &AttentionSelection,
4545
line_sizes: &AttentionLineSizes,

‎crates/cubecl-attention/src/components/stage/dummy/setup.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ impl<
4747
type Config = DummyStageConfig<TA::Config>;
4848

4949
fn setup<AP: crate::components::AttentionPrecision, R: cubecl_core::Runtime>(
50-
client: &ComputeClient<R::Server, R::Channel>,
50+
client: &ComputeClient<R::Server>,
5151
problem: &AttentionProblem,
5252
selection: &AttentionSelection,
5353
line_sizes: &AttentionLineSizes,

‎crates/cubecl-attention/src/components/tile/base.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ pub trait TileAttentionFamily: Send + Sync + 'static {
2929
///
3030
/// This function may return an error if the configuration cannot be supported on the current runtime.
3131
fn setup<AP: AttentionPrecision, R: Runtime>(
32-
client: &ComputeClient<R::Server, R::Channel>,
32+
client: &ComputeClient<R::Server>,
3333
problem: &AttentionProblem,
3434
selection: &AttentionSelection,
3535
line_sizes: &AttentionLineSizes,

‎crates/cubecl-attention/src/components/tile/dummy/attention_matmul/accelerated/setup.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ impl AttentionMatmulFamily for AcceleratedAttentionMatmul {
2222
}
2323

2424
fn setup<AP: AttentionPrecision, R: cubecl_core::Runtime>(
25-
_client: &cubecl_core::prelude::ComputeClient<R::Server, R::Channel>,
25+
_client: &cubecl_core::prelude::ComputeClient<R::Server>,
2626
problem: &crate::components::AttentionProblem,
2727
selection: &crate::components::AttentionSelection,
2828
line_sizes: &crate::components::AttentionLineSizes,

‎crates/cubecl-attention/src/components/tile/dummy/attention_matmul/base.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ pub trait AttentionMatmulFamily: Send + Sync + 'static {
9797
///
9898
/// This function may return an error if the configuration cannot be supported on the current runtime.
9999
fn setup<AP: AttentionPrecision, R: Runtime>(
100-
client: &ComputeClient<R::Server, R::Channel>,
100+
client: &ComputeClient<R::Server>,
101101
problem: &AttentionProblem,
102102
selection: &AttentionSelection,
103103
line_sizes: &AttentionLineSizes,

0 commit comments

Comments
 (0)