[llgo] irgen: generate switch instructions

[axw]

Summary

With this patch, llgo uses ssautil.Switches
to reconstitute (and synthesise) switches,
which can then be lowered to lookup tables,
trees, etc.

We currently only handle integer const case
switches. We erase the comparison blocks (other
than the initial block), and generate a switch
instruction at the end of the block starting
the if-else-if chain. ssautil.Switches does
not remove duplicate const cases (e.g. same
operands for "||"), so we do this in llgo for
now.

[pcc]

I wonder if there is a way to do this without modifying the go/ssa IR. In particular, I don't feel very comfortable about any attempt to define an instruction outside the ssa package. Maybe we can build a mapping to keep track of which ssa.If instructions are to be treated as switch instructions?

I believe LLVM already has a transformation to convert chains of branch instructions into switches, at least in some cases. Have you measured a performance improvement with this change at all? (Asking just out of curiosity. I think it would be worth exploring the performance impact of extending this to support type switches, so I have no objection to this change in principle.)

[axw]

In D6831#105921, @pcc wrote:

I wonder if there is a way to do this without modifying the go/ssa IR. In particular, I don't feel very comfortable about any attempt to define an instruction outside the ssa package. Maybe we can build a mapping to keep track of which ssa.If instructions are to be treated as switch instructions?

It is certainly possible to do that, and I did go down that road initially. I don't really like the kludginess of the current approach either, but the reason I did it this way was to keep it relatively self-contained. Aside from being (IMHO) a bit messy, adding a check to each "ssa.If" would mean we have a performance hit in the frontend for each If instruction, rather than just the ones that we care about. Maybe that's nothing to worry about; I haven't measured it.

I don't know if this will ever be necessary, but another option is for llgo to grow its own IR on top of go/ssa. I'd rather not go there for this, though.

I believe LLVM already has a transformation to convert chains of branch instructions into switches, at least in some cases. Have you measured a performance improvement with this change at all? (Asking just out of curiosity. I think it would be worth exploring the performance impact of extending this to support type switches, so I have no objection to this change in principle.)

Geo-mean improvement of 0.49% (0.23% - 0.75% @ 95% CI)

There seems to be a fair bit of noise in some of the tests, since I'm getting supposed improvements in BenchmarkAppendGrowString, and BenchmarkTCP6ConcurrentReadWrite. One significant improvement that looks genuine is:

BenchmarkCSSEscaper 2983.000000 2476.000000 1.204766

Standard optimisations do not appear to convert branches to switches AFAICT, but perhaps I'm doing something wrong. I did look to see if there was an existing pass, but again I couldn't see anything that looked relevant. There is a pass that goes the opposite direction, converting switches to branches, for targets that don't implement switch.

[pcc]

adding a check to each "ssa.If" would mean we have a performance hit in the frontend for each If instruction

I very much doubt that the performance impact would be measurable. Remember that the switch analysis has to look at each if instruction anyway.

Regardless, I agree that both approaches are a little messy, and I can't see a neater way without modifying go/ssa, so LGTM.

I don't know if this will ever be necessary, but another option is for llgo to grow its own IR on top of go/ssa. I'd rather not go there for this, though.

Agreed. FWIW, the long term direction I'd like to see is attempting to extend go/ssa with what we need (in this particular case, a native switch instruction).

I did look to see if there was an existing pass, but again I couldn't see anything that looked relevant

If you grep for SimplifyCFGOpt::FoldValueComparisonIntoPredecessors you should find something.

Geo-mean improvement of 0.49% (0.23% - 0.75% @ 95% CI)

Great! I'm a little surprised that we can do better than LLVM on its own, but maybe I don't understand what the simplifycfg code is doing.

irgen/ssa.go

func (u *unit) defineFunction(f *ssa.Function) {
   }
 
   fr.blocks = make([]llvm.BasicBlock, len(f.Blocks))
   fr.lastBlocks = make([]llvm.BasicBlock, len(f.Blocks))
   for i, block := range f.Blocks {
     fr.blocks[i] = llvm.AddBasicBlock(fr.function, fmt.Sprintf(".%d.%s", i, block.Comment))
   }
   fr.builder.SetInsertPointAtEnd(fr.blocks[0])
+  fr.transformSwitches(f)
 
   prologueBlock := llvm.InsertBasicBlock(fr.blocks[0], "prologue")
   fr.builder.SetInsertPointAtEnd(prologueBlock)
 
   // Map parameter positions to indices. We use this
   // when processing locals to map back to parameters
   // when generating debug metadata.
   paramPos := make(map[token.Pos]int)
if !isMethod && f.Name() == "init" {
     registerGcBlock = fr.emitInitPrologue()
     fr.isInit = true
   }
 
   fr.builder.CreateBr(fr.blocks[0])
   fr.allocaBuilder.SetInsertPointBefore(prologueBlock.FirstInstruction())
 
   for _, block := range f.DomPreorder() {
-    fr.translateBlock(block, fr.blocks[block.Index])
+    llblock := fr.blocks[block.Index]
+    if llblock.IsNil() {
+      continue
+    }
+    fr.translateBlock(block, llblock)
   }
 
   fr.fixupPhis()
 
   if !fr.unwindBlock.IsNil() {
     fr.setupUnwindBlock(f.Recover, f.Signature.Results())
   }
 
case *ssa.Store:
     // If this is the init function, see if we can simulate the effect
     // of the store in a global's initializer, in which case we can avoid
     // generating code for it.
     if !fr.isInit || !fr.maybeStoreInInitializer(value, addr) {
       fr.nilCheck(instr.Addr, addr)
       fr.builder.CreateStore(value, addr)
     }
 
+  case *switchInstr:
+    fr.emitSwitch(instr)
+
   case *ssa.TypeAssert:
     x := fr.value(instr.X)
     if instr.CommaOk {
       v, ok := fr.interfaceTypeCheck(x, instr.AssertedType)
       fr.tuples[instr] = []*govalue{v, ok}
     } else {
       fr.env[instr] = fr.interfaceTypeAssert(x, instr.AssertedType)
     }

irgen/switches.go

+//===- switches.go - misc utils -------------------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements transformations and IR generation for switches.
+//
+//===----------------------------------------------------------------------===//
+
+package irgen
+
+import (
+  "go/token"
+
+  "llvm.org/llgo/third_party/go.tools/go/exact"
+  "llvm.org/llgo/third_party/go.tools/go/ssa"
+  "llvm.org/llgo/third_party/go.tools/go/ssa/ssautil"
+  "llvm.org/llvm/bindings/go/llvm"
+)
+
+// switchInstr is an instruction representing a switch on constant
+// integer values.
+type switchInstr struct {
+  ssa.Instruction
+  ssautil.Switch
+}
+
+func (sw *switchInstr) String() string {
+  return sw.Switch.String()
+}
+
+func (sw *switchInstr) Parent() *ssa.Function {
+  return sw.Default.Instrs[0].Parent()
+}
+
+func (sw *switchInstr) Block() *ssa.BasicBlock {
+  return sw.Start
+}
+
+func (sw *switchInstr) Operands(rands []*ssa.Value) []*ssa.Value {
+  return nil
+}
+
+func (sw *switchInstr) Pos() token.Pos {
+  return token.NoPos
+}
+
+// emitSwitch emits an LLVM switch instruction.
+func (fr *frame) emitSwitch(instr *switchInstr) {
+  cases, _ := dedupConstCases(fr, instr.ConstCases)
+  ncases := len(cases)
+  elseblock := fr.block(instr.Default)
+  llswitch := fr.builder.CreateSwitch(fr.llvmvalue(instr.X), elseblock, ncases)
+  for _, c := range cases {
+    llswitch.AddCase(fr.llvmvalue(c.Value), fr.block(c.Body))
+  }
+}
+
+// transformSwitches replaces the final If statement in start blocks
+// with a high-level switch instruction, and erases chained condition
+// blocks.
+func (fr *frame) transformSwitches(f *ssa.Function) {
+  for _, sw := range ssautil.Switches(f) {
+    if sw.ConstCases == nil {
+      // TODO(axw) investigate switch
+      // on hashes in type switches.
+      continue
+    }
+    if !isInteger(sw.X.Type()) && !isBoolean(sw.X.Type()) {
+      // LLVM switches can only operate on integers.
+      continue
+    }
+    instr := &switchInstr{Switch: sw}
+    sw.Start.Instrs[len(sw.Start.Instrs)-1] = instr
+    for _, c := range sw.ConstCases[1:] {
+      fr.blocks[c.Block.Index].EraseFromParent()
+      fr.blocks[c.Block.Index] = llvm.BasicBlock{}
+    }
+
+    // Fix predecessors in successor blocks for fixupPhis.
+    cases, duplicates := dedupConstCases(fr, instr.ConstCases)
+    for _, c := range cases {
+      for _, succ := range c.Block.Succs {
+        for i, pred := range succ.Preds {
+          if pred == c.Block {
+            succ.Preds[i] = sw.Start
+            break
+          }
+        }
+      }
+    }
+
+    // Remove redundant edges corresponding to duplicate cases
+    // that will not feature in the LLVM switch instruction.
+    for _, c := range duplicates {
+      for _, succ := range c.Block.Succs {
+        for i, pred := range succ.Preds {
+          if pred == c.Block {
+            head := succ.Preds[:i]
+            tail := succ.Preds[i+1:]
+            succ.Preds = append(head, tail...)
+            removePhiEdge(succ, i)
+            break
+          }
+        }
+      }
+    }
+  }
+}
+
+// dedupConstCases separates duplicate const cases.
+//
+// TODO(axw) fix this in go/ssa/ssautil.
+func dedupConstCases(fr *frame, in []ssautil.ConstCase) (unique, duplicates []ssautil.ConstCase) {
+  unique = make([]ssautil.ConstCase, 0, len(in))
+dedup:
+  for i, c1 := range in {
+    for _, c2 := range in[i+1:] {
+      if exact.Compare(c1.Value.Value, token.EQL, c2.Value.Value) {
+        duplicates = append(duplicates, c1)
+        continue dedup
+      }
+    }
+    unique = append(unique, c1)
+  }
+  return unique, duplicates
+}
+
+// removePhiEdge removes the i'th edge from each PHI
+// instruction in the specified basic block.
+func removePhiEdge(bb *ssa.BasicBlock, i int) {
+  for _, instr := range bb.Instrs {
+    instr, ok := instr.(*ssa.Phi)
+    if !ok {
+      return
+    }
+    head := instr.Edges[:i]
+    tail := instr.Edges[i+1:]
+    instr.Edges = append(head, tail...)
+  }
+}

test/irgen/switch.go

+// RUN: llgo -S -emit-llvm -o - %s | FileCheck %s
+
+package foo
+
+// CHECK: switch i32
+// CHECK-NEXT: i32 0, label %[[L0:.*]]
+// CHECK-NEXT: i32 1, label %[[L1:.*]]
+// CHECK-NEXT: i32 2, label %[[L2:.*]]
+// CHECK-NEXT: ]
+// CHECK: [[L0]]:
+// CHECK-NEXT: ret i32 1
+// CHECK: [[L1]]:
+// CHECK-NEXT: ret i32 2
+// CHECK: [[L2]]:
+// CHECK-NEXT: ret i32 0
+func F1(x int32) int32 {
+  switch x {
+  case 0:
+    return 1
+  case 1:
+    return 2
+  case 2:
+    return 0
+  }
+  panic("unreachable")
+}
+
+// CHECK: switch i64
+// CHECK-NEXT: i64 0
+// CHECK-NEXT: ]
+// CHECK: icmp eq i64 {{.*}}, 1
+func F2(x int64) bool {
+  return x == 0 || x == 0 || x == 1
+}
+
+// CHECK: switch i64
+// CHECK-NEXT: i64 0
+// CHECK-NEXT: ]
+func F3(x int64) bool {
+  return x == 0 || x == 0 || x == 0
+}
+
+// CHECK: switch i64
+// CHECK-NEXT: i64 0
+// CHECK-NEXT: i64 1
+// CHECK-NEXT: i64 2
+// CHECK-NEXT: ]
+// CHECK: icmp eq i64 {{.*}}, 3
+func F4(x int64) bool {
+  return x == 0 || x == 1 || x == 2 || x == 3
+}
+
+// CHECK-NOT: switch double
+func F5(x float64) float64 {
+  switch x {
+  case 0:
+    return 1.0
+  case 1.0:
+    return 0
+  }
+  panic("unreachable")
+}