-
Notifications
You must be signed in to change notification settings - Fork 590
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
[Codegen][GPU] Add pass to expand multi_mma op shapes to intrinsic layout #18139
[Codegen][GPU] Add pass to expand multi_mma op shapes to intrinsic layout #18139
Conversation
…yout Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/ConcretizeMmaShapes.cpp
Outdated
Show resolved
Hide resolved
Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM, a couple quick comments
compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/ConcretizeMmaShapes.cpp
Show resolved
Hide resolved
compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/ConcretizeMmaShapes.cpp
Outdated
Show resolved
Hide resolved
let dependentDialects = [ | ||
"::mlir::tensor::TensorDialect", | ||
"::mlir::iree_compiler::IREE::GPU::IREEGPUDialect", | ||
"::mlir::iree_compiler::IREE::VectorExt::IREEVectorExtDialect", |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why do we need a VectorExt dep here?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
getConcreteMFMALayout
generates some PerDimLayoutAttr
which are unused, but I didn't want to refactor the whole function, so I just added the dependent dialect.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
gah, that's a pain. Can you add a TODO somewhere to drop that dependency?
compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/ConcretizeMmaShapes.cpp
Show resolved
Hide resolved
compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/ConcretizeMmaShapes.cpp
Show resolved
Hide resolved
Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
This PR adds a new pass to explicitly materialize the dimensions of intrinsic layouts for
iree_gpu.multi_mma
ops. This means adding an expand_shape on each of the inputs to go from theOpaqueMmaLayout
shape to theConcreteMmaLayout
shape. This makes it easy to extract the correct data from the tensors when it is time to distribute the multi_mma op to lanes, since the shape will match the number of offsets and sizes needed for the slice.