Skip to content

Commit c6b1a88

Browse files
ConnorGraymodularbot
ConnorGray
authored and
modularbot
committed
[stdlib][max] cleanup: Migrate uses of has_*() to ctx.target.is_*()
This begins the process of removing code that assumes a single global accelerator target that is separate from `current_target()`, and instead scopes down to a locally-available target from the specific `DeviceContext` a kernel is going to be run on. This commit handles just the "easy" cases where a `has_*()` exists in a function that already has a `DeviceContext` instance locally. MODULAR_ORIG_COMMIT_REV_ID: 4fd91541adf46d5daeb257532bc73ee56d27f6e3
1 parent b9605ab commit c6b1a88

76 files changed

Lines changed: 219 additions & 316 deletions

File tree

Some content is hidden

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

‎docs/code/gpu/fundamentals/device_info.mojo‎

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,7 @@
1212
# ===----------------------------------------------------------------------=== #
1313
# DOC: max/gpu/fundamentals.mdx
1414

15-
from std.sys import (
16-
exit,
17-
has_accelerator,
18-
has_amd_gpu_accelerator,
19-
has_apple_gpu_accelerator,
20-
)
15+
from std.sys import exit, has_accelerator
2116

2217
from max.gpu.host import DeviceAttribute, DeviceContext
2318

@@ -48,7 +43,7 @@ def main() raises:
4843
ctx.get_attribute(DeviceAttribute.MAX_SHARED_MEMORY_PER_BLOCK),
4944
)
5045

51-
comptime if not has_apple_gpu_accelerator():
46+
comptime if not ctx.target.is_apple_gpu():
5247
# Not currently defined for Apple GPUs
5348

5449
print(
@@ -76,9 +71,7 @@ def main() raises:
7671
ctx.get_attribute(DeviceAttribute.MAX_REGISTERS_PER_BLOCK),
7772
)
7873

79-
comptime if not (
80-
has_amd_gpu_accelerator() or has_apple_gpu_accelerator()
81-
):
74+
comptime if not ctx.target.is_amd_gpu() or ctx.target.is_apple_gpu():
8275
# Not currently defined for AMD and Apple GPUs
8376

8477
print(

‎docs/code/tile-tensor/tensors/tile_tensor_gpu_examples.mojo‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,7 @@ from max.gpu.host import DeviceContext, DeviceBuffer, get_gpu_target
2525
from layout import Coord, Idx, stack_allocation, TileTensor
2626
from layout.tile_layout import row_major, blocked_product
2727
from std.sys import has_accelerator
28-
from std.sys.info import (
29-
has_apple_gpu_accelerator,
30-
has_nvidia_gpu_accelerator,
31-
is_apple_gpu,
32-
is_nvidia_gpu,
33-
simd_width_of,
34-
)
28+
from std.sys.info import simd_width_of
3529
from std.testing import assert_equal, assert_false, assert_true
3630
from std.sys import exit
3731

‎max/examples/custom_ops/kernels/tensor_core_mma.mojo‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,7 @@ struct TensorCoreMMA[algorithm: StaticString]:
154154
# - "mma_tile_buffers": A matrix multiplication using tile buffers and AMD Tensor Core instructions.
155155

156156
comptime if Self.algorithm == "naive_tensor":
157-
comptime if has_nvidia_gpu_accelerator() or has_amd_gpu_accelerator():
157+
comptime if gpu_ctx.target.is_nvidia_gpu() or gpu_ctx.target.is_amd_gpu():
158158
comptime BM = 64
159159
comptime BN = 64
160160
comptime BK = 8
@@ -186,7 +186,7 @@ struct TensorCoreMMA[algorithm: StaticString]:
186186
block_dim=(NUM_THREADS, 1),
187187
)
188188
elif Self.algorithm == "basic_shared_mem":
189-
comptime if has_nvidia_gpu_accelerator() or has_amd_gpu_accelerator():
189+
comptime if gpu_ctx.target.is_nvidia_gpu() or gpu_ctx.target.is_amd_gpu():
190190
comptime BM = 64
191191
comptime BN = 64
192192
comptime BK = 8
@@ -218,7 +218,7 @@ struct TensorCoreMMA[algorithm: StaticString]:
218218
block_dim=(NUM_THREADS, 1),
219219
)
220220
elif Self.algorithm == "multi_block_tiled":
221-
comptime if has_nvidia_gpu_accelerator() or has_amd_gpu_accelerator():
221+
comptime if gpu_ctx.target.is_nvidia_gpu() or gpu_ctx.target.is_amd_gpu():
222222
comptime BM = 256
223223
comptime BN = 256
224224
comptime BK = 64
@@ -254,7 +254,7 @@ struct TensorCoreMMA[algorithm: StaticString]:
254254
block_dim=(NUM_THREADS, 1),
255255
)
256256
elif Self.algorithm == "scheduler_hints":
257-
comptime if has_nvidia_gpu_accelerator() or has_amd_gpu_accelerator():
257+
comptime if gpu_ctx.target.is_nvidia_gpu() or gpu_ctx.target.is_amd_gpu():
258258
comptime BM = 256
259259
comptime BN = 256
260260
comptime BK = 64
@@ -290,7 +290,7 @@ struct TensorCoreMMA[algorithm: StaticString]:
290290
block_dim=(NUM_THREADS, 1),
291291
)
292292
elif Self.algorithm == "double_buffer":
293-
comptime if has_nvidia_gpu_accelerator() or has_amd_gpu_accelerator():
293+
comptime if gpu_ctx.target.is_nvidia_gpu() or gpu_ctx.target.is_amd_gpu():
294294
comptime BM = 128
295295
comptime BN = 128
296296
comptime BK = 32
@@ -326,7 +326,7 @@ struct TensorCoreMMA[algorithm: StaticString]:
326326
block_dim=(NUM_THREADS, 1),
327327
)
328328
elif Self.algorithm == "mma_tile_buffers":
329-
comptime if has_nvidia_gpu_accelerator() or has_amd_gpu_accelerator():
329+
comptime if gpu_ctx.target.is_nvidia_gpu() or gpu_ctx.target.is_amd_gpu():
330330
comptime BM = 256
331331
comptime BN = 256
332332
comptime BK = 64

‎max/examples/layout_tensor/layout_tensor_gpu_examples.mojo‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,6 @@ from std.memory import Pointer
2828
from std.sys import has_accelerator
2929
from std.sys.info import (
3030
has_apple_gpu_accelerator,
31-
has_nvidia_gpu_accelerator,
32-
is_apple_gpu,
3331
is_nvidia_gpu,
3432
simd_width_of,
3533
)

‎max/kernels/benchmarks/gpu/linalg/bench_bmm.mojo‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ from std.sys import (
1717
get_defined_int,
1818
simd_width_of,
1919
)
20-
from std.sys.info import has_amd_gpu_accelerator
2120

2221
from layout import Coord, TileTensor, row_major, CoordLike, Idx
2322
import linalg.matmul.vendor.blas as vendor_blas
@@ -191,7 +190,7 @@ def bench_bmm[
191190
@inline(.always)
192191
def kernel_launch(ctx: DeviceContext, iteration: Int) raises {imm}:
193192
comptime if use_vendor_blas:
194-
comptime if has_amd_gpu_accelerator():
193+
comptime if ctx.target.is_amd_gpu():
195194
var c_buffer = TileTensor(
196195
c_device.ptr, row_major(Coord(m, n))
197196
)

‎max/kernels/benchmarks/gpu/nn/bench_conv2d.mojo‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -346,7 +346,7 @@ def bench_conv2d[
346346
[ThroughputMeasure(BenchMetric.flops, flops)],
347347
)
348348
comptime if resolved == "cudnn":
349-
comptime assert has_nvidia_gpu_accelerator(), (
349+
comptime assert ctx.target.is_nvidia_gpu(), (
350350
"impl=cudnn requires an NVIDIA target. Build for an NVIDIA"
351351
" accelerator (e.g. cuda:b200) or pick a different impl."
352352
)
@@ -557,7 +557,7 @@ def bench_conv2d[
557557
# is not available on AMD. The AMD `amd_4wave` arm has its own
558558
# dedicated correctness tests under max/kernels/test/gpu/nn/.
559559
var do_verify = verify
560-
comptime if not has_nvidia_gpu_accelerator():
560+
comptime if not ctx.target.is_nvidia_gpu():
561561
if verify:
562562
print("verify: skipped (requires NVIDIA cuDNN as the reference)")
563563
do_verify = False

‎max/kernels/benchmarks/gpu/nn/bench_conv3d.mojo‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@ Usage (kbench):
3131
from std.math import ceildiv
3232
from std.random import rand
3333
from std.sys import get_defined_dtype, get_defined_int
34-
from std.sys.info import has_amd_gpu_accelerator
3534

3635
from max.benchmark import bencher_iter_custom
3736
from std.benchmark import (
@@ -318,7 +317,7 @@ def bench_conv3d[
318317
"dispatch_1x1x1_matmul_conv3d declined: " + bench_input_id
319318
)
320319
elif impl == "qslice":
321-
comptime if not has_amd_gpu_accelerator():
320+
comptime if not ctx.target.is_amd_gpu():
322321
var accepted = dispatch_qslice_conv3d_sm100(
323322
input_tt,
324323
filter_qrscf_tt,
@@ -340,7 +339,7 @@ def bench_conv3d[
340339
+ bench_input_id
341340
)
342341
elif impl == "native_3d":
343-
comptime if has_amd_gpu_accelerator():
342+
comptime if ctx.target.is_amd_gpu():
344343
# Autotune-override knobs (0 = use the dispatcher's
345344
# built-in heuristic). Pass via `-D BM_OVERRIDE=64` etc
346345
# to sweep configs for the per-shape dispatch table.
@@ -426,7 +425,7 @@ def bench_conv3d[
426425
[ThroughputMeasure(BenchMetric.flops, flops)],
427426
)
428427
elif impl == "qslice":
429-
comptime if not has_amd_gpu_accelerator():
428+
comptime if not ctx.target.is_amd_gpu():
430429

431430
@inline(.always)
432431
def qslice_bench(
@@ -453,7 +452,7 @@ def bench_conv3d[
453452
[ThroughputMeasure(BenchMetric.flops, flops)],
454453
)
455454
elif impl == "native_3d":
456-
comptime if has_amd_gpu_accelerator():
455+
comptime if ctx.target.is_amd_gpu():
457456
comptime _BM_OVERRIDE = get_defined_int["BM_OVERRIDE", 0]()
458457
comptime _BN_OVERRIDE = get_defined_int["BN_OVERRIDE", 0]()
459458
comptime _BK_OVERRIDE = get_defined_int["BK_OVERRIDE", 0]()
@@ -491,7 +490,7 @@ def bench_conv3d[
491490
[ThroughputMeasure(BenchMetric.flops, flops)],
492491
)
493492
elif impl == "cudnn":
494-
comptime if has_amd_gpu_accelerator():
493+
comptime if ctx.target.is_amd_gpu():
495494
# Capturing `input_buf` / `output_buf` here would alias the
496495
# `input_tt` / `output_tt` views this arm launches through.
497496
@inline(.always)
@@ -594,7 +593,7 @@ def bench_conv3d[
594593

595594
# Optional correctness cross-check against cuDNN.
596595
if verify:
597-
comptime if has_amd_gpu_accelerator():
596+
comptime if ctx.target.is_amd_gpu():
598597
conv_miopen(
599598
input_tt,
600599
filter_qrscf_tt,

‎max/kernels/benchmarks/gpu/nn/bench_topk.mojo‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -911,7 +911,7 @@ def main() raises:
911911
num_blocks_per_input=num_blocks_per_input,
912912
)
913913

914-
comptime if has_apple_gpu_accelerator():
914+
comptime if ctx.target.is_apple_gpu():
915915
if masked_probs or use_dist:
916916
raise Error(
917917
"the masked_probs and topp_dist benchmarks require"

‎max/kernels/src/algorithm/gpu/rowwise.mojo‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,6 @@ from std.memory import UnsafePointer, stack_allocation
6262
from std.sys import simd_width_of, size_of, get_defined_int
6363
from std.sys.info import (
6464
has_amd_gpu_accelerator,
65-
has_apple_gpu_accelerator,
6665
is_apple_gpu,
6766
)
6867
from std.utils.coord import Coord, DynamicCoord, coord_to_index_list
@@ -1680,7 +1679,7 @@ def launch[
16801679
comptime if (
16811680
supports_splitk
16821681
and effective_simd <= _SPLITK_MAX_SIMD
1683-
and not has_apple_gpu_accelerator()
1682+
and not ctx.target.is_apple_gpu()
16841683
):
16851684
if num_rows < sm_count and row_size >= _SPLITK_MIN_ROW:
16861685
# Cap `blocks_per_row` at `_SPLITK_ROW_BLOCK_SIZE`: the

‎max/kernels/src/algorithm/reductions.mojo‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ single reduce phase plus a per-row `emit` — and runs on both CPU and GPU.
2020
"""
2121

2222
from max.gpu.host import DeviceContext
23-
from std.sys.info import has_apple_gpu_accelerator
2423
from std.utils.coord import Coord
2524

2625
from algorithm import rowwise
@@ -433,7 +432,7 @@ def reduce_mean[
433432
var total = acc.slice[params.emit_tile_width]()
434433

435434
comptime if dtype.is_floating_point():
436-
comptime float_type = DType.float32 if has_apple_gpu_accelerator() else DType.float64
435+
comptime float_type = DType.float32 if context.T.target.is_apple_gpu() else DType.float64
437436
# `axis_size == 0` gives `recip = inf` and `total` (the
438437
# `ReduceSum` identity) is `0`, so `total * recip` is the
439438
# IEEE-754 `0 * inf = NaN` — the same "no data" signal

0 commit comments

Comments
 (0)