[tune][xegpu] infrastructure for tuning, applied to XeGPU matmul example#63
Open
[tune][xegpu] infrastructure for tuning, applied to XeGPU matmul example#63
Conversation
tkarna
reviewed
Mar 5, 2026
tkarna
reviewed
Mar 5, 2026
Comment on lines
+69
to
+73
| assert 64 <= wg_m <= 256 and m % wg_m == 0 and wg_m % DPAS.M == 0 | ||
| assert 64 <= wg_n <= 256 and n % wg_n == 0 and wg_n % DPAS.N == 0 | ||
| assert 32 <= sg_m <= 128 and m % sg_m == 0 and sg_m % DPAS.M == 0 | ||
| assert 32 <= sg_n <= 128 and n % sg_n == 0 and sg_n % DPAS.N == 0 | ||
| assert 16 <= k_tile <= 50 and k % k_tile == 0 and k_tile % DPAS.K == 0 |
Contributor
There was a problem hiding this comment.
This is very nice! We do however need to have some mechanism to change the search space bounds from outside. E.g., sometimes you want to autotune with a larger search space. Maybe the bounds could be set in the constructor of the abstract schedule (i.e. one without concrete chosen param values)?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds ability to extract search spaces from schedules with knobs.
Adds transform interpreter semantics to
constrain_paramsop, which interprets the smt ops in its region so that it is possible to check constraints on int params as well as calculate new int params at interpreter time.Modifies the schedule for the XeGPU matmul (& mlp) example to embed its tuning problem.
lighthouse.tuneimplements:trace-ing of transform schedules and the SMT regions they include:Nodes, eachevaluate-able w.r.t. an environment, with leafs such asConstant, which represents just a constant in IR, and evals to a constantint, andKnob, the representative of atransform.tune.knob, which takes its value from the env while knowing what it's possible values are,while
Applydepends on otherNodes as it modelsValues produced by ops dependent on other valuesPredicatemodels which condition/constraint needs to be true, as a boolean-valued function onNodes, for execution to proceed passed a particular op.rewrite-ing of transform schedules, solely through setting theselectedattr onknobops and theselected_regionattr onalternativesops.enumerate-ing all valid assignments forknobandalternativestuneables.__main__-ing to take a .mlir schedule and derive all valid knob configurations and output the corresponding concrete/transform-interpreter interpretable schedules.Inside
lighthouse.dialects, add (extension) dialects:smt_ext: A wrapper forir.Values of!smt.inttype so we can support python operations on them (e.g. addition) also with integers.transform_tune_ext: A wrapper forir.Values produced directly bytransform.tune.knobops so we can do python operations on them, in particular add constraints, and a camel_caseknob() -> KnobValuehelpertransform_smt_ext: so we we can have a version oftransform.smt.constrain_paramswhich has transform-interpreter semantics: by tracing the body, containing smt ops, we get function we can applied to thetransform.paramwhich were arguments toconstrain_params. That is, this version ofconstrain_paramshas a properTransformOpInterfaceimplementation.