Fix conv3d signature to pass device as arg due to metal 0d3558f

Author: jameszianxuTT

lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp

@@ -1667,10 +1667,19 @@ class Conv3dOpConversionPattern
     auto outputDtype =
         dtypeAttr ? dtypeAttr.getValue() : ttcore::DataType::BFloat16;
 
+    // Emit SSA operands in ODS order (input, weight, bias, device) to
+    // maintain correct index mapping, then arrange args in the C++ API
+    // call order (input, weight, device, bias, ...).
+    auto inputAttr = emitter.emit(srcOp.getInput());
+    auto weightAttr = emitter.emit(srcOp.getWeight());
+    auto biasAttr = emitter.emit(srcOp.getBias());
+    auto deviceAttr = emitter.emit(srcOp.getDevice());
+
     llvm::SmallVector<mlir::Attribute> args{
-        emitter.emit(srcOp.getInput()),
-        emitter.emit(srcOp.getWeight()),
-        emitter.emit(srcOp.getBias()),
+        inputAttr,
+        weightAttr,
+        deviceAttr,
+        biasAttr,
         emitter.emitConv3dConfig(srcOp.getConv3dConfig()),
         emitter.emit(outputDtype),
         emitter.emit(srcOp.getOutChannels()),

lib/OpModel/TTNN/TTNNOpModel.cpp

@@ -5111,6 +5111,7 @@ llvm::Expected<OpConstraints> OpModel<Conv3dOp>::getOpConstraints(
   auto conv3dOpQuery = [=, &specs]() {
     return ::ttnn::graph::query_op_constraints(
         ::ttnn::experimental::conv3d, device, specs.inputSpec, specs.weightSpec,
+        std::optional<::tt::tt_metal::distributed::MeshDevice *>(device),
         specs.biasSpec, specs.config, specs.dtype, specs.outputChannels,
         specs.kernelSize, specs.stride, specs.padding,
         std::array<uint32_t, 3>{1, 1, 1}, specs.paddingMode, specs.groups,
@@ -5156,6 +5157,7 @@ llvm::Expected<size_t> OpModel<Conv3dOp>::getOpRuntime(
   auto conv3dOpRuntime = [=, &specs]() {
     return ::ttnn::graph::query_op_runtime(
         ::ttnn::experimental::conv3d, device, specs.inputSpec, specs.weightSpec,
+        std::optional<::tt::tt_metal::distributed::MeshDevice *>(device),
         specs.biasSpec, specs.config, specs.dtype, specs.outputChannels,
         specs.kernelSize, specs.stride, specs.padding,
         std::array<uint32_t, 3>{1, 1, 1}, specs.paddingMode, specs.groups,

runtime/lib/ttnn/operations/conv/conv3d.cpp

@@ -94,10 +94,10 @@ void run(const ::tt::target::ttnn::Conv3dOp *op, ProgramContext &context) {
              "Memory config must exist for device tensors");
 
   ::ttnn::Tensor out = ::ttnn::experimental::conv3d(
-      input, weight, bias, conv3dConfig, outputDtype, op->out_channels(),
-      kernelSize, stride, padding, std::array<uint32_t, 3>{1, 1, 1},
-      op->padding_mode()->str(), op->groups(), outputMemoryConfig,
-      deviceComputeConfig);
+      input, weight, &targetDevice, bias, conv3dConfig, outputDtype,
+      op->out_channels(), kernelSize, stride, padding,
+      std::array<uint32_t, 3>{1, 1, 1}, op->padding_mode()->str(), op->groups(),
+      outputMemoryConfig, deviceComputeConfig);
 
   tensorPool.insertTTNNTensorAndValidate(op->out(), out);
 }