chore: import upstream snapshot with attribution

test/cpp/cinn/optim/CMakeLists.txt

@@ -0,0 +1,3 @@
+cinn_cc_test(test_cast_simplify SRCS cast_simplify_test.cc DEPS cinncore)
+cinn_cc_test(test_replace_cross_thread_reduction SRCS
+             replace_cross_thread_reduction_test.cc DEPS cinncore)

test/cpp/cinn/optim/cast_simplify_test.cc

@@ -0,0 +1,62 @@
+// 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.
+
+#include <gtest/gtest.h>
+
+#include "paddle/cinn/ir/ir_printer.h"
+#include "paddle/cinn/ir/op/ir_operators.h"
+#include "paddle/cinn/optim/ir_simplify.h"
+#include "paddle/cinn/utils/string.h"
+namespace cinn::optim {
+
+TEST(CastSimplify, same_type) {
+  Var n("n");
+  Expr a = ir::Cast::Make(Int(32), n);
+  LOG(INFO) << n->type();
+  LOG(INFO) << a;
+  SimplifyCast(&a);
+  ASSERT_EQ(utils::GetStreamCnt(a), "n");
+}
+
+TEST(CastSimplify, Imm_int) {
+  Expr a = ir::Cast::Make(Int(64), Expr(1));
+  Expr c = ir::Cast::Make(Int(32), a);
+  LOG(INFO) << c;
+  SimplifyCast(&c);
+  LOG(INFO) << c;
+  ASSERT_EQ(utils::GetStreamCnt(c), "1");
+  ASSERT_EQ(c.type(), Int(32));
+}
+
+TEST(CastSimplify, Imm_double) {
+  Expr a = ir::Cast::Make(Float(64), Expr(2.33));
+  Expr c = ir::Cast::Make(Int(32), a);
+  LOG(INFO) << c;
+  SimplifyCast(&c);
+  LOG(INFO) << c;
+  ASSERT_EQ(utils::GetStreamCnt(c), "2");
+  ASSERT_EQ(c.type(), Int(32));
+}
+
+TEST(CastSimplify, Imm_uint) {
+  Expr a = ir::Cast::Make(UInt(64), Expr(1));
+  Expr c = ir::Cast::Make(UInt(32), a);
+  LOG(INFO) << c;
+  SimplifyCast(&c);
+  LOG(INFO) << c;
+  ASSERT_EQ(utils::GetStreamCnt(c), "1");
+  ASSERT_EQ(c.type(), UInt(32));
+}
+
+}  // namespace cinn::optim

test/cpp/cinn/optim/replace_cross_thread_reduction_test.cc

@@ -0,0 +1,102 @@
+// 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.
+
+#include "paddle/cinn/optim/replace_cross_thread_reduction.h"
+
+#include <gtest/gtest.h>
+
+#include <string>
+#include <vector>
+
+#include "paddle/cinn/cinn.h"
+#include "paddle/cinn/ir/ir.h"
+#include "paddle/cinn/ir/ir_printer.h"
+#include "paddle/cinn/ir/op/ir_operators.h"
+#include "paddle/cinn/ir/schedule/ir_schedule.h"
+#include "paddle/cinn/ir/utils/stmt_converter.h"
+#include "paddle/cinn/utils/string.h"
+
+namespace cinn {
+namespace optim {
+
+TEST(CrossThreadReductionReplacer, basic) {
+#ifdef CINN_WITH_CUDA
+  Context::Global().ResetNameId();
+  Placeholder<float> A("A", {Expr(64), Expr(128)});
+  Target target = cinn::common::DefaultNVGPUTarget();
+  Module::Builder builder("reduce_sum", target);
+  Var reduce_j(128, "reduce_j");
+  ir::Tensor B = Compute(
+      {Expr(64)},
+      [&](Var i) { return lang::ReduceSum(A(i, reduce_j), {reduce_j}); },
+      "B");
+  ast_gen_ius::TensorGroup tensor_group({A, B});
+  auto func = lang::LowerToAst("reduce_sum", {A, B}, &tensor_group);
+  VLOG(6) << "original func\n" << func;
+
+  ir::Expr expr_func_body = ir::ConvertStmtBlockToExprBlock(func->body_block);
+  ir::ModuleExpr mod_expr({expr_func_body});
+  ir::IRSchedule ir_sch(mod_expr);
+
+  ir_sch.Bind(ir_sch.GetLoops("B")[0], "blockIdx.x");
+  ir_sch.Bind(ir_sch.GetLoops("B")[1], "threadIdx.x");
+
+  ir::Expr block = ir_sch.GetBlock("B");
+  block.As<ir::ScheduleBlockRealize>()
+      ->schedule_block.As<ir::ScheduleBlock>()
+      ->reduce_method = ir::BlockReduceMethod();
+
+  ir::Expr func_body = ir_sch.GetModule().GetExprs()[0];
+  std::vector<ir::Argument> args{
+      ir::Argument(ir::Var("A"), ir::Argument::IO::kInput),
+      ir::Argument(ir::Var("B"), ir::Argument::IO::kOutput)};
+  auto new_func = ir::_LoweredFunc_::Make("test_func", args, func_body, {});
+  VLOG(6) << "After Bind: " << new_func->body;
+
+  ReplaceCrossThreadReduction(new_func);
+  VLOG(6) << "After ReplaceCrossThreadReduction: " << new_func->body;
+
+  EXPECT_EQ(utils::GetStreamCnt(new_func->body), utils::Trim(R"ROC({
+  ScheduleBlock(root)
+  {
+    {
+      thread_bind[blockIdx.x] for (i, 0, 64)
+      {
+        ScheduleBlock(B__reduce_init)
+        {
+          i0 = axis.bind(i)
+          {
+            B__reduce_init[i0] = 0.00000000f
+          }
+        }
+        thread_bind[threadIdx.x] for (reduce_j, 0, 128)
+        {
+          ScheduleBlock(B)
+          {
+            i0_0, i1 = axis.bind(i, reduce_j)
+            {
+              B[i0_0] = cinn_block_reduce_sum_fp32(A[i0_0, i1], _Buffer_<cinn_buffer_t*: 32>(shm32__fp32_reduce), false)
+            }
+          }
+        }
+      }
+    }
+  }
+}
+)ROC"));
+#endif
+}
+
+}  // namespace optim
+}  // namespace cinn