wip: xir

include/luisa/ast/constant_data.h

@@ -42,7 +42,6 @@ class LC_AST_API ConstantDecoder {
 };
 
 class LC_AST_API ConstantData {
-
     friend class CallableLibrary;
 
 private:
@@ -59,6 +58,7 @@ class LC_AST_API ConstantData {
     [[nodiscard]] auto raw() const noexcept { return _raw; }
     [[nodiscard]] auto type() const noexcept { return _type; }
     [[nodiscard]] auto hash() const noexcept { return _hash; }
+    [[nodiscard]] explicit operator bool() const noexcept { return _raw != nullptr; }
     [[nodiscard]] bool operator==(const ConstantData &rhs) const noexcept { return _hash == rhs._hash; }
     void decode(ConstantDecoder &d) const noexcept { d.decode(_type, _raw); }
 };

include/luisa/xir/argument.h

@@ -4,13 +4,14 @@
 
 namespace luisa::compute::xir {
 
-class Argument : public Value {
+class LC_XIR_API Argument : public Value {
 
 private:
     bool _by_ref = false;
 
 public:
-    explicit Argument(const Type *type = nullptr,
+    explicit Argument(Pool *pool,
+                      const Type *type = nullptr,
                       bool by_ref = false,
                       const Name *name = nullptr) noexcept;
     void set_by_ref(bool by_ref) noexcept { _by_ref = by_ref; }

include/luisa/xir/basic_block.h

@@ -7,24 +7,19 @@ namespace luisa::compute::xir {
 class Function;
 class Instruction;
 
-class LC_XIR_API BasicBlock : public PooledObject {
+class LC_XIR_API BasicBlock : public Value {
 
 private:
     Function *_function = nullptr;
-    BasicBlock *_parent_block = nullptr;
-    const Name *_name = nullptr;
     InlineInstructionList _instructions;
 
 public:
-    explicit BasicBlock(Function *function = nullptr,
-                        BasicBlock *parent_block = nullptr,
+    explicit BasicBlock(Pool *pool,
+                        Function *function = nullptr,
                         const Name *name = nullptr) noexcept;
     void set_function(Function *function) noexcept;
-    void set_parent_block(BasicBlock *parent_block) noexcept;
-    void set_name(const Name *name) noexcept;
-    [[nodiscard]] auto function() const noexcept { return _function; }
-    [[nodiscard]] auto parent_block() const noexcept { return _parent_block; }
-    [[nodiscard]] auto name() const noexcept { return _name; }
+    [[nodiscard]] auto function() noexcept { return _function; }
+    [[nodiscard]] auto function() const noexcept { return const_cast<const Function *>(_function); }
     [[nodiscard]] auto &instructions() noexcept { return _instructions; }
     [[nodiscard]] auto &instructions() const noexcept { return _instructions; }
 };

include/luisa/xir/constant.h

@@ -11,8 +11,7 @@ class LC_XIR_API Constant : public Value {
     ConstantData _data;
 
 public:
-    Constant() noexcept = default;
-    explicit Constant(ConstantData data, const Name *name = nullptr) noexcept;
+    explicit Constant(Pool *pool, ConstantData data = {}, const Name *name = nullptr) noexcept;
     void set_data(ConstantData data) noexcept;
     [[nodiscard]] auto data() const noexcept { return _data; }
 };

include/luisa/xir/function.h

@@ -4,7 +4,11 @@
 
 namespace luisa::compute::xir {
 
-class Function : public PooledObject {
+class LC_XIR_API Function : public PooledObject {
+
+public:
+    explicit Function(Pool *pool) noexcept : PooledObject{pool} {}
+
 };
 
 }// namespace luisa::compute::xir

include/luisa/xir/ilist.h

@@ -108,9 +108,11 @@ class IntrusiveNode : public Base {
 
 public:
     using Super = IntrusiveNode;
-    using Base::Base;
     static_assert(std::is_base_of_v<PooledObject, Base>);
 
+protected:
+    using Base::Base;
+
 private:
     friend IntrusiveList<T>;
     friend InlineIntrusiveList<T>;
@@ -162,9 +164,9 @@ class IntrusiveList : public detail::IntrusiveListImpl<IntrusiveList<Node>> {
     Node *_tail_sentinel = nullptr;
 
 public:
-    explicit IntrusiveList(Pool &pool) noexcept {
-        _head_sentinel = pool.create<Node>();
-        _tail_sentinel = pool.create<Node>();
+    explicit IntrusiveList(Pool *pool) noexcept {
+        _head_sentinel = pool->create<Node>();
+        _tail_sentinel = pool->create<Node>();
         _head_sentinel->_next = _tail_sentinel;
         _tail_sentinel->_prev = _head_sentinel;
     }
@@ -182,7 +184,8 @@ class InlineIntrusiveList : public detail::IntrusiveListImpl<InlineIntrusiveList
     Node _tail_sentinel;
 
 public:
-    InlineIntrusiveList() noexcept {
+    explicit InlineIntrusiveList(Pool *pool) noexcept
+        : _head_sentinel{pool}, _tail_sentinel{pool} {
         _head_sentinel._next = &_tail_sentinel;
         _tail_sentinel._prev = &_head_sentinel;
     }
@@ -243,9 +246,11 @@ class IntrusiveForwardNode : public Base {
 
 public:
     using Super = IntrusiveForwardNode;
-    using Base::Base;
     static_assert(std::is_base_of_v<PooledObject, Base>);
 
+protected:
+    using Base::Base;
+
 private:
     template<typename>
     friend class IntrusiveForwardList;

include/luisa/xir/instruction.h

@@ -12,7 +12,8 @@ class LC_XIR_API Instruction : public IntrusiveNode<Instruction, User> {
     BasicBlock *_parent_block = nullptr;
 
 public:
-    explicit Instruction(const Type *type = nullptr,
+    explicit Instruction(Pool *pool,
+                         const Type *type = nullptr,
                          BasicBlock *parent_block = nullptr,
                          const Name *name = nullptr) noexcept;
     void remove_self() noexcept override;

include/luisa/xir/instructions/branch.h

@@ -1,19 +1,35 @@
 #pragma once
 
+#include "luisa/xir/basic_block.h"
+
 #include <luisa/xir/instruction.h>
 
 namespace luisa::compute::xir {
 
 class BasicBlock;
 
-class BranchInst : public Instruction {
-
-private:
-    BasicBlock *_true_block = nullptr;
-    BasicBlock *_false_block = nullptr;
+class LC_XIR_API BranchInst : public Instruction {
 
 public:
+    explicit BranchInst(Pool *pool,
+                        Value *cond = nullptr,
+                        BasicBlock *true_block = nullptr,
+                        BasicBlock *false_block = nullptr,
+                        BasicBlock *parent_block = nullptr,
+                        const Name *name = nullptr) noexcept;
+
+    void set_cond(Value *cond) noexcept;
+    void set_true_block(BasicBlock *block) noexcept;
+    void set_false_block(BasicBlock *block) noexcept;
+
+    [[nodiscard]] Value *cond() noexcept;
+    [[nodiscard]] const Value *cond() const noexcept;
+
+    [[nodiscard]] BasicBlock *true_block() noexcept;
+    [[nodiscard]] const BasicBlock *true_block() const noexcept;
 
+    [[nodiscard]] BasicBlock *false_block() noexcept;
+    [[nodiscard]] const BasicBlock *false_block() const noexcept;
 };
 
 }// namespace luisa::compute::xir

include/luisa/xir/metadata.h

@@ -4,7 +4,11 @@
 
 namespace luisa::compute::xir {
 
-struct LC_XIR_API Metadata : IntrusiveForwardNode<Metadata> {};
+class LC_XIR_API Metadata : public IntrusiveForwardNode<Metadata> {
+
+public:
+    explicit Metadata(Pool *pool) noexcept;
+};
 
 using MetadataList = IntrusiveForwardList<Metadata>;
 

include/luisa/xir/metadata/location.h

@@ -1,9 +1,27 @@
 #pragma once
 
+#include <luisa/core/stl/filesystem.h>
 #include <luisa/xir/metadata.h>
 
 namespace luisa::compute::xir {
 
+class LC_XIR_API LocationMD : public Metadata {
 
+private:
+    luisa::filesystem::path _file;
+    int _line;
+    int _column;
 
-}
+public:
+    explicit LocationMD(Pool *pool,
+                        luisa::filesystem::path file = {},
+                        int line = -1, int column = -1) noexcept;
+    void set_file(luisa::filesystem::path file) noexcept { _file = std::move(file); }
+    void set_line(int line) noexcept { _line = line; }
+    void set_column(int column) noexcept { _column = column; }
+    [[nodiscard]] auto &file() const noexcept { return _file; }
+    [[nodiscard]] auto line() const noexcept { return _line; }
+    [[nodiscard]] auto column() const noexcept { return _column; }
+};
+
+}// namespace luisa::compute::xir

include/luisa/xir/name.h

@@ -5,12 +5,14 @@
 
 namespace luisa::compute::xir {
 
-class Name : public PooledObject {
+class LC_XIR_API Name : public PooledObject {
 
 private:
     luisa::string _s;
 
 public:
+    explicit Name(Pool *pool, luisa::string s = {}) noexcept;
+    [[nodiscard]] auto string() const noexcept { return _s; }
     [[nodiscard]] const auto &operator*() const noexcept { return _s; }
     [[nodiscard]] const auto *operator->() const noexcept { return &_s; }
 };

include/luisa/xir/pool.h

@@ -8,10 +8,19 @@
 
 namespace luisa::compute::xir {
 
-struct LC_XIR_API PooledObject {
+class Pool;
 
-    PooledObject() noexcept = default;
+class LC_XIR_API PooledObject {
+
+private:
+    Pool *_pool;
+
+protected:
+    explicit PooledObject(Pool *pool) noexcept : _pool{pool} {}
+
+public:
     virtual ~PooledObject() noexcept = default;
+    [[nodiscard]] auto pool() const noexcept { return _pool; }
 
     // make the object pinned to its memory location
     PooledObject(PooledObject &&) noexcept = delete;
@@ -33,7 +42,8 @@ class LC_XIR_API Pool : public concepts::Noncopyable {
     template<typename T, typename... Args>
         requires std::derived_from<T, PooledObject>
     [[nodiscard]] T *create(Args &&...args) {
-        auto object = luisa::new_with_allocator<T>(std::forward<Args>(args)...);
+        auto object = luisa::new_with_allocator<T>(this, std::forward<Args>(args)...);
+        assert(object->pool() == this && "PooledObject must be created with the correct pool.");
         _objects.emplace_back(object);
         return object;
     }

include/luisa/xir/shared.h

@@ -4,7 +4,7 @@
 
 namespace luisa::compute::xir {
 
-class Shared : public Value {
+class LC_XIR_API Shared : public Value {
 public:
     using Value::Value;
 };

include/luisa/xir/use.h

@@ -14,18 +14,16 @@ class LC_XIR_API Use : public IntrusiveForwardNode<Use> {
     User *_user = nullptr;
 
 public:
-    [[nodiscard]] auto value() noexcept { return _value; }
-    [[nodiscard]] auto value() const noexcept { return _value; }
-    [[nodiscard]] auto user() noexcept { return _user; }
-    [[nodiscard]] auto user() const noexcept { return _user; }
-
-public:
-    Use() noexcept = default;
-    Use(Value *value, User *user) noexcept;
+    explicit Use(Pool *pool, Value *value = nullptr, User *user = nullptr) noexcept;
     // set value, also update the use list of the old and new values
     void set_value(Value *value) noexcept;
     // set user, also update the use list of the old and new users
     void set_user(User *user) noexcept;
+
+    [[nodiscard]] auto value() noexcept { return _value; }
+    [[nodiscard]] auto user() noexcept { return _user; }
+    [[nodiscard]] auto value() const noexcept { return const_cast<const Value *>(_value); }
+    [[nodiscard]] auto user() const noexcept { return const_cast<const User *>(_user); }
 };
 
 using UseList = IntrusiveForwardList<Use>;

include/luisa/xir/user.h

@@ -13,10 +13,15 @@ class LC_XIR_API User : public Value {
     using Value::Value;
     void remove_operand_uses() noexcept;
     void add_operand_uses() noexcept;
-    void set_operands(luisa::vector<Use *> operands) noexcept;
-    void set_operands(Pool &pool, luisa::span<Value *const> operands) noexcept;
+    void set_operand(size_t index, Value *value) noexcept;
+    [[nodiscard]] Use *operand_use(size_t index) noexcept;
+    [[nodiscard]] const Use *operand_use(size_t index) const noexcept;
+    [[nodiscard]] Value *operand(size_t index) noexcept;
+    [[nodiscard]] const Value *operand(size_t index) const noexcept;
+    void set_operand_count(size_t n) noexcept;
+    void set_operands(luisa::span<Value *const> operands) noexcept;
     [[nodiscard]] auto operands() noexcept { return luisa::span{_operands}; }
-    [[nodiscard]] auto operands() const noexcept { return luisa::span{_operands}; }
+    [[nodiscard]] auto operands() const noexcept { return luisa::span<const Use *const>{_operands}; }
 };
 
 }// namespace luisa::compute::xir

include/luisa/xir/value.h

@@ -19,7 +19,7 @@ class LC_XIR_API Value : public PooledObject {
     MetadataList _metadata_list;
 
 public:
-    explicit Value(const Type *type = nullptr, const Name *name = nullptr) noexcept;
+    explicit Value(Pool *pool, const Type *type = nullptr, const Name *name = nullptr) noexcept;
     void set_type(const Type *type) noexcept { _type = type; }
     void set_name(const Name *name) noexcept { _name = name; }
 

src/xir/CMakeLists.txt

@@ -5,6 +5,7 @@ set(LUISA_COMPUTE_XIR_SOURCES
         function.cpp
         instruction.cpp
         metadata.cpp
+        name.cpp
         pool.cpp
         shared.cpp
         use.cpp

src/xir/argument.cpp

@@ -2,7 +2,8 @@
 
 namespace luisa::compute::xir {
 
-Argument::Argument(const Type *type, bool by_ref, const Name *name) noexcept
-    : Value{type, name}, _by_ref{by_ref} {}
+Argument::Argument(Pool *pool, const Type *type,
+                   bool by_ref, const Name *name) noexcept
+    : Value{pool, type, name}, _by_ref{by_ref} {}
 
 }// namespace luisa::compute::xir

src/xir/basic_block.cpp

@@ -2,8 +2,10 @@
 
 namespace luisa::compute::xir {
 
-BasicBlock::BasicBlock(Function *function, BasicBlock *parent_block, const Name *name) noexcept
-    : _function{function}, _parent_block{parent_block}, _name{name} {
+BasicBlock::BasicBlock(Pool *pool, Function *function, const Name *name) noexcept
+    : Value{pool, nullptr, name},
+      _function{function},
+      _instructions{pool} {
     _instructions.head_sentinel()->set_parent_block(this);
     _instructions.tail_sentinel()->set_parent_block(this);
 }
@@ -12,12 +14,4 @@ void BasicBlock::set_function(Function *function) noexcept {
     _function = function;
 }
 
-void BasicBlock::set_parent_block(BasicBlock *parent_block) noexcept {
-    _parent_block = parent_block;
-}
-
-void BasicBlock::set_name(const Name *name) noexcept {
-    _name = name;
-}
-
 }// namespace luisa::compute::xir

src/xir/constant.cpp

@@ -3,11 +3,11 @@
 
 namespace luisa::compute::xir {
 
-Constant::Constant(ConstantData data, const Name *name) noexcept
-    : Value{data.type(), name} { set_data(data); }
+Constant::Constant(Pool *pool, ConstantData data, const Name *name) noexcept
+    : Value{pool, data.type(), name} { set_data(data); }
 
 void Constant::set_data(ConstantData data) noexcept {
-    LUISA_DEBUG_ASSERT(data.type() == type(),
+    LUISA_DEBUG_ASSERT(!data || data.type() == type(),
                        "Constant data type mismatch: {} vs {}",
                        data.type()->description(),
                        type()->description());