llvmでmemsetのコードサイズを小さくする
はじめに
解決したいこと
Rust で配列などをゼロ初期化すると、memset 関数呼び出しが生成されることがあります。このとき呼び出される memset 関数はコードサイズがわりと大きめです。RISC-V の rv32imc ターゲットだと92バイトの memset が生成されます。
組み込み向けのプログラムなどを作っているとできるだけコードサイズを小さくしたくなります。もっとコンパイラにがんばってもらって小さくするにはどうしたら良いでしょうか。
理想と現実
rustc でゼロ埋めを含むコードをコンパイルして RISC-V の rv32imc ターゲットでコード生成すると長い memset が生成されます。（説明のためのコメントは筆者が追加しています。）
rustc が生成する memset
memset:
li a3, 15
bgeu a3, a2, memset_8 ; サイズが15バイト以下なら memset_8 にジャンプ
neg a3, a0
andi a3, a3, 3
add a4, a0, a3
beqz a3, memset_4 ; 先頭アドレスが4の倍数なら memset_4 にジャンプ
mv a5, a0
memset_3:
sb a1, 0(a5) ; 最初の半端分を1バイト単位で埋める
addi a5, a5, 1
bltu a5, a4, memset_3
memset_4:
sub a2, a2, a3
andi a5, a2, -4
add a3, a4, a5
blez a5, memset_7 ; 残りサイズが3バイト以下なら memset_7 にジャンプ
andi a6, a1, 255
lui a5, 4112
addi a5, a5, 257
mul a5, a6, a5 ; 埋める値を4つ並べた4バイト値を作る
memset_6:
sw a5, 0(a4) ; 4バイト単位で埋める
addi a4, a4, 4
bltu a4, a3, memset_6
memset_7:
andi a2, a2, 3
bnez a2, memset_9 ; 4バイトに満たない半端があれば memset_9 にジャンプ
j memset_11
memset_8:
mv a3, a0
beqz a2, memset_11 ; サイズが0なら memset_11 にジャンプ
memset_9:
add a2, a2, a3
memset_10:
sb a1, 0(a3) ; 半端な部分は1バイト単位で埋める
addi a3, a3, 1
bltu a3, a2, memset_10
memset_11:
ret
簡単に言うと、アラインメントが揃っていないかもしれない、埋める値は0でないかもしれない、サイズがワードサイズの倍数でないかもしれない、と様々なケースを想定してコードが書かれているため長くなっています。
しかし、memset 対象領域のアラインメントはワードサイズに揃っている場合が多く、書き込む値も0であるケースが多いです。
もしアラインメントが揃っていてサイズが0より大きくワード境界になっていて埋める値も0である、とわかっていれば次のような12バイトのコードで済みます。
aligned_bzero:
add a2, a2, a0
aligned_bzero_1:
sw zero, (a0)
addi a0, a0, 4
bne a0, a2, aligned_bzero_1
しかもこの程度であればインラインで展開してしまえば、呼び出し側でリンクレジスタを始めとする callee-saved レジスタ退避が不要になりさらにコードが短くなる可能性もあります。
LLVM で memset がどう生成されるか
LLVM IR では memset は intrinsic （組み込み命令）として表現されます。
サイズが固定である程度短い場合はメモリストアループのアンローリングが行われてインライン展開され、そうでない場合は memset 関数呼び出しが生成されるようになっているようです。アーキテクチャによってはより最適な専用命令を使ったコード生成が行われるものもあります。
今回ほしいのは、インライン展開されるけどループのアンローリングまではしないというものです。 LLVM を改造してこういうコード生成をできるようにしていきます。
memset に対する SelectionDAG ノード生成
memset intrinsic を SelectionDAG というグラフ形式に変換しているのが SelectionDAG::getMemset です。
ここで前記のメモリストアループアンローリングまたは memset 関数呼び出しが生成されるのですが、さらにアーキテクチャに依存した最適化のためにフックが可能になっています。これは SelectionDAGTargetInfo というクラスを派生することで実現できます。
今回対象アーキテクチャは RISC-V とします。RISC-V では SelectionDAGTargetInfo を使った特殊化はしていないので実装クラス自体がまだありません。SelectionDAGTargetInfo を派生した RISCVSelectionDAGInfo クラスを作ります。 EmitTargetCodeForMemset という仮想関数が memset の生成時に呼ばれる関数なので、これを実装します。
class RISCVSelectionDAGInfo : public SelectionDAGTargetInfo {
public:
SDValue EmitTargetCodeForMemset(SelectionDAG &DAG, const SDLoc &dl,
SDValue Chain, SDValue Op1, SDValue Op2,
SDValue Op3, Align Alignment, bool isVolatile,
bool AlwaysInline,
MachinePointerInfo DstPtrInfo) const override;
}
条件を満たしたときに最適化された SelectionDAG を生成するようにします。RISC-V には 32bit と 64bit がありますがここでは簡単のため 32bit 決め打ちとします。すなわち、条件は以下を満たした場合とします。
- 最適化オプションが0でない
- 書き込む値が定数であり0である
- サイズが定数であり0でなく4の倍数である
- 書き込み先アドレスのアラインメントが4の倍数である
サイズは定数でなくても良い気がしますが、非定数で4の倍数の判定方法がわからなかったため定数としました。
このとき ALIGNED_BZERO4 というノードを持つ SelectionDAG を生成します。このノードは自分で定義したもので、後ほど詳しく説明しますがアドレスを2つ受け取りそのアドレス範囲を0で埋めてくれる内部的な命令を表すノードです。
一般的な bzero のように先頭アドレスとサイズを受け取るのではなく、メモリ領域を示す2つのアドレスを受け取るようにしている理由は、そのほうが生成コードが最適化されやすいためです。ただし終了アドレスを計算するコードを生成する必要があるため、 DAG.getMemBasePlusOffset を使って DAG を構築します。
static bool shouldGenerateAlignedBzero4(SelectionDAG &DAG,
ConstantSDNode *ConstantSrc, ConstantSDNode *ConstantSize, Align Alignment) {
auto &F = DAG.getMachineFunction().getFunction();
if (F.hasOptNone())
return false;
if (!(ConstantSrc && ConstantSrc->getZExtValue() == 0))
return false;
if (!(ConstantSize && ConstantSize->getZExtValue() > 0 && (ConstantSize->getZExtValue() & 3) == 0))
return false;
if (!((Alignment.value() & 3) == 0))
return false;
return true;
}
SDValue RISCVSelectionDAGInfo::EmitTargetCodeForMemset(
SelectionDAG &DAG, const SDLoc &dl, SDValue Chain, SDValue Dst, SDValue Src,
SDValue Size, Align Alignment, bool isVolatile, bool AlwaysInline,
MachinePointerInfo DstPtrInfo) const {
ConstantSDNode *ConstantSrc = dyn_cast<ConstantSDNode>(Src);
ConstantSDNode *ConstantSize = dyn_cast<ConstantSDNode>(Size);
if (!AlwaysInline && shouldGenerateAlignedBzero4(DAG, ConstantSrc, ConstantSize, Alignment))
return DAG.getNode(RISCVISD::ALIGNED_BZERO4, dl, MVT::Other, Chain, Dst,
DAG.getMemBasePlusOffset(Dst, DAG.getZExtOrTrunc(Size, dl, MVT::i32), dl));
return SDValue();
}
これで RISCVSelectionDAGInfo は完成ですが、デフォルトの SelectionDAGTargetInfo の代わりに今回作った RISCVSelectionDAGInfo を使うように RISCVSubtarget.h にある RISCVSubtarget クラスのメンバを差し替えます。
- SelectionDAGTargetInfo TSInfo;
+ RISCVSelectionDAGInfo TSInfo;
SelectionDAG ノードからコード生成
SelectionDAG は Lowering と呼ばれる処理を経てターゲットアーキテクチャのコードになります。
ALIGNED_BZERO4 ノードを定義します。ノード定義は RISCVInstrInfo.td で行います。
def SDT_RISCVAlignedBzero4 : SDTypeProfile<0, 2, [
SDTCisPtrTy<0>, SDTCisPtrTy<1>
]>;
def riscv_aligned_bzero4 : SDNode<"RISCVISD::ALIGNED_BZERO4",
SDT_RISCVAlignedBzero4,
[SDNPHasChain, SDNPMayStore]>;
let usesCustomInserter = 1, hasNoSchedulingInfo = 1 in
def AlignedBzero4 : Pseudo<(outs), (ins GPR:$dst, GPR:$lst), [(riscv_aligned_bzero4 GPR:$dst, GPR:$lst)]>;
ポインタを2つ受け取り、戻り値がないノードを定義しています。 usesCustomInserter = 1 によって Lowering 処理を自分で書けるようにします。
コード生成の流れです。
ストア中のポインタを保持するために一時レジスタが必要になるため、 RegInfo.createVirtualRegister を使って定義します。
ここでは RISC-V の命令である SW, ADDI, BNE に対応するものが生成可能なので生成したい命令列を作っていきます。
static MachineBasicBlock *emitAlignedBzero4Pseudo(MachineInstr &MI,
MachineBasicBlock *BB) {
assert(MI.getOpcode() == RISCV::AlignedBzero4 && "Unexpected instruction");
MachineFunction &MF = *BB->getParent();
const BasicBlock *LLVM_BB = BB->getBasicBlock();
MachineFunction::iterator It = ++BB->getIterator();
MachineBasicBlock *EntryMBB = BB;
MachineBasicBlock *LoopMBB = MF.CreateMachineBasicBlock(LLVM_BB);
MF.insert(It, LoopMBB);
MachineBasicBlock *DoneMBB = MF.CreateMachineBasicBlock(LLVM_BB);
MF.insert(It, DoneMBB);
// Transfer the remainder of BB and its successor edges to DoneMBB.
DoneMBB->splice(DoneMBB->begin(), EntryMBB,
std::next(MachineBasicBlock::iterator(MI)), EntryMBB->end());
DoneMBB->transferSuccessorsAndUpdatePHIs(EntryMBB);
EntryMBB->addSuccessor(LoopMBB);
MachineRegisterInfo &RegInfo = MF.getRegInfo();
Register CurrDstReg = RegInfo.createVirtualRegister(&RISCV::GPRRegClass);
Register NextDstReg = RegInfo.createVirtualRegister(&RISCV::GPRRegClass);
Register Dst = MI.getOperand(0).getReg();
Register Lst = MI.getOperand(1).getReg();
DebugLoc DL = MI.getDebugLoc();
const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
BuildMI(LoopMBB, DL, TII->get(RISCV::PHI), CurrDstReg)
.addUse(Dst)
.addMBB(EntryMBB)
.addUse(NextDstReg)
.addMBB(LoopMBB);
// sw zero, 0(CurrDstReg)
BuildMI(LoopMBB, DL, TII->get(RISCV::SW))
.addUse(RISCV::X0)
.addUse(CurrDstReg)
.addImm(0);
// addi NextDstReg, CurrDstReg, 4
BuildMI(LoopMBB, DL, TII->get(RISCV::ADDI), NextDstReg)
.addUse(CurrDstReg)
.addImm(4);
// bne NextDstReg, Lst, LoopMBB
BuildMI(LoopMBB, DL, TII->get(RISCV::BNE))
.addUse(NextDstReg)
.addUse(Lst)
.addMBB(LoopMBB);
LoopMBB->addSuccessor(LoopMBB);
LoopMBB->addSuccessor(DoneMBB);
MI.eraseFromParent();
return DoneMBB;
}
ADDI のところは、CurrDstReg を直接更新して addi CurrDstReg, CurrDstReg, 4 でも良さそうな気がしますが、それだと実行時にエラーになってしまいます。LLVM には SSA (Static Single Assignment) という形で中間コードを生成する必要があるためです。
SSA は一つの変数へのアサインは1度だけできるという表現方法で、これにより最適化がしやすくなるようです。
SSA で更新を表現するには、PHI という疑似命令を使います。 PHI は、直前のブロックが何であったかによって変数に値を代入します。上のコードの PHI 命令では、直前が EntryMBB だったときは Dst を、直前が LoopMBB だったときは NextDstReg を CurrDstReg に代入しています。
以上の変更で固定長のゼロ埋めに対して短いコードが生成されるようになります。
変更点のまとめ
固定長の memcpy に対しても同じような変更を加えることで短いコードを生成できるようになります。
bzero と memcpy の両方に対応したものをLLVM 17 へのパッチとしてまとめています。
今後の課題
とりあえず 32 bit 向けでは動くものになっているのですが、実用には 32 bit だけでなく 64 bit のサポートも必要になるでしょう。今のコードだと 64 bit のコードがちゃんと生成できません。他には、浮動小数点命令やベクタ命令を使った効率化も考えられます。
Rust で書いたカーネルをコンパイルするのに使ってみたところ memset や memcpy が生成されず全てインラインに展開されていい感じなのですが、固定長の配列を多用しているためというのもあるかもしれません。普通のプログラムだとアライメントはされているけれど可変長の配列を使用する場合も多そうなので、そのあたりの対応も課題になりそうです。
Discussion