[RFC] Add std.switch and scf.switch

Branching off the discussion in Switch statement missing, I’d like to add switch ops to the standard and structured control flow dialects.

std.switch:
This would be unstructured control flow based on llvm.switch. Given that branch weights are not present on other std branch operations, I do not think we should add them here at this time. In all other aspects, llvm.switch seems like the right way to model this and I anticipate mostly copying over the code (including the custom switch cases printer/parser).

I believe it should be ok to add this new op to std, despite the ongoing effort to break std into multiple dialects because it is closely related to the existing br and cond_br ops and would move along with them to whatever new dialect we might create (say a “control flow” dialect for unstructured control flow).

scf.switch:
This would be the structured version of switch, using regions, as in scf.if and able to yield values, but taking no block arguments.

%x, %y = scf.switch %flag -> (f32, f32) {
    %x_default = ...
    %y_default = ...
    %scf.yield %x_default, %y_default
  } [
    0: {
      %x_0 = ...
      %y_0 = ...
      %scf.yield %x_0, %y_0
    }
    -1: {
      %x_m1 = ...
      %y_m1 = ...
      %scf.yield %x_m1, %y_m1
    }
  ]

+1 to these. switch is as natural of a control flow construct as br, cond_br, etc. It’d be nice to finally get this in tree.

– River

std.switch LGTM!
About scf.switch, I wonder about the type of the condition and how pluggable/extensible we should make it? I.e. have you thought about non-integer matching.

Makes sense to me too

+1. These make sense to add.

Looks great! Minor question: for multiple cases sharing the same block, does llvm.switch avoid block replication (in memory)? The pretty printer can of course be changed to something like:

...  
)} [
1, 2, 3, 4, 5, 6, 7: {
      %x_0 = ...
      %y_0 = ...
      %scf.yield %x_0, %y_0
    }
0: {
  ...

if $blockDestinations could have the same blocks repeating.

I think your question would be about scf.switch, since it’s using regions? llvm.switch (and the proposed std.switch) just record a block identifier and the arguments, so I’m not sure there’s much value deduping for the complexity it adds, but for the proposed scf.switch, which would contain entire regions, that does seem potentially useful. Maybe each case should have an array of attributes of the flag type rather than a single one.

Given the response, I’m inclined to start the patch for std.switch and continue discussing scf.switch, so we can further explore Mehdi and Uday’s comments

1 Like

Sorry for the delay here. ⚙ D99925 [MLIR] Add a switch operation to the standard dialect adds std.switch

Thanks to reviews from Uday and River, this is now landed. It diverges from llvm.switch in a couple other ways now, allowing any signless int to switch on and placing the default case within the list of cases with an explicit keyword:

  switch %flag : i32, [
    default: ^bb1(%a : i32),
    42: ^bb1(%b : i32),
    43: ^bb3(%c : i32)
  ]

For scf.switch, I propose we make it more general, adopting Mehdi and Uday’s suggestions above.

%x, %y = scf.switch %flag : !mydialect.string -> (f32, f32) [
  default: {
    %x_default = ...
    %y_default = ...
    %scf.yield %x_default, %y_default
  },
  "foo": {
    %x_0 = ...
    %y_0 = ...
    %scf.yield %x_0, %y_0
  }
  "bar", "baz": {
    %x_m1 = ...
    %y_m1 = ...
    %scf.yield %x_m1, %y_m1
  }
]

which would be a definition like:

def IfOp : SCF_Op<"switch",
      [DeclareOpInterfaceMethods<RegionBranchOpInterface>,
       SingleBlockImplicitTerminator<"scf::YieldOp">, RecursiveSideEffects,
       NoRegionArguments]> {

  let arguments = (ins AnyType:$flag
                       OptionalAttr<ArrayAttr>:$case_values,
                       OptionalAttr<I32ElementsAttr>:$case_value_offsets);
  let results = (outs Variadic<AnyType>:$results);
  let regions = (region SizedRegion<1>:$defaultRegion,
                        Variadic<SizedRegion<1>>:$caseRegions);
  // ...
}

Some particular points:

  • the default keyword will always be on its own. I don’t think there’s any reason we should allow a declared case value to share a region with the default case. It would never make sense to construct such an op statically and if identifying matching regions in a canonicalization, we would just delete them.
  • We can use offsets to group the case values (as with operands in the llvm and std switch). We could also make this an array of arrays def 2DArrayAttr : TypedArrayAttrBase<ArrayAttr, "array of arrays">. Not sure which will end up being more ergonomic.

PS: discourse syntax highlighting of the basic block args there is weird. I don’t suppose we can fix that?

How is the matching implemented? We need somehow to specify a comparison function or something for the type used to dispatch the switch? Or will this just make the lowering of scf.switch depends on conversion pattern provided by the dialect used? Like in your example !mydialect.string is used so !mydialect should also provide some lowering patterns? How will this all be plugged together? (i.e. we don’t have a clear “scf-to-std” path right now with this I think.

We can: this is a work in progress, it’ll take some time but we’re on it :slight_smile:

Implemented where? In the case of a constant (e.g. in canonicalization), I’d expect to use equality of the Attribute. For lowering, I think something that knows about the type would have to define a lowering. SCF would only have lowerings that use built-in types. For integers it could go straight to std.switch. For non-int, I guess it would have to be a series of ifs.

That said, I don’t have a use case for this generality at the moment. In fact it turns out that the mhlo.case ops I was dealing with all got folded away, so I don’t need any of this for my immediate work :stuck_out_tongue: I’m just following up on this since I proposed it previously and it seems like a missing piece that was of general interest. I think leaving things open to more pluggability is a good idea in general. But I’m fine sticking to built-in types for now. I do worry if we do that then lots of places will assume built in types (because general guidance is not to recheck things checked by the verifier) so extending this will become much harder.

Sure, but even then you’d still need to have an op that does the comparison for these types.

Yeah so this is what I’m asking. I think we should plan for this before adding this to the codebase.

Here I was talking specifically about the case of builtin types, which is the only set of lowerings that would be “part of” the SCF dialect. There we have comparators already.

I would prefer scf.switch sticks to the set of primitive types that std.switch supports. This makes it a reliable, understandable primitive that lowers fairly predictably.

We can generalize it with a TypeInterface later if necessary, but there’s a lot of concerns there, and we want to be demand driven. There are many, many way to lower switches on arbitrary types, (with very different requirements on the type), such as straight line equality comparisons (requires only “==”), balanced (or, with branch probabilities, Huffman-based) tree based lowerings (requires “<” operator), perfect hashing (requires hashing functionality). And what you do as a matter of cost model once you have those primitives varies (e.g. doing two levels of tree-based lowering followed by equality comparisons within a “bucket”).

1 Like

I’m onboard with that, with the caveat I mentioned earlier that I’m concerned that we frequently have something that is guaranteed by the verifier and so not tested anywhere else and then updating it in other places is much trickier when we relax the verifier. That seems like a broader discussion though.

Like I said, I don’t have any immediate use case at the moment, so I’m happy to implement the much-simpler thing :smiley: I just feel kind of obligated to implement this now because people were asking for it and I said I would do it :stuck_out_tongue:

This looks great to me! Reg. the “offsets” vs “array of arrays”: looks like the one that more conveniently gives us the block given a case value will be better?

Can you clarify whether you mean the proposal as originally written above or restricted to integer types?

One concern I have is that the easiest way to represent case values if they’re all integers is with an ElementsAttr, but if we want to later allow other types, this doesn’t extend nicely.

I’m not very sure one way or the other but I tend to agree with the point you make above on “something that is guaranteed by the verifier and so not tested anywhere else and then updating it in other places is much trickier when we relax the verifier”. It’d look like a big update when you need to remove the integer type restriction. So the proposal as originally written above sounds good to me.

Ok it appears we have some disagreement about whether it makes sense to restrict the op to only integers.

My position is: if we think it likely that we will extend this op to accept a wider range of types, we should construct it accordingly to begin with (accept AnyType, use ArrayAttr instead of ElementsAttr for the values) even if we don’t offer all the plumbing for this more general case. This avoids creating an invariant that we’re later going to break. So now someone doing a lowering that relies on the type being an integer (e.g. the lowering to std) will test that in the lowering rather than relying on the definition of the op. This has the disadvantage that everyone who uses the op needs to test this feature, but I think this is a) not very hard to do, and b) actually a good thing if we have any plans to extend the op. I don’t think we should run into issues of underspecification because the op itself remains well defined: it’s pretty clear what it means for any type (and we can easily make canonicalizations operate on any type). How it is lowered is a much more complicated question, as Sean mentions, but that would be up to someone defining a lowering pipeline to define, and it seems totally reasonable that they would use scf.switch for it. We don’t have a guarantee that someone using their weird dialect type will be able to lower scf.switch using passes defined in core, but that’s also true of std.constant today, given that it can create constants of any type.

I’m not dead set on this, so I’m happy to yield if others do not find this argument convincing, but I thought I’d state it a bit more clearly given that at least Uday thought it was worth considering :slightly_smiling_face:

Introducing things as deliberately constrained until further extension points are needed is the typical way we do it – it’s WAY easier to loosen overconstrained things than restrict underconstrained things that users have started to use in their own ways that we didn’t verify/constrain.

Like I said, extending this to non-integer types is going to be a very complex undertaking regardless, and ArrayAttr vs ElementsAttr doesn’t seem like a major stumbling block (C++'s type system will catch most issues at compile time). You can do whatever is most convenient for the patch.