diff --git a/python/quadrants/lang/ast/ast_transformer.py b/python/quadrants/lang/ast/ast_transformer.py index 79c639761e..64acc8d79e 100644 --- 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) + + 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 index edec8978d8..046dfae7f1 100644 --- 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 index a910c1ff69..691acaef8a 100644 --- 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 get_spmd_info() override { auto thread_idx = builder->CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, ArrayRef{}); diff --git a/quadrants/codegen/llvm/codegen_llvm.cpp b/quadrants/codegen/llvm/codegen_llvm.cpp index 50e8a1599b..62ac93db4a 100644 --- 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 args; if (stmt->with_runtime_context) diff --git a/quadrants/inc/internal_ops.inc.h b/quadrants/inc/internal_ops.inc.h index bef62c0d05..1b331a82d5 100644 --- 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 index 97a6a33704..baaaa8522d 100644 --- 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 index 4e46f4b79b..3bf59fe609 100644 --- 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 M, - int max_reg) { +JITModule *JITSessionAMDGPU::add_module(std::unique_ptr 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 index 0ef45aa831..d3ece6fbc1 100644 --- 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 index 0a49e55375..fa108e49da 100644 --- 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 index df145e155e..32c7e73fb0 100644 --- 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] -