diff --git a/include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td b/include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td index 93dc7beb9c..c0e5faff27 100644 --- 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 index 9172b6a6d9..0624336d98 100644 --- a/lib/Conversion/TTKernelToEmitC/TTKernelToEmitC.cpp +++ b/lib/Conversion/TTKernelToEmitC/TTKernelToEmitC.cpp @@ -313,6 +313,13 @@ class TTMetalToEmitCOpaqueRewriter : public OpConversionPattern { 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) { + SmallVector 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, TTMetalToEmitCOpaqueRewriter, TTMetalToEmitCOpaqueRewriter, - TTMetalToEmitCOpaqueRewriter>( + TTMetalToEmitCOpaqueRewriter, + TTMetalToEmitCOpaqueRewriter>( typeConverter, funcOp.getContext()); patterns.add>( diff --git a/test/ttmlir/Conversion/TTKernelToEmitC/ttkernel.mlir b/test/ttmlir/Conversion/TTKernelToEmitC/ttkernel.mlir index 30c6e9cc82..030d09c546 100644 --- 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"() : () -> ()