From 34f1d5e565c0350866d2193d3d0ab10f780a65db Mon Sep 17 00:00:00 2001
From: Philippe Tillet <ptillet@g.harvard.edu>
Date: Mon, 2 Nov 2020 15:05:41 -0500
Subject: [PATCH] [CODEGEN] Fixed bug in 2D reductions

---
 lib/codegen/selection/generator.cc | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/lib/codegen/selection/generator.cc b/lib/codegen/selection/generator.cc
index 5813c1ff1..3de7b2c99 100644
--- a/lib/codegen/selection/generator.cc
+++ b/lib/codegen/selection/generator.cc
@@ -977,11 +977,13 @@ void generator::visit_reduce_inst(ir::reduce_inst* x) {
 
   analysis::data_layout* arg_layout = layouts_->get(arg);
   if(auto* L = dynamic_cast<analysis::scanline_layout*>(arg_layout)){
-    bool can_optimize = true;
+    bool can_optimize = L->get_rank() == 1;
+    /*
     for(size_t r = 0; r < L->get_rank(); r++){
       if(r != axis)
         can_optimize = can_optimize && (L->mts(r) == L->get_shape()[r]);
     }
+    */
     if(can_optimize){
       Value *thread_acc = nullptr;
       // reduce within thread