Skip to content

feat(tensilelite): reorder wmma #12767

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
chengchingwen wants to merge 1 commit into
base: develop
Choose a base branch
from
+42,675 −783
Open

feat(tensilelite): reorder wmma #12767

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 64 additions & 14 deletions projects/hipblaslt/tensilelite/rocisa/rocisa/include/instruction/instruction.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,13 +169,45 @@ namespace rocisa
}
}

// What a GIR action does. Values mirror StinkyTofu's GirActionKind; the converter casts.
enum class GirActionKind : uint8_t
{
Other,
Ds,
Tensor,
Wmma
};

// One shared-memory touch this instruction makes. It is a fact ABOUT THE INSTRUCTION, so it
// rides on it rather than in a side table the frame contract keys by action id.
struct GirAccess
{
bool isWrite = false;
int operand = -1; // id of the (operand, region) pair this touch names
int ring = 1;
int genId = -1; // -1 = not bound to a generation
int gdelta = 0;
int absolute = -1; // -1 = relative, not pinned
bool crossAgent = false;
};

struct GirActionData
{
uint64_t actionId = 0;
uint64_t anchorAction = 0;
GirActionKind kind = GirActionKind::Other;
std::vector<GirAccess> accesses;
};

struct Instruction : public Item
{
InstType instType;
std::string comment;
std::string instStr;
bool outputInlineAsm;
InstType instType;
std::string comment;
std::string instStr;
bool outputInlineAsm;
std::shared_ptr<MemTokenData> m_memToken;
// Complete GIR semantic action record carried with every physical realization.
std::shared_ptr<GirActionData> m_girAction;

Instruction(InstType instType, const std::string& comment = "")
: instType(instType)
Expand All @@ -191,9 +223,10 @@ namespace rocisa
, comment(other.comment)
, instStr(other.instStr)
, outputInlineAsm(other.outputInlineAsm)
, m_memToken(other.m_memToken
? std::make_shared<MemTokenData>(*other.m_memToken)
: nullptr)
, m_memToken(other.m_memToken ? std::make_shared<MemTokenData>(*other.m_memToken)
: nullptr)
, m_girAction(other.m_girAction ? std::make_shared<GirActionData>(*other.m_girAction)
: nullptr)
{
}

Expand All @@ -207,6 +240,20 @@ namespace rocisa
return m_memToken;
}

void setGirActionData(uint64_t actionId,
uint64_t anchorAction,
GirActionKind kind,
const std::vector<GirAccess>& accesses)
{
m_girAction = std::make_shared<GirActionData>(
GirActionData{actionId, anchorAction, kind, accesses});
}

std::shared_ptr<GirActionData> getGirActionData() const
{
return m_girAction;
}

std::shared_ptr<Item> clone() const override
{
std::cout << typeid(*this).name() << std::endl;
Expand Down Expand Up @@ -297,22 +344,25 @@ namespace rocisa
hasVgpr = true;
}
}
if(!hasVgpr){
if(!hasVgpr)
{
if(getVgprMsb() == -1)
// Base layer WA: -2 means no-vgpr inst
rocIsa::getInstance().setVgprMsb(-2);
return;
}
int newVal = msbSrc[0] + (msbSrc[1] << 2) + (msbSrc[2] << 4) + (msbDst << 6);
int oriVal = getVgprMsb();
if(newVal != oriVal && !outputInlineAsm){
if(newVal != oriVal && !outputInlineAsm)
{
// Base layer WA: need to store previous msb value in [15:8] bits.
//int setVal = oriVal < 0? newVal : newVal + (oriVal << 8);
// only set newVal until compiler support it
int setVal = newVal;
std::string msbStr = "s_set_vgpr_msb " + std::to_string(setVal);
std::string msbComment = std::string("src0: " + std::to_string(msbSrc[0]) + ", src1: " + std::to_string(msbSrc[1]) + \
", src2: " + std::to_string(msbSrc[2]) + ", dst: " + std::to_string(msbDst));
// only set newVal until compiler support it
int setVal = newVal;
std::string msbStr = "s_set_vgpr_msb " + std::to_string(setVal);
std::string msbComment = std::string(
"src0: " + std::to_string(msbSrc[0]) + ", src1: " + std::to_string(msbSrc[1])
+ ", src2: " + std::to_string(msbSrc[2]) + ", dst: " + std::to_string(msbDst));
msbStr = formatStr(false, msbStr, msbComment, false);
// Base layer WA: add a no-vgpr inst if oriVal is non-determined and right after label
if(oriVal == -1)
Expand Down
62 changes: 61 additions & 1 deletion projects/hipblaslt/tensilelite/rocisa/rocisa/src/instruction/instruction.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,13 +85,73 @@ void init_inst(nb::module_ m)
m_inst, "InstructionInputVector", "A vector of InstructionInputs.");

// Base class
nb::enum_<rocisa::GirActionKind>(m_inst, "GirActionKind")
.value("Other", rocisa::GirActionKind::Other)
.value("Ds", rocisa::GirActionKind::Ds)
.value("Tensor", rocisa::GirActionKind::Tensor)
.value("Wmma", rocisa::GirActionKind::Wmma);

nb::class_<rocisa::Instruction, rocisa::Item, rocisa::PyInstruction>(m_inst, "Instruction")
.def(nb::init<rocisa::InstType, const std::string&>())
.def_rw("instType", &rocisa::Instruction::instType)
.def_rw("comment", &rocisa::Instruction::comment)
.def_rw("memToken", &rocisa::Instruction::m_memToken)
.def("setMemToken", &rocisa::Instruction::setMemToken)
// None clears it: the member is a shared_ptr whose absence is meaningful, so a caller that
// wants no token must be able to say so.
.def("setMemToken", &rocisa::Instruction::setMemToken, nb::arg("token").none())
.def("getMemToken", &rocisa::Instruction::getMemToken)
.def(
"setGirActionData",
[](rocisa::Instruction& self,
uint64_t actionId,
uint64_t anchorAction,
rocisa::GirActionKind kind,
const nb::sequence& accesses) {
std::vector<rocisa::GirAccess> out;
for(auto item : accesses)
{
auto d = nb::cast<nb::dict>(item);
rocisa::GirAccess a;
a.isWrite = nb::cast<bool>(d["is_write"]);
a.operand = nb::cast<int>(d["operand"]);
a.ring = nb::cast<int>(d["ring"]);
a.genId = nb::cast<int>(d["gen"]);
a.gdelta = nb::cast<int>(d["gdelta"]);
a.absolute = nb::cast<int>(d["absolute"]);
a.crossAgent = nb::cast<bool>(d["cross_agent"]);
out.push_back(std::move(a));
}
self.setGirActionData(actionId, anchorAction, kind, out);
},
nb::arg("action_id"),
nb::arg("anchor_action"),
nb::arg("kind"),
nb::arg("accesses"))
.def("getGirActionData",
[](const rocisa::Instruction& self) -> nb::object {
auto data = self.getGirActionData();
if(!data)
return nb::none();
nb::dict result;
result["action_id"] = nb::int_(data->actionId);
result["anchor_action"] = nb::int_(data->anchorAction);
result["kind"] = nb::int_(static_cast<uint8_t>(data->kind));
nb::list accesses;
for(const rocisa::GirAccess& src : data->accesses)
{
nb::dict access;
access["is_write"] = nb::bool_(src.isWrite);
access["operand"] = nb::int_(src.operand);
access["ring"] = nb::int_(src.ring);
access["gen"] = nb::int_(src.genId);
access["gdelta"] = nb::int_(src.gdelta);
access["absolute"] = nb::int_(src.absolute);
access["cross_agent"] = nb::bool_(src.crossAgent);
accesses.append(std::move(access));
}
result["accesses"] = std::move(accesses);
return std::move(result);
})
.def("setInlineAsm", &rocisa::Instruction::setInlineAsm)
.def("getParams", &rocisa::Instruction::getParams)
.def("preStr", &rocisa::Instruction::preStr)
Expand Down
27 changes: 27 additions & 0 deletions projects/hipblaslt/tensilelite/rocisa/test/test_gir_action_data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
# Copyright Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT

from rocisa.instruction import GirActionKind, SSchedulingFence


def test_gir_action_getter_returns_complete_record():
inst = SSchedulingFence()
assert inst.getGirActionData() is None

accesses = [{
"is_write": False,
"operand": 3,
"ring": 2,
"gen": 1,
"gdelta": -1,
"absolute": -1,
"cross_agent": True,
}]
inst.setGirActionData(17, 9, GirActionKind.Wmma, accesses)

assert inst.getGirActionData() == {
"action_id": 17,
"anchor_action": 9,
"kind": 3,
"accesses": accesses,
}
42 changes: 28 additions & 14 deletions ...cts/hipblaslt/tensilelite/rocisa_stinkytofu_adaptor/rocisa_stinkytofu_adaptor/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,20 +34,8 @@
"instruction", "label", "macro", "register",
})

def __getattr__(name: str):
if name in _LAZY_SUBMODULES:
import importlib, sys as _sys
mod = importlib.import_module(f".{name}", __name__)
globals()[name] = mod
# Wire under ``rocisa.*`` so ``from rocisa.code import Module``
# finds the same object (the adapter is installed as ``rocisa``).
_sys.modules.setdefault(f"rocisa.{name}", mod)
return mod
if name in ("toStinkyTofuModule", "StinkyAsmModule", "CloneSpec"):
_ensure_stinky_types()
val = globals()[name]
return val
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
# ``__getattr__`` is defined once, after ``_ensure_stinky_types``. A second
# definition would replace this lookup and hide ``toStinkyTofuModule``.

_P = "rocisa"

Expand Down Expand Up @@ -460,6 +448,32 @@ def _mfma_types():
# ``hasattr(_rocisa, ...)``.


def __getattr__(name: str):
"""Lazy submodules, the conversion entry points, then stinkytofu re-exports.

``PipelineExtensionPoint`` and friends are defined by the stinkytofu module and surfaced
through ``rocisa`` on the native backend; forward them rather than restate the enum here.
"""
if name in _LAZY_SUBMODULES:
import importlib, sys as _sys
mod = importlib.import_module(f".{name}", __name__)
globals()[name] = mod
# Wire under ``rocisa.*`` so ``from rocisa.code import Module``
# finds the same object (the adapter is installed as ``rocisa``).
_sys.modules.setdefault(f"rocisa.{name}", mod)
return mod
if name in ("toStinkyTofuModule", "StinkyAsmModule", "CloneSpec"):
_ensure_stinky_types()
return globals()[name]
import stinkytofu as _st
try:
return getattr(_st, name)
except AttributeError:
raise AttributeError(
f"module {__name__!r} has no attribute {name!r} "
f"(not provided by the adapter and not found in stinkytofu)") from None


def hasStinkyTofuBackend() -> bool:
"""Return True if the standalone stinkytofu binding exposes arch probes."""
try:
Expand Down
Loading
Loading