masahi commented on code in PR #14817:
URL: https://github.com/apache/tvm/pull/14817#discussion_r1189440294
##########
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()));
+ const BufferLoadNode* load = call_node->args[0].as<BufferLoadNode>();
+ Var src_buffer_var = load->buffer->data;
+ const VarNode* src_buffer_node = src_buffer_var.get();
+ PrimExpr src_index = load->indices[0];
+ DataType src_ele_dtype = GetElementDataType(src_buffer_node);
+ spirv::SType src_ele_stype = builder_->GetSType(src_ele_dtype);
+ spirv::Value src_buffer_val = MakeValue(src_buffer_var);
+ spirv::SType src_ptr_type =
+ builder_->GetPointerType(src_ele_stype,
src_buffer_val.stype.storage_class);
+ ICHECK(var_map_.count(src_buffer_node));
+ spirv::Value src_ptr =
+ builder_->StructArrayAccess(src_ptr_type, var_map_[src_buffer_node],
MakeValue(src_index));
+ spirv::SType type_bool = builder_->GetSType(DataType::UInt(1));
+ spirv::Value t_val = builder_->UIntImm(type_bool, 1);
+ spirv::Value f_val = builder_->UIntImm(type_bool, 0);
+ spirv::Value loaded =
+ builder_->MakeValue(spv::OpCooperativeMatrixLoadNV, fragment_type,
src_ptr, stride_val,
+ (layout != "row_major") ? t_val : f_val);
+ builder_->MakeInst(spv::OpStore, dst_ptr, loaded,
spv::MemoryAccessMaskNone);
+ return spirv::Value();
+ } else if (op->op.same_as(builtin::tvm_mma_sync())) {
+ const VarNode* buffer_d = op->args[0].as<VarNode>();
+ const VarNode* buffer_a = op->args[2].as<VarNode>();
+ const VarNode* buffer_b = op->args[4].as<VarNode>();
+ const VarNode* buffer_c = op->args[6].as<VarNode>();
+ PrimExpr index_d = op->args[1];
+ PrimExpr index_a = op->args[3];
+ PrimExpr index_b = op->args[5];
+ tvm::tir::ExprDeepEqual expr_equal;
+ PrimExpr index_c = op->args[7];
+ bool is_equal = ((buffer_d == buffer_c) && expr_equal(index_d, index_c));
+ spirv::SType& fragment_type_d = fragment_info_[buffer_d].stype;
+ spirv::SType& fragment_type_a = fragment_info_[buffer_a].stype;
+ spirv::SType& fragment_type_b = fragment_info_[buffer_b].stype;
+ spirv::SType& fragment_type_c = fragment_info_[buffer_c].stype;
+ spv::StorageClass storage = fragment_info_[buffer_d].sclass;
+ spirv::SType ptr_type_d = builder_->GetPointerType(fragment_type_d,
storage);
+ spirv::SType ptr_type_a = builder_->GetPointerType(fragment_type_a,
storage);
+ spirv::SType ptr_type_b = builder_->GetPointerType(fragment_type_b,
storage);
+ spirv::SType ptr_type_c = builder_->GetPointerType(fragment_type_c,
storage);
+ spirv::Value ptr_d =
+ builder_->StructArrayAccess(ptr_type_d, var_map_[buffer_d],
MakeValue(index_d));
+ spirv::Value ptr_a =
+ builder_->StructArrayAccess(ptr_type_a, var_map_[buffer_a],
MakeValue(index_a));
+ spirv::Value ptr_b =
+ builder_->StructArrayAccess(ptr_type_b, var_map_[buffer_b],
MakeValue(index_b));
+ spirv::Value ptr_c =
+ is_equal ? ptr_d
+ : builder_->StructArrayAccess(ptr_type_c, var_map_[buffer_c],
MakeValue(index_c));
+ uint32_t mask = spv::MemoryAccessMaskNone;
+ spirv::Value loaded_a = builder_->MakeValue(spv::OpLoad, fragment_type_a,
ptr_a, mask);
+ spirv::Value loaded_b = builder_->MakeValue(spv::OpLoad, fragment_type_b,
ptr_b, mask);
+ spirv::Value loaded_c = builder_->MakeValue(spv::OpLoad, fragment_type_c,
ptr_c, mask);
Review Comment:
See the discussion in
https://github.com/apache/tvm/pull/14770#discussion_r1187486958. I was thinking
that we cannot load / store matrices to a buffer, but looking at this code I'm
realizing that that was a wrong assumption.
But even then, I prefer my solution that materializes individual matrices as
phi value. Can you share your thought to the discussion in #14770? Feel free to
add comments to other aspects as well, I'd really appreciate your feedback.
--
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]
