LLVM最適化パスを使ったゲート最適化

quantum computing

LLVM最適化パスを使ったゲート最適化

LLVMは簡単にPassを作って組み込むことが出来る。このPassというものは何か?というとオフィシャルのドキュメントに記載がある。コンパイラが行う変換と最適化を行うもので、PassManagerがそれらの経路の制御を行うものになる。何かわかりやすい良い図が無いかしら…と探してみたら、この辺のが分かりやすいと思われ。LLVMの魅力というか根幹の一つになる。
このPassの仕組みを量子向けに作っている言語(qlang)、C言語に適用して量子回路の最適化を行ってみることにする。

■ 最適化は何をしているのか?

そもそも最適化は何をしているか?というと、身の回りで分かりやすいのはオプティマイザになる。gccとかだと、-O[0-3]とかで付けるもの。実際にこの最適化、何をやっているのか?というと簡単なコードで見てみると分かりやすい。

– サンプルソースコード
#include <stdio.h>

int main() {
  int i, j = 0;
  for (i = 0; i < 10; i++)
    j += 1;
  return 0;
}
- 最適化無し(LLVM IR)
define dso_local signext i32 @main() #0 {
  %1 = alloca i32, align 4
  %2 = alloca i32, align 4
  %3 = alloca i32, align 4
  store i32 0, i32* %1, align 4
  store i32 0, i32* %3, align 4
  store i32 0, i32* %2, align 4
  br label %4

4:                                                ; preds = %10, %0
  %5 = load i32, i32* %2, align 4
  %6 = icmp slt i32 %5, 10
  br i1 %6, label %7, label %13

7:                                                ; preds = %4
  %8 = load i32, i32* %3, align 4
  %9 = add nsw i32 %8, 1
  store i32 %9, i32* %3, align 4
  br label %10

10:                                               ; preds = %7
  %11 = load i32, i32* %2, align 4
  %12 = add nsw i32 %11, 1
  store i32 %12, i32* %2, align 4
  br label %4

13:                                               ; preds = %4
  ret i32 0
}
- 最適化無し 実行ファイルダンプ
   1019c: 01 11                         addi    sp, sp, -32
   1019e: 06 ec                         sd      ra, 24(sp)
   101a0: 22 e8                         sd      s0, 16(sp)
   101a2: 00 10                         addi    s0, sp, 32
   101a4: 01 45                         mv      a0, zero
   101a6: 23 26 a4 fe                   sw      a0, -20(s0)
   101aa: 23 22 a4 fe                   sw      a0, -28(s0)
   101ae: 23 24 a4 fe                   sw      a0, -24(s0)
   101b2: 6f 00 40 00                   j       4
   101b6: 03 25 84 fe                   lw      a0, -24(s0)
   101ba: a5 45                         addi    a1, zero, 9
   101bc: 63 c2 a5 02                   blt     a1, a0, 36
   101c0: 6f 00 40 00                   j       4
   101c4: 03 25 44 fe                   lw      a0, -28(s0)
   101c8: 05 05                         addi    a0, a0, 1
   101ca: 23 22 a4 fe                   sw      a0, -28(s0)
   101ce: 6f 00 40 00                   j       4
   101d2: 03 25 84 fe                   lw      a0, -24(s0)
   101d6: 05 05                         addi    a0, a0, 1
   101d8: 23 24 a4 fe                   sw      a0, -24(s0)
   101dc: 6f f0 bf fd                   j       -38
   101e0: 01 45                         mv      a0, zero
   101e2: 42 64                         ld      s0, 16(sp)
   101e4: e2 60                         ld      ra, 24(sp)
   101e6: 05 61                         addi    sp, sp, 32
   101e8: 82 80                         ret
- 最適化 -O3(LLVM IR)
define dso_local signext i32 @main() local_unnamed_addr #0 {
  ret i32 0
}
- 最適化 -O3 実行ファイルダンプ
000000000001019c main:
   1019c: 01 45                        	mv	a0, zero
   1019e: 82 80                        	ret

このくらいのコードでも顕著に差が出てくる。最適化のパスを通すと結果だけになるという、非常にシンプルになる。メリットとしては命令(プロセッサに対する命令の回数)が減るためアプリが高速化する・ファイルサイズが減る....etc といった感じ。

■ Xゲートに対する最適化

このPassの仕組みを使って量子回路のゲート最適化をする。ここでは例としてシンプルに1qbitに対するXゲートに対する最適化をしてみる。XゲートはX軸に対する回転で、例えば同じqbitに対して2回連続で行うと360度回転して同じ状態になる。ということは、同じqbitに対して連続して2回あるものは打ち消して回路上から消すことが出来る。利点としては上記のIR・実行コードのように実行回数が減るという利点がある。

■ C/qlangで吐き出されるアセンブリコード

まず最初に量子命令セットを組み込んだソースコード、実行形式はこんな感じになっている。

# ソースコード(test_q.c)
int main(int argc, char* argv[]) {
    int ret1 = 0;

    // call qoox for duplicate optimizer test.
    asm volatile( "qoox.k qzero,qt1,qzero,1");
    asm volatile( "qoox.k qzero,qt1,qzero,1");

    // call qmeas
    asm volatile( "qmeas.k  %0,qt1,qzero,1" :"=r"(ret1));
    printf("%d\n", ret2);
    
    return 0;
}
# コンパイル(pass無し)
clang -target riscv64-unknown-linux-gnu -c test_q.c -emit-llvm -o test_q.bc
llvm-dis test_q.bc
llc -march=riscv64 -relocation-model=pic -filetype=asm test_q.bc  -o test_q.s
riscv64-unknown-elf-gcc test_q.s -o test_q -march=rv64imafdkc
# 実行形式のダンプ(riscv64-unknown-elf-objdump -D test_q)
00000000000101a0 
: 101a0: 7179 addi sp,sp,-48 101a2: f406 sd ra,40(sp) 101a4: f022 sd s0,32(sp) 101a6: 1800 addi s0,sp,48 101a8: fe042623 sw zero,-20(s0) 101ac: fea42423 sw a0,-24(s0) 101b0: feb43023 sd a1,-32(s0) 101b4: fc042e23 sw zero,-36(s0) 101b8: fc042c23 sw zero,-40(s0) 101bc: 0200c00b qoox.k qzero,qt1,qzero,1 101bc: 0200c00b qoox.k qzero,qt1,qzero,1 101c0: 8200850b qmeas.k a0,qt1,qzero,1 101c4: fca42e23 sw a0,-36(s0) 101c8: fd842583 lw a1,-40(s0) 101cc: 0000c517 auipc a0,0xc 101d0: 18450513 addi a0,a0,388 # 1c350 <__clzdi2+0x38> 101d4: 16a000ef jal ra,1033e 101d8: 00000513 li a0,0 101dc: 7402 ld s0,32(sp) 101de: 70a2 ld ra,40(sp) 101e0: 6145 addi sp,sp,48 101e2: 8082 ret

とりあえずテスト的にqoox.kを連続して実行してmeasureしているだけにしている。このqoox.kをclangのPassを使って最適化することで除去するようにする。

■ LLVM Passの実装

PassはLLVMに容易されているサンプルを見ながら実装していく。この辺のPassの実装で説明してあるのはめちゃくちゃ例も少ないから、LLVM Passの実装をする人の参考にでもなればと思ったり。

#include "llvm/Pass.h"
#include "llvm/IR/Function.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
#include <vector>
#include <map>
#include <stack>
using namespace llvm;

namespace {
  struct quantumDupOptimizerPass : public FunctionPass {
    static char ID;
    quantumDupOptimizerPass() : FunctionPass(ID) {}
    std::stack dupErasekList;
    std::multimap dupErasekMap;
    std::vector dupEraseQuantumAsm = {"qoox.k"};

    virtual bool runOnFunction(Function &F) {
      for (llvm::BasicBlock &BB : F) {
        for (llvm::Instruction &II: BB) {
          std::string str;
          llvm::raw_string_ostream rso(str);
          II.print(rso);

          // remove comment(ok?)
          int npos = str.find("#");
          if (npos > 0) str.erase(npos - 1);

          // just simple duplicate gate only now.
          for (std::string &v : dupEraseQuantumAsm) {
            std::size_t found = str.find(v);
            if (std::string::npos != found) {
              if (dupErasekMap.count(str) == 1) {
                auto itr = dupErasekMap.find(str);
                dupErasekList.push(itr->second);
                dupErasekList.push(&II);
                dupErasekMap.erase(str);
              } else 
                dupErasekMap.insert(std::make_pair(str, &II));
            }
          }
        }
      }

      while (!dupErasekList.empty()) {
        Instruction *II = dupErasekList.top();
        II->eraseFromParent();
        dupErasekList.pop();
      }
      return false;
    }
  };
}

char quantumDupOptimizerPass::ID = 0;
static void registerquantumDupOptimizerPass(const PassManagerBuilder &,
                         legacy::PassManagerBase &PM) {
  PM.add(new quantumDupOptimizerPass());
}
static RegisterStandardPasses
  RegisterMyPass(PassManagerBuilder::EP_EarlyAsPossible,
                 registerquantumDupOptimizerPass);

ちょっとダサい実装をしたのが、Instructionからアセンブリコード(文字列)にしてカウント&マッチしたら除去する対象(dupErasekList)に入れて、最後に除去している。もっとカッコいいやり方だと、LLVM本体を大改造して量子拡張命令(Operand)を処理すればいいものの、LLVMのビルドが凄まじく大きいのでビビって手を付けないで泥臭い方法を使った。あとXゲートの他にも何か連続ものを消せるようにdupEraseQuantumAsmに登録できるようにしている(同じようにZ軸に対する回転であるZゲート:qooz.kとかも配列に加えればOK)。

とりあえずPassが出来たので、最初のコードをこのパスを通してビルドする。ビルドにはPassの共有ライブラリを作って、clangのコンパイル時に通してあげる。

# コンパイル
# libqot.so(pass library)
g++ -shared -fPIC -o libqot.so quantumDupOptimizerPass.cpp `llvm-config --cxxflags`

# compile with pass
clang -Xclang -load -Xclang libqot.so -target riscv64-unknown-linux-gnu -c test_q.c -emit-llvm -o test_q.bc
llvm-dis test_q.bc
llc -march=riscv64 -relocation-model=pic -filetype=asm test_q.bc  -o test_q.s
riscv64-unknown-elf-gcc test_q.s -o test_q -march=rv64imafdkc
# 実行形式のダンプ(riscv64-unknown-elf-objdump -D test_q)
00000000000101a0 
: 101a0: 7179 addi sp,sp,-48 101a2: f406 sd ra,40(sp) 101a4: f022 sd s0,32(sp) 101a6: 1800 addi s0,sp,48 101a8: fe042623 sw zero,-20(s0) 101ac: fea42423 sw a0,-24(s0) 101b0: feb43023 sd a1,-32(s0) 101b4: fc042e23 sw zero,-36(s0) 101b8: fc042c23 sw zero,-40(s0) 101bc: 8200850b qmeas.k a0,qt1,qzero,1 101c0: fca42e23 sw a0,-36(s0) 101c4: fd842583 lw a1,-40(s0) 101c8: 0000c517 auipc a0,0xc 101cc: 18850513 addi a0,a0,392 # 1c350 <__clzdi2+0x3c> 101d0: 16a000ef jal ra,1033a 101d4: 00000513 li a0,0 101d8: 7402 ld s0,32(sp) 101da: 70a2 ld ra,40(sp) 101dc: 6145 addi sp,sp,48 101de: 8082 ret

Passを通さないのと比べてみると分かるように、qoox.k 命令を除去することが出来ている。このPassのクラスを量子向け言語qlangに登録(Regisしておくと、もっと簡単に命令セットの最適化をすることが出来る。

といった感じで、LLVMのPassを使ったコンパイル時の最適化方法はこんな感じになる。シンプルな例になるけど、Passを経由することで非常に柔軟で独自の最適化を定義することが出来るのと、今回開発しているRISC-V 命令セット量子拡張のような独自の実装に対しても行うことが出来る。

コメントを残す

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です