masahi commented on code in PR #14817:
URL: https://github.com/apache/tvm/pull/14817#discussion_r1195785617
##########
src/target/spirv/codegen_spirv.cc:
##########
@@ -726,5 +885,51 @@ void CodeGenSPIRV::VisitStmt_(const SeqStmtNode* op) {
void CodeGenSPIRV::VisitStmt_(const EvaluateNode* op) { MakeValue(op->value); }
+int CodeGenSPIRV::Stoi(const std::string& str) {
+ try {
+ return std::stoi(str);
+ } catch (std::invalid_argument& e) {
+ LOG(FATAL) << "Cannot convert \"" << str << "\" to int";
+ throw;
+ }
+}
+
+std::pair<int32_t, int32_t> CodeGenSPIRV::GetWmmaFragmentSize(const VarNode*
variable) {
Review Comment:
Please don't duplicate this with `codegen_cuda.cc`
##########
src/target/spirv/ir_builder.cc:
##########
@@ -500,7 +519,18 @@ SType IRBuilder::DeclareType(const DataType& dtype) {
t.id = id_counter_++;
t.type = dtype;
SType base_type = GetSType(dtype.element_of());
- ib_.Begin(spv::OpTypeVector).AddSeq(t, base_type,
dtype.lanes()).Commit(&global_);
+
+ if (row * col == 0) {
+ ICHECK((row == 0) && (col == 0));
+ ib_.Begin(spv::OpTypeVector).AddSeq(t, base_type,
dtype.lanes()).Commit(&global_);
+ } else {
+ Value v_row = GetSpecConst(GetSType(DataType::UInt(32)), row);
+ Value v_col = GetSpecConst(GetSType(DataType::UInt(32)), col);
Review Comment:
I don't get why we need specialization constants. Aren't they for
configuring values that can change (or values that are only known) at runtime?
##########
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:
Lower `address_of` in its own else if branch (see my PR)
--
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]
