【CINN】add correctness check for bind shedule (#58770)

Courtesy-Xs · web-flow · commit 30e4175ea750 · 2023-11-17T11:08:08.000+08:00
* add correctness check for bind shedule

* fix gramma error

* fix extent as_int32 bug

* fix limits bug

* delete unused comments

* fix ir_schedule test case

* fix test case bug
diff --git a/paddle/cinn/auto_schedule/cost_model/feature_extractor_test.cc b/paddle/cinn/auto_schedule/cost_model/feature_extractor_test.cc
@@ -89,6 +89,7 @@ TEST(FeatureExtractor, SimpleAssign) {
   ASSERT_EQ(to_check[29], slog(3));
 }
 
+#ifdef CINN_WITH_CUDA
 TEST(FeatureExtractor, MatrixMultiply) {
   Context::Global().ResetNameId();
 #ifdef CINN_WITH_CUDA
@@ -162,6 +163,6 @@ TEST(FeatureExtractor, MatrixMultiply) {
   // GpuBind loop
   ASSERT_EQ(to_check[37], slog(out_loop));
 }
-
+#endif
 }  // namespace auto_schedule
 }  // namespace cinn
diff --git a/paddle/cinn/backends/ir_schedule_test.cc b/paddle/cinn/backends/ir_schedule_test.cc
@@ -690,6 +690,7 @@ void test_unroll(void* _args, int32_t num_args)
   ASSERT_EQ(utils::Trim(target_code), utils::Trim(source_code));
 }
 
+#ifdef CINN_WITH_CUDA
 TEST(IrSchedule, bind) {
   Context::Global().ResetNameId();
   Expr M(32);
@@ -733,6 +734,7 @@ function test_bind (_A, _B)
 }
 )ROC"));
 }
+#endif
 
 TEST(IrSchedule, simple_compute_at) {
   Context::Global().ResetNameId();
diff --git a/paddle/cinn/common/target.cc b/paddle/cinn/common/target.cc
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 #ifdef CINN_WITH_CUDA
+#include <cuda.h>
 #include <cuda_runtime_api.h>
 #include <driver_types.h>
 #endif
@@ -20,6 +21,7 @@
 
 #include <sstream>
 
+#include "paddle/cinn/backends/cuda_util.h"
 #include "paddle/cinn/common/target.h"
 #include "paddle/cinn/runtime/cinn_runtime.h"
 
@@ -206,6 +208,18 @@ int GetMaxThreads() {
   return max_threads;
 }
 
+std::array<int, 3> GetCUDAMaxBlockDims() {
+  std::array<int, 3> ret;
+#ifdef CINN_WITH_CUDA
+  cudaDeviceProp prop;
+  CUDA_CALL(cudaGetDeviceProperties(&prop, 0));
+  ret[0] = prop.maxThreadsDim[0];
+  ret[1] = prop.maxThreadsDim[1];
+  ret[2] = prop.maxThreadsDim[2];
+#endif
+  return ret;
+}
+
 int GetMaxBlocks() {
   // cudaDeviceGetAttribute ( int* value, cudaDeviceAttr attr, int  device )
   int max_blocks = 1;
@@ -222,6 +236,18 @@ int GetMaxBlocks() {
   return max_blocks;
 }
 
+std::array<int, 3> GetCUDAMaxGridDims() {
+  std::array<int, 3> ret;
+#ifdef CINN_WITH_CUDA
+  cudaDeviceProp prop;
+  CUDA_CALL(cudaGetDeviceProperties(&prop, 0));
+  ret[0] = prop.maxGridSize[0];
+  ret[1] = prop.maxGridSize[1];
+  ret[2] = prop.maxGridSize[2];
+#endif
+  return ret;
+}
+
 const Target &DefaultTarget() {
 #ifdef CINN_WITH_CUDA
   return DefaultNVGPUTarget();
diff --git a/paddle/cinn/common/target.h b/paddle/cinn/common/target.h
@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <array>
 #include <ostream>
 #include <string>
 #include <vector>
@@ -113,6 +114,10 @@ int GetMaxThreads();
 
 int GetMaxBlocks();
 
+std::array<int, 3> GetCUDAMaxBlockDims();
+
+std::array<int, 3> GetCUDAMaxGridDims();
+
 std::ostream& operator<<(std::ostream& os, Target::Arch arch);
 
 }  // namespace common
diff --git a/paddle/cinn/ir/schedule/ir_schedule.cc b/paddle/cinn/ir/schedule/ir_schedule.cc
@@ -28,6 +28,7 @@
 #include "paddle/cinn/common/cas.h"
 #include "paddle/cinn/common/common.h"
 #include "paddle/cinn/common/ir_util.h"
+#include "paddle/cinn/common/target.h"
 #include "paddle/cinn/ir/dy_schedule/ir_schedule.h"
 #include "paddle/cinn/ir/ir.h"
 #include "paddle/cinn/ir/ir_mutator.h"
@@ -364,9 +365,19 @@ void StScheduleImpl::Bind(const Expr& loop, const std::string& thread_axis) {
   CHECK(thread_axes.count(thread_axis))
       << "thread_axis " << thread_axis << " is not supported";
   int offset = thread_axis.back() - 'x';
+  const std::array<int, 3> kMaxBlockDims = common::GetCUDAMaxBlockDims();
+  const std::array<int, 3> kMaxGridDims = common::GetCUDAMaxGridDims();
+  auto check_offset = [&](const char& c) -> bool {
+    auto extent = loop.As<ir::For>()->extent.as_int32();
+    return extent <= (c == 'b' ? kMaxGridDims[offset] : kMaxBlockDims[offset]);
+  };
   if (thread_axis[0] == 'b') {
+    CHECK(check_offset(thread_axis[0]))
+        << "Invalid Bind! The extent of loop is out of range on grid size!\n";
     MutateForType(loop, ForType::GPUBlock, offset);
   } else {
+    CHECK(check_offset(thread_axis[0]))
+        << "Invalid Bind! The extent of loop is out of range on block size!\n";
     MutateForType(loop, ForType::GPUThread, offset);
   }
 }
diff --git a/paddle/cinn/ir/test/schedule_desc_test.cc b/paddle/cinn/ir/test/schedule_desc_test.cc
@@ -779,6 +779,7 @@ TEST_F(TestScheduleDesc, StepKind_ReverseComputeInline) {
   CheckReplayResult(ir_sch, ir_sch.GetTraceDesc());
 }
 
+#ifdef CINN_WITH_CUDA
 TEST_F(TestScheduleDesc, StepKind_Bind) {
   lowered_funcs = LowerCompute({32, 128}, target);
   ir::IRSchedule ir_sch = MakeIRSchedule(lowered_funcs);
@@ -794,6 +795,7 @@ TEST_F(TestScheduleDesc, StepKind_Bind) {
   CheckReplayResult(ir_sch, trace);
   CheckReplayResult(ir_sch, ir_sch.GetTraceDesc());
 }
+#endif
 
 TEST_F(TestScheduleDesc, StepKind_Rfactor) {
   Expr M(32);

Original file line number	Diff line number	Diff line change
`@@ -89,6 +89,7 @@ TEST(FeatureExtractor, SimpleAssign) {`
`89`	`89`	`ASSERT_EQ(to_check[29], slog(3));`
`90`	`90`	`}`
`91`	`91`
	`92`	`+#ifdef CINN_WITH_CUDA`
`92`	`93`	`TEST(FeatureExtractor, MatrixMultiply) {`
`93`	`94`	`Context::Global().ResetNameId();`
`94`	`95`	`#ifdef CINN_WITH_CUDA`
`@@ -162,6 +163,6 @@ TEST(FeatureExtractor, MatrixMultiply) {`
`162`	`163`	`// GpuBind loop`
`163`	`164`	`ASSERT_EQ(to_check[37], slog(out_loop));`
`164`	`165`	`}`
`165`		`-`
	`166`	`+#endif`
`166`	`167`	`} // namespace auto_schedule`
`167`	`168`	`} // namespace cinn`
Original file line number	Diff line number	Diff line change
`@@ -690,6 +690,7 @@ void test_unroll(void* _args, int32_t num_args)`
`690`	`690`	`ASSERT_EQ(utils::Trim(target_code), utils::Trim(source_code));`
`691`	`691`	`}`
`692`	`692`
	`693`	`+#ifdef CINN_WITH_CUDA`
`693`	`694`	`TEST(IrSchedule, bind) {`
`694`	`695`	`Context::Global().ResetNameId();`
`695`	`696`	`Expr M(32);`
`@@ -733,6 +734,7 @@ function test_bind (_A, _B)`
`733`	`734`	`}`
`734`	`735`	`)ROC"));`
`735`	`736`	`}`
	`737`	`+#endif`
`736`	`738`
`737`	`739`	`TEST(IrSchedule, simple_compute_at) {`
`738`	`740`	`Context::Global().ResetNameId();`
Original file line number	Diff line number	Diff line change
`@@ -779,6 +779,7 @@ TEST_F(TestScheduleDesc, StepKind_ReverseComputeInline) {`
`779`	`779`	`CheckReplayResult(ir_sch, ir_sch.GetTraceDesc());`
`780`	`780`	`}`
`781`	`781`
	`782`	`+#ifdef CINN_WITH_CUDA`
`782`	`783`	`TEST_F(TestScheduleDesc, StepKind_Bind) {`
`783`	`784`	`lowered_funcs = LowerCompute({32, 128}, target);`
`784`	`785`	`ir::IRSchedule ir_sch = MakeIRSchedule(lowered_funcs);`
`@@ -794,6 +795,7 @@ TEST_F(TestScheduleDesc, StepKind_Bind) {`
`794`	`795`	`CheckReplayResult(ir_sch, trace);`
`795`	`796`	`CheckReplayResult(ir_sch, ir_sch.GetTraceDesc());`
`796`	`797`	`}`
	`798`	`+#endif`
`797`	`799`
`798`	`800`	`TEST_F(TestScheduleDesc, StepKind_Rfactor) {`
`799`	`801`	`Expr M(32);`