D2M: cleanup and preparation for bcast (#5375)

### Ticket
None

### Problem description
Break down the bcast PR into smaller ones, this one is mainly for
cleanups.

### What's changed
* Remove redundant `bcast_type` that clutters the IR.
* Extract the logic that generates the compute region of
`linalg.generic` into its own function.
* Remove support for `FuncOp` from the insert DST register pass as it's
getting in the way more and more.
* Update the simple `FuncOp` tests to full `d2m.generic` equivalents
(OMG so painful, but it's a one-time cost).

### Checklist
- [x] New/Existing tests provide coverage for changes

Signed-off-by: wenbinlyuTT <[email protected]>

Author: wenbinlyuTT

include/ttmlir/Dialect/TTKernel/IR/TTKernelOpsEnums.td

@@ -105,22 +105,6 @@ def TTKernel_ReduceType : I32EnumAttr<"ReduceType", "TTKernel Reduce Types",
   let cppNamespace = "::mlir::tt::ttkernel";
 }
 
-def TTKernel_BcastTypeNone : I32EnumAttrCase<"None", 0, "bcast_type_none">;
-def TTKernel_BcastTypeCol : I32EnumAttrCase<"Col", 1, "bcast_type_col">;
-def TTKernel_BcastTypeRow : I32EnumAttrCase<"Row", 2, "bcast_type_row">;
-def TTKernel_BcastTypeScalar : I32EnumAttrCase<"Scalar", 3, "bcast_type_scalar">;
-
-def TTKernel_BcastType : I32EnumAttr<"BcastType", "TTKernel broadcast types",
-                         [
-                           TTKernel_BcastTypeNone,
-                           TTKernel_BcastTypeCol,
-                           TTKernel_BcastTypeRow,
-                           TTKernel_BcastTypeScalar
-                         ]> {
-  let genSpecializedAttr = 0;
-  let cppNamespace = "::mlir::tt::ttkernel";
-}
-
 def TTKernel_ReduceDimRow : I32EnumAttrCase<"Row", 0, "reduce_dim_row">;
 def TTKernel_ReduceDimCol : I32EnumAttrCase<"Col", 1, "reduce_dim_col">;
 def TTKernel_ReduceDimScalar
@@ -134,6 +118,22 @@ def TTKernel_ReduceDim
   let cppNamespace = "::mlir::tt::ttkernel";
 }
 
+def TTKernel_BcastTypeNone :   I32EnumAttrCase<"None",   0, "none">;
+def TTKernel_BcastTypeCol :    I32EnumAttrCase<"Col",    1, "col">;
+def TTKernel_BcastTypeRow :    I32EnumAttrCase<"Row",    2, "row">;
+def TTKernel_BcastTypeScalar : I32EnumAttrCase<"Scalar", 3, "scalar">;
+
+def TTKernel_BcastType : I32EnumAttr<"BcastType", "TTKernel Broadcast Types",
+                         [
+                           TTKernel_BcastTypeNone,
+                           TTKernel_BcastTypeCol,
+                           TTKernel_BcastTypeRow,
+                           TTKernel_BcastTypeScalar
+                         ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::tt::ttkernel";
+}
+
 def TTKernel_ArgTypeCBPort : I32EnumAttrCase<"CBPort", 0, "cb_port">;
 def TTKernel_ArgTypeBufferAddress : I32EnumAttrCase<"BufferAddress", 1, "buffer_address">;
 def TTKernel_ArgTypeSemaphore : I32EnumAttrCase<"Semaphore", 2, "semaphore">;

lib/Conversion/TTIRToD2M/TTIRToD2M.cpp

@@ -16,16 +16,13 @@
 #include "ttmlir/Dialect/D2M/IR/D2MGenericRegionOps.h"
 #include "ttmlir/Dialect/TTCore/IR/TTCoreOpsTypes.h"
 #include "ttmlir/Dialect/TTIR/IR/TTIROps.h"
-#include "ttmlir/Dialect/TTIR/Utils/Utils.h"
 #include "ttmlir/Dialect/TTNN/IR/TTNNOpsAttrs.h"
 
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/IRMapping.h"
 #include "mlir/IR/TypeRange.h"
 #include "mlir/IR/ValueRange.h"
 #include "mlir/Transforms/DialectConversion.h"
@@ -404,7 +401,26 @@ class D2MNamedElementwiseRewriter final
       std::is_same_v<TileOp, d2m::TileLtzOp> ||
       std::is_same_v<TileOp, d2m::TileLezOp>;
 
-private:
+  void createComputeRegion(mlir::OpBuilder &bbBuilder, mlir::Location bbLoc,
+                           mlir::ValueRange bbArgs,
+                           mlir::ConversionPatternRewriter &rewriter,
+                           mlir::Location loc, const size_t numInputs,
+                           const size_t numOutputs) const {
+    mlir::ValueRange operands = bbArgs.take_front(numInputs);
+    mlir::TypeRange resultTypes = bbArgs.take_back(numOutputs);
+
+    mlir::Value yield;
+    if constexpr (isComparisonOp) {
+      // For comparison ops, first subtract then compare with zero.
+      yield = bbBuilder.create<d2m::TileSubOp>(loc, resultTypes, operands);
+      yield = bbBuilder.create<TileOp>(loc, resultTypes, yield);
+    } else {
+      yield = bbBuilder.create<TileOp>(loc, resultTypes, operands);
+    }
+
+    bbBuilder.create<mlir::linalg::YieldOp>(bbLoc, yield);
+  }
+
   LogicalResult
   matchAndRewrite(ConcreteOp op, typename ConcreteOp::Adaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const final {
@@ -467,25 +483,8 @@ class D2MNamedElementwiseRewriter final
             linalgIteratorTypes,
             [&](mlir::OpBuilder &bbBuilder, mlir::Location bbLoc,
                 mlir::ValueRange bbArgs) {
-              mlir::Value yield;
-
-              if constexpr (isComparisonOp) {
-                // For comparison ops, first subtract then compare with zero.
-                mlir::Value subResult = bbBuilder.create<d2m::TileSubOp>(
-                    loc, /*resultTypes=*/bbArgs.take_back(numOutputs),
-                    /*operands=*/bbArgs.take_front(numInputs));
-                yield = bbBuilder.create<TileOp>(
-                    loc, /*resultTypes=*/bbArgs.take_back(numOutputs),
-                    /*operands=*/subResult);
-              } else {
-                // For regular elementwise ops, create TileOp directly.
-                yield = bbBuilder.create<TileOp>(
-                    loc,
-                    /* resultTypes */ bbArgs.take_back(numOutputs).getTypes(),
-                    /* operands */ bbArgs.take_front(numInputs));
-              }
-
-              bbBuilder.create<mlir::linalg::YieldOp>(bbLoc, yield);
+              createComputeRegion(bbBuilder, bbLoc, bbArgs, rewriter, loc,
+                                  numInputs, numOutputs);
             });
 
         rewriter.create<d2m::YieldOp>(loc, linalgGeneric->getResults());

lib/Conversion/TTIRToLinalg/TTIRToLinalg.cpp

@@ -8,10 +8,7 @@
 #include "ttmlir/Dialect/TTIR/Utils/Utils.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
-#include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/IR/Attributes.h"
@@ -20,7 +17,6 @@
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Value.h"
 #include "mlir/IR/ValueRange.h"
-#include "mlir/Pass/Pass.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/SmallVector.h"
@@ -664,7 +660,7 @@ class BroadcastOpConversionPattern
     auto zerosConst =
         rewriter.create<tosa::ConstOp>(op.getLoc(), resultType, zerosAttr);
 
-    // Multiply by ones to implicitly broadcast
+    // Add by zeros to implicitly broadcast.
     auto result = rewriter.create<tosa::AddOp>(op.getLoc(), resultType, input,
                                                zerosConst);