changes

Author: tapspatel

include/ttmlir/Conversion/TTNNToEmitC/EmitCConversion.h

@@ -659,6 +659,9 @@ struct EmitCTypeConverter<::ttnn::TensorMemoryLayout> {
     case ttnn::TensorMemoryLayout::WidthSharded:
       rso << "WIDTH_SHARDED";
       return buf;
+    case ttnn::TensorMemoryLayout::NDSharded:
+      rso << "ND_SHARDED";
+      return buf;
     }
 
     llvm_unreachable("Unknown ttnn::TensorMemoryLayout");

include/ttmlir/Conversion/TTNNToEmitPy/EmitPyConversion.h

@@ -696,6 +696,9 @@ struct EmitPyTypeConverter<::ttnn::TensorMemoryLayout> {
     case ttnn::TensorMemoryLayout::WidthSharded:
       rso << "WIDTH_SHARDED";
       break;
+    case ttnn::TensorMemoryLayout::NDSharded:
+      rso << "ND_SHARDED";
+      break;
     }
 
     return buf;

include/ttmlir/Dialect/TTNN/IR/TTNNOpsEnums.td

@@ -25,13 +25,15 @@ def TTNN_TensorMemoryLayout_Interleaved : I32EnumAttrCase<"Interleaved", 0, "int
 def TTNN_TensorMemoryLayout_HeightSharded : I32EnumAttrCase<"HeightSharded", 2, "height_sharded">;
 def TTNN_TensorMemoryLayout_WidthSharded : I32EnumAttrCase<"WidthSharded", 3, "width_sharded">;
 def TTNN_TensorMemoryLayout_BlockSharded : I32EnumAttrCase<"BlockSharded", 4, "block_sharded">;
+def TTNN_TensorMemoryLayout_NDSharded : I32EnumAttrCase<"NDSharded", 5, "nd_sharded">;
 
 def TTNN_TensorMemoryLayout : I32EnumAttr<"TensorMemoryLayout", "TTNN Tensor Memory Layout",
                            [
                             TTNN_TensorMemoryLayout_Interleaved,
                             TTNN_TensorMemoryLayout_HeightSharded,
                             TTNN_TensorMemoryLayout_WidthSharded,
                             TTNN_TensorMemoryLayout_BlockSharded,
+                            TTNN_TensorMemoryLayout_NDSharded,
                            ]> {
   let genSpecializedAttr = 0;
   let cppNamespace = "::mlir::tt::ttnn";

include/ttmlir/OpModel/TTNN/TTNNOpModel.h

@@ -128,7 +128,7 @@ template <>
 struct OpModel<SiluOp> : UnaryEltwiseOpModel<SiluOp> {};
 
 template <>
-struct OpModel<MishOp> : UnaryEltwiseOpModel<MishOp> {};
+struct OpModel<MishOp> : UnaryEltwiseWithFastApproxModeOpModel<MishOp> {};
 
 template <>
 struct OpModel<RsqrtOp> : UnaryEltwiseWithFastApproxModeOpModel<RsqrtOp> {};

include/ttmlir/Target/TTNN/types.fbs

@@ -7,6 +7,7 @@ enum TensorMemoryLayout: ushort {
   HeightSharded,
   WidthSharded,
   BlockSharded,
+  NDSharded,
 }
 
 enum StorageType: ushort {

include/ttmlir/Target/Utils/MLIRToFlatbuffer.h

@@ -220,6 +220,8 @@ toFlatbuffer(FlatbufferObjectCache &, ttnn::TensorMemoryLayout memLayout) {
     return ::tt::target::ttnn::TensorMemoryLayout::WidthSharded;
   case ttnn::TensorMemoryLayout::BlockSharded:
     return ::tt::target::ttnn::TensorMemoryLayout::BlockSharded;
+  case ttnn::TensorMemoryLayout::NDSharded:
+    return ::tt::target::ttnn::TensorMemoryLayout::NDSharded;
   }
 }
 

lib/OpModel/TTNN/Conversion.cpp

@@ -331,6 +331,8 @@ getTensorMemoryLayout(const TensorMemoryLayout tensorMemoryLayout) {
     return ::tt::tt_metal::TensorMemoryLayout::WIDTH_SHARDED;
   case TensorMemoryLayout::BlockSharded:
     return ::tt::tt_metal::TensorMemoryLayout::BLOCK_SHARDED;
+  case TensorMemoryLayout::NDSharded:
+    return ::tt::tt_metal::TensorMemoryLayout::ND_SHARDED;
   }
 }
 TensorMemoryLayout
@@ -344,6 +346,8 @@ getTensorMemoryLayout(const ::tt::tt_metal::TensorMemoryLayout memLayout) {
     return TensorMemoryLayout::WidthSharded;
   case ::tt::tt_metal::TensorMemoryLayout::BLOCK_SHARDED:
     return TensorMemoryLayout::BlockSharded;
+  case ::tt::tt_metal::TensorMemoryLayout::ND_SHARDED:
+    return TensorMemoryLayout::NDSharded;
   }
 }
 

lib/OpModel/TTNN/TTNNOpModel.cpp

@@ -948,7 +948,7 @@ template struct UnaryEltwiseOpModel<ReciprocalOp>;
 template struct UnaryEltwiseOpModel<CbrtOp>;
 template struct UnaryEltwiseOpModel<BitwiseNotOp>;
 template struct UnaryEltwiseOpModel<SiluOp>;
-template struct UnaryEltwiseOpModel<MishOp>;
+template struct UnaryEltwiseWithFastApproxModeOpModel<MishOp>;
 template struct UnaryEltwiseWithFastApproxModeOpModel<Log1pOp>;
 template struct UnaryEltwiseOpModel<Expm1Op>;
 template struct UnaryEltwiseWithFastApproxModeOpModel<RsqrtOp>;

runtime/lib/ttnn/operations/eltwise/unary/unary.cpp

@@ -182,7 +182,7 @@ void run(const ::tt::target::ttnn::EltwiseUnaryOp *op,
     break;
   }
   case ::tt::target::ttnn::EltwiseUnaryOpType::Mish: {
-    runEltwiseUnaryOp(op, tensorPool, ::ttnn::mish);
+    runEltwiseUnaryWithFastAndApproximateModeOp(op, tensorPool, ::ttnn::mish);
     break;
   }
   case ::tt::target::ttnn::EltwiseUnaryOpType::Sin: {

runtime/lib/ttnn/utils/utils.cpp

@@ -257,6 +257,8 @@ ::ttnn::TensorMemoryLayout toTTNNTensorMemoryLayout(
     return ::ttnn::TensorMemoryLayout::WIDTH_SHARDED;
   case ::tt::target::ttnn::TensorMemoryLayout::BlockSharded:
     return ::ttnn::TensorMemoryLayout::BLOCK_SHARDED;
+  case ::tt::target::ttnn::TensorMemoryLayout::NDSharded:
+    return ::ttnn::TensorMemoryLayout::ND_SHARDED;
   }
 }
 
@@ -271,6 +273,8 @@ fromTTNNTensorMemoryLayout(::ttnn::TensorMemoryLayout tensorMemoryLayout) {
     return ::tt::target::ttnn::TensorMemoryLayout::WidthSharded;
   case ::ttnn::TensorMemoryLayout::BLOCK_SHARDED:
     return ::tt::target::ttnn::TensorMemoryLayout::BlockSharded;
+  case ::ttnn::TensorMemoryLayout::ND_SHARDED:
+    return ::tt::target::ttnn::TensorMemoryLayout::NDSharded;
   }
 }