/* Copyright 2023 The OpenXLA Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#ifndef XLA_SERVICE_CPU_GPU_SHAPE_VERIFIER_H_
#define XLA_SERVICE_CPU_GPU_SHAPE_VERIFIER_H_

#include <memory>
#include <utility>

#include "xla/service/hlo_verifier.h"

namespace xla {

// Verifies that HLO Shapes are supported by the XLA-CPU and XLA-GPU compilers.
class CpuGpuShapeVerifier : public ShapeVerifier {
 public:
  explicit CpuGpuShapeVerifier(const HloVerifierOpts& opts)
      : ShapeVerifier(opts) {}

  absl::Status Preprocess(HloInstruction* hlo) override;
};

// A verifier metadata class that uses the CpuGpuShapeVerifier.
class CpuGpuVerifierMetadata : public TargetVerifierMetadata {
 public:
  explicit CpuGpuVerifierMetadata(HloVerifierOpts&& opts)
      : TargetVerifierMetadata(std::move(opts)) {}

  std::unique_ptr<ShapeVerifier> GetVerifier() const override {
    return std::make_unique<CpuGpuShapeVerifier>(GetVerifierOpts());
  }
};

}  // namespace xla

#endif  // XLA_SERVICE_CPU_GPU_SHAPE_VERIFIER_H_