mei-ye commented on code in PR #14817:
URL: https://github.com/apache/tvm/pull/14817#discussion_r1197165433
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -395,6 +396,135 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const CallNode* op)
{
LOG(FATAL) << "SPIR-V shader cannot make extern calls. Graph contains
extern \""
<< Downcast<StringImm>(op->args[0]) << "\"";
return spirv::Value();
+ } else if (op->op.same_as(builtin::tvm_fill_fragment())) {
+ ICHECK_EQ(op->args.size(), 6U);
+ const VarNode* buffer_node = op->args[0].as<VarNode>();
+ ICHECK(buffer_node && fragment_info_.count(buffer_node));
+ DataType ele_dtype = GetElementDataType(buffer_node);
+ ICHECK(ele_dtype.is_float()) << "Only floating point fragment accumulator
is supported";
+ spirv::SType ele_stype = builder_->GetSType(ele_dtype);
+ spirv::SType& fragment_type = fragment_info_[buffer_node].stype;
+ double init =
static_cast<uint64_t>(Downcast<FloatImm>(op->args[5])->value);
+ PrimExpr prim_index = op->args[4];
+ spirv::Value init_val = builder_->GetCompositeConst(ele_stype,
fragment_type, init);
+ spirv::SType ptr_type =
+ builder_->GetPointerType(fragment_type,
fragment_info_[buffer_node].sclass);
+ spirv::Value index = MakeValue(prim_index);
+ ICHECK(var_map_.count(buffer_node));
+ spirv::Value ptr = builder_->StructArrayAccess(ptr_type,
var_map_[buffer_node], index);
+ builder_->MakeInst(spv::OpStore, ptr, init_val, spv::MemoryAccessMaskNone);
+ return spirv::Value();
+
+ } else if (op->op.same_as(builtin::tvm_load_matrix_sync())) {
+ ICHECK_EQ(op->args.size(), 8U);
+ const VarNode* buffer_node = op->args[0].as<VarNode>();
+ ICHECK(buffer_node && fragment_info_.count(buffer_node));
+ spirv::SType& fragment_type = fragment_info_[buffer_node].stype;
+ PrimExpr dst_index = op->args[4];
+ PrimExpr src_ptr_expr = op->args[5];
+ int stride = static_cast<int>(Downcast<IntImm>(op->args[6])->value);
+ auto type_int = builder_->GetSType(DataType::Int(32));
+ spirv::Value stride_val = builder_->IntImm(type_int, stride);
+ std::string layout = (op->args[7].as<StringImmNode>())->value;
+ spirv::SType dst_ptr_type =
+ builder_->GetPointerType(fragment_type,
fragment_info_[buffer_node].sclass);
+ spirv::Value dst_ptr =
+ builder_->StructArrayAccess(dst_ptr_type, var_map_[buffer_node],
MakeValue(dst_index));
+ const CallNode* call_node = src_ptr_expr.as<CallNode>();
+ ICHECK(call_node && call_node->op.same_as(builtin::address_of()));
Review Comment:
I factor out the lowering of "address_of" in the latest commit. However,
calls to "address_of" is an argument of calls to "tvm_load_matrix_sync", like
below:
@tir.tvm_load_matrix_sync(meta[tir.Var][1], 16, 16, 16, 0,
@tir.address_of(X_shared_2: Buffer(X_shared, float16, [1], [],
scope="shared")[0], dtype=handle), 16, "row_major", dtype=handle)
The call to "address_of" is not recursively visited in: spirv::Value
CodeGenSPIRV::VisitExpr_(const CallNode* op)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
