[dx11] Add underlying DX11 device, memory allocation, and some tests (#…

…3971)

* [dx11] Add ID3D11Device and ID3D11DeviceContext

This change adds the actual D3D11 device class (ID3D11Device and
ID3D11DeviceContext) from <d3d11.h>.

Tested: Tested with unit test and saw the Dx11Device and the D3D11
device it creates can be successfully created. Enabled Graphics Debug
and confirmed there are no remaining live objects when the test is
finished.

* [dx11] Add a class for getting info from the D3D11 info queue

This class fetches information from the D3D11 info queue. This enables
us to count the number of live D3D11 objects. This is most meant for
usage with unit tests and during development.

Tested: Tested with unit tests

* [dx11] Allocate and deallocate memory

This change adds Dx11Device::allocate_memory and
Dx11Device::dealloc_memory.

Tested: Tested with unit tests. The D3D11 objects (Buffer and UAV) are
created and released as expected.

* Auto Format

* Don't build DX11 test when the backend is not built

* Update prtags.json

* [dx11] Make kD3d11DebugEnabled a compile-time constant

This change makes kD3d11DebugEnabled a constant.

Tested: rebuilt-and ran unit test with either kD3d11DebugEnabled set to
ON or OFF.

* Update taichi/backends/dx/dx_info_queue.h

Co-authored-by: Ye Kuang <[email protected]>

* Auto Format

Co-authored-by: Taichi Gardener <[email protected]>
Co-authored-by: Bob Cao <[email protected]>
Co-authored-by: Ye Kuang <[email protected]>

cmake/TaichiCore.cmake

@@ -8,7 +8,6 @@ option(TI_WITH_CC "Build with the C backend" ON)
 option(TI_WITH_VULKAN "Build with the Vulkan backend" OFF)
 option(TI_WITH_DX11 "Build with the DX11 backend" OFF)
 
-
 if(UNIX AND NOT APPLE)
     # Handy helper for Linux
     # https://stackoverflow.com/a/32259072/12003165

cmake/TaichiTests.cmake

@@ -13,6 +13,7 @@ endif()
 file(GLOB_RECURSE TAICHI_TESTS_SOURCE
         "tests/cpp/analysis/*.cpp"
         "tests/cpp/aot/*.cpp"
+        "tests/cpp/backends/*.cpp"
         "tests/cpp/codegen/*.cpp"
         "tests/cpp/common/*.cpp"
         "tests/cpp/ir/*.cpp"

misc/prtags.json

@@ -12,6 +12,8 @@
   "metal"           : "Metal backend",
   "opengl"          : "OpenGL backend",
   "vulkan"          : "Vulkan backend",
+  "dx11"            : "DirectX 11 backend",
+  "spirv"           : "SPIR-V common codegen",
   "wasm"            : "WebAssembly backend",
   "misc"            : "Miscellaneous",
   "std"             : "Standard library",

taichi/backends/dx/dx_api.h

@@ -1,6 +1,7 @@
 #pragma once
 #pragma comment(lib, "d3d11.lib")
 #pragma comment(lib, "d3dcompiler.lib")
+#pragma comment(lib, "dxguid.lib")
 
 #include "taichi/common/core.h"
 

taichi/backends/dx/dx_device.cpp

@@ -4,6 +4,12 @@ namespace taichi {
 namespace lang {
 namespace directx11 {
 
+void check_dx_error(HRESULT hr, const char *msg) {
+  if (!SUCCEEDED(hr)) {
+    TI_ERROR("Error in {}: {}", msg, hr);
+  }
+}
+
 Dx11ResourceBinder::~Dx11ResourceBinder() {
 }
 
@@ -19,19 +25,199 @@ ResourceBinder *Dx11Pipeline::resource_binder() {
   return nullptr;
 }
 
+namespace {
+HRESULT create_compute_device(ID3D11Device **out_device,
+                              ID3D11DeviceContext **out_context,
+                              bool force_ref,
+                              bool debug_enabled) {
+  const D3D_FEATURE_LEVEL levels[] = {
+      D3D_FEATURE_LEVEL_11_1,
+      D3D_FEATURE_LEVEL_11_0,
+      D3D_FEATURE_LEVEL_10_1,
+      D3D_FEATURE_LEVEL_10_0,
+  };
+
+  UINT flags = 0;
+  if (debug_enabled)
+    flags |= D3D11_CREATE_DEVICE_DEBUG;
+
+  ID3D11Device *device = nullptr;
+  ID3D11DeviceContext *context = nullptr;
+  HRESULT hr;
+
+  D3D_DRIVER_TYPE driver_types[] = {
+      D3D_DRIVER_TYPE_HARDWARE, D3D_DRIVER_TYPE_SOFTWARE,
+      D3D_DRIVER_TYPE_REFERENCE, D3D_DRIVER_TYPE_WARP};
+  const char *driver_type_names[] = {
+      "D3D_DRIVER_TYPE_HARDWARE", "D3D_DRIVER_TYPE_SOFTWARE",
+      "D3D_DRIVER_TYPE_REFERENCE", "D3D_DRIVER_TYPE_WARP"};
+
+  const int num_types = sizeof(driver_types) / sizeof(driver_types[0]);
+
+  int attempt_idx = 0;
+  if (force_ref) {
+    attempt_idx = 2;
+  }
+
+  for (; attempt_idx < num_types; attempt_idx++) {
+    D3D_DRIVER_TYPE driver_type = driver_types[attempt_idx];
+    hr = D3D11CreateDevice(nullptr, driver_type, nullptr, flags, levels,
+                           _countof(levels), D3D11_SDK_VERSION, &device,
+                           nullptr, &context);
+
+    if (FAILED(hr) || device == nullptr) {
+      TI_WARN("Failed to create D3D11 device with type {}: {}\n", driver_type,
+              driver_type_names[attempt_idx]);
+      continue;
+    }
+
+    if (device->GetFeatureLevel() < D3D_FEATURE_LEVEL_11_0) {
+      D3D11_FEATURE_DATA_D3D10_X_HARDWARE_OPTIONS hwopts = {0};
+      device->CheckFeatureSupport(D3D11_FEATURE_D3D10_X_HARDWARE_OPTIONS,
+                                  &hwopts, sizeof(hwopts));
+      if (!hwopts.ComputeShaders_Plus_RawAndStructuredBuffers_Via_Shader_4_x) {
+        device->Release();
+        TI_WARN(
+            "DirectCompute not supported via "
+            "ComputeShaders_Plus_RawAndStructuredBuffers_Via_Shader_4");
+      }
+      continue;
+    }
+
+    TI_INFO("Successfully created DX11 device with type {}",
+            driver_type_names[attempt_idx]);
+    *out_device = device;
+    *out_context = context;
+    break;
+  }
+
+  if (*out_device == nullptr || *out_context == nullptr) {
+    TI_ERROR("Failed to create DX11 device using all {} driver types",
+             num_types);
+  }
+
+  return hr;
+}
+
+HRESULT create_raw_buffer(ID3D11Device *device,
+                          UINT size,
+                          void *init_data,
+                          ID3D11Buffer **out_buf) {
+  *out_buf = nullptr;
+  D3D11_BUFFER_DESC desc = {};
+  desc.BindFlags = D3D11_BIND_UNORDERED_ACCESS | D3D11_BIND_SHADER_RESOURCE;
+  desc.ByteWidth = size;
+  desc.MiscFlags = D3D11_RESOURCE_MISC_BUFFER_ALLOW_RAW_VIEWS;
+  if (init_data) {
+    D3D11_SUBRESOURCE_DATA data;
+    data.pSysMem = init_data;
+    return device->CreateBuffer(&desc, &data, out_buf);
+  } else {
+    return device->CreateBuffer(&desc, nullptr, out_buf);
+  }
+}
+
+HRESULT create_buffer_uav(ID3D11Device *device,
+                          ID3D11Buffer *buffer,
+                          ID3D11UnorderedAccessView **out_uav) {
+  D3D11_BUFFER_DESC buf_desc = {};
+  buffer->GetDesc(&buf_desc);
+  D3D11_UNORDERED_ACCESS_VIEW_DESC uav_desc = {};
+  uav_desc.ViewDimension = D3D11_UAV_DIMENSION_BUFFER;
+  uav_desc.Buffer.FirstElement = 0;
+  if (buf_desc.MiscFlags & D3D11_RESOURCE_MISC_BUFFER_ALLOW_RAW_VIEWS) {
+    uav_desc.Format = DXGI_FORMAT_R32_TYPELESS;
+    uav_desc.Buffer.Flags = D3D11_BUFFER_UAV_FLAG_RAW;
+    uav_desc.Buffer.NumElements = buf_desc.ByteWidth / 4;
+  } else if (buf_desc.MiscFlags & D3D11_RESOURCE_MISC_BUFFER_STRUCTURED) {
+    uav_desc.Format = DXGI_FORMAT_UNKNOWN;
+    uav_desc.Buffer.NumElements =
+        buf_desc.ByteWidth / buf_desc.StructureByteStride;
+  } else
+    return E_INVALIDARG;
+  return device->CreateUnorderedAccessView(buffer, &uav_desc, out_uav);
+}
+
+}  // namespace
+
 Dx11Device::Dx11Device() {
+  create_dx11_device();
+  if (kD3d11DebugEnabled) {
+    info_queue_ = std::make_unique<Dx11InfoQueue>(device_);
+  }
   set_cap(DeviceCapability::spirv_version, 0x10300);
 }
 
 Dx11Device::~Dx11Device() {
+  destroy_dx11_device();
+}
+
+void Dx11Device::create_dx11_device() {
+  if (device_ != nullptr && context_ != nullptr) {
+    TI_TRACE("D3D11 device has already been created.");
+    return;
+  }
+  TI_TRACE("Creating D3D11 device");
+  create_compute_device(&device_, &context_, kD3d11ForceRef,
+                        kD3d11DebugEnabled);
+}
+
+void Dx11Device::destroy_dx11_device() {
+  if (device_ != nullptr) {
+    device_->Release();
+    device_ = nullptr;
+  }
+  if (context_ != nullptr) {
+    context_->Release();
+    context_ = nullptr;
+  }
+}
+
+int Dx11Device::live_dx11_object_count() {
+  TI_ASSERT(info_queue_ != nullptr);
+  return info_queue_->live_object_count();
 }
 
 DeviceAllocation Dx11Device::allocate_memory(const AllocParams &params) {
-  TI_NOT_IMPLEMENTED;
+  ID3D11Buffer *buf;
+  HRESULT hr;
+  hr = create_raw_buffer(device_, params.size, nullptr, &buf);
+  check_dx_error(hr, "create raw buffer");
+  alloc_id_to_buffer_[alloc_serial_] = buf;
+
+  ID3D11UnorderedAccessView *uav;
+  hr = create_buffer_uav(device_, buf, &uav);
+  check_dx_error(hr, "create UAV for buffer");
+  alloc_id_to_uav_[alloc_serial_] = uav;
+
+  // Set debug names
+  std::string buf_name = "buffer alloc#" + std::to_string(alloc_serial_) +
+                         " size=" + std::to_string(params.size) + '\0';
+  hr = buf->SetPrivateData(WKPDID_D3DDebugObjectName, buf_name.size(),
+                           buf_name.c_str());
+  check_dx_error(hr, "set name for buffer");
+
+  std::string uav_name = "UAV of " + buf_name;
+  hr = uav->SetPrivateData(WKPDID_D3DDebugObjectName, uav_name.size(),
+                           uav_name.c_str());
+  check_dx_error(hr, "set name for UAV");
+
+  DeviceAllocation alloc;
+  alloc.device = this;
+  alloc.alloc_id = alloc_serial_;
+  ++alloc_serial_;
+
+  return alloc;
 }
 
 void Dx11Device::dealloc_memory(DeviceAllocation handle) {
-  TI_NOT_IMPLEMENTED;
+  uint32_t alloc_id = handle.alloc_id;
+  ID3D11Buffer *buf = alloc_id_to_buffer_[alloc_id];
+  buf->Release();
+  alloc_id_to_buffer_.erase(alloc_id);
+  ID3D11UnorderedAccessView *uav = alloc_id_to_uav_[alloc_id];
+  uav->Release();
+  alloc_id_to_uav_.erase(alloc_id);
 }
 
 std::unique_ptr<Pipeline> Dx11Device::create_pipeline(

taichi/backends/dx/dx_device.h

@@ -1,11 +1,22 @@
 #pragma once
 
 #include "taichi/backends/device.h"
+#include "taichi/backends/dx/dx_info_queue.h"
+#include <d3d11.h>
 
 namespace taichi {
 namespace lang {
 namespace directx11 {
 
+// Only enable debug layer when the corresponding testing facility is enabled
+constexpr bool kD3d11DebugEnabled = true;
+constexpr bool kD3d11ForceRef = false;  // Force REF device. May be used to
+                                        // force software rendering.
+
+void debug_enabled(bool);
+void force_ref(bool);
+void check_dx_error(HRESULT hr, const char *msg);
+
 class Dx11ResourceBinder : public ResourceBinder {
   ~Dx11ResourceBinder() override;
 };
@@ -55,6 +66,25 @@ class Dx11Device : public GraphicsDevice {
                        DeviceAllocation src_img,
                        ImageLayout img_layout,
                        const BufferImageCopyParams &params) override;
+
+  int live_dx11_object_count();
+
+ private:
+  void create_dx11_device();
+  void destroy_dx11_device();
+  ID3D11Buffer *alloc_id_to_buffer(uint32_t alloc_id);
+  ID3D11Buffer *alloc_id_to_buffer_cpu_copy(uint32_t alloc_id);
+  ID3D11UnorderedAccessView *alloc_id_to_uav(uint32_t alloc_id);
+  ID3D11Device *device_{};
+  ID3D11DeviceContext *context_{};
+  std::unique_ptr<Dx11InfoQueue> info_queue_{};
+  std::unordered_map<uint32_t, ID3D11Buffer *>
+      alloc_id_to_buffer_;  // binding ID to buffer
+  std::unordered_map<uint32_t, ID3D11Buffer *>
+      alloc_id_to_cpucopy_;  // binding ID to CPU copy of buffer
+  std::unordered_map<uint32_t, ID3D11UnorderedAccessView *>
+      alloc_id_to_uav_;  // binding ID to UAV
+  int alloc_serial_;
 };
 
 }  // namespace directx11