mei-ye commented on code in PR #14817:
URL: https://github.com/apache/tvm/pull/14817#discussion_r1198332931
##########
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:
OK, fix is checked in now. I recombined the fix of several issues and the
original code into one commit.
--
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]
