Problem
Transform Dialect is a significant part of our production compiler infrastructure.
Recently, we are working on a pass that has transform op execution embedded within it. That is, inside this pass, we would like to first initialize some state variables. These state variables are then passed to the transform ops to be further processed. Afterwards, these state variables are post-processed by the same pass again. We encountered a problem with this interaction between the pass and the transform ops.
Currently, we call applyTransforms to run the transform ops programmatically. Since we also perform analysis on the IR, there are information that we wish to pass into the transform ops as part of the transform state. There seems to be no good way to pass information into the transform ops unless we hack the upstream code.
Current Methodology
We currently create a CustomTransformState class inheriting from the TransformState class (a simplified example is shown below) that includes extra data members to store the additional information.
We then create our customized customApplyTransforms function that will create a CustomTransformState state object with the additional information. As an CustomTransformState object is being constructed, the constructor of TransformState is being invoked. Since the TransformState’s constructor is made private, we had to change it to protected in order to make this to work.
Inside customApplyTransforms, we would then call state.applyTransform which will eventually call the apply function for each transform op. Each transform op is required to downcast the state input argument from TransformState to `CustomTransformState before accessing the additional information.
// schedules and constants are the additional information fields.
class CustomTransformState : public TransformState {
friend LogicalResult customApplyTransforms(
Operation *, TransformOpInterface, const RaggedArray<MappedValue> &,
const TransformOptions &,
DenseMap<Operation *, SmallVector<int>> *schedules,
DenseMap<Operation*, int> *constants);
private:
CustomTransformState(Region *region, Operation *payloadRoot,
const RaggedArray<MappedValue> &extraMappings = {},
const TransformOptions &options = TransformOptions()) :
TransformState(region, payloadRoot, extraMappings, options) {}
public:
DenseMap<Operation *, SmallVector<int>> *schedules;
DenseMap<Operation*, int> *constants;
};
// customApplyTransforms creates a CustomTransformState object.
// schedules and constants are initialized by the pass that invokes customApplyTransforms.
LogicalResult transform::customApplyTransforms(
Operation *payloadRoot, TransformOpInterface transform,
const RaggedArray<MappedValue> &extraMapping,
const TransformOptions &options,
DenseMap<Operation *, SmallVector<int>> *schedules,
DenseMap<Operation *, int> *constants) {
CustomTransformState state(transform->getParentRegion(), payloadRoot,
extraMapping, options);
state.schedules = schedules;
state.constants = constants;
return state.applyTransform(transform).checkAndReport();
}
void CustomPass::runOnOperation() {
// declare variables
DenseMap<Operation *, SmallVector<int>> schedules;
DenseMap<Operation *, int> constants;
// initialize schedules and constants
// ...
// call the transform ops
transform::customApplyTransforms(module, op, extraMappings, options, schedules, constants);
// post process schedules and constants
// ...
}
Other Approaches
We also tried to use the TransformState Extension which seems to be most relevant. However, the extension mechanism allows states to be communicated amongst the transform ops only, while we also want the ability to pass information in and out from the pass level.
The change proposed in this RFC is very simple (i.e. a single word change from private->protected 
) but we’re not sure if there are some design points that we missed. Would be really grateful to hear guidance and feedbacks from upstream!
Thank you!