Merge pull request #7 from openai/broadcastAxis-fix

Fix blocked layout parser
2022-07-15 08:39:49 -07:00
parent 65237f6117 63e6a85901
commit df940aaab0
1 changed files with 3 additions and 0 deletions
--- a/lib/Dialect/TritonGPU/IR/Dialect.cpp
+++ b/lib/Dialect/TritonGPU/IR/Dialect.cpp
@@ -65,6 +65,9 @@ static Attribute parseBlocked(AsmParser &parser, Type type) {
    } else if (attr.getName() == "order") {
      if (parseIntArrayAttr(parser, attr, order, "order").failed())
        return {};
+    } else if (attr.getName() == "broadcastAxis") {
+      if (parseIntArrayAttr(parser, attr, broadcastAxis, "broadcastAxis").failed())
+        return {};
    } else {
      parser.emitError(parser.getNameLoc(), "unexpected key: ")
          << attr.getName().strref();