From e30e35deeb950a593d7030180b61805a80d353db Mon Sep 17 00:00:00 2001
From: Foke Singh <foke.singh@deepkapha.com>
Date: Mon, 20 Mar 2023 15:31:44 +0000
Subject: [PATCH] Update benchmark example to use new user sharding hints and
 drop deprecated module options.

---
 examples/benchmark/main.zig | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/examples/benchmark/main.zig b/examples/benchmark/main.zig
index 8d7604e..cbbe330 100644
--- a/examples/benchmark/main.zig
+++ b/examples/benchmark/main.zig
@@ -5,11 +5,16 @@ const flags = @import("tigerbeetle/flags");
 
 const async_ = asynk.async_;
 
+// set log level to debug to print the generated IR
+pub const std_options = .{
+    .log_level = .debug,
+};
+
 /// Model definition
 const Benchmark = struct {
     pub fn forward(self: Benchmark, a: zml.Tensor, b: zml.Tensor) zml.Tensor {
         _ = self;
-        return a.dot(b, .{.k});
+        return a.withSharding(.{.k}).dot(b.withSharding(.{.k}), .{.k}).withSharding(.{.m});
     }
 };
 
@@ -68,7 +73,7 @@ pub fn asyncMain() !void {
                         deviceKind,
                     });
                     // we only list 1 CPU device
-                    if (target == .cpu) break;
+                    if (target == .cpu and platform.sharding().num_partitions == 1) break;
                 }
             }
         }