Skip to content

Run your first Loom kernel

Example files: loom/docs/examples/getting-started/first-kernel/

This quickstart takes one source file from readable kernel code to checked AMDGPU execution and a real benchmark. The example computes SAXPY, y = alpha * x + y, over a runtime-sized f32 vector.

The same file owns three related things:

  • the kernel and its launch geometry;
  • deterministic correctness samples and expected results; and
  • a named benchmark that selects one of those proven samples.

There is no generated host harness or separate benchmark program. Loom's tools compile the kernel for the selected device, materialize the tensors declared by the case, launch the work, check the result, and measure the same workload when requested.

This page assumes the Loom tools are on PATH and that --device=amdgpu is available. Acquiring Loom describes the current source-build and Bazel setup.

Create one source file

Save the following as saxpy.loom:

saxpy.loom
// Computes y = alpha * x + y for a runtime-sized f32 vector.
kernel.def export("saxpy_f32") @saxpy_f32(%element_count: index) {
  %unit = index.constant 1 : index
  %workgroup_size = index.constant 256 : index
  %rounding = index.sub %workgroup_size, %unit : index
  %rounded_count = index.add %element_count, %rounding : index
  %workgroups = index.div %rounded_count, %workgroup_size : index
  kernel.launch.config workgroups(%workgroups, %unit, %unit) workgroup_size(%workgroup_size, %unit, %unit) : index
} launch(%element_count: index, %alpha: f32, %x: buffer, %y: buffer) {
  %workgroup = kernel.workgroup.id<x> : index
  %lane = kernel.workitem.id<x> : index
  %workgroup_size = kernel.workgroup.size<x> : index
  %element_index = index.madd %workgroup, %workgroup_size, %lane : index
  %in_bounds = index.cmp ult, %element_index, %element_count : index
  %x_noalias, %y_noalias = buffer.assume.noalias %x, %y : buffer, buffer
  %base = index.constant 0 : offset
  %x_view = buffer.view %x_noalias[%base] : buffer -> view<[%element_count]xf32>
  %y_view = buffer.view %y_noalias[%base] : buffer -> view<[%element_count]xf32>
  scf.if %in_bounds {
    %x_value = view.load %x_view[%element_index] : view<[%element_count]xf32> -> f32
    %y_value = view.load %y_view[%element_index] : view<[%element_count]xf32> -> f32
    %result = scalar.fmaf %alpha, %x_value, %y_value : f32
    view.store %result, %y_view[%element_index] : f32, view<[%element_count]xf32>
  }
  kernel.return
}

check.case public @saxpy_f32_case {
  %element_count = check.param.choice values([1, 255, 256, 257, 1009, 65536]) name("element_count") : index
  %alpha = check.literal value(4.0) : f32
  %x = check.generate.fill value(2.0) : tensor<[%element_count]xf32>
  %y = check.generate.fill value(3.0) : tensor<[%element_count]xf32>
  %expected = check.generate.fill value(11.0) : tensor<[%element_count]xf32>
  kernel.launch @saxpy_f32[%element_count](%element_count, %alpha, %x, %y) : [index](index, f32, tensor<[%element_count]xf32>, tensor<[%element_count]xf32>)
  check.expect.equal actual(%y) expected(%expected) : tensor<[%element_count]xf32>
  check.return
}

check.benchmark<@saxpy_f32_case> @saxpy_f32_64k {element_count = 65536}

The kernel.def has two regions. The first maps %element_count to a physical launch using 256 work-items per workgroup. The second is the device body. Each work-item computes one element index, guards the partial final workgroup, loads x and y, performs the fused multiply-add, and stores the result back to y.

%element_count appears in both signatures for different reasons. The value before launch is workload input used by the host-side launch calculation. The value inside launch(...) is an explicit device argument used for bounds and buffer views. Loom does not silently turn one into the other.

The check.case expands to six deterministic samples. Sizes 255, 256, and 257 cover both sides of the workgroup boundary; 1009 exercises a larger tail; and 65536 supplies the benchmark workload. Every sample fills x with 2, initializes y to 3, launches with alpha = 4, and requires every result to equal 11.

The final check.benchmark does not duplicate that setup. It gives the 65536-element assignment a stable workload name and reuses the case's inputs, launch, and expectation.

Check the source before compiling

loom-format --check parses and verifies the complete program and rejects noncanonical text:

loom-format --check saxpy.loom

This is more than whitespace checking. Invalid SSA use, mismatched types, and malformed launch signatures fail at this boundary. The test and benchmark runners then plan concrete parameter samples and reject assignments that do not select a valid case sample.

Execute every correctness sample on AMDGPU

Run the checked program directly:

iree-test-loom saxpy.loom --device=amdgpu >saxpy-test.json

The runner discovers @saxpy_f32_case, plans its six parameter assignments, compiles the reachable kernel for the selected AMDGPU device, allocates and initializes each tensor, executes the launch, and evaluates check.expect.equal. It returns a failing process status if any sample fails and writes the structured report before exiting.

The report distinguishes cases from their concrete samples. This case has one authored definition and six executions; a failure records the exact element_count and stable sample ordinal needed to reproduce it.

Select one sample while editing without changing the source:

iree-test-loom saxpy.loom \
  --device=amdgpu \
  --case=@saxpy_f32_case \
  --sample=3

Sample 3 is the 257-element tail. The ordinal comes from the deterministic parameter plan rather than hidden random state.

Benchmark the proven workload

Measure the named 64K dispatch workload:

iree-benchmark-loom saxpy.loom \
  --device=amdgpu \
  --benchmark=@saxpy_f32_64k \
  --measure=dispatch_complete \
  --batch-size=64 \
  --output=saxpy-benchmark.json

The benchmark runner executes the selected case and verifies its expected result before accepting timing evidence. dispatch_complete measures host submission through device completion for a prepared batch, and the report normalizes the result per logical SAXPY operation. The batch size remains part of the recorded measurement policy instead of becoming source metadata.

Benchmark with an optimized, uninstrumented tool build on an otherwise quiet device. The JSON report carries device identity, parameter values, timing policy, score distribution, and warnings needed to interpret the number.

The checked-in example wraps the same steps and also emits a generic gfx11 deployment artifact:

loom/docs/examples/getting-started/first-kernel/run.sh \
  gfx11-generic build/first-kernel/gfx11-generic

The script prints each public command before executing it and leaves the test, benchmark, plan, and native artifact reports together in the output directory.

Put the source under Bazel

An authoring repository can give the same file automated format, lint, plan, correctness, and benchmark-smoke coverage with one rule:

BUILD.bazel
load(
    "@hrx//loom/build_tools/bazel:defs.bzl",
    "loom_kernel_library",
)
load(
    "@hrx//loom/target/amdgpu:execution_profiles.bzl",
    "AMDGPU_HARDWARE_PROFILE",
)

package(
    default_visibility = ["//visibility:private"],
    licenses = ["notice"],  # Apache 2.0
)

loom_kernel_library(
    name = "saxpy",
    srcs = ["saxpy.loom"],
    execution_profiles = [AMDGPU_HARDWARE_PROFILE],
    targets = ["@hrx//loom/target/amdgpu:gfx11-generic"],
)

The built-in AMDGPU execution profile owns the compiler and driver requirements, --device=amdgpu argument, GPU resource serialization, and stable query tags. The source remains target-independent; the profile chooses where this test instance runs.

Build the reusable bytecode library:

bazel build :saxpy

Run its generated test suite on a configured AMDGPU machine:

bazel test :saxpy_test

saxpy_test performs source lint and canonical-format checks, plans every benchmark without a device, emits and inspects the generic gfx11 deployment artifact, executes all correctness cases on AMDGPU, and runs every benchmark once with no warmups as a smoke test. Full repeated timing remains the explicit iree-benchmark-loom command above; CI should prove that the workload remains correct and executable without turning shared runners into performance laboratories.

loom_kernel_library links the authored file into reusable .loombc and builds a separate test module that retains the root library's checks. Deployment linking selects exported production roots such as @saxpy_f32 and strips check.case and check.benchmark records before target compilation. Test code therefore stays next to the kernel without entering the native artifact.

Follow the boundary that matters next

You now have one portable source file that formats, verifies, executes across tail sizes, names a benchmark workload, participates in automated AMDGPU testing, and remains linkable as a library.