Skip to content
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.

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

[LinalgExt] Generalize attribute setting for attention decomposition #18780

Merged
merged 3 commits into from
Oct 28, 2024

Conversation

Groverkss
Copy link
Contributor

@Groverkss Groverkss commented Oct 15, 2024

This PR teaches attention decomposition to set attributes for attention matmuls by passing attribute dictionaries to iree_linalg_ext.online_attention operation. This allows us to further control codegen of matmuls (generally the root operations) after decomposition (for example, setting lowering config on the decompose matmuls).

Comment on lines 290 to 313
auto qkAttrs = (*this)->getAttrOfType<DictionaryAttr>("qk_attrs");
auto pvAttrs = (*this)->getAttrOfType<DictionaryAttr>("pv_attrs");
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is better to declare constant variables for qk/pv_attrs because they are also used in KernelConfig.cpp. Also, we can get rid of magic numbers/strings. We can declare the constant variables in MarkerUtils.h. What do you think?

Copy link
Contributor Author

@Groverkss Groverkss Oct 21, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think we can do that. MarkerUtils is part of Codegen/, while LinalgExt is part of Dialects/. I can make it part of the attention op definition as well I guess. But it doesnt really make sense to be part of op definition as well? It's codegen specific.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point, I missed that! Can we add some getters to extraClassDeclaration? The concern is that the magic strings are floating in the codebase, which looks brittle to me. It looks better if we can centralize to somewhere, e.g., we can do something like:

// They share the same entry which could be "config".
void setConfig(DictionaryAttr attr) { ... }
DictionaryAttr getConfig() { return (*this)->getAttrOfType<DictionaryAttr>("config"); }
bool hasConfig(StringRef str) {
  auto config = getConfig();
  if (!config) return false;
  return config.getNamed(str);
}

// Marker-like things
static StringRef getQKAttrStr() { return "qk_attrs"; }
static StringRef getPVAttrStr() { return "pv_attrs"; }

let extraClassDeclaration = [{
// Method to implement for specifying output range for
// DestinationStyleOpInterface
MutableOperandRange getDpsInitsMutable();
SmallVector<AffineMap> getIndexingMapsArray();
AffineMap getQueryMap() {
return cast<AffineMap>(getIndexingMapsArray()[0]);
}
AffineMap getKeyMap() {
return cast<AffineMap>(getIndexingMapsArray()[1]);
}
AffineMap getValueMap() {
return cast<AffineMap>(getIndexingMapsArray()[2]);
}
AffineMap getScaleMap() {
return cast<AffineMap>(getIndexingMapsArray()[3]);
}
std::optional<AffineMap> getMaskMap() {
if (getMask()) {
return cast<AffineMap>(getIndexingMapsArray()[4]);
}
return std::nullopt;
}
AffineMap getOutputMap() {
return cast<AffineMap>(getIndexingMapsArray()[getNumDpsInputs()]);
}
int64_t getIterationDomainRank() {
return getQueryMap().getNumDims();
}
}];

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added a decompositon_config attribute to AttentionOp, which acts as the config attribute you asked for above. I also added getQKAttrStr/getPVAttrStr to extraClassDeclarations.

Copy link
Contributor

@hanhanW hanhanW left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, just one final question

Copy link
Contributor

@hanhanW hanhanW left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see, thanks!

@Groverkss Groverkss enabled auto-merge (squash) October 28, 2024 18:46
@Groverkss Groverkss merged commit e66171a into iree-org:main Oct 28, 2024
36 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants