[Misc] Add cross-platform shuffle

[hughperkins]

Issue: #

Brief Summary

copilot:summary

Walkthrough

copilot:walkthrough

[hughperkins]

tests seem reasoanble to me. (got ai to clean these up a lot, and add some missing tests)

[hughperkins]

asking Opus wht this is:

what si InternalFuncStmt for?

⬢ Grepped, searched 1 grep, 1 search
Grepped "class InternalFuncStmt" in .
Searched "...cStmt used for and how is it created?"

InternalFuncStmt is an IR statement node that represents a call to a built-in runtime or intrinsic function — things that aren't user-defined but are part of the
compiler/runtime infrastructure.
It's defined in quadrants/ir/statements.h:

quadrants/ir/statements.h lines 1670-1681

class InternalFuncStmt : public Stmt {
public:
std::string func_name;
std::vector<Stmt *> args;
bool with_runtime_context;
explicit InternalFuncStmt(const std::string &func_name,
const std::vector<Stmt *> &args,
Type *ret_type = nullptr,
bool with_runtime_context = true)
// ...
};

The lifecycle is:

1. Registration — Internal ops are declared via PER_INTERNAL_OP in internal_ops.inc.h (e.g. subgroupShuffle, subgroupBarrier, subgroupAdd, etc.) and given type signatures in
   Operations::init_internals().
2. Flattening — When the frontend encounters a call like subgroup.shuffle(...), it goes through InternalCallOperation::flatten() which pushes an InternalFuncStmt into the IR
   with the op name and args.
3. Default codegen — The base TaskCodeGenLLVM::visit(InternalFuncStmt*) simply forwards it as a runtime function call (looks up func_name in the LLVM module and calls it with
   the args).
4. Backend overrides — Backends like CUDA and AMDGPU override visit(InternalFuncStmt*) to intercept specific op names and emit platform-specific intrinsics instead. That's
   exactly what this branch does — it intercepts "subgroupShuffle" to emit __shfl_sync on CUDA and ds_bpermute on AMDGPU rather than falling through to a generic runtime call.

In short: it's the IR-level representation of "call this named built-in operation," and each codegen backend decides how to lower it — either as a plain runtime call or as
specialized hardware intrinsics.

[hughperkins]

to opus:

why do we override it for shuffle?

Because shuffle is a cross-lane operation — it reads a value from a different thread within the warp/wavefront. That can't be expressed as a regular function call to a C
runtime; it must be a hardware intrinsic instruction.
The default codegen path in TaskCodeGenLLVM::visit just does:

quadrants/codegen/llvm/codegen_llvm.cpp lines 2428-2438

void TaskCodeGenLLVM::visit(InternalFuncStmt *stmt) {
std::vector<llvm::Value *> args;
if (stmt->with_runtime_context)
args.push_back(get_context());
for (auto s : stmt->args) {
args.push_back(llvm_val[s]);
}
llvm_val[stmt] = call(stmt->func_name, std::move(args));
}

That looks up a function by name in the linked runtime module and calls it. This works for things like list management or memory ops that are implemented in C++. But there's
no C function you could write that performs a warp shuffle — it's a single instruction that moves data between registers across lanes in lockstep. Each backend needs to emit
the right intrinsic:
• CUDA: __shfl_sync (PTX shfl.sync.idx)
• AMDGPU: ds_bpermute (LDS-based cross-lane permute)

On top of that, 64-bit types (f64, i64) need the split-into-two-i32-shuffles-and-reassemble pattern because the hardware shuffle instructions only operate on 32-bit values.
That logic is also intrinsic-specific and can't live in a generic runtime function.
So the override exists because shuffle is fundamentally a hardware primitive, not a callable function.

[hughperkins]

to Opus:

why do we need both subgroupShuffle and subgroupBroadcast? Or: what is the difference?

The distinction comes from the Vulkan/SPIR-V spec, where they are genuinely different operations with different requirements:
• subgroupBroadcast (OpGroupNonUniformBroadcast) — the index must be dynamically uniform, meaning every lane in the subgroup must pass the same index value. The GPU can
optimize for this (e.g. a single broadcast wire). This is the "everyone reads from the same lane" operation.
• subgroupShuffle (OpGroupNonUniformShuffle) — the index can be different per lane. Lane 0 can read from lane 3, lane 1 can read from lane 2, etc. This is the general
arbitrary permutation.

You can see the SPIR-V backend emits different opcodes for them:

quadrants/codegen/spirv/spirv_codegen.cpp lines 1346-1348

} else if (stmt->func_name == "subgroupBroadcast") {
  // ...
      spv::OpGroupNonUniformBroadcast, ...

quadrants/codegen/spirv/spirv_codegen.cpp lines 1425-1426

  } else if (ends_with(stmt->func_name, "Shuffle")) {
    spv_op = spv::OpGroupNonUniformShuffle;

On CUDA and AMDGPU, though, the hardware doesn't distinguish — __shfl_sync and ds_bpermute handle both cases identically, which is why the codegen overrides treat them the
same:

quadrants/codegen/cuda/codegen_cuda.cpp lines 788-789

  if (stmt->func_name == "subgroupShuffle" ||
      stmt->func_name == "subgroupBroadcast") {

So both exist because Vulkan requires separate ops (broadcast can be faster when the driver knows all lanes want the same source), but on LLVM backends they collapse to the
same instruction. Keeping them as distinct ops in the IR ensures correct SPIR-V codegen on Vulkan while allowing CUDA/AMDGPU to merge them trivially.

[hughperkins]

Note to reviewers: I'm asking questions on the amdgpu file, so check there for answers first.

[hughperkins]

to Opus:

what is call here? call("thread_idx");

It's a helper on LLVMModuleBuilder (in llvm_codegen_utils.h) that looks up a function by name in the linked runtime module and emits an LLVM CreateCall instruction. The
overload hit by call("thread_idx") is this variadic template:

quadrants/codegen/llvm/llvm_codegen_utils.h lines 162-164

llvm::Value *call(const std::string &func_name, Args &&...args) {
  return call(builder.get(), func_name, std::forward<Args>(args)...);
}

Which resolves to:

quadrants/codegen/llvm/llvm_codegen_utils.h lines 130-135

llvm::Value *call(llvm::IRBuilder<> *builder,
                  const std::string &func_name,
                  std::vector<llvm::Value *> args) {
  auto func = get_runtime_function(func_name);
  return call(builder, func, std::move(args));
}

So call("thread_idx") calls get_runtime_function("thread_idx") to find the function in the LLVM module, then emits a CreateCall with no arguments. It's essentially a shorthand
for "emit an LLVM call to this named runtime/intrinsic wrapper function and return the result as an llvm::Value*."

[hughperkins]

question to Opus:

what is the '4' here

            builder->CreateZExtOrTrunc(llvm_val[stmt->args[1]],
                                       llvm::Type::getInt32Ty(*llvm_context)),
            tlctx->get_constant((int32)4));

It's the dword size in bytes. AMD's ds_bpermute intrinsic takes a byte offset into LDS, not a lane index. Each lane occupies a 4-byte (dword) slot, so to read from lane N you
pass byte offset N * 4.
Compare with CUDA's __shfl_sync, which takes a plain lane index directly — no multiplication needed there.

[hughperkins]

reviewed this function => looks ok

[hughperkins]

wait I should check the args.

[hughperkins]

question to Opus:

what is Intrinsic::amdgcn_mbcnt_lo

It's an LLVM intrinsic for the AMD GCN v_mbcnt_lo_u32_b32 instruction — "masked bit count low." It counts the number of set bits in the low 32 bits of the execution mask for
lanes below the current lane, giving you the lane's position within the active wavefront.
The AMDGPU codegen uses it together with mbcnt_hi to compute the full lane ID:

quadrants/codegen/amdgpu/codegen_amdgpu.cpp lines 472-480

llvm::Value *emit_amdgpu_lane_id() {
  auto i32_ty = llvm::Type::getInt32Ty(*llvm_context);
  auto all_ones = tlctx->get_constant((int32)-1);
  auto zero = tlctx->get_constant((int32)0);
  auto lo = builder->CreateIntrinsic(i32_ty, Intrinsic::amdgcn_mbcnt_lo,
                                     {all_ones, zero});
  return builder->CreateIntrinsic(i32_ty, Intrinsic::amdgcn_mbcnt_hi,
                                  {all_ones, lo});
}

mbcnt_lo counts set bits in the low 32 lanes of the exec mask, mbcnt_hi counts the high 32 and adds the result of lo. Passing -1 (all ones) as the mask means "count all
lanes," so the result is just the thread's lane index within the wavefront (0–63). This is AMD's equivalent of threadIdx.x % warpSize on CUDA.

[hughperkins]

question to Opus:

what are the parametrs of amdgcn_mbcnt_hi ?

amdgcn_mbcnt_hi(mask, src1) takes two i32 parameters:

1. mask — a bitmask over the high 32 lanes (lanes 32–63) of the wavefront. It counts the number of set bits in this mask for lanes below the current one.
2. src1 — an addend that gets added to the bit count result.

So in the codegen:

quadrants/codegen/amdgpu/codegen_amdgpu.cpp lines 477-478

  return builder->CreateIntrinsic(i32_ty, Intrinsic::amdgcn_mbcnt_hi,
                                  {all_ones, lo});

• all_ones (-1) = count all high lanes
• lo = the result of mbcnt_lo, i.e. the count from the low 32 lanes

The result is popcount(hi_mask for lanes below me) + lo, which gives the full lane ID across the 64-lane wavefront. That's why mbcnt_lo is called first with 0 as its addend,
and its result is chained into mbcnt_hi's second parameter.

[hughperkins]

this function looks reasonable to me ✅

[hughperkins]

question to Opus:

is there a way of writing emit_subgroup_shuffle_amdgpu without writing in low level llvm code?

Yes — you could follow the same pattern CUDA uses. CUDA avoids raw CreateIntrinsic calls for shuffle by:

1. Defining stub functions in runtime.cpp (e.g. cuda_shfl_sync_i32 that returns 0)
2. Using patch_intrinsic to rewrite those stubs to the real LLVM intrinsics at module load time
3. In codegen, just calling call("cuda_shfl_sync_i32", mask, value, index, width)

AMDGPU already uses this pattern for a few things:

quadrants/runtime/llvm/llvm_context.cpp lines 578-582

    patch_intrinsic("thread_idx", llvm::Intrinsic::amdgcn_workitem_id_x);
    patch_intrinsic("block_idx", llvm::Intrinsic::amdgcn_workgroup_id_x);
    patch_intrinsic("block_barrier", llvm::Intrinsic::amdgcn_s_barrier,
                    false);
    patch_intrinsic("amdgpu_clock_i64", llvm::Intrinsic::amdgcn_s_memtime);

So you could add stubs like amdgpu_ds_bpermute_i32 and amdgpu_mbcnt to the runtime module, patch them to amdgcn_ds_bpermute / amdgcn_mbcnt_lo / amdgcn_mbcnt_hi, and then the
codegen would simplify to call("amdgpu_ds_bpermute_i32", byte_index, value).
That said, the 64-bit split/reassemble logic (truncate, shift, shuffle both halves, recombine) and the f32 bitcast would still live in the codegen — those are type-level
transformations, not just intrinsic calls. So it would clean up the leaf CreateIntrinsic calls but wouldn't eliminate the function entirely. The same is true for the CUDA side
— notice it also has the manual split logic for f64/i64.

[hughperkins]

question to Opus:

How would this function look if we did this?

The codegen functions would simplify to using call() instead of CreateIntrinsic:

llvm::Value *emit_amdgpu_lane_id() {
auto all_ones = tlctx->get_constant((int32)-1);
auto zero = tlctx->get_constant((int32)0);
auto lo = call("amdgpu_mbcnt_lo", all_ones, zero);
return call("amdgpu_mbcnt_hi", all_ones, lo);
}
llvm::Value *emit_subgroup_shuffle_amdgpu(llvm::Value *value,
DataType dt,
llvm::Value *byte_index) {
auto i32_ty = llvm::Type::getInt32Ty(*llvm_context);
if (dt->is_primitive(PrimitiveTypeID::i32) ||
dt->is_primitive(PrimitiveTypeID::u32)) {
return call("amdgpu_ds_bpermute", byte_index, value);
} else if (dt->is_primitive(PrimitiveTypeID::f32)) {
auto as_i32 = builder->CreateBitCast(value, i32_ty);
auto shuffled = call("amdgpu_ds_bpermute", byte_index, as_i32);
return builder->CreateBitCast(shuffled,
llvm::Type::getFloatTy(*llvm_context));
} else if (dt->is_primitive(PrimitiveTypeID::f64)) {
auto i64_ty = llvm::Type::getInt64Ty(*llvm_context);
auto f64_ty = llvm::Type::getDoubleTy(*llvm_context);
auto as_i64 = builder->CreateBitCast(value, i64_ty);
auto lo = builder->CreateTrunc(as_i64, i32_ty);
auto hi = builder->CreateTrunc(builder->CreateLShr(as_i64, 32), i32_ty);
lo = call("amdgpu_ds_bpermute", byte_index, lo);
hi = call("amdgpu_ds_bpermute", byte_index, hi);
auto lo_64 = builder->CreateZExt(lo, i64_ty);
auto hi_64 = builder->CreateZExt(hi, i64_ty);
auto combined = builder->CreateOr(builder->CreateShl(hi_64, 32), lo_64);
return builder->CreateBitCast(combined, f64_ty);
} else if (dt->is_primitive(PrimitiveTypeID::i64) ||
dt->is_primitive(PrimitiveTypeID::u64)) {
auto i64_ty = llvm::Type::getInt64Ty(*llvm_context);
auto lo = builder->CreateTrunc(value, i32_ty);
auto hi = builder->CreateTrunc(builder->CreateLShr(value, 32), i32_ty);
lo = call("amdgpu_ds_bpermute", byte_index, lo);
hi = call("amdgpu_ds_bpermute", byte_index, hi);
auto lo_64 = builder->CreateZExt(lo, i64_ty);
auto hi_64 = builder->CreateZExt(hi, i64_ty);
return builder->CreateOr(builder->CreateShl(hi_64, 32), lo_64);
} else {
QD_ERROR("subgroup shuffle: unsupported type {} on AMDGPU",
data_type_name(dt));
return nullptr;
}
}

And you'd wire it up with stubs in runtime.cpp:

i32 amdgpu_ds_bpermute(i32 byte_index, i32 value) { return 0; }
i32 amdgpu_mbcnt_lo(i32 mask, i32 base) { return 0; }
i32 amdgpu_mbcnt_hi(i32 mask, i32 base) { return 0; }

Plus patches in llvm_context.cpp alongside the existing AMDGPU ones:

patch_intrinsic("amdgpu_ds_bpermute", llvm::Intrinsic::amdgcn_ds_bpermute);
patch_intrinsic("amdgpu_mbcnt_lo", llvm::Intrinsic::amdgcn_mbcnt_lo);
patch_intrinsic("amdgpu_mbcnt_hi", llvm::Intrinsic::amdgcn_mbcnt_hi);

The change is modest — it just swaps builder->CreateIntrinsic(i32_ty, Intrinsic::amdgcn_ds_bpermute, {byte_index, value}) for call("amdgpu_ds_bpermute", byte_index, value).
The bitcast/trunc/shift/reassemble logic for 64-bit types stays the same either way since that's type manipulation, not intrinsic calls.

[hughperkins]

Question to Opus:

thats about the same number of lines. Anyway we can avoid having to write/see all these createTrunc, createZExt etc? some way we could write as a simple python or c++ program, rather than with builder calls? (or perhaps as ll ir instead of buidler calls?)

Good question. There are two approaches that would let you express this as readable code rather than builder calls.
Option 1: Write it in C in runtime.cpp
The most natural option. Write the split/recombine logic as normal C functions that call the stub amdgpu_ds_bpermute, which gets patched to the intrinsic. After inlining, the
result is identical IR.

// Stub — patched to llvm.amdgcn.ds.bpermute
i32 amdgpu_ds_bpermute(i32 byte_index, i32 value) { return 0; }
QD_ALWAYS_INLINE f32 amdgpu_shuffle_f32(i32 byte_index, f32 value) {
union { f32 f; i32 i; } u;
u.f = value;
u.i = amdgpu_ds_bpermute(byte_index, u.i);
return u.f;
}
QD_ALWAYS_INLINE f64 amdgpu_shuffle_f64(i32 byte_index, f64 value) {
union { f64 d; i64 i; } u;
u.d = value;
i32 lo = (i32)u.i;
i32 hi = (i32)(u.i >> 32);
lo = amdgpu_ds_bpermute(byte_index, lo);
hi = amdgpu_ds_bpermute(byte_index, hi);
u.i = ((i64)(u32)lo) | ((i64)hi << 32);
return u.d;
}

Then the codegen collapses to:

llvm::Value *emit_subgroup_shuffle_amdgpu(llvm::Value *value,
DataType dt,
llvm::Value *byte_index) {
if (dt->is_primitive(PrimitiveTypeID::i32) || dt->is_primitive(PrimitiveTypeID::u32))
return call("amdgpu_ds_bpermute", byte_index, value);
else if (dt->is_primitive(PrimitiveTypeID::f32))
return call("amdgpu_shuffle_f32", byte_index, value);
else if (dt->is_primitive(PrimitiveTypeID::f64))
return call("amdgpu_shuffle_f64", byte_index, value);
else if (dt->is_primitive(PrimitiveTypeID::i64) || dt->is_primitive(PrimitiveTypeID::u64))
return call("amdgpu_shuffle_i64", byte_index, value);
// ...
}

This works because patch_intrinsic replaces the stub's body with an intrinsic wrapper, and then LLVM's inliner propagates it into the callers. The same approach would work for
the CUDA f64/i64 split too — cuda_shfl_sync_f64 in C calling cuda_shfl_sync_i32 twice.
Option 2: Write it as .ll IR
You could write the functions as textual LLVM IR in a .ll file and link it in. More readable than builder calls, but still fairly verbose:

define double @amdgpu_shuffle_f64(i32 %idx, double %val) alwaysinline {
%i64 = bitcast double %val to i64
%lo = trunc i64 %i64 to i32
%hi = trunc i64 lshr(i64 %i64, i64 32) to i32
%lo2 = call i32 @llvm.amdgcn.ds.bpermute(i32 %idx, i32 %lo)
%hi2 = call i32 @llvm.amdgcn.ds.bpermute(i32 %idx, i32 %hi)
; ... reassemble ...
ret double %result
}

Option 1 is the clear winner — it's just C, it uses infrastructure that already exists (patch_intrinsic + always_inline), and the bit manipulation reads like normal code
instead of either builder calls or SSA IR.

[hughperkins]

lets move this code into the runtime too, in C, too, similar to for amdgpu.

[hughperkins]

this function looks ok to me ✅

[hughperkins]

Putting this into draft for now.

Not because it's not mergeable as-is, or that I haven't revieweed it yet (it is; I have).

But beacuse we don't have a clear use-case for it right now.

We can review/merge it if/when we do.

[hughperkins]

Thank you 🎉