# Copyright 2023 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import("//build/buildflag_header.gni")
import("//mojo/public/tools/fuzzers/mojolpm.gni")
import("//services/on_device_model/on_device_model.gni")
import("//services/webnn/features.gni")
import("//third_party/protobuf/proto_library.gni")
import("//third_party/tflite/features.gni")

buildflag_header("buildflags") {
  header = "buildflags.h"
  flags = [
    "WEBNN_USE_TFLITE=$webnn_use_tflite",
    "WEBNN_USE_LITERT=$webnn_use_litert",
    "WEBNN_USE_CHROME_ML_API=$webnn_use_chrome_ml_api",
    "WEBNN_ENABLE_TFLITE_PROFILER=$webnn_enable_tflite_profiler",
    "WEBNN_INSTALL_RUNTIME_IN_CHROME_INSTALLER=$webnn_install_runtime_in_chrome_installer",
  ]
}

source_set("webnn_switches") {
  sources = [
    "webnn_switches.cc",
    "webnn_switches.h",
  ]
  deps = [
    ":buildflags",
    "//base",
  ]
}

source_set("webnn_service_base") {
  visibility = [ ":*" ]

  sources = [
    "error.h",
    "gpu_task_scheduler.cc",
    "gpu_task_scheduler.h",
    "graph_builder_context.cc",
    "graph_builder_context.h",
    "sequence_deleter.cc",
    "sequence_deleter.h",
    "webnn_constant_operand.cc",
    "webnn_constant_operand.h",
    "webnn_context_impl.cc",
    "webnn_context_impl.h",
    "webnn_context_provider_impl.cc",
    "webnn_context_provider_impl.h",
    "webnn_context_provider_in_renderer.cc",
    "webnn_context_provider_in_renderer.h",
    "webnn_graph_builder_impl.cc",
    "webnn_graph_builder_impl.h",
    "webnn_graph_impl.cc",
    "webnn_graph_impl.h",
    "webnn_object_impl.h",
    "webnn_pending_constant_operand.cc",
    "webnn_pending_constant_operand.h",
    "webnn_tensor_impl.cc",
    "webnn_tensor_impl.h",
    "webnn_utils.cc",
    "webnn_utils.h",
  ]

  deps = [
    ":buildflags",
    ":webnn_switches",
    "//base",
    "//gpu/command_buffer/service",
    "//gpu/command_buffer/service:gles2",
    "//gpu/config",
    "//mojo/public/cpp/bindings",
    "//services/viz/privileged/mojom/gl",
    "//services/webnn/public/mojom",
    "//third_party/tflite:buildflags",
  ]

  if (build_tflite_with_xnnpack) {
    deps += [
      "//third_party/pthreadpool",
      "//third_party/xnnpack",
    ]
  }

  if (is_win) {
    deps += [ "//services/webnn/public/cpp:win" ]
  }

  configs += [ "//build/config/compiler:prevent_unsafe_narrowing" ]

  defines = [ "IS_WEBNN_SERVICE_IMPL" ]
}

if (is_apple || webnn_use_tflite || webnn_use_litert) {
  source_set("webnn_service_resource_management") {
    visibility = [ ":*" ]

    sources = [
      "queueable_resource_state.h",
      "queueable_resource_state_base.cc",
      "queueable_resource_state_base.h",
      "resource_task.cc",
      "resource_task.h",
    ]

    deps = [ "//base" ]

    configs += [ "//build/config/compiler:prevent_unsafe_narrowing" ]

    defines = [ "IS_WEBNN_SERVICE_IMPL" ]
  }
}

component("webnn_service") {
  _configs = configs
  configs = []
  configs = [ "//third_party/tflite:tflite_shim_config" ] + _configs
  defines = [ "IS_WEBNN_SERVICE_IMPL" ]

  sources = []

  public_deps = [ ":webnn_service_base" ]

  deps = [
    ":buildflags",
    ":webnn_switches",
    "//base",
    "//gpu/command_buffer/service",
    "//gpu/command_buffer/service:gles2",
    "//gpu/config",
    "//mojo/public/cpp/bindings",
    "//services/viz/privileged/mojom/gl",
    "//services/webnn/public/mojom",
    "//third_party/tflite:buildflags",
  ]

  if (is_apple) {
    sources += [
      "coreml/buffer_content_coreml.h",
      "coreml/buffer_content_coreml.mm",
      "coreml/compiler_context_impl_coreml.h",
      "coreml/compiler_context_impl_coreml.mm",
      "coreml/context_impl_coreml.h",
      "coreml/context_impl_coreml.mm",
      "coreml/graph_builder_coreml.cc",
      "coreml/graph_builder_coreml.h",
      "coreml/graph_impl_coreml.h",
      "coreml/graph_impl_coreml.mm",
      "coreml/tensor_impl_coreml.h",
      "coreml/tensor_impl_coreml.mm",
      "coreml/utils_coreml.h",
      "coreml/utils_coreml.mm",
    ]
    frameworks = [
      "CoreFoundation.framework",
      "CoreML.framework",
      "CoreVideo.framework",
      "Foundation.framework",
      "IOSurface.framework",
    ]
    deps += [
      "//third_party/coremltools:modelformat_proto",
      "//third_party/fp16",
    ]
  }

  if (webnn_use_tflite || webnn_use_litert) {
    sources += [
      "tflite/buffer_content_tflite.cc",
      "tflite/buffer_content_tflite.h",
      "tflite/graph_builder_tflite.cc",
      "tflite/graph_builder_tflite.h",
      "tflite/tensor_impl_tflite.cc",
      "tflite/tensor_impl_tflite.h",
    ]

    if (webnn_use_litert) {
      sources += [
        "tflite/context_impl_litert.cc",
        "tflite/context_impl_litert.h",
        "tflite/graph_impl_litert.cc",
        "tflite/graph_impl_litert.h",
      ]
    }

    deps += [
      "//third_party/flatbuffers",
      "//third_party/fp16",
      "//third_party/tflite",
      "//third_party/tflite:tflite_builtin_op_resolver",
      "//third_party/tflite:tflite_public_headers",
    ]

    if (webnn_use_litert) {
      deps += [
        "//third_party/dawn/src/dawn/native",
        "//third_party/litert",
        "//third_party/litert:buildflags",
      ]
    }

    if (enable_ml_internal) {
      deps += [ "//services/on_device_model/ml" ]
    }
  }

  if (is_win) {
    sources += [
      "ort/compiler_context_impl_ort.cc",
      "ort/compiler_context_impl_ort.h",
      "ort/context_impl_ort.cc",
      "ort/context_impl_ort.h",
      "ort/context_provider_ort.cc",
      "ort/context_provider_ort.h",
      "ort/device_allocator.cc",
      "ort/device_allocator.h",
      "ort/dispatch_context_impl_ort.cc",
      "ort/dispatch_context_impl_ort.h",
      "ort/environment.cc",
      "ort/environment.h",
      "ort/external_weights_manager.cc",
      "ort/external_weights_manager.h",
      "ort/graph_builder_ort.cc",
      "ort/graph_builder_ort.h",
      "ort/graph_impl_ort.cc",
      "ort/graph_impl_ort.h",
      "ort/logging.cc",
      "ort/logging.h",
      "ort/model_editor.cc",
      "ort/model_editor.h",
      "ort/ort_data_type.cc",
      "ort/ort_data_type.h",
      "ort/ort_session_options.cc",
      "ort/ort_session_options.h",
      "ort/ort_status.cc",
      "ort/ort_status.h",
      "ort/platform_functions_ort.cc",
      "ort/platform_functions_ort.h",
      "ort/scoped_ort_types.h",
      "ort/tensor_impl_ort.cc",
      "ort/tensor_impl_ort.h",
      "ort/trivial_model.h",
      "webnn_compiler_service_impl.cc",
      "webnn_compiler_service_impl.h",
    ]
    deps += [
      "//services/webnn/public/cpp:win",
      "//third_party/windows_app_sdk_headers",
    ]

    libs = [ "runtimeobject.lib" ]

    # TODO(crbug.com/40031409): Fix code that adds exit-time destructors and
    # enable the diagnostic by removing this line.
    configs += [ "//build/config/compiler:no_exit_time_destructors" ]
  }

  if (is_apple || webnn_use_tflite || webnn_use_litert) {
    deps += [ ":webnn_service_resource_management" ]
  }
}

source_set("webnn_test_support") {
  testonly = true
  sources = [
    "webnn_test_environment.cc",
    "webnn_test_environment.h",
  ]

  deps = [
    ":webnn_service",
    "//base",
    "//base/test:test_support",
    "//gpu/command_buffer/service",
    "//gpu/config",
    "//mojo/public/cpp/bindings",
    "//services/viz/privileged/mojom/gl",
    "//services/webnn/host",
    "//services/webnn/public/mojom",
  ]

  public_deps = [ ":buildflags" ]

  if (webnn_use_tflite || webnn_use_litert) {
    deps += [ "//services/webnn/public/cpp:in_process_context_provider" ]
  }

  if (is_win) {
    deps += [
      "//services/webnn/public/cpp:win",
      "//ui/gfx/mojom",
    ]
  }
}

source_set("tests") {
  testonly = true

  sources = [
    "webnn_context_provider_impl_unittest.cc",
    "webnn_graph_builder_impl_unittest.cc",
    "webnn_graph_impl_unittest.cc",
    "webnn_test_utils.cc",
    "webnn_test_utils.h",
  ]

  if (webnn_use_tflite || webnn_use_litert || is_mac) {
    sources += [
      "webnn_graph_impl_backend_test.cc",
      "webnn_tensor_impl_backend_test.cc",
    ]
  }

  if (is_win) {
    sources += [
      "ort/model_editor_test.cc",
      "ort/platform_functions_ort_test.cc",
      "ort/test_base_ort.cc",
      "ort/test_base_ort.h",
    ]
  }

  deps = [
    ":buildflags",
    ":webnn_service",
    ":webnn_test_support",
    "//base",
    "//base/test:test_support",
    "//gpu/config",
    "//mojo/public/cpp/bindings",
    "//mojo/public/cpp/test_support:test_utils",
    "//services/webnn/public/mojom",
    "//testing/gtest",
    "//third_party/fp16",
  ]

  if (is_win) {
    deps += [
      "//ui/gl",
      "//ui/gl/init",
    ]
  }
}

# TODO(crbug.com/500747202): Re-enable in debug builds when timeout issues is
# fixed.
if (!is_debug) {
  test("webnn_graph_impl_fuzzer") {
    sources = [
      "webnn_graph_impl_fuzzer.cc",
      "webnn_test_utils.cc",
      "webnn_test_utils.h",
    ]

    fuzztests = [
      "CPU.Activation",
      "CPU.ArgMinMax",
      "CPU.BatchNormalization",
      "CPU.Clamp",
      "CPU.Concat",
      "CPU.Conv2d",
      "CPU.CumulativeSum",
      "CPU.DequantizeLinear",
      "CPU.ElementWiseBinary",
      "CPU.Elu",
      "CPU.Expand",
      "CPU.Gather",
      "CPU.GatherElements",
      "CPU.GatherND",
      "CPU.Gemm",
      "CPU.Gru",
      "CPU.GruCell",
      "CPU.HardSigmoid",
      "CPU.InstanceNormalization",
      "CPU.LayerNormalization",
      "CPU.LeakyRelu",
      "CPU.Linear",
      "CPU.Lstm",
      "CPU.LstmCell",
      "CPU.Matmul",
      "CPU.Pad",
      "CPU.Pool2d",
      "CPU.Prelu",
      "CPU.QuantizeLinear",
      "CPU.Reduce",
      "CPU.Resample2d",
      "CPU.Reshape",
      "CPU.Reverse",
      "CPU.ScatterElements",
      "CPU.ScatterND",
      "CPU.Slice",
      "CPU.Softmax",
      "CPU.Split",
      "CPU.Tile",
      "CPU.Transpose",
      "CPU.Triangular",
      "CPU.Where",
      "CPU.DQActivationQ",
      "CPU.DQArgMax",
      "CPU.DQClampQ",
      "CPU.DQConcatQ",
      "CPU.DQConv2dQ",
      "CPU.DQElementWiseBinaryQ",
      "CPU.DQEluQ",
      "CPU.DQGatherQ",
      "CPU.DQGemmQ",
      "CPU.DQLeakyReluQ",
      "CPU.DQPadQ",
      "CPU.DQPool2dQ",
      "CPU.DQReduceQ",
      "CPU.DQResample2dQ",
      "CPU.DQReshapeQ",
      "CPU.DQSliceQ",
      "CPU.DQSoftmaxQ",
      "CPU.DQSplitQ",
      "CPU.DQTransposeQ",
      "GPU.Activation",
      "GPU.ArgMinMax",
      "GPU.BatchNormalization",
      "GPU.Clamp",
      "GPU.Concat",
      "GPU.Conv2d",
      "GPU.CumulativeSum",
      "GPU.DequantizeLinear",
      "GPU.ElementWiseBinary",
      "GPU.Elu",
      "GPU.Expand",
      "GPU.Gather",
      "GPU.GatherElements",
      "GPU.GatherND",
      "GPU.Gemm",
      "GPU.Gru",
      "GPU.GruCell",
      "GPU.HardSigmoid",
      "GPU.InstanceNormalization",
      "GPU.LayerNormalization",
      "GPU.LeakyRelu",
      "GPU.Linear",
      "GPU.Lstm",
      "GPU.LstmCell",
      "GPU.Matmul",
      "GPU.Pad",
      "GPU.Pool2d",
      "GPU.Prelu",
      "GPU.QuantizeLinear",
      "GPU.Reduce",
      "GPU.Resample2d",
      "GPU.Reshape",
      "GPU.Reverse",
      "GPU.ScatterElements",
      "GPU.ScatterND",
      "GPU.Slice",
      "GPU.Softmax",
      "GPU.Split",
      "GPU.Tile",
      "GPU.Transpose",
      "GPU.Triangular",
      "GPU.Where",
      "GPU.DQActivationQ",
      "GPU.DQArgMax",
      "GPU.DQClampQ",
      "GPU.DQConcatQ",
      "GPU.DQConv2dQ",
      "GPU.DQElementWiseBinaryQ",
      "GPU.DQEluQ",
      "GPU.DQGatherQ",
      "GPU.DQGemmQ",
      "GPU.DQLeakyReluQ",
      "GPU.DQPadQ",
      "GPU.DQPool2dQ",
      "GPU.DQReduceQ",
      "GPU.DQResample2dQ",
      "GPU.DQReshapeQ",
      "GPU.DQSliceQ",
      "GPU.DQSoftmaxQ",
      "GPU.DQSplitQ",
      "GPU.DQTransposeQ",
      "NPU.Activation",
      "NPU.ArgMinMax",
      "NPU.BatchNormalization",
      "NPU.Clamp",
      "NPU.Concat",
      "NPU.Conv2d",
      "NPU.CumulativeSum",
      "NPU.DequantizeLinear",
      "NPU.ElementWiseBinary",
      "NPU.Elu",
      "NPU.Expand",
      "NPU.Gather",
      "NPU.GatherElements",
      "NPU.GatherND",
      "NPU.Gemm",
      "NPU.Gru",
      "NPU.GruCell",
      "NPU.HardSigmoid",
      "NPU.InstanceNormalization",
      "NPU.LayerNormalization",
      "NPU.LeakyRelu",
      "NPU.Linear",
      "NPU.Lstm",
      "NPU.LstmCell",
      "NPU.Matmul",
      "NPU.Pad",
      "NPU.Pool2d",
      "NPU.Prelu",
      "NPU.QuantizeLinear",
      "NPU.Reduce",
      "NPU.Resample2d",
      "NPU.Reshape",
      "NPU.Reverse",
      "NPU.ScatterElements",
      "NPU.ScatterND",
      "NPU.Slice",
      "NPU.Softmax",
      "NPU.Split",
      "NPU.Tile",
      "NPU.Transpose",
      "NPU.Triangular",
      "NPU.Where",
      "NPU.DQActivationQ",
      "NPU.DQArgMax",
      "NPU.DQClampQ",
      "NPU.DQConcatQ",
      "NPU.DQConv2dQ",
      "NPU.DQElementWiseBinaryQ",
      "NPU.DQEluQ",
      "NPU.DQGatherQ",
      "NPU.DQGemmQ",
      "NPU.DQLeakyReluQ",
      "NPU.DQPadQ",
      "NPU.DQPool2dQ",
      "NPU.DQReduceQ",
      "NPU.DQResample2dQ",
      "NPU.DQReshapeQ",
      "NPU.DQSliceQ",
      "NPU.DQSoftmaxQ",
      "NPU.DQSplitQ",
      "NPU.DQTransposeQ",
    ]

    deps = [
      ":webnn_service",
      ":webnn_test_support",
      "//base",
      "//base/test:test_support",
      "//mojo/core/embedder",
      "//mojo/public/cpp/bindings",
      "//services/webnn/public/cpp",
      "//services/webnn/public/mojom",
      "//testing/gtest",
      "//third_party/fp16",
      "//third_party/fuzztest:fuzztest_gtest_main",
    ]

    # ASan inflates stack frames significantly, and the fuzzer deep call chain
    # overflows the default 1MB stack on Windows. Set to 8MB instead.
    if (is_win) {
      ldflags = [ "/STACK:0x800000" ]
    }
  }
}

mojolpm_fuzzer_test("webnn_graph_mojolpm_textproto_fuzzer") {
  sources = [ "webnn_graph_mojolpm_fuzzer.cc" ]

  proto_source = "webnn_graph_mojolpm_fuzzer.proto"

  proto_deps = [ "//services/webnn/public/mojom:mojom_mojolpm" ]
  seed_corpus = "webnn_graph_mojolpm_fuzzer_seed_corpus"
  deps = [
    ":webnn_service",
    ":webnn_test_support",
    "//base",
    "//base/test:test_support",
    "//content/test/fuzzer:mojolpm_fuzzer_support",
    "//services/webnn/public/mojom:mojom_mojolpm",
    "//third_party/libprotobuf-mutator",
  ]
}