LLVMでの効率的なUnion型の扱い

オレオレ言語作り隊アドベントカレンダー12日目です。

自作言語の紹介にしようと思いましたが、~~時間がない~~恥ずかしいので、LLVM使っててはまったunion型の話をしようと思います。

## TL;DR

Union型の表現で迷ったらとりあえずワードの配列で（例えば`[4 x i64]`）。

## 導入

### Union型

Cでおなじみのunion型はメンバが同じメモリ領域を共有している型です。下の例ではパディング等を考えなければ、`foo_t`のサイズは`sizeof int`と`sizeof double`の大きい方になります。`bar`を代入して`baz`の値を取り出すみたいなこともできます。

```c
union foo_t {
  int bar;
  double baz;
};
```

Rust等のモダンな言語では[tagged union](https://doc.rust-lang.org/rust-by-example/custom_types/enum.html)というより型安全な形で目にすることが多いです。また、関数型言語における[代数的データ型](https://wiki.haskell.org/Algebraic_data_type)の表現にも用いられます。

というわけで、union型**大人気**です。

### LLVM IRには無い

LLVMの中間表現にはunion型に直接対応するものはありません。なので構造体等を使って頑張って実装します。基本的にはメンバのうち最大サイズのもののメモリを`alloca`命令や`malloc`関数で確保して、そのメモリのアドレスをビットキャストして各メンバを読みます。

### メモリ使いたくない

自作言語ツクルンジャーとしては、`alloca`命令や`malloc`関数なんてものは忌み嫌うべき存在です。Union型の中身を解釈し直すだけなら、全部レジスタ内の操作で済ませて欲しいものです。また、union型のサイズが小さい場合はunboxedなまま扱いたいです。

## 問題

次のような、あるメンバを取り、対応するunboxedなtagged unionを返すコンストラクタ関数を考えます。

```llvm
define { i64, 何らかの型 } @constructor(double) {
  ...
}
```

とりあえずメンバ間の共通表現はバイト列でいいだろうと思って、

```llvm:foo.ll
define { i64, [8 x i8] } @constructor(double) {
  %2 = alloca { i64, [8 x i8] }
  %3 = bitcast { i64, [8 x i8] }* %2 to { i64, double }*
  %4 = getelementptr { i64, double }, { i64, double }* %3, i32 0, i32 0
  store i64 0, i64* %4
  %5 = getelementptr { i64, double }, { i64, double }* %3, i32 0, i32 1
  store double %0, double* %5
  %6 = load { i64, [8 x i8] }, { i64, [8 x i8] }* %2
  ret { i64, [8 x i8] } %6
}
```

こうしてみました。以下が`llc`でのコンパイル結果です。

```as:foo.s
        .text
        .file   "foo.ll"
        .globl  constructor             # -- Begin function constructor
        .p2align        4, 0x90
        .type   constructor,@function
constructor:                            # @constructor
        .cfi_startproc
# %bb.0:
        movq    $0, -16(%rsp)
        movsd   %xmm0, -8(%rsp)
        movq    -8(%rsp), %rax
        movq    %rax, 8(%rdi)
        movq    $0, (%rdi)
        movq    %rdi, %rax
        retq
.Lfunc_end0:
        .size   constructor, .Lfunc_end0-constructor
        .cfi_endproc
                                        # -- End function

        .section        ".note.GNU-stack","",@progbits
```

悪くなさそう。ここで最適化をかけてみます。

```sh
opt -O3 foo.ll -S > foo.opt.ll
```

LLVM IRでの結果。

```llvm:foo.opt.ll
; ModuleID = 'foo.ll'
source_filename = "foo.ll"

; Function Attrs: norecurse nounwind readnone
define { i64, [8 x i8] } @constructor(double) local_unnamed_addr #0 {
  %2 = bitcast double %0 to i64
  %.sroa.2.8.extract.trunc = trunc i64 %2 to i8
  %.fca.1.0.insert = insertvalue { i64, [8 x i8] } { i64 0, [8 x i8] undef }, i8 %.sroa.2.8.extract.trunc, 1, 0                                                                                              
  %.sroa.2.9.extract.shift = lshr i64 %2, 8
  %.sroa.2.9.extract.trunc = trunc i64 %.sroa.2.9.extract.shift to i8
  %.fca.1.1.insert = insertvalue { i64, [8 x i8] } %.fca.1.0.insert, i8 %.sroa.2.9.extract.trunc, 1, 1
  %.sroa.2.10.extract.shift = lshr i64 %2, 16
  %.sroa.2.10.extract.trunc = trunc i64 %.sroa.2.10.extract.shift to i8
  %.fca.1.2.insert = insertvalue { i64, [8 x i8] } %.fca.1.1.insert, i8 %.sroa.2.10.extract.trunc, 1, 2                                                                                                      
  %.sroa.2.11.extract.shift = lshr i64 %2, 24
  %.sroa.2.11.extract.trunc = trunc i64 %.sroa.2.11.extract.shift to i8
  %.fca.1.3.insert = insertvalue { i64, [8 x i8] } %.fca.1.2.insert, i8 %.sroa.2.11.extract.trunc, 1, 3                                                                                                      
  %.sroa.2.12.extract.shift = lshr i64 %2, 32
  %.sroa.2.12.extract.trunc = trunc i64 %.sroa.2.12.extract.shift to i8
  %.fca.1.4.insert = insertvalue { i64, [8 x i8] } %.fca.1.3.insert, i8 %.sroa.2.12.extract.trunc, 1, 4                                                                                                      
  %.sroa.2.13.extract.shift = lshr i64 %2, 40
  %.sroa.2.13.extract.trunc = trunc i64 %.sroa.2.13.extract.shift to i8
  %.fca.1.5.insert = insertvalue { i64, [8 x i8] } %.fca.1.4.insert, i8 %.sroa.2.13.extract.trunc, 1, 5                                                                                                      
  %.sroa.2.14.extract.shift = lshr i64 %2, 48
  %.sroa.2.14.extract.trunc = trunc i64 %.sroa.2.14.extract.shift to i8
  %.fca.1.6.insert = insertvalue { i64, [8 x i8] } %.fca.1.5.insert, i8 %.sroa.2.14.extract.trunc, 1, 6                                                                                                      
  %.sroa.2.15.extract.shift = lshr i64 %2, 56
  %.sroa.2.15.extract.trunc = trunc i64 %.sroa.2.15.extract.shift to i8
  %.fca.1.7.insert = insertvalue { i64, [8 x i8] } %.fca.1.6.insert, i8 %.sroa.2.15.extract.trunc, 1, 7                                                                                                      
  ret { i64, [8 x i8] } %.fca.1.7.insert
}

attributes #0 = { norecurse nounwind readnone }
```

???

コンパイル後。

```as:foo.opt.s
        .text
        .file   "foo.ll"
        .globl  constructor             # -- Begin function constructor
        .p2align        4, 0x90
        .type   constructor,@function
constructor:                            # @constructor
# %bb.0:
        movq    %xmm0, %rax
        movq    %rax, %r8
        movq    %rax, %r9
        movq    %rax, %rsi
        movq    %rax, %rcx
        movq    %rax, %rdx
        movb    %ah, 9(%rdi)
        movb    %al, 8(%rdi)
        shrq    $16, %rax
        shrq    $24, %r8
        shrq    $32, %r9
        shrq    $40, %rsi
        shrq    $48, %rcx
        shrq    $56, %rdx
        movb    %dl, 15(%rdi)
        movb    %cl, 14(%rdi)
        movb    %sil, 13(%rdi)
        movb    %r9b, 12(%rdi)
        movb    %r8b, 11(%rdi)
        movb    %al, 10(%rdi)
        movq    $0, (%rdi)
        movq    %rdi, %rax
        retq
.Lfunc_end0:
        .size   constructor, .Lfunc_end0-constructor
                                        # -- End function

        .section        ".note.GNU-stack","",@progbits
```

1バイトずつご丁寧に読み書きされてます。

## 解決策

ワード（`i32`や`i64`）の配列をunion型の表現として使います。具体的にはtagged union型の`[8 x i8]`の部分を`[1 x i64]`で置き換えます。

```llvm:bar.ll
define { i64, [1 x i64] } @constructor(double) {
  %2 = alloca { i64, [1 x i64] }
  %3 = bitcast { i64, [1 x i64] }* %2 to { i64, double }*
  %4 = getelementptr { i64, double }, { i64, double }* %3, i32 0, i32 0
  store i64 0, i64* %4
  %5 = getelementptr { i64, double }, { i64, double }* %3, i32 0, i32 1
  store double %0, double* %5
  %6 = load { i64, [1 x i64] }, { i64, [1 x i64] }* %2
  ret { i64, [1 x i64] } %6
}
```

```as:bar.s
        .text
        .file   "bar.ll"
        .globl  constructor             # -- Begin function constructor
        .p2align        4, 0x90
        .type   constructor,@function
constructor:                            # @constructor
        .cfi_startproc
# %bb.0:
        movq    $0, -16(%rsp)
        movsd   %xmm0, -8(%rsp)
        movq    -8(%rsp), %rdx
        xorl    %eax, %eax
        retq
.Lfunc_end0:
        .size   constructor, .Lfunc_end0-constructor
        .cfi_endproc
                                        # -- End function

        .section        ".note.GNU-stack","",@progbits
```

最適化後。

```llvm:bar.opt.ll
; ModuleID = 'bar.ll'
source_filename = "bar.ll"

; Function Attrs: norecurse nounwind readnone
define { i64, [1 x i64] } @constructor(double) local_unnamed_addr #0 {
  %2 = bitcast double %0 to i64
  %.fca.1.0.insert = insertvalue { i64, [1 x i64] } { i64 0, [1 x i64] undef }, i64 %2, 1, 0
  ret { i64, [1 x i64] } %.fca.1.0.insert
}

attributes #0 = { norecurse nounwind readnone }
```

```as:bar.opt.s
        .text
        .file   "bar.ll"
        .globl  constructor             # -- Begin function constructor
        .p2align        4, 0x90
        .type   constructor,@function
constructor:                            # @constructor
# %bb.0:
        movq    %xmm0, %rdx
        xorl    %eax, %eax
        retq
.Lfunc_end0:
        .size   constructor, .Lfunc_end0-constructor
                                        # -- End function

        .section        ".note.GNU-stack","",@progbits
```

大変良い。

## まとめ

Union型の表現はワードの配列にしましょう。生成される命令列が少なく済みます。

ちなみに、LLVMの整数型は`i`の後に任意の自然数を取ることができます。つまり、今回の例では`[16 x i8]`で表されるようなunion型を`i128`にしても良いわけです。この場合についても書こうと思ったのですが、~~時間がなかった~~疲れたため皆さんの宿題にします。答えは、どうでもいいメモリ領域をゼロ埋めする命令が無駄に入るでした。

## おまけ（広告）

[Ein言語](https://github.com/ein-lang/ein)という自作言語を自作中です。名前は昨日決めました。Haskellのような遅延評価による関数型言語です。LLVMで書いていて、最近、関数定義ができるようになりました。[ついった](https://twitter.com/raviqqe)。