ROCm · michaelselehov · Apr 24, 2026 · Apr 24, 2026 · Apr 24, 2026 · Apr 24, 2026
diff --git a/python/quadrants/lang/ast/ast_transformer.py b/python/quadrants/lang/ast/ast_transformer.py
@@ -905,13 +905,15 @@ def build_static_for(ctx: ASTTransformerFuncContext, node: ast.For, is_grouped:
 
     @staticmethod
     def build_range_for(ctx: ASTTransformerFuncContext, node: ast.For) -> None:
+        if len(node.iter.args) not in [1, 2, 3]:
+            raise QuadrantsSyntaxError(f"Range should have 1, 2, or 3 arguments, found {len(node.iter.args)}")
+        if len(node.iter.args) == 3:
+            return ASTTransformer.build_strided_range_for(ctx, node)
         with ctx.variable_scope_guard():
             loop_name = node.target.id
             ctx.check_loop_var(loop_name)
             loop_var = expr.Expr(ctx.ast_builder.make_id_expr(""))
             ctx.create_variable(loop_name, loop_var)
-            if len(node.iter.args) not in [1, 2]:
-                raise QuadrantsSyntaxError(f"Range should have 1 or 2 arguments, found {len(node.iter.args)}")
             if len(node.iter.args) == 2:
                 begin_expr = expr.Expr(build_stmt(ctx, node.iter.args[0]))
                 end_expr = expr.Expr(build_stmt(ctx, node.iter.args[1]))
@@ -940,6 +942,50 @@ def build_range_for(ctx: ASTTransformerFuncContext, node: ast.For) -> None:
             ctx.ast_builder.end_frontend_range_for()
         return None
 
+    @staticmethod
+    def build_strided_range_for(ctx, node):
+        """Desugar `for i in range(start, stop, step)` into a while loop.
+
+        The Quadrants IR does not natively support a step parameter in
+        range-for loops.  We lower `range(start, stop, step)` into::
+
+            i = start
+            while i < stop:   # (or i > stop when step < 0)
+                <body>
+                i = i + step
+        """
+        with ctx.variable_scope_guard():
+            loop_name = node.target.id
+
+            begin_expr = expr.Expr(build_stmt(ctx, node.iter.args[0]))
+            end_expr = expr.Expr(build_stmt(ctx, node.iter.args[1]))
+            step_expr = expr.Expr(build_stmt(ctx, node.iter.args[2]))
+
+            begin = qd_ops.cast(begin_expr, primitive_types.i32)
+            end = qd_ops.cast(end_expr, primitive_types.i32)
+            step = qd_ops.cast(step_expr, primitive_types.i32)
+
+            loop_var = impl.expr_init(begin)
+            ctx.create_variable(loop_name, loop_var)
+
+            with ctx.loop_scope_guard():
+                stmt_dbg_info = _qd_core.DebugInfo(ctx.get_pos_info(node))
+                ctx.ast_builder.begin_frontend_while(expr.Expr(1, dtype=primitive_types.i32).ptr, stmt_dbg_info)
+
+                cond = loop_var < end
+                impl.begin_frontend_if(ctx.ast_builder, cond, stmt_dbg_info)
+                ctx.ast_builder.begin_frontend_if_true()
+                ctx.ast_builder.pop_scope()
+                ctx.ast_builder.begin_frontend_if_false()
+                ctx.ast_builder.insert_break_stmt(stmt_dbg_info)
+                ctx.ast_builder.pop_scope()
+
+                build_stmts(ctx, node.body)
+
+                loop_var._assign(loop_var + step)
+                ctx.ast_builder.pop_scope()
+        return None
+
     @staticmethod
     def build_ndrange_for(ctx: ASTTransformerFuncContext, node: ast.For) -> None:
         with ctx.variable_scope_guard():

diff --git a/python/quadrants/lang/simt/subgroup.py b/python/quadrants/lang/simt/subgroup.py
@@ -147,6 +147,10 @@ def shuffle_xor(value, mask):
     pass
 
 
+def dpp_swap_pairs(value):
+    return impl.call_internal("subgroupDppSwapPairs", value, with_runtime_context=False)
+
+
 def shuffle_up(value, offset):
     return impl.call_internal("subgroupShuffleUp", value, offset, with_runtime_context=False)
 
@@ -188,4 +192,5 @@ def shuffle_down(value, offset):
     "shuffle_xor",
     "shuffle_up",
     "shuffle_down",
+    "dpp_swap_pairs",
 ]
diff --git a/quadrants/codegen/amdgpu/codegen_amdgpu.cpp b/quadrants/codegen/amdgpu/codegen_amdgpu.cpp
@@ -543,6 +543,15 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM {
     }
   }
 
+  void visit(InternalFuncStmt *stmt) override {
+    if (stmt->func_name == "subgroupDppSwapPairs") {
+      llvm_val[stmt] = emit_amdgpu_dpp_swap_pairs(llvm_val[stmt->args[0]],
+                                                  stmt->args[0]->ret_type);
+    } else {
+      TaskCodeGenLLVM::visit(stmt);
+    }
+  }
+
   void visit(BinaryOpStmt *stmt) override {
     auto op = stmt->op_type;
     auto ret_quadrants_type = stmt->ret_type;
@@ -584,6 +593,45 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM {
   }
 
  private:
+  llvm::Value *emit_amdgpu_dpp_swap_pairs(llvm::Value *value, DataType dt) {
+    auto *i32_ty = llvm::Type::getInt32Ty(*llvm_context);
+    auto *i1_ty = llvm::Type::getInt1Ty(*llvm_context);
+    auto *ctrl = llvm::ConstantInt::get(i32_ty, 0xB1);
+    auto *rmask = llvm::ConstantInt::get(i32_ty, 0xF);
+    auto *bmask = llvm::ConstantInt::get(i32_ty, 0xF);
+    auto *bctrl = llvm::ConstantInt::getFalse(i1_ty);
+
+    auto emit_dpp_32 = [&](llvm::Value *v) -> llvm::Value * {
+      auto *ty = v->getType();
+      return builder->CreateIntrinsic(
+          Intrinsic::amdgcn_update_dpp, {ty},
+          {llvm::Constant::getNullValue(ty), v, ctrl, rmask, bmask, bctrl});
+    };
+
+    if (dt->is_primitive(PrimitiveTypeID::i32) ||
+        dt->is_primitive(PrimitiveTypeID::u32) ||
+        dt->is_primitive(PrimitiveTypeID::f32)) {
+      return emit_dpp_32(value);
+    }
+    if (dt->is_primitive(PrimitiveTypeID::f64) ||
+        dt->is_primitive(PrimitiveTypeID::i64) ||
+        dt->is_primitive(PrimitiveTypeID::u64)) {
+      auto *i64_ty = llvm::Type::getInt64Ty(*llvm_context);
+      auto *i64_val = builder->CreateBitCast(value, i64_ty);
+      auto *lo = builder->CreateTrunc(i64_val, i32_ty);
+      auto *hi = builder->CreateTrunc(builder->CreateLShr(i64_val, 32), i32_ty);
+      lo = emit_dpp_32(lo);
+      hi = emit_dpp_32(hi);
+      auto *result = builder->CreateOr(
+          builder->CreateZExt(lo, i64_ty),
+          builder->CreateShl(builder->CreateZExt(hi, i64_ty), 32));
+      return builder->CreateBitCast(result, value->getType());
+    }
+    QD_ERROR("subgroupDppSwapPairs: unsupported type {} on AMDGPU",
+             data_type_name(dt));
+    return nullptr;
+  }
+
   std::tuple<llvm::Value *, llvm::Value *> get_spmd_info() override {
     auto thread_idx = builder->CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x,
                                                ArrayRef<llvm::Value *>{});

diff --git a/quadrants/codegen/llvm/codegen_llvm.cpp b/quadrants/codegen/llvm/codegen_llvm.cpp
@@ -2446,6 +2446,16 @@ void TaskCodeGenLLVM::visit(ClearListStmt *stmt) {
 }
 
 void TaskCodeGenLLVM::visit(InternalFuncStmt *stmt) {
+  if (stmt->func_name == "subgroupDppSwapPairs") {
+    QD_ERROR(
+        "Internal op \"{}\" requires a GPU backend (AMDGPU or CUDA). "
+        "Wrap the call site with a backend guard such as "
+        "qd.static(backend == gs.amdgpu) so the CPU path never reaches "
+        "it.",
+        stmt->func_name);
+    return;
+  }
+
   std::vector<llvm::Value *> args;
 
   if (stmt->with_runtime_context)

diff --git a/quadrants/inc/internal_ops.inc.h b/quadrants/inc/internal_ops.inc.h
@@ -27,6 +27,7 @@ PER_INTERNAL_OP(subgroupBarrier)
 PER_INTERNAL_OP(subgroupMemoryBarrier)
 PER_INTERNAL_OP(subgroupElect)
 PER_INTERNAL_OP(subgroupBroadcast)
+PER_INTERNAL_OP(subgroupDppSwapPairs)
 PER_INTERNAL_OP(subgroupSize)
 PER_INTERNAL_OP(subgroupInvocationId)
 PER_INTERNAL_OP(subgroupAdd)

diff --git a/quadrants/ir/type_system.cpp b/quadrants/ir/type_system.cpp
@@ -379,6 +379,7 @@ void Operations::init_internals() {
   PLAIN_OP(subgroupMemoryBarrier, i32_void, false);
   PLAIN_OP(subgroupElect, i32, false);
   POLY_OP(subgroupBroadcast, false, Signature({}, {ValueT, !u32}, ValueT));
+  POLY_OP(subgroupDppSwapPairs, false, Signature({}, {ValueT}, ValueT));
   PLAIN_OP(subgroupSize, i32, false);
   PLAIN_OP(subgroupInvocationId, i32, false);
   POLY_OP(subgroupAdd, false, Signature({}, {ValueT}, ValueT));

diff --git a/quadrants/runtime/amdgpu/jit_amdgpu.cpp b/quadrants/runtime/amdgpu/jit_amdgpu.cpp
@@ -17,8 +17,8 @@
 namespace quadrants {
 namespace lang {
 #if defined(QD_WITH_AMDGPU)
-JITModule *JITSessionAMDGPU ::add_module(std::unique_ptr<llvm::Module> M,
-                                         int max_reg) {
+JITModule *JITSessionAMDGPU::add_module(std::unique_ptr<llvm::Module> M,
+                                        int max_reg) {
   // HSACo caching
   auto cache_key = compute_module_cache_key(M.get());
   auto cache_it = hsaco_cache_.find(cache_key);

diff --git a/quadrants/runtime/llvm/llvm_runtime_executor.cpp b/quadrants/runtime/llvm/llvm_runtime_executor.cpp
@@ -143,7 +143,7 @@ LlvmRuntimeExecutor::LlvmRuntimeExecutor(CompileConfig &config,
     // magic number 32
     // I didn't find the relevant parameter to limit the max block num per CU
     // So ....
-    int query_max_block_per_cu{32};
+    int query_max_block_per_cu{8};
     if (config.max_block_dim == 0) {
       config.max_block_dim = query_max_block_dim;
     }

diff --git a/tests/python/test_ast_refactor.py b/tests/python/test_ast_refactor.py
@@ -399,15 +399,18 @@ def foo(x: qd.i32):
 def test_range_for_three_arguments():
     a = qd.field(qd.i32, shape=(10,))
 
-    with pytest.raises(qd.QuadrantsCompilationError, match="Range should have 1 or 2 arguments, found 3"):
-
-        @qd.kernel
-        def foo(x: qd.i32):
-            for i in range(3, 7, 2):
-                a[i] = x
+    @qd.kernel
+    def foo(x: qd.i32):
+        for i in range(3, 7, 2):
+            a[i] = x
 
-        x = 5
-        foo(x)
+    a.fill(0)
+    foo(5)
+    for i in range(10):
+        if i in (3, 5):
+            assert a[i] == 5
+        else:
+            assert a[i] == 0
 
 
 @test_utils.test(print_preprocessed_ir=True)

diff --git a/tests/python/test_exception.py b/tests/python/test_exception.py
@@ -108,9 +108,10 @@ def foo():
 @test_utils.test(print_full_traceback=False)
 def test_exception_in_node_with_body():
     frameinfo = getframeinfo(currentframe())
+
     @qd.kernel
     def foo():
-        for i in range(1, 2, 3):
+        for i in range():
             a = 1
             b = 1
             c = 1
@@ -121,9 +122,8 @@ def foo():
     lineno = frameinfo.lineno
     file = frameinfo.filename
     msg = f"""
-File "{file}", line {lineno + 3}, in foo:
-        for i in range(1, 2, 3):
-        ^^^^^^^^^^^^^^^^^^^^^^^^
-Range should have 1 or 2 arguments, found 3"""
+File "{file}", line {lineno + 4}, in foo:
+        for i in range():
+        ^^^^^^^^^^^^^^^^^
+Range should have 1, 2, or 3 arguments, found 0"""
     assert msg in e.value.args[0]
-