tensorflow/compiler/mlir/tfrt/transforms/lmhlo_to_gpu/convolution_pattern.cc
// Create op to build a convolution plan, which can be used to run the
// convolution. This is the variant with fused activation.
Create op to build a convolution plan, which can be used to run the
convolution. This is the variant with fused activation.
Create op to build a convolution plan which can be used to run the convolution . This is the variant with fused activation .
Value CreateBuildFusedConvPlanOp(Value input, Value output, Value bias,
                                 Value handle, mlir::Location loc,
                                 const xla::gpu::GpuConvConfig& config,
                                 cudnnBackendDescriptorType_t backend_type,
                                 ConversionPatternRewriter& rewriter) {
  se::dnn::BatchDescriptor bias_descriptor = GetBiasDescriptor(config);
  bias_descriptor.set_layout(config.output_descriptor.layout());
  tfrt::gpu::wrapper::DnnDataType input_type = MlirTypeToDnnDataType(
      GetMemRefElementType(input), config.input_descriptor.layout());
  tfrt::gpu::wrapper::DnnDataType output_type = MlirTypeToDnnDataType(
      GetMemRefElementType(output), config.output_descriptor.layout());
  tfrt::gpu::wrapper::DnnDataType bias_type = MlirTypeToDnnDataType(
      GetMemRefElementType(bias), bias_descriptor.layout());
  int vector_size, vector_dim;
  std::tie(vector_size, vector_dim) =
      tfrt::gpu::wrapper::GetTensorVectorizedSizeAndDim(input_type);
  std::vector<int64_t> input_dims = config.input_descriptor.vectorized_dims(
      se::dnn::DataLayout::kBatchDepthYX, vector_size, vector_dim);
  std::vector<int64_t> input_strides =
      config.input_descriptor.vectorized_strides(
          se::dnn::DataLayout::kBatchDepthYX, vector_size, vector_dim);
  std::tie(vector_size, vector_dim) =
      tfrt::gpu::wrapper::GetTensorVectorizedSizeAndDim(output_type);
  std::vector<int64_t> output_dims = config.output_descriptor.vectorized_dims(
      se::dnn::DataLayout::kBatchDepthYX, vector_size, vector_dim);
  std::vector<int64_t> output_strides =
      config.output_descriptor.vectorized_strides(
          se::dnn::DataLayout::kBatchDepthYX, vector_size, vector_dim);
  std::tie(vector_size, vector_dim) =
      tfrt::gpu::wrapper::GetTensorVectorizedSizeAndDim(input_type);
  std::vector<int64_t> filter_dims = config.filter_descriptor.vectorized_dims(
      se::dnn::FilterLayout::kOutputInputYX, vector_size, vector_dim);
  std::vector<int64_t> filter_strides =
      config.filter_descriptor.vectorized_strides(
          se::dnn::FilterLayout::kOutputInputYX, vector_size, vector_dim);
  std::tie(vector_size, vector_dim) =
      tfrt::gpu::wrapper::GetTensorVectorizedSizeAndDim(input_type);
  std::vector<int64_t> bias_dims = bias_descriptor.vectorized_dims(
      se::dnn::DataLayout::kBatchDepthYX, vector_size, vector_dim);
  std::vector<int64_t> bias_strides = bias_descriptor.vectorized_strides(
      se::dnn::DataLayout::kBatchDepthYX, vector_size, vector_dim);
  const auto* conv_desc = &config.conv_desc;
  cudnnConvolutionMode_t conv_mode =
      config.conv_desc.convolution_not_crosscorr() ? CUDNN_CONVOLUTION
                                                   : CUDNN_CROSS_CORRELATION;
  int conv_dim = config.conv_desc.ndims();
  auto conv_dialations = xla::llvm_ir::AsArrayRef(conv_desc->dilations());
  auto conv_padding = xla::llvm_ir::AsArrayRef(conv_desc->padding());
  auto conv_strides = xla::llvm_ir::AsArrayRef(conv_desc->strides());
  std::vector<int64_t> tuning_knob_ids, tuning_knob_values;
  tuning_knob_ids.reserve(config.algorithm.TuningKnobs().size());
  tuning_knob_values.reserve(config.algorithm.TuningKnobs().size());
  for (auto iter : config.algorithm.TuningKnobs()) {
    tuning_knob_ids.push_back(iter.first);
    tuning_knob_values.push_back(iter.second);
  }
  auto alpha = rewriter.create<tfrt::compiler::ConstantF64Op>(
      loc, llvm::APFloat(config.conv_result_scale));
  auto alpha2 = rewriter.create<tfrt::compiler::ConstantF64Op>(
      loc, llvm::APFloat(config.fusion->side_input_scale));
  cudnnActivationMode_t activaton_mode = config.fusion->mode == se::dnn::kRelu
                                             ? CUDNN_ACTIVATION_RELU
                                             : CUDNN_ACTIVATION_IDENTITY;
  return rewriter.create<tfrt::gpu::DnnBuildFusedConvolutionOp>(
      loc, handle, input_type, output_type, bias_type,
      rewriter.getI64ArrayAttr(input_dims),
      rewriter.getI64ArrayAttr(input_strides),
      rewriter.getI64ArrayAttr(output_dims),
      rewriter.getI64ArrayAttr(output_strides),
      rewriter.getI64ArrayAttr(filter_dims),
      rewriter.getI64ArrayAttr(filter_strides),
      rewriter.getI64ArrayAttr(bias_dims),
      rewriter.getI64ArrayAttr(bias_strides), conv_mode, conv_dim,
      rewriter.getI64ArrayAttr(conv_dialations),
      rewriter.getI64ArrayAttr(conv_padding),
      rewriter.getI64ArrayAttr(conv_strides), backend_type, alpha, alpha2,
      activaton_mode, config.algorithm.algo_id(),
      rewriter.getI64ArrayAttr(tuning_knob_ids),
      rewriter.getI64ArrayAttr(tuning_knob_values));
}
Value CreateBuildFusedConvPlanOp ( Value input , Value output , Value bias , Value handle , mlir :: Location loc , const xla :: gpu :: GpuConvConfig & config , cudnnBackendDescriptorType_t backend_type , ConversionPatternRewriter & rewriter ) { se :: dnn :: BatchDescriptor bias_descriptor = GetBiasDescriptor ( config ) ; bias_descriptor . set_layout ( config . output_descriptor . layout ( ) ) ; tfrt :: gpu :: wrapper :: DnnDataType input_type = MlirTypeToDnnDataType ( GetMemRefElementType ( input ) , config . input_descriptor . layout ( ) ) ; tfrt :: gpu :: wrapper :: DnnDataType output_type = MlirTypeToDnnDataType ( GetMemRefElementType ( output ) , config . output_descriptor . layout ( ) ) ; tfrt :: gpu :: wrapper :: DnnDataType bias_type = MlirTypeToDnnDataType ( GetMemRefElementType ( bias ) , bias_descriptor . layout ( ) ) ; int vector_size , vector_dim ; std :: tie ( vector_size , vector_dim ) = tfrt :: gpu :: wrapper :: GetTensorVectorizedSizeAndDim ( input_type ) ; std :: vector < int64_t > input_dims = config . input_descriptor . vectorized_dims ( se :: dnn :: DataLayout :: kBatchDepthYX , vector_size , vector_dim ) ; std :: vector < int64_t > input_strides = config . input_descriptor . vectorized_strides ( se :: dnn :: DataLayout :: kBatchDepthYX , vector_size , vector_dim ) ; std :: tie ( vector_size , vector_dim ) = tfrt :: gpu :: wrapper :: GetTensorVectorizedSizeAndDim ( output_type ) ; std :: vector < int64_t > output_dims = config . output_descriptor . vectorized_dims ( se :: dnn :: DataLayout :: kBatchDepthYX , vector_size , vector_dim ) ; std :: vector < int64_t > output_strides = config . output_descriptor . vectorized_strides ( se :: dnn :: DataLayout :: kBatchDepthYX , vector_size , vector_dim ) ; std :: tie ( vector_size , vector_dim ) = tfrt :: gpu :: wrapper :: GetTensorVectorizedSizeAndDim ( input_type ) ; std :: vector < int64_t > filter_dims = config . filter_descriptor . vectorized_dims ( se :: dnn :: FilterLayout :: kOutputInputYX , vector_size , vector_dim ) ; std :: vector < int64_t > filter_strides = config . filter_descriptor . vectorized_strides ( se :: dnn :: FilterLayout :: kOutputInputYX , vector_size , vector_dim ) ; std :: tie ( vector_size , vector_dim ) = tfrt :: gpu :: wrapper :: GetTensorVectorizedSizeAndDim ( input_type ) ; std :: vector < int64_t > bias_dims = bias_descriptor . vectorized_dims ( se :: dnn :: DataLayout :: kBatchDepthYX , vector_size , vector_dim ) ; std :: vector < int64_t > bias_strides = bias_descriptor . vectorized_strides ( se :: dnn :: DataLayout :: kBatchDepthYX , vector_size , vector_dim ) ; const auto * conv_desc = & config . conv_desc ; cudnnConvolutionMode_t conv_mode = config . conv_desc . convolution_not_crosscorr ( ) ? CUDNN_CONVOLUTION : CUDNN_CROSS_CORRELATION ; int conv_dim = config . conv_desc . ndims ( ) ; auto conv_dialations = xla :: llvm_ir :: AsArrayRef ( conv_desc -> dilations ( ) ) ; auto conv_padding = xla :: llvm_ir :: AsArrayRef ( conv_desc -> padding ( ) ) ; auto conv_strides = xla :: llvm_ir :: AsArrayRef ( conv_desc -> strides ( ) ) ; std :: vector < int64_t > tuning_knob_ids , tuning_knob_values ; tuning_knob_ids . reserve ( config . algorithm . TuningKnobs ( ) . size ( ) ) ; tuning_knob_values . reserve ( config . algorithm . TuningKnobs ( ) . size ( ) ) ; for ( auto iter : config . algorithm . TuningKnobs ( ) ) { tuning_knob_ids . push_back ( iter . first ) ; tuning_knob_values . push_back ( iter . second ) ; } auto alpha = rewriter . create < tfrt :: compiler :: ConstantF64Op > ( loc , llvm :: APFloat ( config . conv_result_scale ) ) ; auto alpha2 = rewriter . create < tfrt :: compiler :: ConstantF64Op > ( loc , llvm :: APFloat ( config . fusion -> side_input_scale ) ) ; cudnnActivationMode_t activaton_mode = config . fusion -> mode == se :: dnn :: kRelu ? CUDNN_ACTIVATION_RELU : CUDNN_ACTIVATION_IDENTITY ; return rewriter . create < tfrt :: gpu :: DnnBuildFusedConvolutionOp > ( loc , handle , input_type , output_type , bias_type , rewriter . getI64ArrayAttr ( input_dims ) , rewriter . getI64ArrayAttr ( input_strides ) , rewriter . getI64ArrayAttr ( output_dims ) , rewriter . getI64ArrayAttr ( output_strides ) , rewriter . getI64ArrayAttr ( filter_dims ) , rewriter . getI64ArrayAttr ( filter_strides ) , rewriter . getI64ArrayAttr ( bias_dims ) , rewriter . getI64ArrayAttr ( bias_strides ) , conv_mode , conv_dim , rewriter . getI64ArrayAttr ( conv_dialations ) , rewriter . getI64ArrayAttr ( conv_padding ) , rewriter . getI64ArrayAttr ( conv_strides ) , backend_type , alpha , alpha2 , activaton_mode , config . algorithm . algo_id ( ) , rewriter . getI64ArrayAttr ( tuning_knob_ids ) , rewriter . getI64ArrayAttr ( tuning_knob_values ) ) ; }
Create op to build a convolution plan, which can be used to run the
convolution.
Create op to build a convolution plan which can be used to run the convolution .
// For the purposes of the cudnn graph, say that the bias tensor has the same // layout as the output tensor.  It doesn't actually matter, because bias is a // 1D array.  But we need to get the correct vectorization, otherwise the // cudnn graph API rejects this tensor, even though vectorized float tensors // aren't even a thing in cuDNN.