Merge fixes from directml branch (#124)

Merges some of the recent changes from the directml branch:
* Use compute queue for AMD devices (#102)
* Register List Kernels for DML (#95)
* Update DirectMLX to latest (#104)
* Remove extra rows from test email (#106)
* Fix DML's Select kernel for int64 (#113)
* Fix list kernels and tensor array ops registration (#114)
* Simplify CI scripts (#112)
* Fix StridedSlice's input size coalescing (#115)
* Disable int64 image test (#116)
* Fix network share copy path (#117)
* Pipeline should continue if a test job fails (#118)
* Switch network share path to use build number instead of build ID
* Add missing HostMemory int32 registrations for _Arg and _RetVal (#122)
* Implement all the arithmetic Scatter and ResourceScatter operators (#121)
* Register emulated kernel implementations for RandomStandardNormal and TruncatedNormal (#120)

Author: jstoecker

tensorflow/core/common_runtime/dml/dml_device_state.cc

@@ -24,6 +24,7 @@ limitations under the License.
 #include "dml_upload_heap.h"
 #include "dml_util.h"
 #include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/util/env_var.h"
 #include "tensorflow/stream_executor/platform/default/dso_loader.h"
 
 using Microsoft::WRL::ComPtr;
@@ -70,8 +71,20 @@ namespace tensorflow {
   ComPtr<IDMLDevice> dml_device;
   dml_device = CreateDmlDevice(d3d_device.Get(), dml_flags);
 
+  // Default to using compute queues for AMD since it seems to mitigate TDRs and
+  // improve performance
+  const bool use_compute_queue_default = adapter.VendorID() == VendorID::kAmd;
+
+  bool use_compute_queue;
+  Status s = ReadBoolFromEnvVar("TF_DIRECTML_USE_COMPUTE_QUEUE",
+                                use_compute_queue_default, &use_compute_queue);
+
+  D3D12_COMMAND_LIST_TYPE queue_type = use_compute_queue
+                                           ? D3D12_COMMAND_LIST_TYPE_COMPUTE
+                                           : D3D12_COMMAND_LIST_TYPE_DIRECT;
+
   D3D12_COMMAND_QUEUE_DESC command_queue_desc = {};
-  command_queue_desc.Type = D3D12_COMMAND_LIST_TYPE_DIRECT;
+  command_queue_desc.Type = queue_type;
   command_queue_desc.Priority = D3D12_COMMAND_QUEUE_PRIORITY_NORMAL;
   command_queue_desc.Flags = D3D12_COMMAND_QUEUE_FLAG_DISABLE_GPU_TIMEOUT;
   command_queue_desc.NodeMask = 0;

tensorflow/core/common_runtime/gpu/gpu_process_state.cc

@@ -26,6 +26,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h"
 #include "tensorflow/core/common_runtime/gpu/gpu_id_utils.h"
 #include "tensorflow/core/common_runtime/gpu/gpu_init.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_mem_allocator.h"
 #include "tensorflow/core/common_runtime/pool_allocator.h"
 #include "tensorflow/core/common_runtime/shared_counter.h"
 #include "tensorflow/core/framework/allocator.h"

tensorflow/core/kernels/BUILD

@@ -2727,10 +2727,17 @@ tf_kernel_library(
 tf_kernel_library(
     name = "list_kernels",
     srcs = ["list_kernels.cc"],
-    hdrs = ["list_kernels.h"],
+    hdrs = [
+        "list_kernels.h",
+        "tensor_array.h",
+        "aggregate_ops.h",
+        "split_lib.h"] + if_dml(["dml_tensor_array.h"]),
     gpu_srcs = [
         "list_kernels.cu.cc",
         "list_kernels.h",
+        "tensor_array.h",
+        "aggregate_ops.h",
+        "split_lib.h",
     ],
     deps = [
         ":concat_lib",
@@ -8094,6 +8101,7 @@ tf_kernel_library(
         "dml_gather_op.cc",
         "dml_gather_nd_op.cc",
         "dml_scatter_nd_op.cc",
+        "dml_scatter_update_ops.cc",
         "dml_tensor_scatter_ops.cc",
         "dml_scan_ops.cc",
         "dml_dynamic_stitch_op.cc",
@@ -8145,6 +8153,7 @@ tf_kernel_library(
         "dml_kernel_wrapper.h",
         "dml_ops_common.h",
         "assign_op.h",
+        "random_op.h",
         "stateless_random_ops.h",
         "tensor_array.h",
         "concat_lib.h",

tensorflow/core/kernels/dml_addn_op.cc

@@ -66,7 +66,7 @@ class DmlAddNKernel : public DmlKernel {
                                    &identity_desc};
       Initialize(ctx, std::move(tensors), op_desc);
     } else {
-      auto scope = dml::Scope(ctx->GetDmlDevice());
+      auto scope = dml::Graph(ctx->GetDmlDevice());
       auto result = dml::InputTensor(scope, 0, inputs[0]);
 
       for (uint32_t i = 1; i < inputs.size(); ++i) {

tensorflow/core/kernels/dml_batch_norm_ops.cc

@@ -321,7 +321,7 @@ class DmlFusedBatchNormKernel : public DmlKernel {
     // the mean/variance tensors back to TF.
 
     auto scope =
-        dml::Scope(ctx->GetDmlDevice(), GetDmlXTensorPolicy(tensor_format));
+        dml::Graph(ctx->GetDmlDevice(), GetDmlXTensorPolicy(tensor_format));
     auto x = dml::InputTensor(scope, 0, input_descs[0]);
     auto scale = dml::InputTensor(scope, 1, input_descs[1]);
     auto offset = dml::InputTensor(scope, 2, input_descs[2]);
@@ -451,7 +451,7 @@ class DmlFusedBatchNormKernel : public DmlKernel {
     auto output_descs = GetDmlTensorDescs(tensors.outputs);
 
     auto scope =
-        dml::Scope(ctx->GetDmlDevice(), GetDmlXTensorPolicy(tensor_format));
+        dml::Graph(ctx->GetDmlDevice(), GetDmlXTensorPolicy(tensor_format));
     auto x = dml::InputTensor(scope, 0, input_descs[0]);
     auto mean = dml::InputTensor(scope, 1, input_descs[1]);
     auto variance = dml::InputTensor(scope, 2, input_descs[2]);
@@ -574,7 +574,7 @@ class DmlBatchNormWithGlobalNormalizationKernel : public DmlKernel {
     auto output_descs = GetDmlTensorDescs(tensors.outputs);
 
     const uint32_t beta_index = scale_after_normalization ? 4 : 3;
-    auto scope = dml::Scope(ctx->GetDmlDevice());
+    auto scope = dml::Graph(ctx->GetDmlDevice());
     auto t = dml::InputTensor(scope, 0, input_descs[0]);
     auto m = dml::InputTensor(scope, 1, input_descs[1]);
     auto v = dml::InputTensor(scope, 2, input_descs[2]);
@@ -684,7 +684,7 @@ class DmlFusedBatchNormGradKernel : public DmlKernel {
     auto output_descs = GetDmlTensorDescs(tensors.outputs);
 
     auto scope =
-        dml::Scope(ctx->GetDmlDevice(), GetDmlXTensorPolicy(tensor_format));
+        dml::Graph(ctx->GetDmlDevice(), GetDmlXTensorPolicy(tensor_format));
 
     auto y_backprop =
         dml::InputTensor(scope, kYBackprop, input_descs[kYBackprop]);
@@ -885,7 +885,7 @@ class DmlBatchGlobalNormGradKernel : public DmlKernel {
     auto input_descs = GetDmlTensorDescs(tensors.inputs);
     auto output_descs = GetDmlTensorDescs(tensors.outputs);
 
-    auto scope = dml::Scope(ctx->GetDmlDevice());
+    auto scope = dml::Graph(ctx->GetDmlDevice());
 
     const uint32_t back_prop_index =
         scale_after_normalization ? kBackProp : kBackProp - 1;

tensorflow/core/kernels/dml_batch_to_space_op.cc

@@ -310,7 +310,7 @@ class DmlBatchToSpaceKernel : public DmlKernel {
       return;
     }
 
-    auto scope = dml::Scope(ctx->GetDmlDevice());
+    auto scope = dml::Graph(ctx->GetDmlDevice());
     auto input = dml::InputTensor(scope, 0, inputs[0]);
 
     absl::Span<const int64> internal_block_sizes =
@@ -393,7 +393,7 @@ class DmlBatchToSpaceKernel : public DmlKernel {
     // Finally, slice the appropriate dimensions
     dml::TensorDesc::Dimensions slice_offsets(perm_reshaped_sizes.size());
     dml::TensorDesc::Dimensions slice_sizes = perm_reshaped_sizes;
-    dml::TensorDesc::Dimensions slice_strides(perm_reshaped_sizes.size(), 1);
+    absl::InlinedVector<int32_t, 4> slice_strides(perm_reshaped_sizes.size(), 1);
 
     absl::Span<const int64> internal_crops = init_helper->GetInternalCrops();
 

tensorflow/core/kernels/dml_check_numerics_op.cc

@@ -73,7 +73,7 @@ class DmlCheckNumericsKernel : public DmlKernel {
     tensors.outputs = {output};
 
     auto inputs = GetDmlTensorDescs(tensors.inputs);
-    auto scope = dml::Scope(ctx->GetDmlDevice());
+    auto scope = dml::Graph(ctx->GetDmlDevice());
     auto input_tensor = dml::InputTensor(scope, 0, inputs[0]);
 
     // Reduce doesn't support less than 32bit integer datatypes, so we need to

tensorflow/core/kernels/dml_conv_ops.cc

@@ -539,7 +539,7 @@ class DmlFusedConv2DKernel : public DmlKernel {
     auto input_descs = GetDmlTensorDescs(tensors.inputs);
     auto output_descs = GetDmlTensorDescs(tensors.outputs);
 
-    auto scope = dml::Scope(ctx->GetDmlDevice(), GetDmlXTensorPolicy(conv_params.data_format));
+    auto scope = dml::Graph(ctx->GetDmlDevice(), GetDmlXTensorPolicy(conv_params.data_format));
     auto input = dml::InputTensor(scope, 0, input_descs[0]);
     auto filter = dml::InputTensor(scope, 1, input_descs[1]);
 

tensorflow/core/kernels/dml_cwise_ops.cc

@@ -184,7 +184,7 @@ class DmlCompositeBinaryKernel : public DmlKernel {
     auto inputs = GetDmlTensorDescs(tensors.inputs);
     auto outputs = GetDmlTensorDescs(tensors.outputs);
 
-    auto scope = dml::Scope(ctx->GetDmlDevice());
+    auto scope = dml::Graph(ctx->GetDmlDevice());
     auto x = dml::InputTensor(scope, 0, inputs[0]);
     auto y = dml::InputTensor(scope, 1, inputs[1]);
 
@@ -319,7 +319,7 @@ class DmlCompositeUnaryKernel : public DmlKernel {
     auto inputs = GetDmlTensorDescs(tensors.inputs);
     auto outputs = GetDmlTensorDescs(tensors.outputs);
 
-    auto scope = dml::Scope(ctx->GetDmlDevice());
+    auto scope = dml::Graph(ctx->GetDmlDevice());
     auto x = dml::InputTensor(scope, 0, inputs[0]);
 
     ExpressionFunctor expression;
@@ -803,7 +803,7 @@ class DmlBinaryWithZeroKernel : public DmlKernel {
     auto inputs = GetDmlTensorDescs(tensors.inputs);
     auto outputs = GetDmlTensorDescs(tensors.outputs);
 
-    auto scope = dml::Scope(ctx->GetDmlDevice());
+    auto scope = dml::Graph(ctx->GetDmlDevice());
     auto x = dml::InputTensor(scope, 0, inputs[0]);
     auto y = dml::InputTensor(scope, 1, inputs[1]);
     auto zero = dml::ZeroTensor(scope, x.GetOutputDesc().dataType,
@@ -914,7 +914,7 @@ class DmlSquaredDifferenceKernel : public DmlKernel {
     auto inputs = GetDmlTensorDescs(tensors.inputs);
     auto outputs = GetDmlTensorDescs(tensors.outputs);
 
-    auto scope = dml::Scope(ctx->GetDmlDevice());
+    auto scope = dml::Graph(ctx->GetDmlDevice());
     auto x = dml::InputTensor(scope, 0, inputs[0]);
     auto y = dml::InputTensor(scope, 1, inputs[1]);
     auto diff = x - y;
@@ -1043,7 +1043,7 @@ class DmlApproximateEqualKernel : public DmlKernel {
     auto inputs = GetDmlTensorDescs(tensors.inputs);
     auto outputs = GetDmlTensorDescs(tensors.outputs);
 
-    auto scope = dml::Scope(ctx->GetDmlDevice());
+    auto scope = dml::Graph(ctx->GetDmlDevice());
     auto x = dml::InputTensor(scope, 0, inputs[0]);
     auto y = dml::InputTensor(scope, 1, inputs[1]);
 
@@ -1196,7 +1196,7 @@ class DmlBitCountKernel : public DmlKernel {
       // 2D so that we can reduce each adjacent pair of counts.
       dml::TensorDesc::Dimensions double_sizes = {1, 1, num_elements, 2};
 
-      auto scope = dml::Scope(ctx->GetDmlDevice());
+      auto scope = dml::Graph(ctx->GetDmlDevice());
       auto in_64_bit = dml::InputTensor(scope, 0, in_desc);
       auto in_32_bit = dml::Reinterpret(in_64_bit, DML_TENSOR_DATA_TYPE_UINT32,
                                         double_sizes, dml::NullOpt);

tensorflow/core/kernels/dml_data_format_dim_map.cc

@@ -82,7 +82,7 @@ class DmlDataFormaDimMapKernel : public DmlKernel {
       }
     }
 
-    auto scope = dml::Scope(ctx->GetDmlDevice());
+    auto scope = dml::Graph(ctx->GetDmlDevice());
 
     DmlKernelTensors tensors = GetTensorInfos(ctx, {});
     auto inputs = GetDmlTensorDescs(tensors.inputs);