[Frontend] Fix ops conversion

Author: YWHyuk

PyTorchSimFrontend/mlir/mlir_codegen_backend.py

@@ -803,12 +803,11 @@ def where(condition, operand1, operand2, *args, var_info=None, **kwargs):
         cond_type = var_info[condition]
         operand_type = var_info[operand1]
         if cond_type[0] < tile_size:
-            condition = ops.broadcast(condition, operand_type[0])
+            condition = ops.broadcast(condition, tile_size)
         elif cond_type[0] > tile_size:
-            operand1 = ops.broadcast(operand1, operand_type[0])
-            operand2 = ops.broadcast(operand2, operand_type[0])
+            operand1 = ops.broadcast(operand1, cond_type[0])
+            operand2 = ops.broadcast(operand2, cond_type[0])
         tile_size, ret_type = var_info[operand1]
-
         shape = f"vector<{tile_size}x{ret_type}>" if tile_size > 1 else ret_type
         cond_shape = f"vector<{tile_size}xi1>," if tile_size > 1 else ""
         return f"arith.select %{condition}, %{operand1}, %{operand2} : {cond_shape} {shape}", [tile_size, ret_type]
@@ -1164,10 +1163,6 @@ def store(self, name: str, index: sympy.Expr, value, *args, **kwargs):
             # Todo. If tile_size is not same (i.e., view operation), we can't apply peephole optimization easily
             require_store = self.spad_buffer_dict[str(value)][1] != tile_size
 
-        if compute_vec_size < self.var_info[value][0]:
-            value = self.cse.generate(self.stores, f"vector.extract_strided_slice  %{value} {{offsets = [0], sizes = [{compute_vec_size}], strides = [1]}}: vector<{self.var_info[value][0]}x{self.var_info[value][1]}> to {vshape}")
-            self.register_var_info(value, [compute_vec_size, mlir_dtype])
-
         if require_store:
             # Define scratch pad buffer
             sram_var, sram_index_var = self.get_scratchpad_buffer(dtype, name, local_tile_desc, index)
@@ -1176,6 +1171,11 @@ def store(self, name: str, index: sympy.Expr, value, *args, **kwargs):
             _, operand_type = self.var_info[value]
             if mlir_dtype != operand_type:
                 value = ops.custom_cast(value, mlir_dtype)
+
+            if compute_vec_size < self.var_info[value][0]:
+                value = self.cse.generate(self.stores, f"vector.extract_strided_slice  %{value} {{offsets = [0], sizes = [{compute_vec_size}], strides = [1]}}: vector<{self.var_info[value][0]}x{self.var_info[value][1]}> to {vshape}")
+                self.register_var_info(value, [compute_vec_size, mlir_dtype])
+
             with self.override_buffer_cse(buffer=self.stores):
                 ops._store(value, sram_var, compute_index_var, tile_shape, buffer_name=name)
         else:

PyTorchSimFrontend/mlir/mlir_template.py

@@ -919,24 +919,23 @@ def store_reduction_epilogue(self, name, index, value):
                 compute_index_var = ",".join(partial_zero_var_list)
 
             with self.override_buffer_cse(buffer=self.reductions_suffix):
-                out = ops._load(partial_vec_size, mlir_dtype, sram_var, compute_index_var, partial_tile_shape)
+                out = ops._load(partial_vec_size, mlir_dtype, value, compute_index_var, partial_tile_shape)
                 ops._store(init_vec, value, compute_index_var, partial_tile_shape) # Clear the partial buffer to zero
 
                 # 2 step reduction
                 new_vec_size = 2
-                new_reduced_shape = f"<{new_vec_size}x{mlir_dtype}>"
+                new_reduced_shape = f"vector<{new_vec_size}x{mlir_dtype}>"
                 reduction_type = self.reduction_info[value][0]
-                out = ops.multi_reduction(out, init_vec, partial_vec_size, new_vec_size, reduction_type, partial_vshape, self.reduction_info[value][0], mlir_dtype)
+                out = ops.multi_reduction(out, init_vec2, partial_vec_size, new_vec_size, partial_vshape, reduction_type, mlir_dtype)
 
             out2 = self.cse.generate(self.reductions_suffix, f"vector.shuffle %{out}, %{out} [1, 0] : {new_reduced_shape}, {new_reduced_shape}")
             self.register_var_info(out2, [new_vec_size, mlir_dtype])
 
             with self.override_buffer_cse(buffer=self.reductions_suffix):
                 out = reduction_partial_combine_vec(self.reduction_info[value][0], out, out2)
 
-            if self.welford_reduce_out is not None:
-                # NOTE: It not a real welford algorithm... We just used E(X^2) - E(X)^2
-                with self.override_buffer_cse(buffer=self.reductions_suffix):
+                if self.welford_reduce_out is not None:
+                    # NOTE: It not a real welford algorithm... We just used E(X^2) - E(X)^2
                     divider = ops.constant(float(self.reduction_axis_size), "f32")
                     if self.buffer_types[name][1] > 1:
                         divider_vec = ops.broadcast(divider, new_vec_size)
@@ -955,9 +954,9 @@ def store_reduction_epilogue(self, name, index, value):
                         m2 = ops.mul(variance, divider_vec)
                         out = m2
 
-            final_zero_var_list[-1] = f"%{body_index_var}"
-            final_compute_index_var = ",".join(final_zero_var_list)
-            ops._store(out, sram_var, final_compute_index_var, final_tile_shape, buffer_name=name)
+                final_zero_var_list[-1] = f"%{body_index_var}"
+                final_compute_index_var = ",".join(final_zero_var_list)
+                ops._store(out, sram_var, final_compute_index_var, final_tile_shape, buffer_name=name)
 
         # MVOUT Encoding
         # Generate DMA instruction