text: /*******************************************************************************
  * Copyright 2019-2021 Intel Corporation
  *
  * 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 GPU_JIT_GEMM_GEN_GEMM_KERNEL_HPP
  #define GPU_JIT_GEMM_GEN_GEMM_KERNEL_HPP
  
  #include "common/c_types_map.hpp"
  #include "common/type_helpers.hpp"
  #include "gen_gemm_kernel_generator.hpp"
  #include "gpu/compute/compute.hpp"
  #include "gpu/compute/device_info.hpp"
  #include "gpu/jit/jit_generator_base.hpp"
  #include "gpu/jit/ngen_type_bridge.hpp"
  #include "gpu/primitive_conf.hpp"
  
  namespace dnnl {
  namespace impl {
  namespace gpu {
  namespace jit {
  
  struct gen_gemm_kernel_t : public jit_generator_base {
  
      status_t init_gemm(compute::gpu_arch_t arch) {
          hw_ = convert_dnnl_arch_to_hw(arch);
          if (hw_ == ngen::HW::Unknown) return status::unimplemented;
  
          auto status = complete_strategy();
          if (status != status::success) return status;
  
          return init_interface();
      }
  
      const char *kernel_name() const override { return "gemm_kernel"; }
      cl_kernel get_kernel(cl_context context, cl_device_id device) override;
  
      CommonDriverInfo driver_info() const;
  
  protected:
      static Type convert_dnnl_to_kernel_type(data_type_t type) {
          switch (type) {
              default: assert(!"Unknown type");
              case data_type::f32: return Type::f32;
              case data_type::f16: return Type::f16;
              case data_type::s32: return Type::s32;
              case data_type::u8: return Type::u8;
              case data_type::s8: return Type::s8;
          }
      }
  
      static ngen::HW convert_dnnl_arch_to_hw(compute::gpu_arch_t arch) {
          switch (arch) {
              case compute::gpu_arch_t::gen9: return ngen::HW::Gen9;
              case compute::gpu_arch_t::xe_lp: return ngen::HW::Xe_LP;
              default: return ngen::HW::Unknown;
          }
      }
  
      ngen::HW hw_ = ngen::HW::Unknown;
      GEMMProblem problem_;
      GEMMStrategy strategy_;
      ngen::NEOInterfaceHandler interface_ {ngen::HW::Unknown};
  
  private:
      static bool matching_hw(ngen::HW hw, ngen::HW hw_ref);
      status_t complete_strategy();
      status_t read_strategy(const char *str);
      status_t init_interface();
  };
  
  struct gen_gemm_nocopy_kernel_t : public gen_gemm_kernel_t {
      status_t init(compute::gpu_arch_t arch, int batch_dims, bool trans_a,
              bool trans_b, bool c_offset, bool bias, const post_ops_t &post_ops,
              data_type_t a_type, data_type_t b_type, data_type_t c_type,
              int unroll_m, int unroll_n) {
  
          problem_.Ta = convert_dnnl_to_kernel_type(a_type);
          problem_.Tb = convert_dnnl_to_kernel_type(b_type);
          problem_.Tc = convert_dnnl_to_kernel_type(c_type);
          problem_.Tc_ext = problem_.Ts = problem_.Tc;
          problem_.A.layout = trans_a ? MatrixLayout::T : MatrixLayout::N;
          problem_.B.layout = trans_b ? MatrixLayout::T : MatrixLayout::N;
          problem_.C.layout = MatrixLayout::N;
          problem_.A.crosspack = problem_.B.crosspack = problem_.C.crosspack = 1;
          problem_.A.packSize = problem_.B.packSize = problem_.C.packSize = 0;
          problem_.A.padded = problem_.B.padded = problem_.C.padded = false;
          problem_.A.setAlignment(int(types::data_type_size(a_type)));
          problem_.B.setAlignment(int(types::data_type_size(b_type)));
          problem_.C.setAlignment(int(types::data_type_size(c_type)));
          problem_.A.base = ngen::AddressBase::createA64(true);
          problem_.B.base = ngen::AddressBase::createA64(true);
          problem_.C.base = ngen::AddressBase::createA64(true);
          if (batch_dims > 0) {
              problem_.batch = BatchMode::Strided;
              problem_.batchDims = batch_dims;
          }
          if (c_type == data_type::s32) {
              problem_.abOffset = ABOffset::Calc;
              problem_.Ts = Type::f32;
          }
          if (post_ops.len() > 0) {
              problem_.post_ops = post_ops;
              if (a_type == data_type::f16) problem_.Ts = Type::f32;
          }
          if (c_offset || bias) {
              assert(!(c_offset && bias));
              problem_.cOffset = bias ? COffset::Pre : COffset::Post;
              problem_.CO.base = ngen::AddressBase::createBTS(0);
              problem_.CO.crosspack = 1;
              problem_.CO.padded = false;
              problem_.CO.alignment = problem_.C.alignment;
          }
  
          strategy_.unroll[LoopM] = unroll_m;
          strategy_.unroll[LoopN] = unroll_n;
  
          return init_gemm(arch);
      }
  
      static void choose_unrolls(compute::gpu_arch_t arch, int hw_threads,
              bool trans_a, bool trans_b, data_type_t a_type, data_type_t b_type,
              data_type_t c_type, dim_t m, dim_t n, dim_t k, dim_t batch,
              int &unroll_m, int &unroll_n);
  };
  
  } // namespace jit
  } // namespace gpu
  } // namespace impl
  } // namespace dnnl
  
  #endif
  