[llgo] Elide alloca for unused received values in select

[axw]

Summary

If a receive case in a select statement is not assigned to a named variable, then we can eliminate the alloca and copy at runtime.

[pcc]

irgen/channels.go

63–64

It looks like we don't need this struct definition any more.

64–66

Was there any particular reason to make this a top-level function rather than a method?

129

There could in principle be multiple extract referrers. With this code you end up checking only the first one.

[axw]

irgen/channels.go

63–64

Thanks, removed.

64–66

It feels a little odd for frame to own methods; the frame is the subject here, so it feels more natural to have it as a top-level function operating on a frame. I'll change it back for now, for consistency's sake.

129

The loop and index check should take care of multiple extract instructions. I don't think there would ever be multiple extract referrers with the same index, would there?

[pcc]

LGTM. Comment is mostly nitpicking, this is fine as is.

irgen/channels.go

64–66

Right, if we decide this is worth changing we should do it globally rather than be inconsistent.

129

Yes, I meant multiple extract referrers with the same index. I agree that we probably won't ever see this, but it wouldn't hurt to check for them.

[axw]

Check all Extract referrers with a matching index, not just the first one.

irgen/channels.go

129

Okay, done.

[axw]

Add missing test file.

[axw]

Merge.

irgen/channels.go

 //===- channels.go - IR generation for channels ---------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements IR generation for channels.
 //
 //===----------------------------------------------------------------------===//
 
 package irgen
 
 import (
+  "llvm.org/llgo/third_party/go.tools/go/ssa"
   "llvm.org/llgo/third_party/go.tools/go/types"
   "llvm.org/llvm/bindings/go/llvm"
 )
 
 // makeChan implements make(chantype[, size])
 func (fr *frame) makeChan(chantyp types.Type, size *govalue) *govalue {
   // TODO(pcc): call __go_new_channel_big here if needed
   dyntyp := fr.types.ToRuntime(chantyp)
func (fr *frame) chanRecv(ch *govalue, commaOk bool) (x, ok *govalue) {
   x = newValue(fr.builder.CreateLoad(ptr, ""), elemtyp)
   return
 }
 
 // chanClose implements close(ch)
 func (fr *frame) chanClose(ch *govalue) {
   fr.runtime.builtinClose.call(fr, ch.value)
 }
 
-// selectState is equivalent to ssa.SelectState
-type selectState struct {
-  Dir  types.ChanDir
-  Chan *govalue
-  Send *govalue
-}
-
-func (fr *frame) chanSelect(states []selectState, blocking bool) (index, recvOk *govalue, recvElems []*govalue) {
-  n := uint64(len(states))
-  if !blocking {
+func (fr *frame) chanSelect(sel *ssa.Select) (index, recvOk *govalue, recvElems []*govalue) {
+  n := uint64(len(sel.States))
+  if !sel.Blocking {
     // non-blocking means there's a default case
     n++
   }
   size := llvm.ConstInt(llvm.Int32Type(), n, false)
   selectp := fr.runtime.newSelect.call(fr, size)[0]
 
   // Allocate stack for the values to send and receive.
-  //
-  // TODO(axw) check if received elements have any users, and
-  // elide stack allocation if not (pass nil to recv2 instead.)
-  ptrs := make([]llvm.Value, len(states))
-  for i, state := range states {
+  ptrs := make([]llvm.Value, len(sel.States))
+  for i, state := range sel.States {
     chantyp := state.Chan.Type().Underlying().(*types.Chan)
     elemtyp := fr.types.ToLLVM(chantyp.Elem())
-    ptrs[i] = fr.allocaBuilder.CreateAlloca(elemtyp, "")
     if state.Dir == types.SendOnly {
-      fr.builder.CreateStore(state.Send.value, ptrs[i])
+      ptrs[i] = fr.allocaBuilder.CreateAlloca(elemtyp, "")
+      fr.builder.CreateStore(fr.llvmvalue(state.Send), ptrs[i])
+    } else {
+      // Only allocate stack space if the received value is used.
+      used := chanSelectStateUsed(sel, len(recvElems))
+      if used {
+        ptrs[i] = fr.allocaBuilder.CreateAlloca(elemtyp, "")
       } else {
+        ptrs[i] = llvm.ConstNull(llvm.PointerType(llvm.Int8Type(), 0))
+      }
       recvElems = append(recvElems, newValue(ptrs[i], chantyp.Elem()))
     }
   }
 
   // Create select{send,recv2} calls.
   var receivedp llvm.Value
   if len(recvElems) > 0 {
     receivedp = fr.allocaBuilder.CreateAlloca(fr.types.ToLLVM(types.Typ[types.Bool]), "")
   }
-  if !blocking {
+  if !sel.Blocking {
     // If the default case is chosen, the index must be -1.
     fr.runtime.selectdefault.call(fr, selectp, llvm.ConstAllOnes(llvm.Int32Type()))
   }
-  for i, state := range states {
-    ch := state.Chan.value
+  for i, state := range sel.States {
+    ch := fr.llvmvalue(state.Chan)
     index := llvm.ConstInt(llvm.Int32Type(), uint64(i), false)
     if state.Dir == types.SendOnly {
       fr.runtime.selectsend.call(fr, selectp, ch, ptrs[i], index)
     } else {
       fr.runtime.selectrecv2.call(fr, selectp, ch, ptrs[i], receivedp, index)
     }
   }
 
   // Fire off the select.
   index = newValue(fr.runtime.selectgo.call(fr, selectp)[0], types.Typ[types.Int])
   if len(recvElems) > 0 {
     recvOk = newValue(fr.builder.CreateLoad(receivedp, ""), types.Typ[types.Bool])
     for _, recvElem := range recvElems {
       recvElem.value = fr.builder.CreateLoad(recvElem.value, "")
     }
   }
   return index, recvOk, recvElems
 }
+
+func chanSelectStateUsed(sel *ssa.Select, recvIndex int) bool {
+  for _, instr := range *sel.Referrers() {
+    extract, ok := instr.(*ssa.Extract)
+    if !ok || extract.Index != (recvIndex+2) {
+      continue
+    }
+    if len(*extract.Referrers()) > 0 {
+      return true
+    }
+  }
+  return false
+}

irgen/ssa.go

for i, res := range instr.Results {
       vals[i] = fr.llvmvalue(res)
     }
     fr.retInf.encode(llvm.GlobalContext(), fr.allocaBuilder, fr.builder, vals)
 
   case *ssa.RunDefers:
     fr.runDefers()
 
   case *ssa.Select:
-    states := make([]selectState, len(instr.States))
-    for i, state := range instr.States {
-      states[i] = selectState{
-        Dir:  state.Dir,
-        Chan: fr.value(state.Chan),
-        Send: fr.value(state.Send),
-      }
-    }
-    index, recvOk, recvElems := fr.chanSelect(states, instr.Blocking)
+    index, recvOk, recvElems := fr.chanSelect(instr)
     tuple := append([]*govalue{index, recvOk}, recvElems...)
     fr.tuples[instr] = tuple
 
   case *ssa.Send:
     fr.chanSend(fr.value(instr.Chan), fr.value(instr.X))
 
   case *ssa.Slice:
     x := fr.llvmvalue(instr.X)

test/irgen/select.go

+// RUN: llgo -S -emit-llvm -o - %s | FileCheck %s
+
+package foo
+
+// CHECK-NOT: alloca [1024 x i8]
+// CHECK-NOT: alloca [2048 x i8]
+// CHECK: alloca [4096 x i8]
+func F() {
+  ch1 := make(chan [1024]byte)
+  ch2 := make(chan [2048]byte)
+  ch3 := make(chan [4096]byte)
+  select {
+  case <-ch1:
+  case _ = <-ch2:
+  case x := <-ch3:
+    _ = x[0]
+  }
+}