.

docs/source/conf.py

@@ -31,7 +31,6 @@
 
 import hidet
 hidet.option.cache_dir(os.path.join(hidet.option.get_cache_dir(), 'docs-cache'))
-hidet.utils.clear_op_cache()
 print('Build docs with under cache: {}'.format(hidet.option.get_cache_dir()))
 
 # -- Project information -----------------------------------------------------

docs/source/python_api/driver.rst

@@ -1,6 +1,6 @@
-hidet.driver
-------------
+hidet.drivers
+-------------
 
-.. automodule:: hidet.driver
+.. automodule:: hidet.drivers
   :members:
   :autosummary:

gallery/how-to-guides/add-new-operator-compute-definition.py

@@ -270,16 +270,15 @@ def demo_task():
 from hidet.ir.task import Task
 
 
-def run_task(task: Task, inputs: List[hidet.Tensor], outputs: List[hidet.Tensor]):
+def run_task(task: Task, inputs: List[hidet.Tensor]):
     """Run given task and print inputs and outputs"""
-    from hidet.runtime import CompiledFunction
+    from hidet.runtime import CompiledTask
 
     # build the task
-    func: CompiledFunction = hidet.driver.build_task(task, target='cpu')
-    params = inputs + outputs
+    func: CompiledTask = hidet.drivers.build_task(task, target='cpu')
 
     # run the compiled task
-    func(*params)
+    outputs = func.run_async(inputs)
 
     print('Task:', task.name)
     print('Inputs:')
@@ -317,7 +316,7 @@ def add_example():
     b: TensorNode = tensor_input(name='b', dtype='float32', shape=[5])
     c: TensorNode = compute(name='c', shape=[5], fcompute=lambda i: a[i] + b[i])
     task = Task(name='add', inputs=[a, b], outputs=[c])
-    run_task(task, [hidet.randn([5]), hidet.randn([5])], [hidet.empty([5])])
+    run_task(task, [hidet.randn([5]), hidet.randn([5])])
 
 
 add_example()
@@ -350,7 +349,7 @@ def reduce_sum_example():
         ),
     )
     task = Task('reduce_sum', inputs=[a], outputs=[b])
-    run_task(task, [hidet.randn([4, 3])], [hidet.empty([4])])
+    run_task(task, [hidet.randn([4, 3])])
 
 
 reduce_sum_example()
@@ -371,7 +370,7 @@ def arg_max_example():
         ),
     )
     task = Task('arg_max', inputs=[a], outputs=[b])
-    run_task(task, [hidet.randn([4, 3])], [hidet.empty([4], dtype='int64')])
+    run_task(task, [hidet.randn([4, 3])])
 
 
 arg_max_example()
@@ -391,7 +390,7 @@ def matmul_example():
         ),
     )
     task = Task('matmul', inputs=[a, b], outputs=[c])
-    run_task(task, [hidet.randn([3, 3]), hidet.randn([3, 3])], [hidet.empty([3, 3])])
+    run_task(task, [hidet.randn([3, 3]), hidet.randn([3, 3])])
 
 
 matmul_example()
@@ -411,7 +410,7 @@ def softmax_example():
     softmax = compute('softmax', shape=[3], fcompute=lambda i: exp_a[i] / exp_sum)
 
     task = Task('softmax', inputs=[a], outputs=[softmax])
-    run_task(task, [hidet.randn([3])], [hidet.empty([3])])
+    run_task(task, [hidet.randn([3])])
 
 
 softmax_example()

gallery/how-to-guides/add-new-operator-template-based.py

@@ -17,7 +17,7 @@
 import hidet
 from hidet.ir.compute import TensorNode, compute, reduce
 from hidet.ir.task import Task
-from hidet.ir.func import IRModule
+from hidet.ir.module import IRModule
 
 
 class BatchMatmulFp16Task(Task):
@@ -240,16 +240,6 @@ def demo_usage():
 
 demo_usage()
 
-# %%
-# Generated Source Code
-# ---------------------
-# If you are interested in the generated source code, here it is:
-
-a = hidet.randn([1, 2, 2], dtype='float16', device='cuda')
-b = hidet.randn([1, 2, 2], dtype='float16', device='cuda')
-op = BatchMatmulFp16Op(a, b)
-print(op.task_func.source(color=True))
-
 # %%
 # Summary
 # -------

python/hidet/backend/build.py

@@ -68,7 +68,7 @@ def run_compile_command(self, command: str, src_path, out_lib_path: str):
 
                 # write the compilation log
                 log_name = self.__class__.__name__.lower() + '_output.txt'
-                with open(os.path.join(out_lib_dir, log_name), 'w') as f:
+                with open(os.path.join(out_lib_dir, log_name), 'w', encoding='utf-8') as f:
                     output = '\n'.join([result.stdout.decode('utf-8').strip(), result.stderr.decode('utf-8').strip()])
                     f.write(output.strip())
 

python/hidet/graph/flow_graph.py

@@ -363,7 +363,7 @@ def build(self, *, space=0, allow_hook=False):
 
         Returns
         -------
-        ret: hidet.runtime.model.CompiledModel
+        ret: hidet.runtime.CompiledGraph
             The compiled model.
         """
         from hidet.drivers.build_graph import build_flow_graph

python/hidet/testing/models/resnet.py

@@ -10,8 +10,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from typing import Type, Union, List
-
-import hidet
 from hidet.graph import nn
 
 
@@ -112,31 +110,21 @@ def forward(self, x):
         return x
 
 
-def resnet18(batch_size=1, channels=3, height=224, width=224):
-    model = ResNet(block=BasicBlock, layers=[2, 2, 2, 2])
-    inputs = [hidet.randn([batch_size, channels, height, width])]
-    return model, inputs
+def resnet18():
+    return ResNet(block=BasicBlock, layers=[2, 2, 2, 2])
 
 
-def resnet34(batch_size=1, channels=3, height=224, width=224):
-    model = ResNet(block=BasicBlock, layers=[3, 4, 6, 3])
-    inputs = [hidet.randn([batch_size, channels, height, width])]
-    return model, inputs
+def resnet34():
+    return ResNet(block=BasicBlock, layers=[3, 4, 6, 3])
 
 
-def resnet50(batch_size=1, channels=3, height=224, width=224):
-    model = ResNet(block=Bottleneck, layers=[3, 4, 6, 3])
-    inputs = [hidet.randn([batch_size, channels, height, width])]
-    return model, inputs
+def resnet50():
+    return ResNet(block=Bottleneck, layers=[3, 4, 6, 3])
 
 
-def resnet101(batch_size=1, channels=3, height=224, width=224):
-    model = ResNet(block=Bottleneck, layers=[3, 4, 23, 3])
-    inputs = [hidet.randn([batch_size, channels, height, width])]
-    return model, inputs
+def resnet101():
+    return ResNet(block=Bottleneck, layers=[3, 4, 23, 3])
 
 
-def resnet152(batch_size=1, channels=3, height=224, width=224):
-    model = ResNet(block=Bottleneck, layers=[3, 8, 36, 3])
-    inputs = [hidet.randn([batch_size, channels, height, width])]
-    return model, inputs
+def resnet152():
+    return ResNet(block=Bottleneck, layers=[3, 8, 36, 3])

tests/models/test_gpt2.py

@@ -43,6 +43,7 @@ def test_gpt2(device: str, opt: bool):
         graph = hidet.graph.optimize(graph)
 
     compiled_model = graph.build()
+    compiled_model.save('./outs/compiled.hidet')
 
     generated_text = generate(
         compiled_model,

tests/unit_tests/test_dynamic_shape.py

@@ -55,7 +55,7 @@ def get_graph(seq: Union[int, str]) -> FlowGraph:
 @pytest.mark.parametrize('device', ['cuda'])
 @pytest.mark.parametrize('bs,h,w', [(1, 224, 224), (2, 224, 224), (1, 256, 256)])
 def test_resnet50(device, bs, h, w):
-    model, _ = hidet.testing.models.resnet50()
+    model = hidet.testing.models.resnet50()
     if device == 'cuda':
         model = model.cuda()
     else: