support partial reduction

Author: zhczhong

include/gc/Dialect/Arith/Utils/EasyBuild.h

@@ -1,17 +1,10 @@
-//===- EasyBuild.h - Easy Arith IR Builder utilities ------------*- C++ -*-===//
+//===-- EasyBuild.h - DESC --------------------------------------*- C++ -*-===//
 //
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
-//
-// This header file defines the easy-build utilities for arith dialects. It
-// provides the utility functions, classes and operators to make it easir to
-// program arith dialect operations in C++
-//
-//===----------------------------------------------------------------------===//
-
 #ifndef MLIR_DIALECT_ARITH_UTILS_EASYBUILD_H
 #define MLIR_DIALECT_ARITH_UTILS_EASYBUILD_H
 #include "gc/IR/EasyBuild.h"
@@ -28,12 +21,8 @@ namespace impl {
 
 template <std::size_t size> struct ToFloatType {};
 
-template <> struct ToFloatType<4> {
-  using type = Float32Type;
-};
-template <> struct ToFloatType<8> {
-  using type = Float64Type;
-};
+template <> struct ToFloatType<4> { using type = Float32Type; };
+template <> struct ToFloatType<8> { using type = Float64Type; };
 
 inline Type getElementType(Value v) {
   auto type = v.getType();

include/gc/IR/EasyBuild.h

@@ -1,16 +1,10 @@
-//===- EasyBuild.h - Easy IR Builder utilities ------------------*- C++ -*-===//
+//===-- EasyBuild.h - DESC --------------------------------------*- C++ -*-===//
 //
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
-//
-// This header file defines the easy-build utilities core data structures for
-// building IR.
-//
-//===----------------------------------------------------------------------===//
-
 #ifndef MLIR_IR_EASYBUILD_H
 #define MLIR_IR_EASYBUILD_H
 #include "mlir/IR/Builders.h"

include/gc/IR/EasyBuildSCF.h

@@ -1,10 +1,11 @@
-//===- EasyBuildSCF.h - Easy IR Builder for general control flow *- C++ -*-===//
+//===-- EasyBuildSCF.h - DESC -----------------------------------*- C++ -*-===//
 //
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+
 //
 // This header file defines the helper classes, functions and macros to help to
 // build general structured control flow. Developers can use the utilities in

lib/gc/Transforms/DeepTileContractionNamedOp.cpp

@@ -1,9 +1,9 @@
 //===-- DeepTileContractionNamedOp.cpp - DESC -------------------*- C++ -*-===//
-// 
+//
 // This file is licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-// 
+//
 //===----------------------------------------------------------------------===//
 
 #include "./Tiling.hpp"
@@ -273,9 +273,19 @@ generateOuterLoop(RewriterBase &b, linalg::LinalgOp linalgOp,
       b.setInsertionPoint(currentOp);
       if (auto partialInterface =
               dyn_cast<PartialReductionOpInterface>(currentOp.getOperation())) {
+        for (auto [idx, tile] : llvm::enumerate(tileSizes)) {
+          if (isConstantIntValue(tile, 0)) {
+            tileSizes[idx] = loopRanges[idx].size;
+          }
+        }
+
+        SmallVector<OpFoldResult> newParallelDims;
+        for (auto i = 0UL; i < reductionDims.size(); i++) {
+          newParallelDims.push_back(getAsIndexOpFoldResult(b.getContext(), i));
+        }
         auto tilingResult = linalgX::tileAllUsingForall(
-            b, cast<PartialReductionOpInterface>(currentOp.getOperation()),
-            numThreads, tileSizes, std::nullopt);
+            b, cast<PartialReductionOpInterface>(currentOp.getOperation()), {},
+            tileSizes, newParallelDims, std::nullopt);
         if (failed(tilingResult))
           return failure();
         currentOp = dyn_cast<linalg::LinalgOp>(tilingResult->parallelTiledOp);