PyKernel: get_noc_addr_from_bank_id needed for first eltwise PyKern…

…el (#2050) This PR implements the last TTKernel op needed by [eltwise kernel example](#1984)
tenstorrent · Feb 3, 2025 · 32f72e6 · 32f72e6
1 parent b654322
commit 32f72e6
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 1 deletion.
diff --git a/include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td b/include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td
@@ -479,6 +479,16 @@ def TTKernel_GetNocAddrXYOp : TTKernel_Op<"get_noc_addr_xy"> {
     let results = (outs TTKernel_NocAddr:$nocAddr);
 }
 
+def TTKernel_GetNocAddrFromBankIDOp: TTKernel_Op<"get_noc_addr_from_bank_id"> {
+    let summary = "GetNocAddrFromBankID";
+    let description = [{
+      GetNocAddrFromBankID api
+    }];
+
+    let arguments = (ins I32:$bank_id, I32:$bankAddressOffset);
+    let results = (outs TTKernel_NocAddr:$nocAddr);
+}
+
 def TTKernel_GetNocAddrOp : TTKernel_Op<"get_noc_addr"> {
     let summary = "GetNocAddr";
     let description = [{

diff --git a/lib/Conversion/TTKernelToEmitC/TTKernelToEmitC.cpp b/lib/Conversion/TTKernelToEmitC/TTKernelToEmitC.cpp
@@ -313,6 +313,13 @@ class TTMetalToEmitCOpaqueRewriter : public OpConversionPattern<SourceOp> {
       template_args.push_back(
           emitc::OpaqueAttr::get(op.getContext(), "uint32_t"));
       return ArrayAttr::get(op.getContext(), template_args);
+    } else if constexpr (std::is_same_v<SourceOp,
+                                        ttkernel::GetNocAddrFromBankIDOp>) {
+      SmallVector<Attribute, 1> template_args;
+
+      template_args.push_back(
+          emitc::OpaqueAttr::get(op.getContext(), "true")); // default to DRAM
+      return ArrayAttr::get(op.getContext(), template_args);
     }
     return ArrayAttr();
   }
@@ -613,7 +620,8 @@ class ConvertTTKernelToEmitCPass
           TTMetalToEmitCOpaqueRewriter<ttkernel::GetWritePtrOp>,
           TTMetalToEmitCOpaqueRewriter<ttkernel::GetReadPtrOp>,
           TTMetalToEmitCOpaqueRewriter<ttkernel::GetTileSizeOp>,
-          TTMetalToEmitCOpaqueRewriter<ttkernel::GetCompileArgValOp>>(
+          TTMetalToEmitCOpaqueRewriter<ttkernel::GetCompileArgValOp>,
+          TTMetalToEmitCOpaqueRewriter<ttkernel::GetNocAddrFromBankIDOp>>(
           typeConverter, funcOp.getContext());
 
       patterns.add<TTMetalToEmitCOpaqueRewriter<ttkernel::GetNocAddrXYOp>>(

diff --git a/test/ttmlir/Conversion/TTKernelToEmitC/ttkernel.mlir b/test/ttmlir/Conversion/TTKernelToEmitC/ttkernel.mlir
@@ -22,6 +22,10 @@ module attributes {} {
     %4 = "ttkernel.get_noc_addr_xy"(%c0_i32, %c0_i32, %c262208_i32) : (i32, i32, i32) -> !ttkernel.noc_addr
     // CHECK: emitc.call_opaque "noc_async_read"[[C:.*]]
     "ttkernel.noc_async_read"(%4, %c262432_i32, %c32_i32) : (!ttkernel.noc_addr, i32, i32) -> ()
+    %bank_id = arith.constant 1 : i32
+    %addr_offset = arith.constant 262400 : i32
+    %noc_addr = "ttkernel.get_noc_addr_from_bank_id"(%bank_id, %addr_offset) : (i32, i32) -> !ttkernel.noc_addr
+    // CHECK: = emitc.call_opaque "get_noc_addr_from_bank_id"({{.*}}) {template_args = [#emitc.opaque<"true">]}
     // CHECK: emitc.call_opaque "noc_async_read_barrier"[[C:.*]]
     "ttkernel.noc_async_read_barrier"() : () -> ()
     "ttkernel.return"() : () -> ()