wehub-resource-sync
117 lines
4.1 KiB
C++
117 lines
4.1 KiB
C++
// Copyright (c) 2021 CINN Authors. All Rights Reserved.
|
|
//
|
|
// 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.
|
|
|
|
#pragma once
|
|
|
|
#include "paddle/cinn/ir/ir.h"
|
|
#include "paddle/cinn/ir/lowered_func.h"
|
|
#include "paddle/cinn/pass/pass.h"
|
|
|
|
namespace cinn {
|
|
namespace optim {
|
|
/**
|
|
* Optimizes GPU expressions by transforming variables, buffer indices, and
|
|
* memory access patterns for efficient GPU execution.
|
|
*
|
|
* This pass is applicable in scenarios where GPU-specific expressions need to
|
|
* be optimized for execution on GPU backends. This pass is essential in
|
|
* compiler pipelines that generate or transform GPU code, ensuring that
|
|
* variables and memory accesses are correctly mapped and optimized for GPU
|
|
* architecture.
|
|
*
|
|
* When applied, this pass performs a series of transformations on the IR to
|
|
* optimize expressions for GPU execution:
|
|
* 1) Variable and Loop Transformation
|
|
* 2) Buffer and Memory Access Optimization
|
|
* 3) Expression Simplification and Type Casting
|
|
*/
|
|
void OptimizeExprGPU(ir::stmt::BlockRef func_body);
|
|
|
|
/**
|
|
* Remove the GPU block/thread-bound For loops, add IfThenElse guards if needed.
|
|
*
|
|
* It's usually safe to remove bound loops, because when launching the kernel,
|
|
* we are expected to choose dim sizes that match the extents of these loops.
|
|
* However, there are cases where we cannot simply remove a loop, but need to
|
|
* add an IfThenElse as guard:
|
|
* 1) if the loop doesn't start from 0.
|
|
* 2) if we cannot prove that the loop's extent is always equal to or greater
|
|
* than the corresponding dim size.
|
|
*
|
|
* Example 1:
|
|
* # assume blockDim.x == 256
|
|
* thread_bind[threadIdx.x] for (k, 0, 256):
|
|
* ScheduleBlock(A)
|
|
* =>
|
|
* ScheduleBlock(A)
|
|
*
|
|
* Example 2:
|
|
* # assume gridDim.x == 8
|
|
* thread_bind[blockIdx.x] for (k, 2, min(S0, 8)):
|
|
* ScheduleBlock(A)
|
|
* =>
|
|
* if (blockIdx.x >= 2 && blockIdx.x < min(S0, 8)):
|
|
* ScheduleBlock(A)
|
|
*
|
|
* @param fn The LoweredFunc to process.
|
|
*/
|
|
class RemoveGpuForLoopsPass : public FuncPass {
|
|
public:
|
|
RemoveGpuForLoopsPass() : FuncPass("remove_gpu_for_loops") {}
|
|
|
|
LogicalResult Run(ir::LoweredFunc fn) override;
|
|
};
|
|
std::unique_ptr<FuncPass> CreateRemoveGpuForLoopsPass();
|
|
|
|
/**
|
|
* Removes conditional wrappers around CUDA thread synchronization calls.
|
|
*
|
|
* This pass is applicable in scenarios where CUDA synchronization functions,
|
|
* such as `cuda_sync_threads`, are enclosed within conditional statements
|
|
* (`IfThenElse`) that check if a certain variable equals zero. Such scenarios
|
|
* are common in auto-generated code or optimized code paths where
|
|
* synchronization is conditionally performed based on loop iterations or
|
|
* specific flags.
|
|
*
|
|
* When applied, this pass traverses the Intermediate Representation (IR) of a
|
|
* lowered function to identify `IfThenElse` nodes that contain
|
|
* `cuda_sync_threads` calls with conditions checking for equality to zero. For
|
|
* each identified conditional synchronization:
|
|
* 1) It verifies that the `IfThenElse` condition is an equality (`EQ`)
|
|
* comparison where the second operand is zero.
|
|
* 2) It replaces the entire `IfThenElse` node with the `cuda_sync_threads`
|
|
* call, effectively removing the conditional check.
|
|
*
|
|
* Example 1:
|
|
* if (xxxx == 0) { __syncthreads(); }
|
|
* =>
|
|
* __syncthreads();
|
|
*
|
|
* Example 2:
|
|
* if (xxxx > 0) { __syncthreads(); }
|
|
* =>
|
|
* if (xxxx > 0) { __syncthreads(); }
|
|
*/
|
|
class CudaSyncThreadsDropIfThenElsePass : public BlockPass {
|
|
public:
|
|
CudaSyncThreadsDropIfThenElsePass()
|
|
: BlockPass("cuda_sync_threads_drop_ifthenelse") {}
|
|
|
|
LogicalResult Run(ir::stmt::BlockRef block) override;
|
|
};
|
|
std::unique_ptr<BlockPass> CreateCudaSyncThreadsDropIfThenElsePass();
|
|
|
|
} // namespace optim
|
|
} // namespace cinn
|